def _handle_read_packet(self, packet): """Takes a packet from the client and advances the state machine depending on that packet. This should only be invoked from the READING state. Returns an appropriate DataPacket containing the next block of data. Args: packet: A packet object that has already been unpacked. Returns: a packet object with which to send back to the client. """ assert self.state == READING if not isinstance(packet, packets.AcknowledgementPacket): return packets.ErrorPacket( 0, "Illegal packet type given" " current state of conversation. Host: %s, Port: %s." % (self.client_host, self.client_port)) if self.current_block_num != packet.block_num: return packets.NoOpPacket() previous_block_num = packet.block_num if previous_block_num == self.read_buffer.get_block_count(): self.state = COMPLETED self.log("READREQUEST", "Success") return packets.NoOpPacket() else: self.current_block_num += 1 data = self.read_buffer.get_block(self.current_block_num) return packets.DataPacket(self.current_block_num, data)
def _handle_initial_read_packet(self, packet): """Check if there is an application action to respond to this request If so, then send the first block and move the state to READING. Otherwise, send back an error packet and move the state to COMPLETED. Args: packet: An unpacked ReadRequestPacket. Returns: A Data packet if the request's filename matches any possible read rule. The data packet includes the first block of data from the output of the read action. Otherwise, an ErrorPacket with a file not found error code and message. """ assert isinstance(packet, packets.ReadRequestPacket) self.filename = packet.filename self.mode = packet.mode self.read_buffer = self.response_router.initialize_read( self.filename, self.client_host, self.client_port) if self.read_buffer: self.state = READING data = self.read_buffer.get_block(1) self.current_block_num = 1 return packets.DataPacket(1, data) else: self.log("READREQUEST", "File not found") self.state = COMPLETED return packets.ErrorPacket( 1, "File not found. Host: %s, Port: %s" % (self.client_host, self.client_port))
def test_out_of_lock_step_block_num(self): packet = packets.DataPacket(2, "") conversation = TFTPConversation(self.client_host, self.client_port, StubResponseRouterTwo()) conversation.cached_packet = "stub packet" conversation.state = tftp_conversation.WRITING conversation.current_block_num = 3 response_packet = conversation.handle_packet(packet) self.assertEqual(conversation.state, tftp_conversation.WRITING) self.assertEqual(response_packet.__class__, packets.NoOpPacket)
def test_get_conversation_old_with_data_packet(self): conversation_table = ConversationTable() packet = packets.DataPacket('stub block number', 'stub data') old_conversation = TFTPConversation('10.26.0.1', 3942, 'stub_router') conversation_table.add_conversation('10.26.0.1', 3942, old_conversation) reactor = Reactor('stub_socket', 'stub_router', conversation_table) conversation = reactor.get_conversation('10.26.0.1', 3942, packet) self.assertEqual(len(conversation_table), 1) self.assertTrue(isinstance(conversation, TFTPConversation)) self.assertEqual(conversation, old_conversation)
def test_illegal_packet_type_during_reading_state(self): packet = packets.DataPacket(2, "") conversation = TFTPConversation(self.client_host, self.client_port, StubResponseRouterTwo()) conversation.cached_packet = "stub packet" conversation.state = tftp_conversation.READING conversation.read_buffer = StubReadBufferTwo() response_packet = conversation.handle_packet(packet) self.assertEqual(conversation.state, tftp_conversation.READING) self.assertEqual(response_packet.__class__, packets.ErrorPacket) self.assertEqual(response_packet.error_code, 0) self.assertEqual(conversation.cached_packet, "stub packet")
def test_continue_writing(self): packet = packets.DataPacket(2, "X" * 512) conversation = TFTPConversation(self.client_host, self.client_port, StubResponseRouterTwo()) conversation.state = tftp_conversation.WRITING conversation.write_buffer = StubWriteBufferTwo() conversation.current_block_num = 1 response_packet = conversation.handle_packet(packet) self.assertEqual(conversation.state, tftp_conversation.WRITING) self.assertEqual(conversation.current_block_num, 2) self.assertEqual(conversation.write_buffer.data, "X" * 512) self.assertEqual(conversation.cached_packet, response_packet) self.assertEqual(response_packet.__class__, packets.AcknowledgementPacket) self.assertEqual(response_packet.block_num, 2)
def test_finish_writing(self): packet = packets.DataPacket(3, "O" * 511) conversation = TFTPConversation(self.client_host, self.client_port, StubResponseRouterTwo()) conversation.state = tftp_conversation.WRITING conversation.write_buffer = WriteBuffer() conversation.write_buffer.data = "X" * 512 conversation.filename = "stub_filename" conversation.current_block_num = 2 write_action_wrapper = StubWriteActionWrapper() conversation.write_action = write_action_wrapper.stub_action response_packet = conversation.handle_packet(packet) self.assertEqual(conversation.state, tftp_conversation.COMPLETED) self.assertEqual(conversation.current_block_num, 3) self.assertEqual(conversation.cached_packet, response_packet) self.assertEqual(response_packet.__class__, packets.AcknowledgementPacket) self.assertEqual(response_packet.block_num, 3) # action should get invoked, saving this state in the wrapper class self.assertEqual(write_action_wrapper.received_state, ("10.26.0.3", 12345, "stub_filename", "X" * 512 + "O" * 511))
final_attempts = 0 # number of attempts for final packets print("Sending %s to %s:%s" % (args.filename, args.address, args.port)) with open(args.filename, 'rb') as f: while True: # Send the data we read from the file read_data = f.read(packets.DataPacket.DATA_SIZE) if not read_data: # EOF if args.verbose: print("end of file") break last_packet = packets.DataPacket.DATA_SIZE * (packet_number + 1) >= total_bytes ack = 1 if packet_number == next_ack or last_packet else 0 data_packet = packets.DataPacket(connection_id, total_bytes, packet_number, ack, read_data) if args.verbose: print(data_packet.header_as_string()) sock.sendto(data_packet.as_bytes(), addr) # send the packet # Receive an ACK from the receiver if data_packet.get_ack(): if args.verbose: print("ACKING PACKET %d" % packet_number) if last_packet: print("LAST PACKET") try: success = False while not success: # Wait for successful ACK recv_data, addr = sock.recvfrom(