1
0
Fork 0

A lot of comments and add encryption in option

This commit is contained in:
Ethanell 2019-07-14 19:28:00 +02:00
parent 756be54cca
commit 3fde3f747a

View file

@ -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:
self.get_secure_connexion(self.socket) if encryption: # Set secure connexion is asked
self.send(sock, f"Client service ID: {service_id}") 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}": 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)
self.set_secure_connexion(connexion) Connect a socket server to a socket"""
if self.receive(connexion) == f"Client service ID: {service_id}": 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}") self.send(connexion, f"Server service ID: {service_id}")
return connexion return connexion
else: else: