From b7c30d14910ce9993d255d198b3f855615562f17 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=D0=A0=D0=BE=D0=BC=D0=B0=D0=BD=20=D0=91=D0=BE=D1=80=D0=BE?=
 =?UTF-8?q?=D0=B4=D0=B8=D0=BD?= <roman@ukamnya.ru>
Date: Mon, 8 Apr 2019 09:58:19 +0300
Subject: [PATCH] close() and disconnect()

---
 mods/utcp.py | 58 ++++++++++++++++++++++++++--------------------------
 1 file changed, 29 insertions(+), 29 deletions(-)

diff --git a/mods/utcp.py b/mods/utcp.py
index 94b4017..cba1c83 100644
--- a/mods/utcp.py
+++ b/mods/utcp.py
@@ -208,7 +208,6 @@ class TCP(object):
                     raise EOFError('Connection not in connected devices')
             data_parts = TCP.data_divider(data)
             for data_part in data_parts:
-                data_not_received = True
                 data_chunk = data_part if not self.encrypted else self.peer_keypair[connection].peer_pub.encrypt_raw(data_part)
                 checksum_of_data = TCP.checksum(data_chunk)
                 self.connections[connection].checksum = checksum_of_data
@@ -217,26 +216,28 @@ class TCP(object):
                 packet_to_send = pickle.dumps(self.connections[connection])
                 self.connections[connection].checksum = 0
                 self.connections[connection].data = b''
-                retransmit_count = 0
-                while data_not_received and retransmit_count < 3:
-                    data_not_received = False
-                    try:
-                        self.own_socket.sendto(packet_to_send, connection)
-                        answer = self.find_correct_packet('ACK', connection)
-                        if not answer:
-                            data_not_received = True
-                            retransmit_count += 1
-                    except socket.timeout:
-                        #print('timeout')
-                        data_not_received = True
-                if not answer:
-                    self.drop_connection(connection)
-                    raise EOFError('Connection lost')
+                answer = self.retransmit(connection, packet_to_send)
                 self.connections[connection].seq += len(data_part)
             return len(data)
         except socket.error as error:
             raise EOFError('Socket was closed before executing command. Error is: %s.' % error)
-
+    def retransmit(self, peer_addr, pickled_packet, condition='ACK'):
+        data_not_received = True
+        retransmit_count = 0
+        while data_not_received and retransmit_count < 3:
+            data_not_received = False
+            try:
+                self.own_socket.sendto(pickled_packet, peer_addr)
+                answer = self.find_correct_packet(condition, peer_addr)
+                if not answer:
+                    data_not_received = True
+                    retransmit_count += 1
+            except socket.timeout:
+                data_not_received = True
+        if not answer:
+            self.drop_connection(peer_addr)
+            raise EOFError('Connection lost')
+        return answer
     def recv(self, size, connection=None):
         try:
             if connection not in list(self.connections.keys()):
@@ -379,13 +380,13 @@ class TCP(object):
             fin_packet.set_flags(fin=True)
             packet_to_send = pickle.dumps(fin_packet)
             self.own_socket.sendto(packet_to_send, connection)
-            answer = self.find_correct_packet('ACK', connection)  # change cause may get a None value
+            answer = self.retransmit(connection, packet_to_send)
             conn.ack += 1
             answer = self.find_correct_packet('FIN-ACK', connection)
             if answer.flag_fin != 1:
                 raise Exception('The receiver didn\'t send the fin packet')
             else:
-                self.send_ack(connection, self.connections[connection].ack + 1)
+                self.send_ack(connection, conn.ack + 1)
                 self.drop_connection(connection)
                 if len(self.connections) == 0 and self.client:
                     self.own_socket.close()
@@ -395,16 +396,15 @@ class TCP(object):
             
     def disconnect(self, connection):
         try:
-            self.connections[connection].ack += 1
-            self.connections[connection].seq += 1
-            self.connections[connection].set_flags(ack=True)
-            packet_to_send = pickle.dumps(self.connections[connection])
-            self.own_socket.sendto(packet_to_send, connection)
-            self.connections[connection].set_flags(fin=True, ack=True)
-            self.connections[connection].seq += 1
-            packet_to_send = pickle.dumps(self.connections[connection])
-            self.own_socket.sendto(packet_to_send, connection)
-            answer = self.find_correct_packet('ACK', connection)
+            conn = self.connections[connection]
+            self.send_ack(connection, conn.set_ack(conn.ack + 1))
+            finack_packet = TCPPacket(conn.seq_inc())
+            finack_packet.set_flags(fin=True, ack=True)
+            packet_to_send = pickle.dumps(finack_packet)
+            try:
+                answer = self.retransmit(connection, packet_to_send)
+            except:
+                pass
             with self.connection_lock:
                 self.connections.pop(connection)
         except Exception as error: