flifloo/SecureSocketService
Archived
1
0
Fork 0
Code Activity

Set dynamic buffer size on message and little fix

Browse source
This commit is contained in:
Ethanell 2019-07-22 18:49:37 +02:00
parent 345ceeb8b8
commit 42fc2f09df
2 changed files with 12 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

17
SecureSocketService/__init__.py
View file

@ -26,6 +26,7 @@ class Socket:
if self.fernet and encryption: # Encrypt if available if self.fernet and encryption: # Encrypt if available
message = self.fernet.encrypt(message) message = self.fernet.encrypt(message)
try: # Try to send, if fail raise a custom error try: # Try to send, if fail raise a custom error
sock.send(f"Buffer size: {len(message)}".encode("Utf8"))
sock.send(message) sock.send(message)
except socket.error: except socket.error:
raise ConnectionError("Fail to send") raise ConnectionError("Fail to send")
@ -33,11 +34,12 @@ class Socket:
return True return True
def receive(self, sock: socket.socket, decode="Utf8", encryption=True): def receive(self, sock: socket.socket, decode="Utf8", encryption=True):
"""self, sock (socket), utf8 (str/bool): Utf8, encryption (bool) """self, sock (socket), decode (str): Utf8, encryption (bool)
Receive a message from a socket""" Receive a message from a socket"""
try: # Try to receive, else raise a custom error try: # Try to receive, else raise a custom error
response = sock.recv(1028) buffer_size = int(sock.recv(1028).decode("Utf8")[13:])
except socket.error: response = sock.recv(buffer_size)
except (socket.error, ValueError):
raise ConnectionError("Fail to receive") raise ConnectionError("Fail to receive")
else: else:
if self.fernet and encryption and response != b"": # Encrypt if available if self.fernet and encryption and response != b"": # Encrypt if available
@ -55,12 +57,13 @@ class Socket:
# Receive the public key of the target socket # Receive the public key of the target socket
public_key = serialization.load_pem_public_key( public_key = serialization.load_pem_public_key(
self.receive(sock, False, False), self.receive(sock, "", False),
backend=default_backend() backend=default_backend()
) )
# Encrypt the key and send back with private key encryption to the target socket # Encrypt the key and send back with private key encryption to the target socket
sock.send( self.send(
sock,
public_key.encrypt( public_key.encrypt(
self.key, self.key,
padding.OAEP( padding.OAEP(
@ -68,7 +71,9 @@ class Socket:
algorithm=hashes.SHA256(), algorithm=hashes.SHA256(),
label=None label=None
) )
) ),
"",
False
) )
if self.receive(sock) == "Key receive": # Check if the kay is valid if self.receive(sock) == "Key receive": # Check if the kay is valid

2
setup.py
View file

@ -7,7 +7,7 @@ def readme():
setup(name="SecureSocketService", setup(name="SecureSocketService",
version="1.0.1", version="1.0.3",
description="A socket service with secure connexion", description="A socket service with secure connexion",
long_description=readme(), long_description=readme(),
classifiers=[ classifiers=[