Queer European MD passionate about IT
Browse Source

Allow end-to-end encryption

Davte 5 years ago
parent
commit
0069f9e5ea
2 changed files with 77 additions and 5 deletions
  1. 1 0
      requirements.txt
  2. 76 5
      src/client.py

+ 1 - 0
requirements.txt

@@ -0,0 +1 @@
+cryptography

+ 76 - 5
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':