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)