def test_add_dynamic_size(self): d = Defragmenter() d.add_dynamic_size(10, 2, 2) ret = d.get_message() self.assertIsNone(ret) d.add_data(10, bytearray( b'\xee\xee' + # header bytes b'\x00\x00' + # remaining length # next message b'\xff\xff' + # header bytes b'\x00\x01' + # remaining length b'\xf0')) ret = d.get_message() self.assertIsNotNone(ret) msg_type, data = ret self.assertEqual(10, msg_type) self.assertEqual(bytearray(b'\xee\xee\x00\x00'), data) ret = d.get_message() self.assertIsNotNone(ret) msg_type, data = ret self.assertEqual(10, msg_type) self.assertEqual(bytearray(b'\xff\xff\x00\x01\xf0'), data) ret = d.get_message() self.assertIsNone(ret)
def test_add_dynamic_size(self): d = Defragmenter() d.add_dynamic_size(10, 2, 2) ret = d.get_message() self.assertIsNone(ret) d.add_data( 10, bytearray(b'\xee\xee' + # header bytes b'\x00\x00' + # remaining length # next message b'\xff\xff' + # header bytes b'\x00\x01' + # remaining length b'\xf0')) ret = d.get_message() self.assertIsNotNone(ret) msg_type, data = ret self.assertEqual(10, msg_type) self.assertEqual(bytearray(b'\xee\xee\x00\x00'), data) ret = d.get_message() self.assertIsNotNone(ret) msg_type, data = ret self.assertEqual(10, msg_type) self.assertEqual(bytearray(b'\xff\xff\x00\x01\xf0'), data) ret = d.get_message() self.assertIsNone(ret)
def test_add_dynamic_size_with_incomplete_payload(self): d = Defragmenter() d.add_dynamic_size(10, 2, 2) d.add_data(10, bytearray(b'\xee\xee\x00\x01')) self.assertIsNone(d.get_message()) d.add_data(10, bytearray(b'\x99')) msg_type, data = d.get_message() self.assertEqual(10, msg_type) self.assertEqual(bytearray(b'\xee\xee\x00\x01\x99'), data)
def test_add_dynamic_size_with_incomplete_payload(self): d = Defragmenter() d.add_dynamic_size(10, 2, 2) d.add_data(10, bytearray(b'\xee\xee\x00\x01')) self.assertIsNone(d.get_message()) d.add_data(10, bytearray(b'\x99')) msg_type, data = d.get_message() self.assertEqual(10, msg_type) self.assertEqual(bytearray(b'\xee\xee\x00\x01\x99'), data)
def process(self, state): """Connect to a server.""" sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) sock.settimeout(5) sock.connect((self.hostname, self.port)) # disable Nagle - we handle buffering and flushing ourselves sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) # allow for later buffering of writes to the socket sock = BufferedSocket(sock) defragmenter = Defragmenter() defragmenter.add_static_size(ContentType.alert, 2) defragmenter.add_static_size(ContentType.change_cipher_spec, 1) defragmenter.add_dynamic_size(ContentType.handshake, 1, 3) state.msg_sock = MessageSocket(sock, defragmenter) state.msg_sock.version = self.version
def test_add_dynamic_size_with_two_streams(self): d = Defragmenter() d.add_dynamic_size(9, 0, 3) d.add_dynamic_size(10, 2, 2) d.add_data(10, bytearray(b'\x44\x44\x00\x04')) d.add_data(9, bytearray(b'\x00\x00\x02')) self.assertIsNone(d.get_message()) d.add_data(9, bytearray(b'\x09'*2)) d.add_data(10, bytearray(b'\x10'*4)) msg_type, data = d.get_message() self.assertEqual(msg_type, 9) self.assertEqual(data, bytearray(b'\x00\x00\x02\x09\x09')) msg_type, data = d.get_message() self.assertEqual(msg_type, 10) self.assertEqual(data, bytearray(b'\x44'*2 + b'\x00\x04' + b'\x10'*4))
def test_add_dynamic_size_with_two_streams(self): d = Defragmenter() d.add_dynamic_size(9, 0, 3) d.add_dynamic_size(10, 2, 2) d.add_data(10, bytearray(b'\x44\x44\x00\x04')) d.add_data(9, bytearray(b'\x00\x00\x02')) self.assertIsNone(d.get_message()) d.add_data(9, bytearray(b'\x09' * 2)) d.add_data(10, bytearray(b'\x10' * 4)) msg_type, data = d.get_message() self.assertEqual(msg_type, 9) self.assertEqual(data, bytearray(b'\x00\x00\x02\x09\x09')) msg_type, data = d.get_message() self.assertEqual(msg_type, 10) self.assertEqual(data, bytearray(b'\x44' * 2 + b'\x00\x04' + b'\x10' * 4))
def test_add_dynamic_size_with_incomplete_header(self): d = Defragmenter() d.add_dynamic_size(10, 2, 2) d.add_data(10, bytearray(b'\xee')) self.assertIsNone(d.get_message()) d.add_data(10, bytearray(b'\xee')) self.assertIsNone(d.get_message()) d.add_data(10, bytearray(b'\x00')) self.assertIsNone(d.get_message()) d.add_data(10, bytearray(b'\x00')) ret = d.get_message() self.assertIsNotNone(ret) msg_type, data = ret self.assertEqual(10, msg_type) self.assertEqual(bytearray(b'\xee\xee\x00\x00'), data)
def test_add_dynamic_size_with_incomplete_header(self): d = Defragmenter() d.add_dynamic_size(10, 2, 2) d.add_data(10, bytearray(b'\xee')) self.assertIsNone(d.get_message()) d.add_data(10, bytearray(b'\xee')) self.assertIsNone(d.get_message()) d.add_data(10, bytearray(b'\x00')) self.assertIsNone(d.get_message()) d.add_data(10, bytearray(b'\x00')) ret = d.get_message() self.assertIsNotNone(ret) msg_type, data = ret self.assertEqual(10, msg_type) self.assertEqual(bytearray(b'\xee\xee\x00\x00'), data)
def test_add_dynamic_size_with_invalid_offset(self): d = Defragmenter() with self.assertRaises(ValueError): d.add_dynamic_size(1, -1, 2)
def test_add_dynamic_size_with_double_type(self): d = Defragmenter() d.add_dynamic_size(1, 0, 1) with self.assertRaises(ValueError): d.add_dynamic_size(1, 2, 2)
def test_add_dynamic_size_with_invalid_offset(self): d = Defragmenter() with self.assertRaises(ValueError): d.add_dynamic_size(1, -1, 2)
def test_add_dynamic_size_with_double_type(self): d = Defragmenter() d.add_dynamic_size(1, 0, 1) with self.assertRaises(ValueError): d.add_dynamic_size(1, 2, 2)