mirror of
https://github.com/Grasscutters/mitmproxy.git
synced 2024-11-23 00:01:36 +00:00
socks: add assert_socks5 method
This commit is contained in:
parent
9131b96b63
commit
9aaf10120d
@ -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))
|
||||
|
@ -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")
|
||||
|
Loading…
Reference in New Issue
Block a user