def test_custom_error(self):
code = 0xFF
message = 'Custom error message'
self._check_error(socks.SocksError(message), socks.SocksError, None,
message)
self._check_error(socks.SocksError(message=message), socks.SocksError,
None, message)
self._check_error(socks.SocksError(code=code), socks.SocksError, code,
'')
self._check_error(socks.SocksError(message, code=code),
socks.SocksError, code, message)
self._check_error(socks.SocksError(message=message, code=code),
socks.SocksError, code, message)
def test_end_to_end_successful_relay(self):
class Proto(object):
data = b''
lost = []
def dataReceived(self, d):
self.data = self.data + d
def connectionLost(self, reason):
self.lost.append(reason)
the_proto = Proto()
dis = []
def on_disconnect(error_message):
dis.append(error_message)
sm = socks._SocksMachine(
'CONNECT',
u'1.2.3.4',
443,
on_disconnect=on_disconnect,
create_connection=lambda a, p: the_proto,
)
sm.connection()
sm.feed_data(b'\x05')
sm.feed_data(b'\x00')
# reply with success, port 0x1234
sm.feed_data(b'\x05\x00\x00\x01\x00\x00\x00\x00\x12\x34')
# now some data that should get relayed
sm.feed_data(b'this is some relayed data')
# should *not* have disconnected
self.assertEqual(0, len(dis))
self.assertTrue(the_proto.data, b"this is some relayed data")
sm.disconnected(socks.SocksError("it's fine"))
self.assertEqual(1, len(Proto.lost))
self.assertTrue("it's fine" in str(Proto.lost[0]))