119 lines
3.8 KiB
Python
119 lines
3.8 KiB
Python
|
import socket
|
||
|
|
from cryptography.hazmat.backends import default_backend
|
||
|
|
from cryptography.hazmat.primitives.asymmetric import rsa
|
||
|
|
from cryptography.hazmat.primitives import serialization
|
||
|
|
from cryptography.hazmat.primitives import hashes
|
||
|
|
from cryptography.hazmat.primitives.asymmetric import padding
|
||
|
|
from cryptography.fernet import Fernet
|
||
|
|
|
||
|
|
|
||
|
|
class Socket:
|
||
|
|
def __init__(self):
|
||
|
|
self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||
|
|
self.key = None
|
||
|
|
self.fernet = None
|
||
|
|
|
||
|
|
def send(self, sock: socket.socket, message: str):
|
||
|
|
if type(message) != bytes:
|
||
|
|
message = message.encode("Utf8")
|
||
|
|
if self.fernet:
|
||
|
|
message = self.fernet.encrypt(message)
|
||
|
|
try:
|
||
|
|
sock.send(message)
|
||
|
|
except socket.error:
|
||
|
|
raise ConnectionError("Fail to send")
|
||
|
|
else:
|
||
|
|
return True
|
||
|
|
|
||
|
|
def receive(self, sock: socket.socket, utf8=True, fernet=True):
|
||
|
|
try:
|
||
|
|
response = sock.recv(1028)
|
||
|
|
except socket.error:
|
||
|
|
raise ConnectionError("Fail to receive")
|
||
|
|
else:
|
||
|
|
if self.fernet and fernet:
|
||
|
|
response = self.fernet.decrypt(response)
|
||
|
|
if utf8:
|
||
|
|
response = response.decode("Utf8")
|
||
|
|
return response
|
||
|
|
|
||
|
|
def set_secure_connexion(self, sock: socket.socket):
|
||
|
|
if not self.key:
|
||
|
|
self.key = Fernet.generate_key()
|
||
|
|
self.fernet = Fernet(self.key)
|
||
|
|
|
||
|
|
public_key = serialization.load_pem_public_key(
|
||
|
|
self.receive(sock, False, False),
|
||
|
|
backend=default_backend()
|
||
|
|
)
|
||
|
|
|
||
|
|
sock.send(
|
||
|
|
public_key.encrypt(
|
||
|
|
self.key,
|
||
|
|
padding.OAEP(
|
||
|
|
mgf=padding.MGF1(algorithm=hashes.SHA256()),
|
||
|
|
algorithm=hashes.SHA256(),
|
||
|
|
label=None
|
||
|
|
)
|
||
|
|
)
|
||
|
|
)
|
||
|
|
|
||
|
|
if self.receive(sock) == "Key receive":
|
||
|
|
self.send(sock, "Key valid")
|
||
|
|
return True
|
||
|
|
else:
|
||
|
|
raise ConnectionError("Invalid key !")
|
||
|
|
|
||
|
|
def get_secure_connexion(self, sock: socket.socket):
|
||
|
|
private_key = rsa.generate_private_key(
|
||
|
|
public_exponent=65537,
|
||
|
|
key_size=2048,
|
||
|
|
backend=default_backend()
|
||
|
|
)
|
||
|
|
|
||
|
|
public_key = private_key.public_key()
|
||
|
|
|
||
|
|
pem = public_key.public_bytes(
|
||
|
|
encoding=serialization.Encoding.PEM,
|
||
|
|
format=serialization.PublicFormat.SubjectPublicKeyInfo
|
||
|
|
)
|
||
|
|
|
||
|
|
self.send(sock, pem)
|
||
|
|
|
||
|
|
self.key = private_key.decrypt(
|
||
|
|
self.receive(sock, False, False),
|
||
|
|
padding.OAEP(
|
||
|
|
mgf=padding.MGF1(algorithm=hashes.SHA256()),
|
||
|
|
algorithm=hashes.SHA256(),
|
||
|
|
label=None
|
||
|
|
)
|
||
|
|
)
|
||
|
|
|
||
|
|
self.fernet = Fernet(self.key)
|
||
|
|
|
||
|
|
self.send(sock, "Key receive")
|
||
|
|
if self.receive(sock) != "Key valid":
|
||
|
|
raise ConnectionError("Invalid key !")
|
||
|
|
return True
|
||
|
|
|
||
|
|
def connect_server(self, sock: socket.socket, host: str, port: int, service_id: int):
|
||
|
|
try:
|
||
|
|
sock.connect((host, port))
|
||
|
|
except socket.error:
|
||
|
|
raise ConnectionError("Can't connect to server !")
|
||
|
|
else:
|
||
|
|
self.get_secure_connexion(self.socket)
|
||
|
|
self.send(sock, f"Client service ID: {service_id}")
|
||
|
|
if self.receive(sock) != f"Server service ID: {service_id}":
|
||
|
|
raise ConnectionError("Server service ID")
|
||
|
|
return True
|
||
|
|
|
||
|
|
def connect_client(self, sock: socket, service_id: int):
|
||
|
|
connexion, address = sock.accept()
|
||
|
|
self.set_secure_connexion(connexion)
|
||
|
|
if self.receive(connexion) == f"Client service ID: {service_id}":
|
||
|
|
self.send(connexion, f"Server service ID: {service_id}")
|
||
|
|
return connexion
|
||
|
|
else:
|
||
|
|
raise ConnectionError("Invalid client service ID !")
|