diff options
| -rw-r--r-- | netlib/socks.py | 41 | ||||
| -rw-r--r-- | test/test_socks.py | 57 |
2 files changed, 91 insertions, 7 deletions
diff --git a/netlib/socks.py b/netlib/socks.py index 5a73c61a..eef98f5c 100644 --- a/netlib/socks.py +++ b/netlib/socks.py @@ -6,7 +6,6 @@ from . import tcp, utils class SocksError(Exception): - def __init__(self, code, message): super(SocksError, self).__init__(message) self.code = code @@ -17,21 +16,18 @@ VERSION = utils.BiDi( SOCKS5=0x05 ) - CMD = utils.BiDi( CONNECT=0x01, BIND=0x02, UDP_ASSOCIATE=0x03 ) - ATYP = utils.BiDi( IPV4_ADDRESS=0x01, DOMAINNAME=0x03, IPV6_ADDRESS=0x04 ) - REP = utils.BiDi( SUCCEEDED=0x00, GENERAL_SOCKS_SERVER_FAILURE=0x01, @@ -44,7 +40,6 @@ REP = utils.BiDi( ADDRESS_TYPE_NOT_SUPPORTED=0x08, ) - METHOD = utils.BiDi( NO_AUTHENTICATION_REQUIRED=0x00, GSSAPI=0x01, @@ -58,14 +53,27 @@ class ClientGreeting(object): def __init__(self, ver, methods): self.ver = ver - self.methods = methods + self.methods = array.array("B") + self.methods.extend(methods) + + def assert_socks5(self): + if self.ver != VERSION.SOCKS5: + if self.ver == ord("G") and len(self.methods) == ord("E"): + guess = "Probably not a SOCKS request but a regular HTTP request. " + else: + guess = "" + + raise SocksError( + REP.GENERAL_SOCKS_SERVER_FAILURE, + guess + "Invalid SOCKS version. Expected 0x05, got 0x%x" % self.ver + ) @classmethod def from_file(cls, f): ver, nmethods = struct.unpack("!BB", f.safe_read(2)) methods = array.array("B") methods.fromstring(f.safe_read(nmethods)) - return cls(ver, methods) + return cls(ver, methods.tolist()) def to_file(self, f): f.write(struct.pack("!BB", self.ver, len(self.methods))) @@ -79,6 +87,18 @@ class ServerGreeting(object): self.ver = ver self.method = method + def assert_socks5(self): + if self.ver != VERSION.SOCKS5: + if self.ver == ord("H") and self.method == ord("T"): + guess = "Probably not a SOCKS request but a regular HTTP response. " + else: + guess = "" + + raise SocksError( + REP.GENERAL_SOCKS_SERVER_FAILURE, + guess + "Invalid SOCKS version. Expected 0x05, got 0x%x" % self.ver + ) + @classmethod def from_file(cls, f): ver, method = struct.unpack("!BB", f.safe_read(2)) @@ -97,6 +117,13 @@ class Message(object): self.atyp = atyp self.addr = addr + def assert_socks5(self): + if self.ver != VERSION.SOCKS5: + raise SocksError( + REP.GENERAL_SOCKS_SERVER_FAILURE, + "Invalid SOCKS version. Expected 0x05, got 0x%x" % self.ver + ) + @classmethod def from_file(cls, f): ver, msg, rsv, atyp = struct.unpack("!BBBB", f.safe_read(4)) diff --git a/test/test_socks.py b/test/test_socks.py index a9db4706..eb5d55f9 100644 --- a/test/test_socks.py +++ b/test/test_socks.py @@ -9,6 +9,7 @@ def test_client_greeting(): raw = tutils.treader("\x05\x02\x00\xBE\xEF") out = StringIO() msg = socks.ClientGreeting.from_file(raw) + msg.assert_socks5() msg.to_file(out) assert out.getvalue() == raw.getvalue()[:-1] @@ -18,10 +19,37 @@ def test_client_greeting(): assert 0xEF not in msg.methods +def test_client_greeting_assert_socks5(): + raw = tutils.treader("\x00\x00") + msg = socks.ClientGreeting.from_file(raw) + tutils.raises(socks.SocksError, msg.assert_socks5) + + raw = tutils.treader("HTTP/1.1 200 OK" + " " * 100) + msg = socks.ClientGreeting.from_file(raw) + try: + msg.assert_socks5() + except socks.SocksError as e: + assert "Invalid SOCKS version" in str(e) + assert "HTTP" not in str(e) + else: + assert False + + raw = tutils.treader("GET / HTTP/1.1" + " " * 100) + msg = socks.ClientGreeting.from_file(raw) + try: + msg.assert_socks5() + except socks.SocksError as e: + assert "Invalid SOCKS version" in str(e) + assert "HTTP" in str(e) + else: + assert False + + def test_server_greeting(): raw = tutils.treader("\x05\x02") out = StringIO() msg = socks.ServerGreeting.from_file(raw) + msg.assert_socks5() msg.to_file(out) assert out.getvalue() == raw.getvalue() @@ -29,10 +57,33 @@ def test_server_greeting(): assert msg.method == 0x02 +def test_server_greeting_assert_socks5(): + raw = tutils.treader("HTTP/1.1 200 OK" + " " * 100) + msg = socks.ServerGreeting.from_file(raw) + try: + msg.assert_socks5() + except socks.SocksError as e: + assert "Invalid SOCKS version" in str(e) + assert "HTTP" in str(e) + else: + assert False + + raw = tutils.treader("GET / HTTP/1.1" + " " * 100) + msg = socks.ServerGreeting.from_file(raw) + try: + msg.assert_socks5() + except socks.SocksError as e: + assert "Invalid SOCKS version" in str(e) + assert "HTTP" not in str(e) + else: + assert False + + def test_message(): raw = tutils.treader("\x05\x01\x00\x03\x0bexample.com\xDE\xAD\xBE\xEF") out = StringIO() msg = socks.Message.from_file(raw) + msg.assert_socks5() assert raw.read(2) == "\xBE\xEF" msg.to_file(out) @@ -43,6 +94,12 @@ def test_message(): assert msg.addr == ("example.com", 0xDEAD) +def test_message_assert_socks5(): + raw = tutils.treader("\xEE\x01\x00\x03\x0bexample.com\xDE\xAD\xBE\xEF") + msg = socks.Message.from_file(raw) + tutils.raises(socks.SocksError, msg.assert_socks5) + + def test_message_ipv4(): # Test ATYP=0x01 (IPV4) raw = tutils.treader("\x05\x01\x00\x01\x7f\x00\x00\x01\xDE\xAD\xBE\xEF") |