def test_get_data_fragment_with_padding(): udp_payload = get_random_bytes(DATA_FRAG_PAYLOAD_SIZE - 1) f = FragmentGenerator(udp_payload) fragment = f.get_data_fragment() _, frag_byte, _, payload = cut(fragment, FRAG_ID_SIZE, FRAG_FLAG_SIZE, 1) assert b2i(frag_byte) & FragmentGenerator.LAST_FRAG_FLAG assert b2i(frag_byte) & FragmentGenerator.PADDING_FLAG assert b2i(frag_byte) == 0b1100_0000 assert payload[0:DATA_FRAG_PAYLOAD_SIZE - 1] == udp_payload assert len( fragment) == DATA_FRAG_PAYLOAD_SIZE + FRAG_ID_SIZE + FRAG_FLAG_SIZE udp_payload = get_random_bytes(1) f = FragmentGenerator(udp_payload) fragment = f.get_data_fragment() _, frag_byte, _, payload = cut(fragment, FRAG_ID_SIZE, FRAG_FLAG_SIZE, 2) assert b2i(frag_byte) & FragmentGenerator.LAST_FRAG_FLAG assert b2i(frag_byte) & FragmentGenerator.PADDING_FLAG assert b2i(frag_byte) == 0b1100_0000 assert payload[0:1] == udp_payload assert len( fragment) == DATA_FRAG_PAYLOAD_SIZE + FRAG_ID_SIZE + FRAG_FLAG_SIZE
def test_get_data_fragment_2_fragments(): udp_payload = get_random_bytes(400) fragmented_payload = bytes() f = FragmentGenerator(udp_payload) fragment = f.get_data_fragment() msg_id, is_last, fragment_id, payload = parse_fragment(fragment) fragmented_payload += payload assert msg_id == f.message_id assert not is_last assert fragment_id == 0 fragment = f.get_data_fragment() msg_id, is_last, fragment_id, payload = parse_fragment(fragment) fragmented_payload += payload assert msg_id == f.message_id assert is_last assert fragment_id == 1 assert udp_payload == fragmented_payload
def test_get_data_fragment_too_many(): udp_payload = get_random_bytes(400) f = FragmentGenerator(udp_payload) try: for _ in range(1000): f.get_data_fragment() assert False except ValueError: assert True
def test_parse_fragment(): udp_payload = get_random_bytes(120) f = FragmentGenerator(udp_payload) fragment = f.get_data_fragment() msg_id, is_last, fragment_id, payload = parse_fragment(fragment) assert msg_id == f.message_id assert is_last assert fragment_id == 0 assert udp_payload == payload
def recv_response(self, response): self.last_interaction = time() """Turns the response into a MixMessage and saves its fragments for later sending. """ frag_gen = FragmentGenerator(response) while frag_gen: print(self, "Data", "<-", len(frag_gen.udp_payload)) fragment = frag_gen.get_data_fragment() packet = fragment + get_random_bytes(MIX_COUNT * CTR_PREFIX_LEN) ChannelExit.to_mix.append(DATA_MSG_FLAG + i2b(self.in_chan_id, CHAN_ID_SIZE) + packet)
def test_get_data_fragment(): udp_payload = get_random_bytes(DATA_FRAG_PAYLOAD_SIZE) f = FragmentGenerator(udp_payload) fragment = f.get_data_fragment() _, frag_byte, payload = cut(fragment, FRAG_ID_SIZE, FRAG_FLAG_SIZE) assert b2i(frag_byte) & FragmentGenerator.LAST_FRAG_FLAG assert b2i(frag_byte) == 0b0100_0000 assert payload[0:DATA_FRAG_PAYLOAD_SIZE] == udp_payload assert len( fragment) == DATA_FRAG_PAYLOAD_SIZE + FRAG_ID_SIZE + FRAG_FLAG_SIZE