def test_protocol_basic(self): packets_received = [] def callback(packet): packets_received.append(packet) f = UDPFrame() f.data = b"Hello, world" f.checksum = 0x01 f.destination_port = 9000 f.source_port = 8000 phandler = StreamProtocolHandler(UDPFrame, callback) packet_bytes = f.pack() rem_bytes = packet_bytes while len(rem_bytes) > 0: chunk = rem_bytes[:5] rem_bytes = rem_bytes[5:] phandler.feed(chunk) assert (len(packets_received) == 1) rx = packets_received[0] assert rx.data == b"Hello, world" assert rx.checksum == 0x01 assert rx.destination_port == 9000 assert rx.source_port == 8000
def client(): def message_received(frame): print(frame) s = socket.socket() s.connect(('127.0.0.1', 7070)) proto_handler = StreamProtocolHandler(EchoProtocolFrame, message_received) while True: input = raw_input("Data to send: ") request = EchoProtocolFrame() request.frame_type = FRAME_TYPE_ECHO_REQUEST request.payload = input s.sendall(request.pack()) # to handle asynchronous messages, this would be done asynchronously proto_handler.feed(s.recv(1024))
def test_protocol_magic_scan_full_buffer(self): rx = [] def cb(packet): rx.append(packet) protocol_handler = StreamProtocolHandler(MagicSchema, cb) test_sequence = MagicSchema() test_sequence.value = -29939 pack = test_sequence.pack() # garbage with our bytes in the middle test_bytes = b'\x1A\x3fadbsfkasdf;aslkfjasd;f' + pack + b'\x00\x00asdfn234r' protocol_handler.feed(test_bytes) self.assertEqual(len(rx), 1) self.assertEqual(rx[0].value, -29939)
def test_protocol_violation(self): # verify that when there is a protocol violation, that state gets # cleaned up properly. For this, we have a dispatch field that # has an unmapped type byte rx = [] def bad_cb(packet): self.fail("Callback should not have been executed") def good_cb(packet): rx.append(packet) self.assertEqual(packet.body.value, -31415926) good_msg = ErrorCaseSchema() good_msg.body = MagicSchema() good_msg.body.value = -31415926 good_packet = good_msg.pack() bad_packet = b'\x01\x02\xAA\xAA' + good_packet phandler = StreamProtocolHandler(ErrorCaseSchema, bad_cb) phandler.feed(bad_packet) phandler.packet_callback = good_cb phandler.feed(good_packet) self.assertEqual(len(rx), 1)
class EchoTCPHandler(socketserver.BaseRequestHandler): def _frame_received(self, request_frame): # frame is an instance of EchoProtocolFrame print("Received %r" % request_frame) if request_frame.frame_type == FRAME_TYPE_ECHO_REQUEST: response = EchoProtocolFrame() response.frame_type = FRAME_TYPE_ECHO_RESPONSE response.payload = "You sent %r" % request_frame.payload self.request.sendall(response.pack()) else: print("Unexpected frame: %r" % request_frame) def setup(self): self.handler = StreamProtocolHandler(EchoProtocolFrame, self._frame_received) def handle(self): while True: self.handler.feed(self.request.recv(1024))
def setup(self): self.handler = StreamProtocolHandler(EchoProtocolFrame, self._frame_received)