filebridging/src/client.py

import argparse
import asyncio
import collections
import logging
# import signal
import os
import ssl


class Client:
    def __init__(self, host='localhost', port=3001,
                 buffer_chunk_size=10**4, buffer_length_limit=10**4,
                 password=None):
        self._password = password
        self._host = host
        self._port = port
        self._stopping = False
        # Shared queue of bytes
        self.buffer = collections.deque()
        # How many bytes per chunk
        self._buffer_chunk_size = buffer_chunk_size
        # How many chunks in buffer
        self._buffer_length_limit = buffer_length_limit
        self._file_path = None
        self._working = False
        self._ssl_context = None
        self._encryption_complete = False

    @property
    def host(self) -> str:
        return self._host

    @property
    def port(self) -> int:
        return self._port

    @property
    def stopping(self) -> bool:
        return self._stopping

    @property
    def buffer_length_limit(self) -> int:
        return self._buffer_length_limit

    @property
    def buffer_chunk_size(self) -> int:
        return self._buffer_chunk_size

    @property
    def file_path(self) -> str:
        return self._file_path

    @property
    def working(self) -> bool:
        return self._working

    @property
    def ssl_context(self) -> ssl.SSLContext:
        return self._ssl_context

    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

    @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,
                                                       port=self.port,
                                                       ssl=self.ssl_context)
        writer.write("sender\n".encode('utf-8'))
        await writer.drain()
        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:
            file_path = self.file_path + '.enc'
            # Remove already-encrypted file if present (salt would differ)
            if os.path.isfile(file_path):
                os.remove(file_path)
            asyncio.ensure_future(
                self.encrypt_file(
                    input_file=self.file_path,
                    output_file=file_path
                )
            )
            # Give encryption an edge
            while not os.path.isfile(file_path):
                await asyncio.sleep(.5)
        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)
                    await writer.drain()
                except ConnectionResetError:
                    logging.info('Server closed the connection.')
                    self.stop()
                    break
        writer.close()
        return

    async def run_receiving_client(self, file_path='~/input.txt'):
        self._file_path = file_path
        reader, writer = await asyncio.open_connection(host=self.host,
                                                       port=self.port,
                                                       ssl=self.ssl_context)
        writer.write("receiver\n".encode('utf-8'))
        await writer.drain()
        await reader.readline()  # Wait for server start signal
        await self.receive(reader=reader)

    async def receive(self, reader: asyncio.StreamReader):
        self._working = True
        file_path = self.file_path
        logging.info("Receiving file...")
        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)
        logging.info("File received.")
        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:
            logging.info("Received interruption signal, stopping...")
            self._stopping = True
        else:
            raise KeyboardInterrupt("Not working yet...")


def get_action(action):
    """Parse abbreviations for `action`."""
    if not isinstance(action, str):
        return
    elif action.lower().startswith('r'):
        return 'receive'
    elif action.lower().startswith('s'):
        return 'send'


def get_file_path(path, action='receive'):
    """Check that file `path` is correct and return it."""
    if (
            isinstance(path, str)
            and action == 'send'
            and os.path.isfile(path)
    ):
        return path
    elif (
            isinstance(path, str)
            and action == 'receive'
            and os.access(os.path.dirname(os.path.abspath(path)), os.W_OK)
    ):
        return path
    elif path is not None:
        logging.error(f"Invalid file: `{path}`")


if __name__ == '__main__':
    # noinspection SpellCheckingInspection
    log_formatter = logging.Formatter(
        "%(asctime)s [%(module)-15s %(levelname)-8s]     %(message)s",
        style='%'
    )
    root_logger = logging.getLogger()
    root_logger.setLevel(logging.DEBUG)

    console_handler = logging.StreamHandler()
    console_handler.setFormatter(log_formatter)
    console_handler.setLevel(logging.DEBUG)
    root_logger.addHandler(console_handler)

    # Parse command-line arguments
    cli_parser = argparse.ArgumentParser(description='Run client',
                                         allow_abbrev=False)
    cli_parser.add_argument('--_host', type=str,
                            default=None,
                            required=False,
                            help='server address')
    cli_parser.add_argument('--_port', type=int,
                            default=None,
                            required=False,
                            help='server _port')
    cli_parser.add_argument('--action', type=str,
                            default=None,
                            required=False,
                            help='[S]end or [R]eceive')
    cli_parser.add_argument('--path', type=str,
                            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='*',
                            help='[S]end or [R]eceive (see `action`)')
    args = vars(cli_parser.parse_args())
    _host = args['_host']
    _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:
        try:
            from config import host as _host
        except ImportError:
            _host = None
    if _port is None:
        try:
            from config import port as _port
        except ImportError:
            _port = None
    # Take `s`, `r` etc. from command line as `_action`
    if _action is None:
        for arg in args['others']:
            _action = get_action(arg)
            if _action:
                break
    if _action is None:
        try:
            from config import action as _action
            _action = get_action(_action)
        except ImportError:
            _action = None
    if _file_path is None:
        try:
            from config import file_path as _file_path
            _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:
        _host = input("Enter _host:\t\t\t\t\t\t")
    while _port is None:
        try:
            _port = int(input("Enter _port:\t\t\t\t\t\t"))
        except ValueError:
            logging.info("Invalid _port. Enter a valid _port number!")
            _port = None
    while _action is None:
        _action = get_action(
            input("Do you want to (R)eceive or (S)end a file?\t\t")
        )
    while _file_path is None:
        _file_path = get_file_path(
            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
        _ssl_context = ssl.create_default_context(ssl.Purpose.SERVER_AUTH)
        _ssl_context.check_hostname = False
        _ssl_context.load_verify_locations(certificate)
        client.set_ssl_context(_ssl_context)
    except ImportError:
        logging.warning("Please consider using SSL.")
        certificate, key = None, None
    logging.info("Starting client...")
    if _action == 'send':
        loop.run_until_complete(
            client.run_sending_client(
                file_path=_file_path
            )
        )
    else:
        loop.run_until_complete(
            client.run_receiving_client(
                file_path=_file_path
            )
        )
    loop.close()
    logging.info("Stopped client")