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
75
Socket.py
75
Socket.py
|
@ -8,45 +8,57 @@ from cryptography.fernet import Fernet
|
|||
|
||||
|
||||
class Socket:
|
||||
"""Self
|
||||
Socket service with security system"""
|
||||
def __init__(self):
|
||||
self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||
self.key = None
|
||||
self.fernet = None
|
||||
"""Self
|
||||
Socket service with security system"""
|
||||
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):
|
||||
if type(message) != bytes:
|
||||
def send(self, sock: socket.socket, message: str, encryption=True):
|
||||
"""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")
|
||||
if self.fernet:
|
||||
if self.fernet and encryption: # Encrypt if available
|
||||
message = self.fernet.encrypt(message)
|
||||
try:
|
||||
try: # Try to send, if fail raise a custom error
|
||||
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:
|
||||
def receive(self, sock: socket.socket, utf8=True, encryption=True):
|
||||
"""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)
|
||||
except socket.error:
|
||||
raise ConnectionError("Fail to receive")
|
||||
else:
|
||||
if self.fernet and fernet:
|
||||
if self.fernet and encryption: # Encrypt if available
|
||||
response = self.fernet.decrypt(response)
|
||||
if utf8:
|
||||
if utf8: # decode the utf8 by default
|
||||
response = response.decode("Utf8")
|
||||
return response
|
||||
|
||||
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.fernet = Fernet(self.key)
|
||||
|
||||
# Receive the public key of the targed socket
|
||||
public_key = serialization.load_pem_public_key(
|
||||
self.receive(sock, False, False),
|
||||
backend=default_backend()
|
||||
)
|
||||
|
||||
# Encrypt the key and send back with private key encryption to the target socket
|
||||
sock.send(
|
||||
public_key.encrypt(
|
||||
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")
|
||||
return True
|
||||
else:
|
||||
raise ConnectionError("Invalid key !")
|
||||
|
||||
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(
|
||||
public_exponent=65537,
|
||||
key_size=2048,
|
||||
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(
|
||||
encoding=serialization.Encoding.PEM,
|
||||
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.receive(sock, False, False),
|
||||
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":
|
||||
raise ConnectionError("Invalid key !")
|
||||
return True
|
||||
|
||||
def connect_server(self, sock: socket.socket, host: str, port: int, service_id: int):
|
||||
try:
|
||||
def connect_server(self, sock: socket.socket, host: str, port: int, service_id: int, encryption=True):
|
||||
"""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))
|
||||
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 encryption: # Set secure connexion is asked
|
||||
self.get_secure_connexion(self.socket)
|
||||
|
||||
self.send(sock, f"Client service ID: {service_id}") # Check if service id match
|
||||
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}":
|
||||
def connect_client(self, sock: socket, service_id: int, encryption=True):
|
||||
"""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)
|
||||
|
||||
if self.receive(connexion) == f"Client service ID: {service_id}": # Check if service id match
|
||||
self.send(connexion, f"Server service ID: {service_id}")
|
||||
return connexion
|
||||
else:
|
||||
|
|
Reference in a new issue