From 5f1c8fdf47695d8bb4ad6f3a16c426cc4d874d46 Mon Sep 17 00:00:00 2001
From: Davte <davte@libero.it>
Date: Fri, 10 Apr 2020 15:15:45 +0200
Subject: [PATCH] Send file while encrypting it

---
 src/client.py | 59 ++++++++++++++++++++++++++++++++++-----------------
 1 file changed, 40 insertions(+), 19 deletions(-)

diff --git a/src/client.py b/src/client.py
index bdc06a4..5a1bf4f 100644
--- a/src/client.py
+++ b/src/client.py
@@ -21,6 +21,7 @@ class Client:
         self._file_path = None
         self._working = False
         self._ssl_context = None
+        self._encryption_complete = False
 
     @property
     def host(self) -> str:
@@ -62,6 +63,10 @@ class Client:
         """Password for file encryption or decryption."""
         return self._password
 
+    @property
+    def encryption_complete(self):
+        return self._encryption_complete
+
     async def run_sending_client(self, file_path='~/output.txt'):
         self._file_path = file_path
         reader, writer = await asyncio.open_connection(host=self.host,
@@ -72,35 +77,51 @@ class Client:
         await reader.readline()  # Wait for server start signal
         await self.send(writer=writer)
 
+    async def encrypt_file(self, input_file, output_file):
+        self._encryption_complete = False
+        logging.info("Encrypting file...")
+        stdout, stderr = ''.encode(), ''.encode()
+        try:
+            _subprocess = await asyncio.create_subprocess_shell(
+                "openssl enc -aes-256-cbc "
+                "-md sha512 -pbkdf2 -iter 100000 -salt "
+                f"-in {input_file} -out {output_file} "
+                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.")
+        self._encryption_complete = True
+
     async def send(self, writer: asyncio.StreamWriter):
         self._working = True
         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}"
+            asyncio.ensure_future(
+                self.encrypt_file(
+                    input_file=self.file_path,
+                    output_file=file_path
                 )
-                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...")
+            )
+            while not os.path.isfile(file_path):
+                await asyncio.sleep(0.1)  # Let the encryption begin
+        logging.info("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:
+                    # If encryption is in progress, wait and read again later
+                    if self.password and not self.encryption_complete:
+                        await asyncio.sleep(1)
+                        continue
                     break
                 try:
                     writer.write(output_data)