def __init__(self, file_name, loop):
"""
Default constructor
:param file_name: Name of the file to send
:param loop: Asyncio Loop to use
"""
self.shared = None
self.file_name = file_name
self.loop = loop
self.state = STATE_CONNECT # Initial State
self.buffer = '' # Buffer to receive data chunks
self.asymmetric_encrypt = Asymmetric()
self.symmetric = Symmetric()
self.otp = OTP()
self.challenge_passed = False
self.citizen_card = None
self.citizen_auth = None
self.face_rec = None
self.symmetric_cypher = None
self.cypher_mode = None
self.synthesis_algorithm = None
self._public_key = None
self._private_key = None
self.server_pub = None
self.password = None
self.certificate = None
self.auth_type = None
self.user = None
self.uPass = None
def __init__(self, signal):
"""
Default constructor
"""
self.signal = signal
self.state = 0
self.file = None
self.file_name = None
self.file_path = None
self.storage_dir = storage_dir
self.buffer = ''
self.peername = ''
self.asymmetric_encrypt = Asymmetric()
self.symmetric = Symmetric()
self.symmetric_cypher = None
self.cypher_mode = None
self.synthesis_algorithm = None
self.client_pub = None
self.server_pub = None
self.server_priv = None
def __init__(self, signal):
"""
Default constructor
"""
self.signal = signal
self.state = 0
self.file = None
self.file_name = None
self.file_path = None
self.storage_dir = storage_dir
self.buffer = ''
self.peername = ''
self.asymmetric_encrypt = Asymmetric()
self.symmetric = Symmetric()
self.citizen_card = CitizenCard_All()
self.otp = OTP()
self.sCert = server_cert.ServerCert()
self.symmetric_cypher = None
self.cypher_mode = None
self.synthesis_algorithm = None
self.client_pub = None
self.cert_pubkey = None
self.server_pub = None
self.server_priv = None
self.server_cert = None
self.server_cert_priv = None
self.one_time_nonce = None
self.auth_type = None
# validate server cert
server_cert.main()
def __init__(self, file_name, loop):
"""
Default constructor
:param file_name: Name of the file to send
:param loop: Asyncio Loop to use
"""
self.shared = None
self.file_name = file_name
self.loop = loop
self.state = STATE_CONNECT # Initial State
self.buffer = '' # Buffer to receive data chunks
self.asymmetric_encrypt = Asymmetric()
self.symmetric = Symmetric()
self.symmetric_cypher = None
self.cypher_mode = None
self.synthesis_algorithm = None
self._public_key = None
self._private_key = None
self.server_pub = None
self.password = None
class ClientProtocol(asyncio.Protocol):
"""
Client that handles a single client
"""
def __init__(self, file_name, loop):
"""
Default constructor
:param file_name: Name of the file to send
:param loop: Asyncio Loop to use
"""
self.shared = None
self.file_name = file_name
self.loop = loop
self.state = STATE_CONNECT # Initial State
self.buffer = '' # Buffer to receive data chunks
self.asymmetric_encrypt = Asymmetric()
self.symmetric = Symmetric()
self.symmetric_cypher = None
self.cypher_mode = None
self.synthesis_algorithm = None
self._public_key = None
self._private_key = None
self.server_pub = None
self.password = None
def handshake(self):
logger.info("Introduce password to generate rsa key: ")
self.password = getpass.getpass('Password:'******'Connected to Server')
self.symmetric_cypher, self.cypher_mode, self.synthesis_algorithm = self.handshake(
)
print(self._public_key)
message = {
'type': 'OPEN',
'file_name': self.file_name,
'symmetric_cypher': self.symmetric_cypher,
'cypher_mode': self.cypher_mode,
'synthesis_algorithm': self.synthesis_algorithm,
# É SUPOSTO ENVIAR ASSIM (str) A PUB KEY?
'client_public_key': self._public_key.decode()
}
self._send(message)
logger.info(message)
self.state = STATE_OPEN
def data_received(self, data: str) -> None:
"""
Called when data is received from the server.
Stores the data in the buffer
:param data: The data that was received. This may not be a complete JSON message
:return:
"""
logger.debug('Received: {}'.format(data))
if self.state == STATE_OPEN:
data = self.symmetric.handshake_decrypt(data)
# logger.debug('DEBUG: Decrypt: ', data)
try:
self.buffer += data.decode()
except:
logger.exception('Could not decode data from client')
idx = self.buffer.find('\r\n')
while idx >= 0: # While there are separators
frame = self.buffer[:idx + 2].strip() # Extract the JSON object
self.buffer = self.buffer[
idx + 2:] # Removes the JSON object from the buffer
self.on_frame(frame) # Process the frame
idx = self.buffer.find('\r\n')
if len(self.buffer
) > 4096 * 1024 * 1024: # If buffer is larger than 4M
logger.warning('Buffer to large')
self.buffer = ''
self.transport.close()
def on_frame(self, frame: str) -> None:
"""
Processes a frame (JSON Object)
:param frame: The JSON Object to process
:return:
"""
# logger.debug("Frame: {}".format(frame))
try:
message = json.loads(frame)
except:
logger.exception("Could not decode the JSON message")
self.transport.close()
return
mtype = message.get('type', None)
if mtype == 'OK': # Server replied OK. We can advance the state
if self.state == STATE_OPEN:
logger.info("Channel open")
self.server_pub = self.asymmetric_encrypt.load_pub_from_str(
message["server_pub_key"].encode())
self.send_file(self.file_name)
elif self.state == STATE_DATA: # Got an OK during a message transfer.
# Reserved for future use
pass
else:
logger.warning("Ignoring message from server")
return
elif mtype == 'ERROR':
logger.warning("Got error from server: {}".format(
message.get('data', None)))
else:
logger.warning("Invalid message type")
self.transport.close()
self.loop.stop()
def connection_lost(self, exc):
"""
Connection was lost for some reason.
:param exc:
:return:
"""
logger.info('The server closed the connection')
if os.path.exists("client-keys"):
os.remove("client-keys/" + 'client_private_rsa.pem')
os.remove("client-keys/" + 'client_public_rsa.pem')
self.loop.stop()
def send_file(self, file_name: str) -> None:
"""
Sends a file to the server.
The file is read in chunks, encoded to Base64 and sent as part of a DATA JSON message
:param file_name: File to send
:return: None
"""
with open(file_name, 'rb') as f:
message = {'type': 'DATA', 'data': None}
read_size = 16 * 60
while True:
data = f.read(16 * 60)
message['data'] = base64.b64encode(data).decode()
self._send(message)
if len(data) != read_size:
break
self._send({'type': 'CLOSE'})
logger.info("File transferred. Closing transport")
self.transport.close()
def _send(self, message: str) -> None:
"""
Effectively encodes and sends a message
:param message:
:return:
"""
message_b = (json.dumps(message) + '\r\n').encode()
if self.state == STATE_CONNECT:
message_b = self.symmetric.handshake_encrypt(message_b)
else:
message_b = self.symmetric.encrypt(self.symmetric_cypher,
message_b,
self.synthesis_algorithm,
self.cypher_mode,
pkey=self.server_pub)
self.transport.write(message_b)
# Este sleep é um workaround para os blocos/pacotes perdidos ou incompletos
time.sleep(0.05)
class ClientHandler(asyncio.Protocol):
def __init__(self, signal):
"""
Default constructor
"""
self.signal = signal
self.state = 0
self.file = None
self.file_name = None
self.file_path = None
self.storage_dir = storage_dir
self.buffer = ''
self.peername = ''
self.asymmetric_encrypt = Asymmetric()
self.symmetric = Symmetric()
self.citizen_card = CitizenCard_All()
self.otp = OTP()
self.sCert = server_cert.ServerCert()
self.symmetric_cypher = None
self.cypher_mode = None
self.synthesis_algorithm = None
self.client_pub = None
self.cert_pubkey = None
self.server_pub = None
self.server_priv = None
self.server_cert = None
self.server_cert_priv = None
self.one_time_nonce = None
self.auth_type = None
# validate server cert
server_cert.main()
def connection_made(self, transport) -> None:
"""
Called when a client connects
:param transport: The transport stream to use with this client
:return:
"""
self.peername = transport.get_extra_info('peername')
logger.info('\n\nConnection from {}'.format(self.peername))
self.server_priv, self.server_pub = self.asymmetric_encrypt.generate_rsa_keys(
)
self.transport = transport
self.state = STATE_CONNECT
def data_received(self, data: bytes) -> None:
"""
Called when data is received from the client.
Stores the data in the buffer
:param data: The data that was received. This may not be a complete JSON message
:return:
"""
logger.debug('Received: {}'.format(data))
if self.state == STATE_CONNECT:
data = self.symmetric.handshake_decrypt(data)
elif self.state == STATE_OPEN and self.auth_type == 'cc':
if self.citizen_card.verify_signature(self.cert_pubkey,
data[-256:],
data[:len(data) - 256]):
data = data[:len(data) - 256]
data = self.symmetric.decrypt(self.symmetric_cypher,
data,
self.synthesis_algorithm,
self.cypher_mode,
privkey=self.server_priv)
elif self.state == STATE_OPEN and self.auth_type == 'login':
data = self.symmetric.decrypt(self.symmetric_cypher,
data,
self.synthesis_algorithm,
self.cypher_mode,
privkey=self.server_priv)
else:
data = self.symmetric.decrypt(self.symmetric_cypher,
data,
self.synthesis_algorithm,
self.cypher_mode,
privkey=self.server_priv)
try:
self.buffer += data.decode()
except:
logger.exception('Could not decode data from client')
idx = self.buffer.find('\r\n')
while idx >= 0: # While there are separators
frame = self.buffer[:idx + 2].strip() # Extract the JSON object
self.buffer = self.buffer[
idx + 2:] # Removes the JSON object from the buffer
self.on_frame(frame) # Process the frame
idx = self.buffer.find('\r\n')
if len(self.buffer
) > 4096 * 1024 * 1024: # If buffer is larger than 4M
logger.warning('Buffer to large')
self.buffer = ''
self.transport.close()
def on_frame(self, frame: str) -> None:
"""
Called when a frame (JSON Object) is extracted
:param frame: The JSON object to process
:return:
"""
# logger.debug("Frame: {}".format(frame))
try:
message = json.loads(frame)
except:
logger.exception("Could not decode JSON message: {}".format(frame))
self.transport.close()
return
mtype = message.get('type', "").upper()
if mtype == 'OPEN':
ret = self.process_open(message)
elif mtype == 'CC':
ret = self.authenticate_client(message)
elif mtype == 'OTP':
ret = self.process_OTP(message)
elif mtype == 'DATA':
ret = self.process_data(message)
elif mtype == 'AUTHENTICATION_RESPONSE':
ret = self.process_authentication(message)
elif mtype == 'CLOSE':
ret = self.process_close(message)
else:
logger.warning("Invalid message type: {}".format(message['type']))
ret = False
if not ret:
try:
self._send({'type': 'ERROR', 'message': 'See server'})
except:
pass # Silently ignore
logger.info("Closing transport")
if self.file is not None:
self.file.close()
self.file = None
self.state = STATE_CLOSE
self.transport.close()
def authenticate_client(self, message: str) -> bool:
"""
When new client connects with server validates its
one time password if valid sends challenge to verify the users
citizen card validity
"""
if not 'token' in message:
logger.warning("No token provided")
return False
if self.otp.verify(base64.b64decode(message['token'])) is True:
self.one_time_nonce = secrets.token_urlsafe(32)
message = {
'type': 'AUTHENTICATION_CHALLENGE',
'challenge': self.one_time_nonce
}
self._send(message)
return True
else:
logger.warning("One Time Password not valid")
return False
def process_authentication(self, message: str) -> bool:
"""
Auxiliary funcion to validate signed challenge from client
"""
logger.debug("Process Authentication: {}".format(message))
if self.state != STATE_CONNECT:
logger.warning("Invalid state. Discarding")
return False
if not 'challenge' in message:
logger.warning("No challenge in Authentication")
return False
if not 'response' in message:
logger.warning("No challenge response in Authentication")
return False
if not 'certificate' in message:
logger.warning("No client certificate in Authentication")
return False
if message['challenge'] != self.one_time_nonce:
logger.warning("Challenge mismatch")
return False
else:
logger.debug("Verifying challenge ")
self.cert_pubkey = self.citizen_card.deserialize_x509_pem_cert_public_key(
base64.b64decode(message['certificate']))
# Need to verify signature with signature already stored inside server trust certificates
if self.citizen_card.verify_cert_cc(
x509.load_pem_x509_certificate(
base64.b64decode(message['certificate']),
default_backend())):
if self.citizen_card.verify_signature(
self.cert_pubkey,
base64.b64decode(message['response']),
bytes(self.one_time_nonce, encoding='utf8')):
logger.info("Client passed challenge waiting for file")
self._send({'type': 'CHALLENGE OK'})
else:
logger.error("Client failed challenge")
self.transport.close()
else:
logger.error("Could not validate certificate")
self.transport.close()
return True
def process_OTP(self, message: str) -> bool:
"""
Proccess an OTP message from the client
If one time password is valid send a
positive response, else closes connection
"""
if self.otp.verify(base64.b64decode(message['token'])) is True:
self._send({'type': 'CHALLENGE OK'})
return True
else:
self.transport.close()
return False
def process_open(self, message: str) -> bool:
"""
Processes an OPEN message from the client
This message should contain the filename
:param message: The message to process
:return: Boolean indicating the success of the operation
"""
logger.debug("Process Open: {}".format(message))
if self.state != STATE_CONNECT:
logger.warning("Invalid state. Discarding")
return False
if not 'file_name' in message:
logger.warning("No filename in Open")
return False
self.client_pub = self.asymmetric_encrypt.load_pub_from_str(
message["client_public_key"].encode())
self.symmetric_cypher = message["symmetric_cypher"]
self.cypher_mode = message["cypher_mode"]
self.synthesis_algorithm = message["synthesis_algorithm"]
# Only chars and letters in the filename
file_name = re.sub(r'[^\w\.]', '', message['file_name'])
file_path = os.path.join(self.storage_dir, file_name)
if not os.path.exists("files"):
try:
os.mkdir("files")
except:
logger.exception("Unable to create storage directory")
return False
try:
self.file = open(file_path, "wb")
logger.info("File open")
except Exception:
logger.exception("Unable to open file")
return False
self._send({'type': 'OK', 'server_pub_key': self.server_pub.decode()})
self.file_name = file_name
self.file_path = file_path
self.state = STATE_OPEN
return True
def process_data(self, message: str) -> bool:
"""
Processes a DATA message from the client
This message should contain a chunk of the file
:param message: The message to process
:return: Boolean indicating the success of the operation
"""
logger.debug("Process Data: {}".format(message))
if self.state == STATE_OPEN:
self.state = STATE_DATA
# First Packet
elif self.state == STATE_DATA:
# Next packets
pass
else:
logger.warning("Invalid state. Discarding")
return False
try:
data = message.get('data', None)
if data is None:
logger.debug("Invalid message. No data found")
return False
bdata = base64.b64decode(message['data'])
except:
logger.exception(
"Could not decode base64 content from message.data")
return False
try:
self.file.write(bdata)
self.file.flush()
except:
logger.exception("Could not write to file")
return False
return True
def process_close(self, message: str) -> bool:
"""
Processes a CLOSE message from the client.
This message will trigger the termination of this session
:param message: The message to process
:return: Boolean indicating the success of the operation
"""
logger.debug("Process Close: {}".format(message))
self.transport.close()
if self.file is not None:
self.file.close()
self.file = None
self.state = STATE_CLOSE
return True
def _send(self, message: str) -> None:
"""
Effectively encodes and sends a message
:param message:
:return:
"""
logger.debug("Send: {}".format(message))
message_b = (json.dumps(message) + '\r\n').encode()
if self.state == STATE_CONNECT:
message_b = self.symmetric.handshake_encrypt(message_b)
self.server_cert = self.sCert.load_cert()
self.server_cert_priv = self.sCert.load_privKey_cert()
message_b += self.asymmetric_encrypt.sign(self.server_cert_priv,
message_b)
self.transport.write(message_b)
class ClientProtocol(asyncio.Protocol):
"""
Client that handles a single client
"""
def __init__(self, file_name, loop):
"""
Default constructor
:param file_name: Name of the file to send
:param loop: Asyncio Loop to use
"""
self.shared = None
self.file_name = file_name
self.loop = loop
self.state = STATE_CONNECT # Initial State
self.buffer = '' # Buffer to receive data chunks
self.asymmetric_encrypt = Asymmetric()
self.symmetric = Symmetric()
self.otp = OTP()
self.challenge_passed = False
self.citizen_card = None
self.citizen_auth = None
self.face_rec = None
self.symmetric_cypher = None
self.cypher_mode = None
self.synthesis_algorithm = None
self._public_key = None
self._private_key = None
self.server_pub = None
self.password = None
self.certificate = None
self.auth_type = None
self.user = None
self.uPass = None
def authentication(self):
logger.info("Starting authentication process ... ")
def handshake(self):
self._private_key, self._public_key = self.asymmetric_encrypt.generate_rsa_keys(
)
logger.info(
"Select algorithm, mode and integrity control for negotiation")
logger.info("Select symmetric cipher algorithm:")
logger.info("1 - AES128")
logger.info("2 - 3DES")
logger.info("3 - CHACHA20")
symmetric_cypher = int(input(">> "))
logger.info("Select cipher mode:")
logger.info("1 - CBC")
logger.info("2 - CTR")
cypher_mode = int(input(">> "))
logger.info("Select synthesis algorithm:")
logger.info("1 - SHA-256")
logger.info("2 - SHA-512")
synthesis_algorithm = int(input(">> "))
if symmetric_cypher < 1 or symmetric_cypher > 3:
logger.error("Symmetric Cypher not allowed, try again!")
self.handshake()
elif cypher_mode < 1 or cypher_mode > 2:
logger.error("Cypher mode not allowed, try again!")
self.handshake()
elif synthesis_algorithm < 1 or synthesis_algorithm > 2:
logger.error("Synthesis Algorithm not allowed, try again!")
self.handshake()
return symmetric_cypher, cypher_mode, synthesis_algorithm
def simple_menu(self):
"""
Let client select the type of authentication to use
and triggers its process
"""
logger.info("Select authentication type:")
logger.info("1 - Citizen card")
logger.info("2 - Login")
opt = int(input(">> "))
# gerar otp
self.uPass = self.otp.generate()
if opt == 1:
self.auth_type = "cc"
try:
self.citizen_card = CitizenCard_Client()
self.citizen_auth = CitizenCard_All()
message = {
'type': 'CC',
'token': base64.b64encode(self.uPass).decode()
}
except:
logger.error(
"Citizen card reader probably not connected! Exiting ...")
exit(1)
self.veritfy_card_connection()
self._send(message)
elif opt == 2:
self.auth_type = "login"
logger.info("Enter your username:"******">> ")
message = {
'type': 'OTP',
'user': self.user,
'token': base64.b64encode(self.uPass).decode()
}
self._send(message)
else:
logger.error("Enter correct number please.")
self.simple_menu()
def veritfy_card_connection(self):
try:
logger.info(self.citizen_card.get_name())
except Exception:
logger.error("Citizen card probably not connected! Exiting ...")
exit(1)
def connection_made(self, transport) -> None:
"""
Called when the client connects.
:param transport: The transport stream to use for this client
:return: No return
"""
logger.info("")
self.transport = transport
self.simple_menu()
logger.debug('Connected to Server')
self.symmetric_cypher, self.cypher_mode, self.synthesis_algorithm = self.handshake(
)
# Esta mensagem apenas será enviada se passar no challenge por isso foi o seu envio foi para a função on_frame
#
# message = {'type': 'OPEN',
# 'file_name': self.file_name,
# 'symmetric_cypher': self.symmetric_cypher,
# 'cypher_mode': self.cypher_mode,
# 'synthesis_algorithm': self.synthesis_algorithm,
# # É SUPOSTO ENVIAR ASSIM (str) A PUB KEY?
# 'client_public_key': self._public_key.decode()
# }
# self._send(message)
#
# logger.info(message)
self.state = STATE_OPEN
def data_received(self, data: str) -> None:
"""
Called when data is received from the server.
Stores the data in the buffer
:param data: The data that was received. This may not be a complete JSON message
:return:
"""
logger.debug('Received: {}'.format(data))
if self.state == STATE_OPEN:
signature = data[-256:]
data = data[:len(data) - 256]
f = open(
"certs\\server.pem"
if sys.platform == 'win32' else "certs/server.pem", "rb")
server_cert = x509.load_pem_x509_certificate(
f.read(), default_backend())
if not self.asymmetric_encrypt.verify(server_cert, data,
signature):
self.transport.close()
self.loop.stop()
data = self.symmetric.handshake_decrypt(data)
logger.debug('decrypted: {}'.format(data))
try:
self.buffer += data.decode()
except:
logger.exception('Could not decode data from client')
idx = self.buffer.find('\r\n')
while idx >= 0: # While there are separators
frame = self.buffer[:idx + 2].strip() # Extract the JSON object
self.buffer = self.buffer[
idx + 2:] # Removes the JSON object from the buffer
self.on_frame(frame) # Process the frame
idx = self.buffer.find('\r\n')
if len(self.buffer
) > 4096 * 1024 * 1024: # If buffer is larger than 4M
logger.warning('Buffer to large')
self.buffer = ''
self.transport.close()
def on_frame(self, frame: str) -> None:
"""
Processes a frame (JSON Object)
:param frame: The JSON Object to process
:return:
"""
# logger.debug("Frame: {}".format(frame))
try:
message = json.loads(frame)
except:
logger.exception("Could not decode the JSON message")
self.transport.close()
return
mtype = message.get('type', None)
if mtype == 'OK': # Server replied OK. We can advance the state
if self.state == STATE_OPEN:
logger.info("Channel open")
self.server_pub = self.asymmetric_encrypt.load_pub_from_str(
message["server_pub_key"].encode())
self.send_file(self.file_name)
elif self.state == STATE_DATA: # Got an OK during a message transfer.
# Reserved for future use
pass
else:
logger.warning("Ignoring message from server")
return
elif mtype == 'AUTHENTICATION_CHALLENGE': # Server replied with a challenge to authenticate clients
if self.state == STATE_OPEN and self.auth_type == "cc":
logger.info(
"Authentication process, signing challenge from server")
challenge_response = self.citizen_card.sign_with_cc(
message["challenge"])
# Only need the digital signature certificate from cc
self.certificate = self.citizen_card.get_x509_certificates(
KEY_USAGE=lambda x: x.value.digital_signature)
self.certificate = self.certificate[0]
# Convert certificate to bytes
bytes_cert = self.citizen_auth.serialize(self.certificate)
self._send({
'type':
'AUTHENTICATION_RESPONSE',
'challenge':
message['challenge'],
'response':
base64.b64encode(bytes(challenge_response)).decode(),
'certificate':
base64.b64encode(bytes_cert).decode()
})
elif mtype == 'CHALLENGE OK':
message = {
'type': 'OPEN',
'file_name': self.file_name,
'symmetric_cypher': self.symmetric_cypher,
'cypher_mode': self.cypher_mode,
'synthesis_algorithm': self.synthesis_algorithm,
'client_public_key': self._public_key.decode()
}
self._send(message)
self.challenge_passed = True
elif mtype == 'ERROR':
logger.warning("Got error from server: {}".format(
message.get('data', None)))
else:
logger.warning("Invalid message type")
# self.transport.close()
# self.loop.stop()
def connection_lost(self, exc):
"""
Connection was lost for some reason.
:param exc:
:return:
"""
logger.info('The server closed the connection')
self.loop.stop()
def send_file(self, file_name: str) -> None:
"""
Sends a file to the server.
The file is read in chunks, encoded to Base64 and sent as part of a DATA JSON message
:param file_name: File to send
:return: None
"""
with open(file_name, 'rb') as f:
message = {'type': 'DATA', 'data': None}
read_size = 16 * 60
while True:
data = f.read(16 * 60)
message['data'] = base64.b64encode(data).decode()
self._send(message)
if len(data) != read_size:
break
self._send({'type': 'CLOSE'})
logger.info("File transferred. Closing transport")
self.transport.close()
def _send(self, message: str) -> None:
"""
Effectively encodes and sends a message
:param message:
:return:
"""
logger.debug("Send {}".format(message))
message_b = (json.dumps(message) + '\r\n').encode()
if self.state == STATE_CONNECT:
message_b = self.symmetric.handshake_encrypt(message_b)
elif self.state == STATE_OPEN and self.challenge_passed is False:
message_b = self.symmetric.handshake_encrypt(message_b)
elif self.state == STATE_OPEN and self.challenge_passed is True:
message_b = self.symmetric.encrypt(self.symmetric_cypher,
message_b,
self.synthesis_algorithm,
self.cypher_mode,
pkey=self.server_pub)
self.state = STATE_DATA
elif self.state == STATE_DATA:
message_b = self.symmetric.encrypt(self.symmetric_cypher,
message_b,
self.synthesis_algorithm,
self.cypher_mode,
pkey=self.server_pub)
self.transport.write(message_b)
# Este sleep é um workaround para os blocos/pacotes perdidos ou incompletos
time.sleep(0.05)
class ClientHandler(asyncio.Protocol):
def __init__(self, signal):
"""
Default constructor
"""
self.signal = signal
self.state = 0
self.file = None
self.file_name = None
self.file_path = None
self.storage_dir = storage_dir
self.buffer = ''
self.peername = ''
self.asymmetric_encrypt = Asymmetric()
self.symmetric = Symmetric()
self.symmetric_cypher = None
self.cypher_mode = None
self.synthesis_algorithm = None
self.client_pub = None
self.server_pub = None
self.server_priv = None
def connection_made(self, transport) -> None:
"""
Called when a client connects
:param transport: The transport stream to use with this client
:return:
"""
self.peername = transport.get_extra_info('peername')
logger.info('\n\nConnection from {}'.format(self.peername))
logger.info(
"New client connected, introduce password to generate rsa key")
password = getpass.getpass('Password:'******'Received: {}'.format(data))
if self.state == STATE_CONNECT:
data = self.symmetric.handshake_decrypt(data)
else:
data = self.symmetric.decrypt(self.symmetric_cypher,
data,
self.synthesis_algorithm,
self.cypher_mode,
privkey=self.server_priv)
try:
self.buffer += data.decode()
except:
logger.exception('Could not decode data from client')
idx = self.buffer.find('\r\n')
while idx >= 0: # While there are separators
frame = self.buffer[:idx + 2].strip() # Extract the JSON object
self.buffer = self.buffer[
idx + 2:] # Removes the JSON object from the buffer
self.on_frame(frame) # Process the frame
idx = self.buffer.find('\r\n')
if len(self.buffer
) > 4096 * 1024 * 1024: # If buffer is larger than 4M
logger.warning('Buffer to large')
self.buffer = ''
self.transport.close()
def on_frame(self, frame: str) -> None:
"""
Called when a frame (JSON Object) is extracted
:param frame: The JSON object to process
:return:
"""
# logger.debug("Frame: {}".format(frame))
try:
message = json.loads(frame)
except:
logger.exception("Could not decode JSON message: {}".format(frame))
self.transport.close()
return
mtype = message.get('type', "").upper()
if mtype == 'OPEN':
ret = self.process_open(message)
elif mtype == 'DATA':
ret = self.process_data(message)
elif mtype == 'CLOSE':
ret = self.process_close(message)
else:
logger.warning("Invalid message type: {}".format(message['type']))
ret = False
if not ret:
try:
self._send({'type': 'ERROR', 'message': 'See server'})
except:
pass # Silently ignore
logger.info("Closing transport")
if self.file is not None:
self.file.close()
self.file = None
self.state = STATE_CLOSE
self.transport.close()
def process_open(self, message: str) -> bool:
"""
Processes an OPEN message from the client
This message should contain the filename
:param message: The message to process
:return: Boolean indicating the success of the operation
"""
logger.debug("Process Open: {}".format(message))
if self.state != STATE_CONNECT:
logger.warning("Invalid state. Discarding")
return False
if not 'file_name' in message:
logger.warning("No filename in Open")
return False
# print(message["client_public_key"])
self.client_pub = self.asymmetric_encrypt.load_pub_from_str(
message["client_public_key"].encode())
self.symmetric_cypher = message["symmetric_cypher"]
self.cypher_mode = message["cypher_mode"]
self.synthesis_algorithm = message["synthesis_algorithm"]
# Only chars and letters in the filename
file_name = re.sub(r'[^\w\.]', '', message['file_name'])
file_path = os.path.join(self.storage_dir, file_name)
if not os.path.exists("files"):
try:
os.mkdir("files")
except:
logger.exception("Unable to create storage directory")
return False
try:
self.file = open(file_path, "wb")
logger.info("File open")
except Exception:
logger.exception("Unable to open file")
return False
self._send({'type': 'OK', 'server_pub_key': self.server_pub.decode()})
self.file_name = file_name
self.file_path = file_path
self.state = STATE_OPEN
return True
def process_data(self, message: str) -> bool:
"""
Processes a DATA message from the client
This message should contain a chunk of the file
:param message: The message to process
:return: Boolean indicating the success of the operation
"""
logger.debug("Process Data: {}".format(message))
if self.state == STATE_OPEN:
self.state = STATE_DATA
# First Packet
elif self.state == STATE_DATA:
# Next packets
pass
else:
logger.warning("Invalid state. Discarding")
return False
try:
data = message.get('data', None)
if data is None:
logger.debug("Invalid message. No data found")
return False
bdata = base64.b64decode(message['data'])
except:
logger.exception(
"Could not decode base64 content from message.data")
return False
try:
self.file.write(bdata)
self.file.flush()
except:
logger.exception("Could not write to file")
return False
return True
def process_close(self, message: str) -> bool:
"""
Processes a CLOSE message from the client.
This message will trigger the termination of this session
:param message: The message to process
:return: Boolean indicating the success of the operation
"""
logger.debug("Process Close: {}".format(message))
self.transport.close()
if self.file is not None:
self.file.close()
self.file = None
if os.path.exists("server-keys"):
os.remove("server-keys/" + str(self.peername[1]) +
'_private_rsa.pem')
os.remove("server-keys/" + str(self.peername[1]) +
'_public_rsa.pem')
self.state = STATE_CLOSE
return True
def _send(self, message: str) -> None:
"""
Effectively encodes and sends a message
:param message:
:return:
"""
logger.debug("Send: {}".format(message))
message_b = (json.dumps(message) + '\r\n').encode()
if self.state == STATE_CONNECT:
message_b = self.symmetric.handshake_encrypt(message_b)
self.transport.write(message_b)