commit 4485657e3840f0ea9d9fc96623dd7f50a476a79d Author: Sean Robinson seankrobinson@gmail.com Date: Mon Dec 17 20:02:22 2012 -0700
Add Socks class to handle a socket connection through a SOCKS5 proxy
I attempted to keep this generic and open enough that IPV6 support can be easily added at a later date.
Signed-off-by: Sean Robinson seankrobinson@gmail.com --- test/network.py | 172 +++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 files changed, 172 insertions(+), 0 deletions(-)
diff --git a/test/network.py b/test/network.py index 72bfca7..6ab17f5 100644 --- a/test/network.py +++ b/test/network.py @@ -6,8 +6,20 @@ the tor network.
ProxyError - Base error for proxy issues. +- SocksError - Reports problems returned by the SOCKS proxy. + + Socks - Communicate through a SOCKS5 proxy with a socket interface """
+import socket +import struct + +import stem.util.connection + +SOCKS5_NOAUTH_GREETING = (0x05, 0x01, 0x00) +SOCKS5_NOAUTH_RESPONSE = (0x05, 0x00) +SOCKS5_CONN_BY_IPV4 = (0x05, 0x01, 0x00, 0x01) +SOCKS5_CONN_BY_NAME = (0x05, 0x01, 0x00, 0x03) + class ProxyError(Exception): """ Base error for proxy issues. """
@@ -39,3 +51,163 @@ class SocksError(ProxyError): if self.code in self._ERROR_MESSAGE: code = self.code return "[%s] %s" % (code, self._ERROR_MESSAGE[code]) + +class Socks(socket.socket): + """ + A **socket.socket**-like interface through a SOCKS5 proxy connection. + Tor does not support proxy authentication, so neither does this class. + + This class supports the context manager protocol. When used this way, the + socket will automatically close when leaving the context. An example: + + :: + + from test.network import Socks + + with Socks(('127.0.0.1', 9050)) as socks: + socks.settimeout(2) + socks.connect(('www.torproject.org', 443)) + """ + + def __init__(self, proxy_addr, family = socket.AF_INET, + type_ = socket.SOCK_STREAM, proto = 0, _sock = None): + """ + Creates a SOCKS5-aware socket which will route connections through the + proxy_addr SOCKS5 proxy. Currently, only IPv4 TCP connections are + supported, so the defaults for family and type_ are your best option. + + :param tuple proxy_addr: address of the SOCKS5 proxy, for IPv4 this + contains (host, port) + :param int family: address family of the socket + :param int type_: address type of the socket (see **socket.socket** for + more information about family and type_) + + :returns: :class:`~test.network.Socks` + """ + + _socket_socket.__init__(self, family, type_, proto, _sock) + self._proxy_addr = proxy_addr + + def __enter__(self, *args, **kwargs): + return self + + def __exit__(self, exit_type, value, traceback): + self.close() + return False + + def _recvall(self, expected_size): + """ + Returns expected number bytes from the socket, or dies trying. + + :param int expected_size: number of bytes to return + + :returns: + * **str** in Python 2 (bytes is str) + * **bytes** in Python 3 + + :raises: + * :class:`socket.error` for socket errors + * :class:`test.SocksError` if the received data was more that expected + """ + + while True: + response = self.recv(expected_size * 2) + + if len(response) == 0: + raise socket.error("socket closed unexpectedly?") + elif len(response) == expected_size: + return response + elif len(response) > expected_size: + raise SocksError(0x01) + + def _ints_to_bytes(self, integers): + """ + Returns a byte string converted from integers. + + :param list integers: list of ints to convert + + :returns: + * **str** in Python 2 (bytes is str) + * **bytes** in Python 3 + """ + + if bytes is str: + bytes_ = ''.join([chr(x) for x in integers]) # Python 2 + else: + bytes_ = bytes(integers) # Python 3 + return bytes_ + + def _bytes_to_ints(self, bytes_): + """ + Returns a tuple of integers converted from a string (Python 2) or + bytes (Python 3). + + :param str,bytes bytes_: byte string to convert + + :returns: **list** of ints + """ + + try: + integers = [ord(x) for x in bytes_] # Python 2 + except TypeError: + integers = [x for x in bytes_] # Python 3 + return tuple(integers) + + def _pack_string(self, string_): + """ + Returns a packed string for sending over a socket. + + :param str string_: string to convert + + :returns: + * **str** in Python 2 (bytes is str) + * **bytes** in Python 3 + """ + + try: + return struct.pack(">%ss" % len(string_), string_) + except struct.error: + # Python 3: encode str to bytes + return struct.pack(">%ss" % len(string_), string_.encode()) + + def connect(self, address): + + """ + Establishes a connection to address through the SOCKS5 proxy. + + :param tuple address: target address, for IPv4 this contains + (host, port) + + :raises: :class:`test.SocksError` for any errors + """ + + socket.socket.connect(self, (self._proxy_addr[0], self._proxy_addr[1])) + # ask for non-authenticated connection + self.sendall(self._ints_to_bytes(SOCKS5_NOAUTH_GREETING)) + response = self._bytes_to_ints(self._recvall(2)) + if response != SOCKS5_NOAUTH_RESPONSE: + raise SocksError(0x01) + + if stem.util.connection.is_valid_ip_address(address[0]): + header = self._ints_to_bytes(SOCKS5_CONN_BY_IPV4) + header = header + socket.inet_aton(address[0]) + else: + # As a last gasp, try connecting by name + header = self._ints_to_bytes(SOCKS5_CONN_BY_NAME) + header = header + self._ints_to_bytes([len(address[0])]) + header = header + self._pack_string(address[0]) + + header = header + struct.pack(">H", address[1]) + self.sendall(header) + response = self._bytes_to_ints(self._recvall(10)) + # check the status byte + if response[1] != 0x00: + raise SocksError(response[1]) + + def connect_ex(self, address): + """ + Not Implemented. + """ + + raise NotImplementedError +