From 0069f9e5eaa490d8401cbed7357b66e182adf5f0 Mon Sep 17 00:00:00 2001 From: Davte Date: Fri, 10 Apr 2020 13:41:36 +0200 Subject: [PATCH] Allow end-to-end encryption --- requirements.txt | 1 + src/client.py | 81 +++++++++++++++++++++++++++++++++++++++++++++--- 2 files changed, 77 insertions(+), 5 deletions(-) create mode 100644 requirements.txt diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..0d38bc5 --- /dev/null +++ b/requirements.txt @@ -0,0 +1 @@ +cryptography diff --git a/src/client.py b/src/client.py index 3d3779f..bdc06a4 100644 --- a/src/client.py +++ b/src/client.py @@ -9,7 +9,9 @@ import ssl class Client: def __init__(self, host='localhost', port=3001, - buffer_chunk_size=10**4, buffer_length_limit=10**4): + buffer_chunk_size=10**4, buffer_length_limit=10**4, + password=None): + self._password = password self._host = host self._port = port self._stopping = False @@ -55,6 +57,11 @@ class Client: def set_ssl_context(self, ssl_context: ssl.SSLContext): self._ssl_context = ssl_context + @property + def password(self): + """Password for file encryption or decryption.""" + return self._password + async def run_sending_client(self, file_path='~/output.txt'): self._file_path = file_path reader, writer = await asyncio.open_connection(host=self.host, @@ -67,7 +74,30 @@ class Client: async def send(self, writer: asyncio.StreamWriter): self._working = True - with open(self.file_path, 'rb') as file_to_send: + file_path = self.file_path + if self.password: + logging.info("Encrypting file...") + file_path = self.file_path + '.enc' + stdout, stderr = ''.encode(), ''.encode() + try: + _subprocess = await asyncio.create_subprocess_shell( + "openssl enc -aes-256-cbc " + "-md sha512 -pbkdf2 -iter 100000 -salt " + f"-in {self.file_path} -out {file_path} " + f"-pass pass:{self.password}" + ) + stdout, stderr = await _subprocess.communicate() + except Exception as e: + logging.error( + "Exception {e}:\n{o}\n{er}".format( + e=e, + o=stdout.decode().strip(), + er=stderr.decode().strip() + ) + ) + + logging.info("Encryption completed. Sending file...") + with open(file_path, 'rb') as file_to_send: while not self.stopping: output_data = file_to_send.read(self.buffer_chunk_size) if not output_data: @@ -94,12 +124,36 @@ class Client: async def receive(self, reader: asyncio.StreamReader): self._working = True - with open(self.file_path, 'wb') as file_to_receive: + file_path = self.file_path + if self.password: + file_path += '.enc' + with open(file_path, 'wb') as file_to_receive: while not self.stopping: input_data = await reader.read(self.buffer_chunk_size) if not input_data: break file_to_receive.write(input_data) + if self.password: + logging.info("Decrypting file...") + stdout, stderr = ''.encode(), ''.encode() + try: + _subprocess = await asyncio.create_subprocess_shell( + "openssl enc -aes-256-cbc " + "-md sha512 -pbkdf2 -iter 100000 -salt -d " + f"-in {file_path} -out {self.file_path} " + f"-pass pass:{self.password}" + ) + stdout, stderr = await _subprocess.communicate() + logging.info("Decryption completed.") + except Exception as e: + logging.error( + "Exception {e}:\n{o}\n{er}".format( + e=e, + o=stdout.decode().strip(), + er=stderr.decode().strip() + ) + ) + logging.info("Decryption failed", exc_info=True) def stop(self, *_): if self.working: @@ -138,6 +192,7 @@ def get_file_path(path, action='receive'): if __name__ == '__main__': + # noinspection SpellCheckingInspection log_formatter = logging.Formatter( "%(asctime)s [%(module)-15s %(levelname)-8s] %(message)s", style='%' @@ -169,6 +224,10 @@ if __name__ == '__main__': default=None, required=False, help='File path') + cli_parser.add_argument('--password', '--p', '--pass', type=str, + default=None, + required=False, + help='Password for file encryption or decryption') cli_parser.add_argument('others', metavar='R or S', nargs='*', @@ -178,6 +237,7 @@ if __name__ == '__main__': _port = args['_port'] _action = get_action(args['action']) _file_path = args['path'] + _password = args['password'] # If _host and _port are not provided from command-line, try to import them if _host is None: @@ -199,7 +259,6 @@ if __name__ == '__main__': if _action is None: try: from config import action as _action - _action = get_action(_action) except ImportError: _action = None @@ -209,6 +268,11 @@ if __name__ == '__main__': _file_path = get_action(_file_path) except ImportError: _file_path = None + if _password is None: + try: + from config import password as _password + except ImportError: + _password = None # If import fails, prompt user for _host or _port while _host is None: @@ -228,10 +292,17 @@ if __name__ == '__main__': path=input(f"Enter file to {_action}:\t\t\t\t\t\t"), action=_action ) + if _password is None: + logging.warning( + "You have provided no password for file encryption.\n" + "Your file will be unencoded unless you provide a password in " + "config file." + ) loop = asyncio.get_event_loop() client = Client( host=_host, port=_port, + password=_password ) try: from config import certificate @@ -240,7 +311,7 @@ if __name__ == '__main__': _ssl_context.load_verify_locations(certificate) client.set_ssl_context(_ssl_context) except ImportError: - logging.info("Please consider using SSL.") + logging.warning("Please consider using SSL.") certificate, key = None, None logging.info("Starting client...") if _action == 'send':