A lot of comments and add encryption in option
This commit is contained in:
parent
756be54cca
commit
3fde3f747a
1 changed files with 50 additions and 25 deletions
71
Socket.py
71
Socket.py
|
@ -8,45 +8,57 @@ from cryptography.fernet import Fernet
|
||||||
|
|
||||||
|
|
||||||
class Socket:
|
class Socket:
|
||||||
|
"""Self
|
||||||
|
Socket service with security system"""
|
||||||
def __init__(self):
|
def __init__(self):
|
||||||
self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
"""Self
|
||||||
self.key = None
|
Socket service with security system"""
|
||||||
self.fernet = None
|
self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) # Set the socket object
|
||||||
|
self.key = None # Set the key encryption
|
||||||
|
self.fernet = None # Set the encryption object
|
||||||
|
|
||||||
def send(self, sock: socket.socket, message: str):
|
def send(self, sock: socket.socket, message: str, encryption=True):
|
||||||
if type(message) != bytes:
|
"""self, sock (socket), message (str), encryption (bool)
|
||||||
|
Send a message with a specific socket, encrypted if available"""
|
||||||
|
if type(message) != bytes: # If the message are not a bytes type, encode in Utf8
|
||||||
message = message.encode("Utf8")
|
message = message.encode("Utf8")
|
||||||
if self.fernet:
|
if self.fernet and encryption: # Encrypt if available
|
||||||
message = self.fernet.encrypt(message)
|
message = self.fernet.encrypt(message)
|
||||||
try:
|
try: # Try to send, if fail raise a custom error
|
||||||
sock.send(message)
|
sock.send(message)
|
||||||
except socket.error:
|
except socket.error:
|
||||||
raise ConnectionError("Fail to send")
|
raise ConnectionError("Fail to send")
|
||||||
else:
|
else:
|
||||||
return True
|
return True
|
||||||
|
|
||||||
def receive(self, sock: socket.socket, utf8=True, fernet=True):
|
def receive(self, sock: socket.socket, utf8=True, encryption=True):
|
||||||
try:
|
"""self, sock (socket), utf8 (bool), encryption (bool)
|
||||||
|
Receive a message from a socket"""
|
||||||
|
try: # Try to receive, else raise a custom error
|
||||||
response = sock.recv(1028)
|
response = sock.recv(1028)
|
||||||
except socket.error:
|
except socket.error:
|
||||||
raise ConnectionError("Fail to receive")
|
raise ConnectionError("Fail to receive")
|
||||||
else:
|
else:
|
||||||
if self.fernet and fernet:
|
if self.fernet and encryption: # Encrypt if available
|
||||||
response = self.fernet.decrypt(response)
|
response = self.fernet.decrypt(response)
|
||||||
if utf8:
|
if utf8: # decode the utf8 by default
|
||||||
response = response.decode("Utf8")
|
response = response.decode("Utf8")
|
||||||
return response
|
return response
|
||||||
|
|
||||||
def set_secure_connexion(self, sock: socket.socket):
|
def set_secure_connexion(self, sock: socket.socket):
|
||||||
if not self.key:
|
"""self, sock (socket)
|
||||||
|
Set encryption for a connected socket"""
|
||||||
|
if not self.key: # If any key generated, generate one
|
||||||
self.key = Fernet.generate_key()
|
self.key = Fernet.generate_key()
|
||||||
self.fernet = Fernet(self.key)
|
self.fernet = Fernet(self.key)
|
||||||
|
|
||||||
|
# Receive the public key of the targed socket
|
||||||
public_key = serialization.load_pem_public_key(
|
public_key = serialization.load_pem_public_key(
|
||||||
self.receive(sock, False, False),
|
self.receive(sock, False, False),
|
||||||
backend=default_backend()
|
backend=default_backend()
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# Encrypt the key and send back with private key encryption to the target socket
|
||||||
sock.send(
|
sock.send(
|
||||||
public_key.encrypt(
|
public_key.encrypt(
|
||||||
self.key,
|
self.key,
|
||||||
|
@ -58,28 +70,33 @@ class Socket:
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|
||||||
if self.receive(sock) == "Key receive":
|
if self.receive(sock) == "Key receive": # Check if the kay is valid
|
||||||
self.send(sock, "Key valid")
|
self.send(sock, "Key valid")
|
||||||
return True
|
return True
|
||||||
else:
|
else:
|
||||||
raise ConnectionError("Invalid key !")
|
raise ConnectionError("Invalid key !")
|
||||||
|
|
||||||
def get_secure_connexion(self, sock: socket.socket):
|
def get_secure_connexion(self, sock: socket.socket):
|
||||||
|
"""self, sock (socket)
|
||||||
|
Get encryption for a connected socket"""
|
||||||
|
# Generate private key
|
||||||
private_key = rsa.generate_private_key(
|
private_key = rsa.generate_private_key(
|
||||||
public_exponent=65537,
|
public_exponent=65537,
|
||||||
key_size=2048,
|
key_size=2048,
|
||||||
backend=default_backend()
|
backend=default_backend()
|
||||||
)
|
)
|
||||||
|
|
||||||
public_key = private_key.public_key()
|
public_key = private_key.public_key() # Create pubic kay
|
||||||
|
|
||||||
|
# Make a sendable public key
|
||||||
pem = public_key.public_bytes(
|
pem = public_key.public_bytes(
|
||||||
encoding=serialization.Encoding.PEM,
|
encoding=serialization.Encoding.PEM,
|
||||||
format=serialization.PublicFormat.SubjectPublicKeyInfo
|
format=serialization.PublicFormat.SubjectPublicKeyInfo
|
||||||
)
|
)
|
||||||
|
|
||||||
self.send(sock, pem)
|
self.send(sock, pem) # Send the public key
|
||||||
|
|
||||||
|
# Decrypt the receive key with the private key and set it
|
||||||
self.key = private_key.decrypt(
|
self.key = private_key.decrypt(
|
||||||
self.receive(sock, False, False),
|
self.receive(sock, False, False),
|
||||||
padding.OAEP(
|
padding.OAEP(
|
||||||
|
@ -89,29 +106,37 @@ class Socket:
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|
||||||
self.fernet = Fernet(self.key)
|
self.fernet = Fernet(self.key) # Set the encryption object
|
||||||
|
|
||||||
self.send(sock, "Key receive")
|
self.send(sock, "Key receive") # Check if the key is valid
|
||||||
if self.receive(sock) != "Key valid":
|
if self.receive(sock) != "Key valid":
|
||||||
raise ConnectionError("Invalid key !")
|
raise ConnectionError("Invalid key !")
|
||||||
return True
|
return True
|
||||||
|
|
||||||
def connect_server(self, sock: socket.socket, host: str, port: int, service_id: int):
|
def connect_server(self, sock: socket.socket, host: str, port: int, service_id: int, encryption=True):
|
||||||
try:
|
"""self, sock (socket), host (str), port (int), service_id (int), encryption (bool)
|
||||||
|
Connect a socket to a socket server"""
|
||||||
|
try: # Try to connect, else raise a custom error
|
||||||
sock.connect((host, port))
|
sock.connect((host, port))
|
||||||
except socket.error:
|
except socket.error:
|
||||||
raise ConnectionError("Can't connect to server !")
|
raise ConnectionError("Can't connect to server !")
|
||||||
else:
|
else:
|
||||||
|
if encryption: # Set secure connexion is asked
|
||||||
self.get_secure_connexion(self.socket)
|
self.get_secure_connexion(self.socket)
|
||||||
self.send(sock, f"Client service ID: {service_id}")
|
|
||||||
|
self.send(sock, f"Client service ID: {service_id}") # Check if service id match
|
||||||
if self.receive(sock) != f"Server service ID: {service_id}":
|
if self.receive(sock) != f"Server service ID: {service_id}":
|
||||||
raise ConnectionError("Server service ID")
|
raise ConnectionError("Server service ID")
|
||||||
return True
|
return True
|
||||||
|
|
||||||
def connect_client(self, sock: socket, service_id: int):
|
def connect_client(self, sock: socket, service_id: int, encryption=True):
|
||||||
connexion, address = sock.accept()
|
"""self, sock (socket), service_id (int), encryption (bool)
|
||||||
|
Connect a socket server to a socket"""
|
||||||
|
connexion, address = sock.accept() # Await for connexion
|
||||||
|
if encryption:
|
||||||
self.set_secure_connexion(connexion)
|
self.set_secure_connexion(connexion)
|
||||||
if self.receive(connexion) == f"Client service ID: {service_id}":
|
|
||||||
|
if self.receive(connexion) == f"Client service ID: {service_id}": # Check if service id match
|
||||||
self.send(connexion, f"Server service ID: {service_id}")
|
self.send(connexion, f"Server service ID: {service_id}")
|
||||||
return connexion
|
return connexion
|
||||||
else:
|
else:
|
||||||
|
|
Reference in a new issue