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

.

Browse Source
This commit is contained in:
Роман Бородин 2019-04-08 09:22:52 +03:00
parent e2ae47d056
commit 586e69552b
1 changed files with 38 additions and 44 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

82
mods/utcp.py
View File

@ -30,8 +30,15 @@ class Connection:
HIGHEST_STARTING_SEQ = 4294967295 HIGHEST_STARTING_SEQ = 4294967295
def __init__(self, remote, encrypted=False): def __init__(self, remote, encrypted=False):
self.peer_addr = remote self.peer_addr = remote
self.ack = 0
self.seq = Connection.gen_starting_seq_num() self.seq = Connection.gen_starting_seq_num()
self.my_key self.my_key = None
if encrypted:
self.my_key = KeyPair(simplecrypto.RsaKeypair())
self.pubkey = self.my_key.publickey.serialize()
self.peer_pub = None
self.recv_lock = threading.Lock()
self.send_lock = threading.Lock()
@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)
@ -86,6 +93,9 @@ class TCPPacket(object):
self.flag_fin = 1 self.flag_fin = 1
else: else:
self.flag_fin = 0 self.flag_fin = 0
def set_data(self, data):
self.checksum = TCP.checksum(data)
self.data = data
@ -145,7 +155,6 @@ class TCP(object):
self.settimeout() self.settimeout()
self.connection_lock = threading.Lock() self.connection_lock = threading.Lock()
self.queue_lock = threading.Lock() self.queue_lock = threading.Lock()
self.peer_keypair = {}
self.connections = {} self.connections = {}
self.connection_queue = [] self.connection_queue = []
self.packets_received = {'SYN': {}, 'ACK': {}, 'SYN-ACK': {}, 'DATA or FIN': {}, 'FIN-ACK': {}} self.packets_received = {'SYN': {}, 'ACK': {}, 'SYN-ACK': {}, 'DATA or FIN': {}, 'FIN-ACK': {}}
@ -254,15 +263,14 @@ class TCP(object):
answer, address = self.find_correct_packet('SYN') answer, address = self.find_correct_packet('SYN')
with self.queue_lock: with self.queue_lock:
if len(self.connection_queue) < max_connections: if len(self.connection_queue) < max_connections:
conn = Connection(address, self.encrypted)
if self.encrypted: if self.encrypted:
try: try:
peer_pub = answer.data conn.peer_pub = answer.data
self.peer_keypair[address] = KeyPair(simplecrypto.RsaKeypair()) # FIXME: for some reason slowly creates a key (~5 sec)
self.peer_keypair[address].peer_pub = simplecrypto.RsaPublicKey(peer_pub)
except: except:
self.peer_keypair.pop(address) self.peer_keypair.pop(address)
raise socket.error('Init peer public key error') raise socket.error('Init peer public key error')
self.connection_queue.append((answer, address)) 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) self.own_socket.sendto('Connections full', address)
@ -287,34 +295,28 @@ class TCP(object):
self.new_conn_event.wait(0.1) self.new_conn_event.wait(0.1)
if self.connection_queue: if self.connection_queue:
with self.queue_lock: with self.queue_lock:
answer, address = self.connection_queue.pop() answer, conn = self.connection_queue.pop()
self.connections[conn.peer_addr] = conn
self.connections[address] = TCPPacket() packet = TCPPacket(conn.seq)
self.connections[address].ack = answer.seq + 1 packet.ack = answer.seq + 1
self.connections[address].seq += 1 packet.seq = conn.seq_inc()
self.connections[address].set_flags(ack=True, syn=True) packet.set_flags(ack=True, syn=True)
if self.encrypted: if self.encrypted:
pubkey = self.peer_keypair[address].my_key.publickey.serialize() packet.set_data(conn.peer_pub.encrypt_raw(conn.pubkey))
self.connections[address].data = self.peer_keypair[address].peer_pub.encrypt_raw(pubkey) packet_to_send = pickle.dumps(packet)
self.connections[address].checksum = TCP.checksum(self.connections[address].data) #On packet lost retransmit
packet_to_send = pickle.dumps(self.connections[address])
if self.encrypted:
self.connections[address].data = b''
self.connections[address].checksum = 0
#lock address, connections dictionary?
packet_not_sent_correctly = True packet_not_sent_correctly = True
while packet_not_sent_correctly or answer is None: while packet_not_sent_correctly or answer is None:
try: try:
packet_not_sent_correctly = False packet_not_sent_correctly = False
self.own_socket.sendto(packet_to_send, address) self.own_socket.sendto(packet_to_send, conn.peer_addr)
answer = self.find_correct_packet('ACK', address) answer = self.find_correct_packet('ACK', conn.peer_addr)
except socket.timeout: except socket.timeout:
packet_not_sent_correctly = True packet_not_sent_correctly = True
self.connections[address].set_flags() conn.ack = answer.seq + 1
self.connections[address].ack = answer.seq + 1 return ConnectedSOCK(self, conn.peer_addr), conn.peer_addr
return ConnectedSOCK(self, address), address
except Exception as error: except Exception as error:
self.close(address) self.close(conn.peer_addr)
raise EOFError('Something went wrong in accept func: ' + str(error)) raise EOFError('Something went wrong in accept func: ' + str(error))
def connect(self, server_address=('127.0.0.1', 10000)): def connect(self, server_address=('127.0.0.1', 10000)):
@ -323,32 +325,24 @@ class TCP(object):
self.status = 1 self.status = 1
self.client = True self.client = True
self.central_receive() self.central_receive()
self.connections[server_address] = TCPPacket() conn = Connection(server_address, self.encrypted)
self.connections[server_address].set_flags(syn=True) self.connections[server_address] = Connection(server_address, self.encrypted)
syn_packet = TCPPacket(conn.seq)
syn_packet.set_flags(syn=True)
if self.encrypted: if self.encrypted:
self.peer_keypair[server_address] = KeyPair(simplecrypto.RsaKeypair()) # FIXME: but here it creates the key quickly syn_packet.set_data(conn.pubkey)
pubkey = self.peer_keypair[server_address].my_key.publickey.serialize() first_packet_to_send = pickle.dumps(syn_packet)
pub_checksum = TCP.checksum(pubkey) self.own_socket.sendto(first_packet_to_send, server_address)
self.connections[server_address].checksum = pub_checksum
self.connections[server_address].data = pubkey
first_packet_to_send = pickle.dumps(self.connections[server_address])
self.connections[server_address].data = b''
self.connections[server_address].checksum = 0
self.own_socket.sendto(first_packet_to_send, list(self.connections.keys())[FIRST])
self.connections[server_address].set_flags()
answer = self.find_correct_packet('SYN-ACK', server_address) answer = self.find_correct_packet('SYN-ACK', server_address)
if type(answer) == str: # == 'Connections full': if type(answer) == str: # == 'Connections full':
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 = self.peer_keypair[server_address].my_key.decrypt_raw(answer.data) peer_pub = conn.my_key.decrypt_raw(answer.data)
self.peer_keypair[server_address].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')
if not peer_pub or answer.checksum != TCP.checksum(answer.data): ack_packet =
raise socket.error('Invalid peer public key')
self.connections[server_address].ack = answer.seq + 1 self.connections[server_address].ack = answer.seq + 1
self.connections[server_address].seq += 1 self.connections[server_address].seq += 1
self.connections[server_address].set_flags(ack=True) self.connections[server_address].set_flags(ack=True)