#!/usr/bin/env python3

import sys
import os
# import ssl

import re
import logging

import psycopg2

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, TLS_FTPHandler
from pyftpdlib.servers import FTPServer
from pyftpdlib.authorizers import DummyAuthorizer, AuthenticationFailed

def conn_db(cfg):
    return  psycopg2.connect(dbname=cfg.dbname, 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):
    def __init__(self, cfg):
        # Initialize the DummyAuthorizer and add the admin user
        super().__init__()
        self.add_user(
            cfg.adminuser[0], cfg.adminuser[1], cfg.adminuser[2], perm=cfg.adminuser[3])

        # Definisci la connessione al database
        conn = conn_db(cfg)

        # Crea un cursore
        cur = conn.cursor()
        cur.execute(f'SELECT ftpuser, hash, virtpath, perm FROM {cfg.dbschema}.{cfg.dbusertable} WHERE deleted_at IS NULL')

        for ftpuser, hash, virtpath, perm in cur.fetchall():
            self.add_user(ftpuser, hash, virtpath, perm)
            try:
                Path(cfg.virtpath + ftpuser).mkdir(parents=True, exist_ok=True)
            except:
                self.responde('551 Error in create virtual user path.')

    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):

    def __init__(self, conn, server, ioloop=None):
        # Initialize the FTPHandler and add custom commands
        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 DELU': dict(perm='M', auth=True, arg=True,
                               help='Syntax: SITE <SP> DELU USERNAME (remove virtual user).')}
        )
        self.proto_cmds.update(
            {'SITE RESU': dict(perm='M', auth=True, arg=True,
                               help='Syntax: SITE <SP> RESU USERNAME (restore 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):
        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)

                    # Crea un cursore
                    cur = conn.cursor()
                    try:
                        cur.execute(f"INSERT INTO {cfg.dbschema}.{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(), lines))
                        conn.commit()
                        conn.close()

                    except psycopg2.Error 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):
        # Remove partially uploaded files
        os.remove(file)

    def ftp_SITE_ADDU(self, line):
        """
            Add a virtual user and save the virtuser configuration file.
            Create a directory for the virtual user in the specified virtpath.
        """
        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:
            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:
                self.respond('551 Error in create virtual user path.')
            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
                    # Definisci la connessione al database
                    conn = conn_db(cfg)

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

    def ftp_SITE_DELU(self, line):
        """
            Remove a virtual user and save the virtuser configuration file.
        """
        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.dbschema}.{cfg.dbusertable} SET deleted_at = now() WHERE ftpuser = '{user}'")
            conn.commit()
            conn.close()

            logging.info("User {} deleted.".format(user))
            self.respond('200 SITE DELU successful.')
        except psycopg2.Error as e:
            self.respond('501 SITE DELU failed.')
            print(e)

    def ftp_SITE_RESU(self, line):
        """
            Restore a virtual user and save the virtuser configuration file.
        """
        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.dbschema}.{cfg.dbusertable} SET deleted_at = null WHERE ftpuser = '{user}'")
                conn.commit()
            except psycopg2.Error as e:
                logging.error("Update DB failed: {}".format(e))

            cur.execute(f"SELECT ftpuser, hash, virtpath, perm FROM {cfg.dbschema}.{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:
                self.responde('551 Error in create virtual user path.')

            conn.close()

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

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

    def ftp_SITE_LSTU(self, line):
        """
            List all virtual users.
        """
        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.dbschema}.{cfg.dbusertable} WHERE deleted_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:
            self.respond('501 list users failed.')

def main():
    # 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: {}.".format("Shutdown requested...exiting")
        )

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

if __name__ == "__main__":
    main()