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

.

utcp_uuid
Роман Бородин 2019-04-08 17:48:13 +03:00
parent 544a98edd3
commit 4c65194dbc
1 changed files with 29 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

38
mods/utcp.py
View File

@ -35,6 +35,13 @@ class Connection:
self.peer_pub = None self.peer_pub = None
self.recv_lock = threading.Lock() self.recv_lock = threading.Lock()
self.send_lock = threading.Lock() self.send_lock = threading.Lock()
self.packet_buffer = {
'SYN': [],
'ACK': [],
'SYN-ACK': [],
'DATA': [],
'FIN-ACK': []
}
@staticmethod @staticmethod
def gen_starting_seq_num(): def gen_starting_seq_num():
return random.randint(Connection.SMALLEST_STARTING_SEQ, Connection.HIGHEST_STARTING_SEQ) return random.randint(Connection.SMALLEST_STARTING_SEQ, Connection.HIGHEST_STARTING_SEQ)
@ -46,15 +53,19 @@ class Connection:
return ack return ack
class Ack: class Ack:
type = 'ACK'
def __init__(self, id_): def __init__(self, id_):
self.id = id_ self.id = id_
class Fin: class Fin:
type = 'FIN'
def __init__(self): def __init__(self):
self.id = uuid.uuid4().bytes self.id = uuid.uuid4().bytes
class FinAck: class FinAck:
type = 'FIN-ACK'
def __init__(self, id_): def __init__(self, id_):
self.id = id_ self.id = id_
class Syn: class Syn:
type = 'SYN'
checksum = None checksum = None
def __init__(self): def __init__(self):
self.id = uuid.uuid4().bytes self.id = uuid.uuid4().bytes
@ -62,6 +73,7 @@ class Syn:
self.checksum = TCP.checksum(pubkey) self.checksum = TCP.checksum(pubkey)
self.pubkey = pubkey self.pubkey = pubkey
class SynAck: class SynAck:
type = 'SYN-ACK'
checksum = None checksum = None
def __init__(self, id_): def __init__(self, id_):
self.id = id_ self.id = id_
@ -87,7 +99,7 @@ class Packet:
self.checksum = TCP.checksum(data) self.checksum = TCP.checksum(data)
self.data = data self.data = data
class RestrictedUnpickler(pickle.Unpickler): class RestrictedUnpickler(pickle.Unpickler):
def find_class(self, module, name): def find_class(self, module, name):
@ -233,7 +245,7 @@ class TCP(object):
packet = Packet() packet = Packet()
packet.set_data(data_chunk) packet.set_data(data_chunk)
packet_to_send = pickle.dumps(packet) packet_to_send = pickle.dumps(packet)
answer = self.retransmit(connection, packet_to_send, wnat_id=packet.id) 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)
@ -281,7 +293,7 @@ class TCP(object):
conn = Connection(address, self.encrypted) conn = Connection(address, self.encrypted)
if self.encrypted: if self.encrypted:
try: try:
conn.peer_pub = simplecrypto.RsaPublicKey(answer.data) conn.peer_pub = simplecrypto.RsaPublicKey(answer.pubkey)
except: except:
raise socket.error('Init peer public key error') raise socket.error('Init peer public key error')
self.connection_queue.append((answer, conn)) self.connection_queue.append((answer, conn))
@ -341,7 +353,7 @@ class TCP(object):
raise socket.error('Server cant receive any connections right now.') raise socket.error('Server cant receive any connections right now.')
if self.encrypted: if self.encrypted:
try: try:
peer_pub = conn.my_key.decrypt_raw(answer.data) peer_pub = conn.my_key.decrypt_raw(answer.pubkey)
conn.peer_pub = simplecrypto.RsaPublicKey(peer_pub) conn.peer_pub = simplecrypto.RsaPublicKey(peer_pub)
except: except:
raise socket.error('Decrypt peer public key error') raise socket.error('Decrypt peer public key error')
@ -363,6 +375,8 @@ class TCP(object):
with self.connection_lock: with self.connection_lock:
if len(self.connections): if len(self.connections):
self.connections.pop(connection) self.connections.pop(connection)
for k in list(self.packets_received.keys()):
self.packets_received[k].pop(connection)
def close(self, connection=None): def close(self, connection=None):
try: try:
if connection not in list(self.connections.keys()): if connection not in list(self.connections.keys()):
@ -432,7 +446,8 @@ class TCP(object):
if not len(self.packets_received[condition][address]): if not len(self.packets_received[condition][address]):
del self.packets_received[condition][address] del self.packets_received[condition][address]
else: else:
packet = self.packets_received[condition].pop(address) with conn.recv_lock:
packet = self.packets_received[condition].pop(address)
if want_id and packet.id != want_id: if want_id and packet.id != want_id:
raise KeyError raise KeyError
return packet return packet
@ -448,6 +463,9 @@ 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):
with self.queue_lock:
pass
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, Packet):
@ -462,11 +480,13 @@ class TCP(object):
ack = Ack(packet.id) ack = Ack(packet.id)
self.send_ack(address, ack) self.send_ack(address, ack)
self.blink_incoming_packet_event() self.blink_incoming_packet_event()
elif packet.packet_type() == '':
#print('redundant packet found', packet)
pass
else: else:
self.packets_received[packet.packet_type()][address] = packet if address in self.packets_received[packet.type]:
conn = self.connections[address]
with conn.recv_lock:
if packet.id == self.packets_received[packet.type][address].id:
return
self.packets_received[packet.type][address] = packet
self.blink_incoming_packet_event() self.blink_incoming_packet_event()
def central_receive_handler(self): def central_receive_handler(self):