#!.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 ADDU USERNAME PASSWORD (add virtual user).')} ) self.proto_cmds.update( {'SITE DISU': dict(perm='M', auth=True, arg=True, help='Syntax: SITE DISU USERNAME (disable virtual user).')} ) self.proto_cmds.update( {'SITE ENAU': dict(perm='M', auth=True, arg=True, help='Syntax: SITE ENAU USERNAME (enable virtual user).')} ) self.proto_cmds.update( {'SITE LSTU': dict(perm='M', auth=True, arg=None, help='Syntax: SITE 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()