Added libs
This commit is contained in:
Lucas
264
fspn/protocol/connection.py
Normal file
264
fspn/protocol/connection.py
Normal file
@@ -0,0 +1,264 @@
|
||||
from ..utils.observable import Observable
|
||||
from ..utils.wrapper_util import threaded
|
||||
from .security import Security
|
||||
|
||||
from enum import Enum
|
||||
import socket
|
||||
import ipaddress
|
||||
import struct
|
||||
import time, datetime
|
||||
import logging, traceback
|
||||
import json
|
||||
import uuid
|
||||
|
||||
# TODO Impossible: Hiding ip in a p2p connection hahahahahaha
|
||||
|
||||
HEADER_STRUCTURE = '!I?IIb' # size, encrypted, nonce_len, mac_len, is_binary
|
||||
HEADER_SIZE = struct.calcsize(HEADER_STRUCTURE) # 14 bytes
|
||||
|
||||
MAX_CONNECTION_TRIES = 3
|
||||
|
||||
class EVENTS(Enum):
|
||||
ON_CONNECTION = 0
|
||||
ON_CONNECTION_ERROR = 1
|
||||
ON_DISCONNECTION = 2
|
||||
ON_MESSAGE = 3
|
||||
|
||||
events = [EVENTS.ON_CONNECTION, EVENTS.ON_CONNECTION_ERROR, EVENTS.ON_DISCONNECTION, EVENTS.ON_MESSAGE]
|
||||
|
||||
class STATUS(Enum):
|
||||
DISCONNECTED = 0
|
||||
CONNECTED = 1
|
||||
CONNECTING = 2
|
||||
HANDSHAKING = 3
|
||||
ERROR = -1
|
||||
|
||||
class Connection(Observable):
|
||||
def __init__(self, user, pmc, conn=None):
|
||||
super().__init__(events)
|
||||
self.security = Security(user, pmc)
|
||||
self.status = STATUS.DISCONNECTED
|
||||
self.address = None
|
||||
self.hostname = None
|
||||
self.bind_address = None
|
||||
self.conn = conn
|
||||
if conn:
|
||||
self.set_addresses()
|
||||
self.id = str(uuid.uuid4())
|
||||
self.handshake_payload = None
|
||||
|
||||
def set_addresses(self, address=None):
|
||||
if address:
|
||||
self.address = address
|
||||
host, port = address
|
||||
try:
|
||||
ipaddress.ip_address(host)
|
||||
except ValueError:
|
||||
self.hostname = host
|
||||
else:
|
||||
self.address = self.conn.getpeername()
|
||||
self.bind_address = self.conn.getsockname()
|
||||
|
||||
|
||||
@threaded
|
||||
def connect(self, address, bind_address=('0.0.0.0', 0)):
|
||||
self.set_addresses(address)
|
||||
self.bind_address = bind_address
|
||||
try:
|
||||
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||
s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
||||
#s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)
|
||||
s.bind(bind_address)
|
||||
s.settimeout(10)
|
||||
self.conn = s
|
||||
except Exception:
|
||||
self.status = STATUS.ERROR
|
||||
self.fire_event(EVENTS.ON_CONNECTION_ERROR, error="Error to setup connection")
|
||||
raise
|
||||
|
||||
self.status = STATUS.CONNECTING
|
||||
logging.info(f'Socket trying to connect: {self.bind_address} -> {address}')
|
||||
|
||||
for i in range(MAX_CONNECTION_TRIES):
|
||||
try:
|
||||
s.settimeout(None)
|
||||
self.conn = s
|
||||
self.bind_address = s.getsockname()
|
||||
s.connect(address)
|
||||
self.address = s.getpeername()
|
||||
break
|
||||
except Exception as e:
|
||||
logging.exception("ERROR")
|
||||
if i < MAX_CONNECTION_TRIES - 1:
|
||||
continue
|
||||
else:
|
||||
self.status = STATUS.ERROR
|
||||
self.fire_event(EVENTS.ON_CONNECTION_ERROR, error=f"No connection could be made in {MAX_CONNECTION_TRIES} retries: {e}")
|
||||
raise
|
||||
|
||||
self.new_connection()
|
||||
|
||||
def new_connection(self):
|
||||
self.handshake_create_payload()
|
||||
|
||||
def handshake_create_payload(self):
|
||||
self.status = STATUS.HANDSHAKING
|
||||
logging.info(f'Socket handshaking: {self.bind_address} -> {self.address}')
|
||||
my_ecdsa_str = self.security.user
|
||||
my_proof_of_work = self.security.proof_of_work
|
||||
my_ecdh_pk = self.security.ecdh.public_key_to_str()
|
||||
date = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
|
||||
self.handshake_payload = {'ecdsa':my_ecdsa_str, 'pof':my_proof_of_work, 'ecdh':my_ecdh_pk, 'date':date}
|
||||
payload = {'ecdsa':my_ecdsa_str, 'pof':my_proof_of_work, 'ecdh':my_ecdh_pk, 'date':date}
|
||||
request_id = self.security.sign_message_ecdsa(json.dumps(payload), self.payload_signature_callback, "Handshake Connection")
|
||||
self.wait_signature()
|
||||
|
||||
@threaded
|
||||
def wait_signature(self):
|
||||
# TODO Config this time as security time validation
|
||||
time.sleep(120)
|
||||
if self.status == STATUS.HANDSHAKING:
|
||||
logging.info(f"Closing connection, payload signature time exceded")
|
||||
self.status = STATUS.DISCONNECTED
|
||||
self.fire_event(EVENTS.ON_CONNECTION_ERROR, error="Payload signature time exceded")
|
||||
self.close_connection()
|
||||
|
||||
def payload_signature_callback(self, request_id, signature):
|
||||
if self.status != STATUS.HANDSHAKING:
|
||||
logging.info(f"Not using signature, connection is closed")
|
||||
return None
|
||||
if signature:
|
||||
self.handshake_payload['signature'] = signature
|
||||
self.send_message(json.dumps(self.handshake_payload), False)
|
||||
|
||||
while self.status == STATUS.HANDSHAKING:
|
||||
message = self.message_reader()["data"]
|
||||
logging.error
|
||||
if(message):
|
||||
payload = json.loads(message)
|
||||
payload_signed = payload.copy()
|
||||
payload_signed.pop('signature')
|
||||
self.security.handshake_validation(payload['ecdsa'], json.dumps(payload_signed), payload['signature'], payload['ecdh'], payload['pof'], payload['date'])
|
||||
|
||||
self.status = STATUS.CONNECTED
|
||||
logging.info(f'Ready: {self.bind_address} -> {self.address}')
|
||||
self.fire_event(EVENTS.ON_CONNECTION)
|
||||
self.wait_message()
|
||||
else:
|
||||
self.status = STATUS.DISCONNECTED
|
||||
self.fire_event(EVENTS.ON_CONNECTION_ERROR, error="Payload signature is None")
|
||||
self.close_connection()
|
||||
|
||||
@threaded
|
||||
def wait_message(self):
|
||||
while self.status == STATUS.CONNECTED:
|
||||
try:
|
||||
message = self.message_reader()
|
||||
except socket.timeout as to:
|
||||
logging.exception("ERROR")
|
||||
continue
|
||||
except (ConnectionAbortedError, EOFError, ConnectionResetError, OSError):
|
||||
# TODO Maybe try to reconnect
|
||||
self.status = STATUS.DISCONNECTED
|
||||
self.fire_event(EVENTS.ON_DISCONNECTION)
|
||||
break
|
||||
except Exception as e:
|
||||
logging.exception("ERROR")
|
||||
self.status = STATUS.DISCONNECTED
|
||||
self.fire_event(EVENTS.ON_DISCONNECTION)
|
||||
break
|
||||
|
||||
try:
|
||||
if(message):
|
||||
# logging.debug(message)
|
||||
self.fire_event(EVENTS.ON_MESSAGE, message=message)
|
||||
except Exception as e:
|
||||
logging.error(traceback.format_exc())
|
||||
|
||||
def close_connection(self):
|
||||
logging.info(f'Socket closed: {self.address}')
|
||||
self.conn.close()
|
||||
self.status = STATUS.DISCONNECTED
|
||||
self.fire_event(EVENTS.ON_DISCONNECTION)
|
||||
|
||||
@threaded
|
||||
def send_binary(self, data: bytes, meta: dict = None, encrypted=True):
|
||||
meta_json = json.dumps(meta or {})
|
||||
meta_encoded = meta_json.encode('utf-8')
|
||||
meta_len_bytes = struct.pack('!I', len(meta_encoded))
|
||||
full_payload = meta_len_bytes + meta_encoded + data
|
||||
|
||||
nonce = b''
|
||||
mac = b''
|
||||
|
||||
if encrypted:
|
||||
nonce, full_payload, mac = self.security.encrypt_message(full_payload)
|
||||
|
||||
message_header = struct.pack(
|
||||
HEADER_STRUCTURE,
|
||||
len(nonce) + len(mac) + len(full_payload),
|
||||
encrypted,
|
||||
len(nonce),
|
||||
len(mac),
|
||||
1 # is_binary = True
|
||||
)
|
||||
|
||||
self.conn.sendall(message_header + nonce + mac + full_payload)
|
||||
|
||||
@threaded
|
||||
def send_message(self, message: str, encrypted=True):
|
||||
nonce = b''
|
||||
mac = b''
|
||||
meta = b''
|
||||
encoded_msg = message.encode('utf-8')
|
||||
|
||||
if encrypted:
|
||||
nonce, encoded_msg, mac = self.security.encrypt_message(encoded_msg)
|
||||
|
||||
message_header = struct.pack(
|
||||
HEADER_STRUCTURE,
|
||||
len(nonce) + len(mac) + len(meta) + len(encoded_msg),
|
||||
encrypted,
|
||||
len(nonce),
|
||||
len(mac),
|
||||
0 # is_binary = False
|
||||
)
|
||||
|
||||
self.conn.sendall(message_header + nonce + mac + encoded_msg)
|
||||
|
||||
|
||||
def message_reader(self):
|
||||
message_header = self.recv_all(HEADER_SIZE)
|
||||
if not message_header:
|
||||
self.close_connection()
|
||||
return None
|
||||
|
||||
total_len, encrypted, nonce_len, mac_len, is_binary = struct.unpack(HEADER_STRUCTURE, message_header)
|
||||
data = self.recv_all(total_len)
|
||||
|
||||
if encrypted:
|
||||
nonce = data[:nonce_len]
|
||||
mac = data[nonce_len:nonce_len + mac_len]
|
||||
payload = data[nonce_len + mac_len:]
|
||||
decrypted = self.security.decrypt_message(nonce, payload, mac)
|
||||
else:
|
||||
decrypted = data
|
||||
|
||||
if is_binary:
|
||||
meta_length = struct.unpack('!I', decrypted[:4])[0]
|
||||
meta_raw = decrypted[4:4 + meta_length]
|
||||
meta = json.loads(meta_raw)
|
||||
file_data = decrypted[4 + meta_length:]
|
||||
return {"meta": meta, "data": file_data}
|
||||
else:
|
||||
return {"data": decrypted.decode('utf-8')}
|
||||
|
||||
|
||||
def recv_all(self, n: int) -> bytes:
|
||||
buffer = b''
|
||||
while len(buffer) < n:
|
||||
chunk = self.conn.recv(n - len(buffer))
|
||||
if not chunk:
|
||||
raise ConnectionError("Connection closed before receive all bytes")
|
||||
buffer += chunk
|
||||
return buffer
|
||||
Reference in New Issue
Block a user