ASE/FtpCsvReceiver.py

#!.venv/bin/python
"""This module implements an FTP server with custom commands for managing virtual users and handling CSV file uploads."""

import sys
import os
# import ssl

import re
import logging

import mysql.connector

from hashlib import sha256
from pathlib import Path

from utils.time import timestamp_fmt as ts
from utils.config import set_config as setting

from pyftpdlib.handlers import FTPHandler
from pyftpdlib.servers import FTPServer
from pyftpdlib.authorizers import DummyAuthorizer, AuthenticationFailed


def conn_db(cfg):
    """Establishes a connection to the MySQL database.

    Args:
        cfg: The configuration object containing database connection details.

    Returns:
        A MySQL database connection object.
    """
    return mysql.connector.connect(user=cfg.dbuser, password=cfg.dbpass, host=cfg.dbhost, port=cfg.dbport)


def extract_value(patterns, primary_source, secondary_source, default='Not Defined'):
    """Extracts the first match for a list of patterns from the primary source.
    Falls back to the secondary source if no match is found.
    """
    for source in (primary_source, secondary_source):
        for pattern in patterns:
            matches = re.findall(pattern, source, re.IGNORECASE)
            if matches:
                return matches[0]  # Return the first match immediately
    return default  # Return default if no matches are found



class DummySha256Authorizer(DummyAuthorizer):
    """Custom authorizer that uses SHA256 for password hashing and manages users from a database."""

    def __init__(self, cfg):
        """Initializes the authorizer, adds the admin user, and loads users from the database.

        Args:
            cfg: The configuration object.
        """
        super().__init__()
        self.add_user(
            cfg.adminuser[0], cfg.adminuser[1], cfg.adminuser[2], perm=cfg.adminuser[3])

        # Define the database connection
        conn = conn_db(cfg)

        # Create a cursor
        cur = conn.cursor()
        cur.execute(f'SELECT ftpuser, hash, virtpath, perm FROM {cfg.dbname}.{cfg.dbusertable} WHERE disabled_at IS NULL')

        for ftpuser, hash, virtpath, perm in cur.fetchall():
            self.add_user(ftpuser, hash, virtpath, perm)
            """
            Create the user's directory if it does not exist.
            """
            try:
                Path(cfg.virtpath + ftpuser).mkdir(parents=True, exist_ok=True)
            except  Exception as e:
                self.responde(f'551 Error in create virtual user path: {e}')

    def validate_authentication(self, username, password, handler):
        # Validate the user's password against the stored hash
        hash = sha256(password.encode("UTF-8")).hexdigest()
        try:
            if self.user_table[username]["pwd"] != hash:
                raise KeyError
        except KeyError:
            raise AuthenticationFailed

class ASEHandler(FTPHandler):
    """Custom FTP handler that extends FTPHandler with custom commands and file handling."""

    def __init__(self, conn, server, ioloop=None):
        """Initializes the handler, adds custom commands, and sets up command permissions.

        Args:
            conn: The connection object.
            server: The FTP server object.
            ioloop: The I/O loop object.
        """
        super().__init__(conn, server, ioloop)
        self.proto_cmds = FTPHandler.proto_cmds.copy()
        # Add custom FTP commands for managing virtual users - command in lowercase
        self.proto_cmds.update(
            {'SITE ADDU': dict(perm='M', auth=True, arg=True,
                               help='Syntax: SITE <SP> ADDU USERNAME PASSWORD (add virtual user).')}
        )
        self.proto_cmds.update(
            {'SITE DISU': dict(perm='M', auth=True, arg=True,
                               help='Syntax: SITE <SP> DISU USERNAME (disable virtual user).')}
        )
        self.proto_cmds.update(
            {'SITE ENAU': dict(perm='M', auth=True, arg=True,
                               help='Syntax: SITE <SP> ENAU USERNAME (enable virtual user).')}
        )
        self.proto_cmds.update(
            {'SITE LSTU': dict(perm='M', auth=True, arg=None,
                               help='Syntax: SITE <SP> LSTU (list virtual users).')}
        )

    def on_file_received(self, file):
        """Handles the event when a file is successfully received.

        Args:
            file: The path to the received file.
        """
        if not os.stat(file).st_size:
            os.remove(file)
            logging.info(f'File {file} was empty: removed.')
        else:
            cfg = self.cfg
            path, filenameExt = os.path.split(file)
            filename, fileExtension = os.path.splitext(filenameExt)
            if (fileExtension.upper() in (cfg.fileext)):
                with open(file, 'r') as csvfile:
                    lines = csvfile.readlines()

                    unit_name = extract_value(cfg.units_name, filename, str(lines[0:9]))
                    unit_type = extract_value(cfg.units_type, filename, str(lines[0:9]))
                    tool_name = extract_value(cfg.tools_name, filename, str(lines[0:9]))
                    tool_type = extract_value(cfg.tools_type, filename, str(lines[0:9]))

                    conn = conn_db(cfg)

                    # Create a cursor
                    cur = conn.cursor()
                    try:
                        cur.execute(f"INSERT INTO {cfg.dbname}.{cfg.dbrectable} (filename, unit_name, unit_type, tool_name, tool_type, tool_data) VALUES (%s, %s, %s, %s, %s, %s)", (filename, unit_name.upper(), unit_type.upper(), tool_name.upper(), tool_type.upper(), ''.join(lines)))
                        conn.commit()
                        conn.close()

                    except Exception as e:
                        logging.error(f'File {file} not loaded. Held in user path.')
                        logging.error(f'{e}')
                    else:
                        os.remove(file)
                        logging.info(f'File {file} loaded: removed.')

    def on_incomplete_file_received(self, file):
        """Removes partially uploaded files.
        Args:
            file: The path to the incomplete file.
        """
        os.remove(file)

    def ftp_SITE_ADDU(self, line):
        """Adds a virtual user, creates their directory, and saves their details to the database.
        """
        cfg = self.cfg
        try:
            parms = line.split()
            user = os.path.basename(parms[0])  # Extract the username
            password = parms[1]  # Get the password
            hash = sha256(password.encode("UTF-8")).hexdigest()  # Hash the password
        except  IndexError:
            self.respond('501 SITE ADDU failed. Command needs 2 arguments')
        else:
            try:
                # Create the user's directory
                Path(cfg.virtpath + user).mkdir(parents=True, exist_ok=True)
            except Exception as e:
                self.respond(f'551 Error in create virtual user path: {e}')
            else:
                try:
                    # Add the user to the authorizer
                    self.authorizer.add_user(str(user),
                                            hash, cfg.virtpath + "/" + user, perm=cfg.defperm)
                    # Save the user to the database
                    # Define the database connection
                    conn = conn_db(cfg)

                    # Create a cursor
                    cur = conn.cursor()
                    cur.execute(f"INSERT INTO {cfg.dbname}.{cfg.dbusertable} (ftpuser, hash, virtpath, perm) VALUES ('{user}', '{hash}', '{cfg.virtpath + user}', '{cfg.defperm}')")
                    conn.commit()
                    conn.close()
                    logging.info(f"User {user} created.")
                    self.respond('200 SITE ADDU successful.')
                except Exception as e:
                    self.respond(f'501 SITE ADDU failed: {e}.')
                    print(e)

    def ftp_SITE_DISU(self, line):
        """Removes a virtual user from the authorizer and marks them as deleted in the database."""
        cfg = self.cfg
        parms = line.split()
        user = os.path.basename(parms[0])  # Extract the username
        try:
            # Remove the user from the authorizer
            self.authorizer.remove_user(str(user))
            # Delete the user from database
            conn = conn_db(cfg)

            # Crea un cursore
            cur = conn.cursor()
            cur.execute(f"UPDATE {cfg.dbname}.{cfg.dbusertable} SET disabled_at = now() WHERE ftpuser = '{user}'")
            conn.commit()
            conn.close()

            logging.info(f"User {user} deleted.")
            self.respond('200 SITE DISU successful.')
        except Exception as e:
            self.respond('501 SITE DISU failed.')
            print(e)

    def ftp_SITE_ENAU(self, line):
        """Restores a virtual user by updating their status in the database and adding them back to the authorizer."""
        cfg = self.cfg
        parms = line.split()
        user = os.path.basename(parms[0])  # Extract the username
        try:
            # Restore the user into database
            conn = conn_db(cfg)

            # Crea un cursore
            cur = conn.cursor()
            try:
                cur.execute(f"UPDATE {cfg.dbname}.{cfg.dbusertable} SET disabled_at = null WHERE ftpuser = '{user}'")
                conn.commit()
            except Exception as e:
                logging.error(f"Update DB failed: {e}")

            cur.execute(f"SELECT ftpuser, hash, virtpath, perm FROM {cfg.dbname}.{cfg.dbusertable} WHERE ftpuser = '{user}'")

            ftpuser, hash, virtpath, perm = cur.fetchone()
            self.authorizer.add_user(ftpuser, hash, virtpath, perm)
            try:
                Path(cfg.virtpath + ftpuser).mkdir(parents=True, exist_ok=True)
            except  Exception as e:
                self.responde(f'551 Error in create virtual user path: {e}')

            conn.close()

            logging.info(f"User {user} restored.")
            self.respond('200 SITE ENAU successful.')

        except Exception as e:
            self.respond('501 SITE ENAU failed.')
            print(e)

    def ftp_SITE_LSTU(self, line):
        """Lists all virtual users from the database."""
        cfg = self.cfg
        users_list = []
        try:
            # Connect to the SQLite database to fetch users
            conn = conn_db(cfg)

            # Crea un cursore
            cur = conn.cursor()
            self.push("214-The following virtual users are defined:\r\n")
            cur.execute(f'SELECT ftpuser, perm FROM {cfg.dbname}.{cfg.dbusertable} WHERE disabled_at IS NULL ')
            [users_list.append(f'Username: {ftpuser}\tPerms: {perm}\r\n') for ftpuser, perm in cur.fetchall()]
            self.push(''.join(users_list))
            self.respond("214 LSTU SITE command successful.")

        except  Exception as e:
            self.respond(f'501 list users failed: {e}')

def main():
    """Main function to start the FTP server."""
    # Load the configuration settings
    cfg = setting.config()

    try:
        # Initialize the authorizer and handler
        authorizer = DummySha256Authorizer(cfg)
        handler = ASEHandler
        handler.cfg = cfg
        handler.authorizer = authorizer
        handler.masquerade_address = cfg.proxyaddr
        # Set the range of passive ports for the FTP server
        _range = list(range(cfg.firstport, cfg.firstport + cfg.portrangewidth))
        handler.passive_ports = _range

        # Configure logging
        logging.basicConfig(
            format="%(asctime)s %(message)s",
            filename=cfg.logfilename,
            level=logging.INFO,
        )

        # Create and start the FTP server
        server = FTPServer(("0.0.0.0", 2121), handler)
        server.serve_forever()
    except KeyboardInterrupt:
        logging.info(
            "Info: Shutdown requested...exiting"
        )

    except Exception:
        print(
            f"{ts.timestamp("log")} - PID {os.getpid():>5} >> Error: {sys.exc_info()[1]}."
        )


if __name__ == "__main__":
    main()