Grasscutters/mitmproxy
3
0
Fork 0
mirror of https://github.com/Grasscutters/mitmproxy.git synced 2024-11-26 18:18:25 +00:00
Code Issues Projects Releases Wiki Activity

socks: optionally fail early

Browse Source
This commit is contained in:
Maximilian Hils 2015-07-03 02:45:12 +02:00
parent 9aaf10120d
commit 880c66fe48
2 changed files with 13 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
netlib/socks.py
View File

@ -69,11 +69,16 @@ class ClientGreeting(object):
)
@classmethod
def from_file(cls, f):
def from_file(cls, f, fail_early=False):
"""
:param fail_early: If true, a SocksError will be raised if the first byte does not indicate socks5.
"""
ver, nmethods = struct.unpack("!BB", f.safe_read(2))
methods = array.array("B")
methods.fromstring(f.safe_read(nmethods))
return cls(ver, methods.tolist())
client_greeting = cls(ver, [])
if fail_early:
client_greeting.assert_socks5()
client_greeting.methods.fromstring(f.safe_read(nmethods))
return client_greeting
def to_file(self, f):
f.write(struct.pack("!BB", self.ver, len(self.methods)))
@ -115,7 +120,7 @@ class Message(object):
self.ver = ver
self.msg = msg
self.atyp = atyp
self.addr = addr
self.addr = tcp.Address.wrap(addr)
def assert_socks5(self):
if self.ver != VERSION.SOCKS5:

3
test/test_socks.py
View File

@ -44,6 +44,9 @@ def test_client_greeting_assert_socks5():
else:
assert False
raw = tutils.treader("XX")
tutils.raises(socks.SocksError, socks.ClientGreeting.from_file, raw, fail_early=True)
def test_server_greeting():
raw = tutils.treader("\x05\x02")