roman
/
black-mamba
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

working version

utcp_uuid
Роман Бородин 2019-04-08 22:15:04 +03:00
parent 9932525bd2
commit 11b7ec1ee2
1 changed files with 103 additions and 79 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

174
mods/utcp.py
View File

@ -24,6 +24,7 @@ class Connection:
def __init__(self, remote, encrypted=False): def __init__(self, remote, encrypted=False):
self.peer_addr = remote self.peer_addr = remote
self.seq = Connection.gen_starting_seq_num() self.seq = Connection.gen_starting_seq_num()
self.recv_seq = 0
self.my_key = None self.my_key = None
if encrypted: if encrypted:
self.my_key = simplecrypto.RsaKeypair() self.my_key = simplecrypto.RsaKeypair()
@ -50,49 +51,56 @@ class Connection:
class UTCPPacket: class UTCPPacket:
def __cmp__(self, other): def __cmp__(self, other):
return (self.seq > other.seq) - (self.seq < other.seq) return (self.seq > other.seq) - (self.seq < other.seq)
def __lt__(self, other):
return self.seq < other.seq
def __gt__(self, other):
return self.seq > other.seq
def __eq__(self, other):
return self.seq == other.seq
class Ack(UTCPPacket): class Ack(UTCPPacket):
type = 'ACK' type = 'ACK'
def __init__(self, id_, seq): def __init__(self, id_):
self.id = id_ self.id = id_
self.seq = seq self.seq = 0
class Fin(UTCPPacket): class Fin(UTCPPacket):
type = 'FIN' type = 'FIN'
def __init__(self, seq): def __init__(self):
self.id = uuid.uuid4().bytes self.id = uuid.uuid4().bytes
self.seq = seq self.seq = 0
class FinAck(UTCPPacket): class FinAck(UTCPPacket):
type = 'FIN-ACK' type = 'FIN-ACK'
def __init__(self, id_, seq): def __init__(self, id_):
self.id = id_ self.id = id_
self.seq = seq self.seq = 0
class Syn(UTCPPacket): class Syn(UTCPPacket):
type = 'SYN' type = 'SYN'
checksum = None checksum = None
def __init__(self, seq): def __init__(self):
self.id = uuid.uuid4().bytes self.id = uuid.uuid4().bytes
self.seq = seq self.seq = 0
def set_pub(self, pubkey): def set_pub(self, pubkey):
self.checksum = TCP.checksum(pubkey) self.checksum = TCP.checksum(pubkey)
self.pubkey = pubkey self.pubkey = pubkey
class SynAck(UTCPPacket): class SynAck(UTCPPacket):
type = 'SYN-ACK' type = 'SYN-ACK'
checksum = None checksum = None
def __init__(self, id_, seq): def __init__(self, id_):
self.id = id_ self.id = id_
self.seq = seq self.seq = 0
def set_pub(self, pubkey): def set_pub(self, pubkey):
self.checksum = TCP.checksum(pubkey) self.checksum = TCP.checksum(pubkey)
self.pubkey = pubkey self.pubkey = pubkey
class Packet(UTCPPacket): class Packet(UTCPPacket):
def __init__(self, seq): type = 'DATA'
def __init__(self):
self.id = uuid.uuid4().bytes self.id = uuid.uuid4().bytes
self.checksum = 0 self.checksum = 0
self.data = b'' self.data = b''
self.seq = seq self.seq = 0
def __repr__(self): def __repr__(self):
return f'TCPpacket(type={self.packet_type()})' return f'Packet(seq={self.seq})'
def __str__(self): def __str__(self):
return 'UUID: %d, DATA:%s' % (uuid.UUID(bytes=self.id), self.data) return 'UUID: %d, DATA:%s' % (uuid.UUID(bytes=self.id), self.data)
@ -246,19 +254,37 @@ class TCP(object):
data_chunk = data_part if not self.encrypted else conn.peer_pub.encrypt_raw(data_part) data_chunk = data_part if not self.encrypted else conn.peer_pub.encrypt_raw(data_part)
packet = Packet() packet = Packet()
packet.set_data(data_chunk) packet.set_data(data_chunk)
packet_to_send = pickle.dumps(packet) answer = self.__send_packet(connection, packet, retransmit=True)
answer = self.retransmit(connection, packet_to_send, want_id=packet.id)
return len(data) return len(data)
except socket.error as error: except socket.error as error:
raise EOFError('Socket was closed before executing command. Error is: %s.' % error) raise EOFError('Socket was closed before executing command. Error is: %s.' % error)
def retransmit(self, peer_addr, pickled_packet, condition='ACK', want_id=None): def recv(self, size, connection=None):
try:
if connection not in list(self.connections.keys()):
if connection is None:
connection = list(self.connections.keys())[0]
else:
raise EOFError('Connection not in connected devices')
data = self.find_correct_packet('DATA', connection, size)
if not self.status:
raise EOFError('Disconnecting')
return data
except socket.error as error:
raise EOFError('Socket was closed before executing command. Error is: %s.' % error)
def __send_packet(self, peer_addr, packet, retransmit=False, wait_cond='ACK'):
conn = self.connections[peer_addr]
packet.seq = conn.seq_inc()
packet_to_send = pickle.dumps(packet)
if not retransmit:
self.own_socket.sendto(packet_to_send, peer_addr)
else:
data_not_received = True data_not_received = True
retransmit_count = 0 retransmit_count = 0
while data_not_received and retransmit_count < 3: while data_not_received and retransmit_count < 3:
data_not_received = False data_not_received = False
try: try:
self.own_socket.sendto(pickled_packet, peer_addr) self.own_socket.sendto(packet_to_send, peer_addr)
answer = self.find_correct_packet(condition, peer_addr, want_id=want_id) answer = self.find_correct_packet(wait_cond, peer_addr, want_id=packet.id)
if not answer: if not answer:
data_not_received = True data_not_received = True
retransmit_count += 1 retransmit_count += 1
@ -268,23 +294,6 @@ class TCP(object):
self.drop_connection(peer_addr) self.drop_connection(peer_addr)
raise EOFError('Connection lost') raise EOFError('Connection lost')
return answer return answer
def recv(self, size, connection=None):
try:
if connection not in list(self.connections.keys()):
if connection is None:
connection = list(self.connections.keys())[0]
else:
raise EOFError('Connection not in connected devices')
data = self.find_correct_packet('DATA', connection, size)
if not self.status:
raise EOFError('Disconnecting')
return data
except socket.error as error:
raise EOFError('Socket was closed before executing command. Error is: %s.' % error)
def send_ack(self, connection, ack):
packet_to_send = pickle.dumps(ack)
self.own_socket.sendto(packet_to_send, connection)
def listen_handler(self, max_connections): def listen_handler(self, max_connections):
try: try:
while True and self.status: while True and self.status:
@ -301,7 +310,9 @@ class TCP(object):
self.connection_queue.append((answer, conn)) self.connection_queue.append((answer, conn))
self.blink_new_conn_event() self.blink_new_conn_event()
else: else:
self.own_socket.sendto('Connections full', address) if answer.id in map(lambda x: x[0].id, self.connection_queue):
continue
self.own_socket.sendto(b'Connections full', address)
except KeyError: except KeyError:
continue continue
except socket.error as error: except socket.error as error:
@ -328,8 +339,7 @@ class TCP(object):
syn_ack = SynAck(answer.id) syn_ack = SynAck(answer.id)
if self.encrypted: if self.encrypted:
syn_ack.set_pub(conn.peer_pub.encrypt_raw(conn.pubkey)) syn_ack.set_pub(conn.peer_pub.encrypt_raw(conn.pubkey))
packet_to_send = pickle.dumps(syn_ack) answer = self.__send_packet(conn.peer_addr, syn_ack, retransmit=True)
self.retransmit(conn.peer_addr, packet_to_send, 'ACK', want_id=syn_ack.id)
return ConnectedSOCK(self, conn.peer_addr), conn.peer_addr return ConnectedSOCK(self, conn.peer_addr), conn.peer_addr
except Exception as error: except Exception as error:
self.close(conn.peer_addr) self.close(conn.peer_addr)
@ -346,9 +356,8 @@ class TCP(object):
syn = Syn() syn = Syn()
if self.encrypted: if self.encrypted:
syn.set_pub(conn.pubkey) syn.set_pub(conn.pubkey)
first_packet_to_send = pickle.dumps(syn)
try: try:
answer = self.retransmit(server_address, first_packet_to_send, 'SYN-ACK', want_id=syn.id) answer = self.__send_packet(server_address, syn, retransmit=True, wait_cond='SYN-ACK')
except EOFError: except EOFError:
raise EOFError('Remote peer unreachable') raise EOFError('Remote peer unreachable')
if type(answer) == str: # == 'Connections full': if type(answer) == str: # == 'Connections full':
@ -360,8 +369,7 @@ class TCP(object):
except: except:
raise socket.error('Decrypt peer public key error') raise socket.error('Decrypt peer public key error')
ack = Ack(answer.id) ack = Ack(answer.id)
second_packet_to_send = pickle.dumps(ack) self.__send_packet(server_address, ack)
self.own_socket.sendto(second_packet_to_send, server_address)
self.channel = UTCPChannel(self) self.channel = UTCPChannel(self)
except socket.error as error: except socket.error as error:
self.own_socket.close() self.own_socket.close()
@ -385,15 +393,13 @@ class TCP(object):
else: else:
raise EOFError('Connection not in connected devices') raise EOFError('Connection not in connected devices')
fin = Fin() fin = Fin()
packet_to_send = pickle.dumps(fin) answer = self.__send_packet(connection, fin, retransmit=True)
self.own_socket.sendto(packet_to_send, connection)
answer = self.retransmit(connection, packet_to_send, want_id=fin.id)
answer = self.find_correct_packet('FIN-ACK', connection, want_id=fin.id) answer = self.find_correct_packet('FIN-ACK', connection, want_id=fin.id)
if not answer: if not answer:
raise Exception('The receiver didn\'t send the fin packet') raise Exception('The receiver didn\'t send the fin packet')
else: else:
ack = Ack(fin.id) ack = Ack(fin.id)
self.send_ack(connection, ack) self.__send_packet(connection, ack)
self.drop_connection(connection) self.drop_connection(connection)
if len(self.connections) == 0 and self.client: if len(self.connections) == 0 and self.client:
self.own_socket.close() self.own_socket.close()
@ -404,11 +410,10 @@ class TCP(object):
def disconnect(self, connection, fin_id): def disconnect(self, connection, fin_id):
try: try:
ack = Ack(fin_id) ack = Ack(fin_id)
self.send_ack(connection, ack) self.__send_packet(connection, ack)
fin_ack = FinAck(fin_id) fin_ack = FinAck(fin_id)
packet_to_send = pickle.dumps(fin_ack)
try: try:
answer = self.retransmit(connection, packet_to_send, want_id=fin_id) answer = self.__send_packet(connection, fin_ack, retransmit=True)
except: except:
pass pass
with self.connection_lock: with self.connection_lock:
@ -420,7 +425,6 @@ class TCP(object):
def data_divider(data): def data_divider(data):
'''Divides the data into a list where each element's length is 1024''' '''Divides the data into a list where each element's length is 1024'''
data = [data[i:i + DATA_DIVIDE_LENGTH] for i in range(0, len(data), DATA_DIVIDE_LENGTH)] data = [data[i:i + DATA_DIVIDE_LENGTH] for i in range(0, len(data), DATA_DIVIDE_LENGTH)]
data.append(PACKET_END)
return data return data
@staticmethod @staticmethod
@ -440,14 +444,30 @@ class TCP(object):
if condition in ['ACK', 'SYN-ACK', 'FIN-ACK']: if condition in ['ACK', 'SYN-ACK', 'FIN-ACK']:
tries += 1 tries += 1
if condition == 'DATA': if condition == 'DATA':
if len(conn.packet_buffer[condition]):
data = b''
while size and conn.packet_buffer[condition]:
with conn.recv_lock: with conn.recv_lock:
packet = self.packets_received[condition][address][:size] packet = conn.packet_buffer[condition][0]
self.packets_received[condition][address] = self.packets_received[condition][address][size:] chunk = packet.data[:size]
if not len(self.packets_received[condition][address]): chunk_len = len(chunk)
del self.packets_received[condition][address] data += chunk
packet.data = packet.data[size:]
size -= chunk_len
if not len(packet.data):
try:
conn.packet_buffer[condition].pop(0)
except IndexError:
size = 0
return data
else:
raise KeyError
else: else:
with conn.recv_lock: with conn.recv_lock:
packet = self.packets_received[condition].pop(address) if len(conn.packet_buffer[condition]):
packet = conn.packet_buffer[condition].pop()
else:
raise KeyError
if want_id and packet.id != want_id: if want_id and packet.id != want_id:
raise KeyError raise KeyError
return packet return packet
@ -463,30 +483,34 @@ class TCP(object):
def sort_answers(self, packet, address): def sort_answers(self, packet, address):
if address not in self.connections and not isinstance(packet, Syn): if address not in self.connections and not isinstance(packet, Syn):
return return
if isinstance(packet, Syn): if not isinstance(packet, Syn):
with self.queue_lock: conn = self.connections[address]
pass if conn.recv_seq == packet.seq:
ack = Ack(packet.id)
self.__send_packet(address, ack)
return
elif conn.recv_seq > packet.seq:
return
else:
conn.recv_seq = packet.seq
if isinstance(packet, Packet):
if packet.checksum == TCP.checksum(packet.data):
if self.encrypted:
packet.data = conn.my_key.decrypt_raw(packet.data)
else:
return
if isinstance(packet, Fin): if isinstance(packet, Fin):
self.disconnect(address, packet.id) self.disconnect(address, packet.id)
elif isinstance(packet, Packet): elif isinstance(packet, Syn):
if packet.checksum == TCP.checksum(packet.data): if packet.id not in map(lambda x: x.id, self.syn_received.values()):
conn = self.connections[address] self.syn_received[address] = packet
data_chunk = packet.data if not self.encrypted else conn.my_key.decrypt_raw(packet.data)
if data_chunk != PACKET_END:
with conn.recv_lock:
if address not in self.packets_received['DATA']:
self.packets_received['DATA'][address] = b''
self.packets_received['DATA'][address] += data_chunk
ack = Ack(packet.id)
self.send_ack(address, ack)
self.blink_incoming_packet_event()
else: else:
if address in self.packets_received[packet.type]:
conn = self.connections[address]
with conn.recv_lock: with conn.recv_lock:
if packet.id == self.packets_received[packet.type][address].id: if packet.id not in map(lambda x: x.id, conn.packet_buffer[packet.type]):
return bisect.insort(conn.packet_buffer[packet.type], packet)
self.packets_received[packet.type][address] = packet if isinstance(packet, Packet):
ack = Ack(packet.id)
self.__send_packet(address, ack)
self.blink_incoming_packet_event() self.blink_incoming_packet_event()
def central_receive_handler(self): def central_receive_handler(self):