diff options
Diffstat (limited to 'netlib')
| -rw-r--r-- | netlib/socks.py | 58 | ||||
| -rw-r--r-- | netlib/tcp.py | 3 |
2 files changed, 60 insertions, 1 deletions
diff --git a/netlib/socks.py b/netlib/socks.py index 51ad1c63..57ccd1be 100644 --- a/netlib/socks.py +++ b/netlib/socks.py @@ -10,7 +10,6 @@ class SocksError(Exception): super(SocksError, self).__init__(message) self.code = code - VERSION = utils.BiDi( SOCKS4=0x04, SOCKS5=0x05 @@ -47,6 +46,10 @@ METHOD = utils.BiDi( NO_ACCEPTABLE_METHODS=0xFF ) +USERNAME_PASSWORD_VERSION = utils.BiDi( + DEFAULT=0x01 +) + class ClientGreeting(object): __slots__ = ("ver", "methods") @@ -113,6 +116,59 @@ class ServerGreeting(object): f.write(struct.pack("!BB", self.ver, self.method)) +class UsernamePasswordAuth(object): + __slots__ = ("ver", "username", "password") + + def __init__(self, ver, username, password): + self.ver = ver + self.username = username + self.password = password + + def assert_authver1(self): + if self.ver != USERNAME_PASSWORD_VERSION.DEFAULT: + raise SocksError( + 0, + "Invalid auth version. Expected 0x01, got 0x%x" % self.ver + ) + + @classmethod + def from_file(cls, f): + ver, ulen = struct.unpack("!BB", f.safe_read(2)) + username = f.safe_read(ulen) + plen, = struct.unpack("!B", f.safe_read(1)) + password = f.safe_read(plen) + return cls(ver, username.decode(), password.decode()) + + def to_file(self, f): + f.write(struct.pack("!BB", self.ver, len(self.username))) + f.write(self.username.encode()) + f.write(struct.pack("!B", len(self.password))) + f.write(self.password.encode()) + + +class UsernamePasswordAuthResponse(object): + __slots__ = ("ver", "status") + + def __init__(self, ver, status): + self.ver = ver + self.status = status + + def assert_authver1(self): + if self.ver != USERNAME_PASSWORD_VERSION.DEFAULT: + raise SocksError( + 0, + "Invalid auth version. Expected 0x01, got 0x%x" % self.ver + ) + + @classmethod + def from_file(cls, f): + ver, status = struct.unpack("!BB", f.safe_read(2)) + return cls(ver, status) + + def to_file(self, f): + f.write(struct.pack("!BB", self.ver, self.status)) + + class Message(object): __slots__ = ("ver", "msg", "atyp", "addr") diff --git a/netlib/tcp.py b/netlib/tcp.py index 6423888a..574f3845 100644 --- a/netlib/tcp.py +++ b/netlib/tcp.py @@ -458,9 +458,11 @@ class _Connection(object): def __init__(self, connection): if connection: self.connection = connection + self.peer_address = Address(connection.getpeername()) self._makefile() else: self.connection = None + self.peer_address = None self.rfile = None self.wfile = None @@ -701,6 +703,7 @@ class TCPClient(_Connection): 'Error connecting to "%s": %s' % (self.address.host, err)) self.connection = connection + self.peer_address = Address(connection.getpeername()) self._makefile() def settimeout(self, n): |