diff --git a/vm1/src/utils/authorizers/database_authorizer.py b/vm1/src/utils/authorizers/database_authorizer.py index 080acd3..75f4f9f 100644 --- a/vm1/src/utils/authorizers/database_authorizer.py +++ b/vm1/src/utils/authorizers/database_authorizer.py @@ -113,8 +113,23 @@ class DatabaseAuthorizer(DummyAuthorizer): # Temporarily add user to in-memory table for this session # This allows pyftpdlib to work correctly for the duration of the session - if username not in self.user_table: - self.add_user(ftpuser, stored_hash, virtpath, perm) + # We add/update directly to avoid issues with add_user() checking if user exists + if username in self.user_table: + # User already exists, just update credentials + self.user_table[username]['pwd'] = stored_hash + self.user_table[username]['home'] = virtpath + self.user_table[username]['perm'] = perm + self.user_table[username]['operms'] = {} + else: + # User doesn't exist, add to table directly with all required fields + self.user_table[username] = { + 'pwd': stored_hash, + 'home': virtpath, + 'perm': perm, + 'operms': {}, # Optional per-directory permissions + 'msg_login': '230 Login successful.', + 'msg_quit': '221 Goodbye.' + } logger.info(f"Successful login for user: {username}") diff --git a/vm1/src/utils/servers/sftp_server.py b/vm1/src/utils/servers/sftp_server.py index 002c3f0..8743559 100644 --- a/vm1/src/utils/servers/sftp_server.py +++ b/vm1/src/utils/servers/sftp_server.py @@ -36,12 +36,14 @@ class ASESSHServer(asyncssh.SSHServer): def __init__(self, cfg): """Initialize SSH server with configuration.""" self.cfg = cfg + self.user_home_dirs = {} # Store user home directories after authentication super().__init__() def connection_made(self, conn): """Called when connection is established.""" # Store config in connection for later use conn._cfg = self.cfg + conn._ssh_server = self # Store reference to server for accessing user_home_dirs logger.info(f"SSH connection from {conn.get_extra_info('peername')[0]}") def connection_lost(self, exc): @@ -62,7 +64,9 @@ class ASESSHServer(asyncssh.SSHServer): # Check if user is admin if username == self.cfg.adminuser[0]: if self.cfg.adminuser[1] == password_hash: - logger.info(f"Admin user '{username}' authenticated successfully") + # Store admin home directory + self.user_home_dirs[username] = self.cfg.adminuser[2] + logger.info(f"Admin user '{username}' authenticated successfully (home: {self.cfg.adminuser[2]})") return True else: logger.warning(f"Failed admin login attempt for user: {username}") @@ -106,7 +110,10 @@ class ASESSHServer(asyncssh.SSHServer): logger.error(f"Failed to create directory for user {username}: {e}") return False - logger.info(f"Successful SFTP login for user: {username}") + # Store the user's home directory for chroot + self.user_home_dirs[username] = virtpath + + logger.info(f"Successful SFTP login for user: {username} (home: {virtpath})") return True except Exception as e: @@ -127,39 +134,79 @@ class SFTPFileHandler(asyncssh.SFTPServer): """Extended SFTP server with file upload handling.""" def __init__(self, chan): - super().__init__(chan) + super().__init__(chan, chroot=self._get_user_home(chan)) self.cfg = chan.get_connection()._cfg + self._open_files = {} # Track open files for processing - async def close(self): - """Handle session close.""" - await super().close() + @staticmethod + def _get_user_home(chan): + """Get the home directory for the authenticated user.""" + conn = chan.get_connection() + username = conn.get_extra_info('username') + ssh_server = getattr(conn, '_ssh_server', None) - # Override file operations to add custom handling - async def rename(self, oldpath, newpath): - """ - Handle file rename/move - called when upload completes. - This is where we trigger the CSV processing like in FTP. - """ - result = await super().rename(oldpath, newpath) + if ssh_server and username in ssh_server.user_home_dirs: + return ssh_server.user_home_dirs[username] - # Check if it's a CSV file that was uploaded - if newpath.lower().endswith('.csv'): - try: - # Trigger file processing (same as FTP on_file_received) - logger.info(f"CSV file uploaded via SFTP: {newpath}") - # Create a mock handler object with required attributes - mock_handler = type('obj', (object,), { - 'cfg': self.cfg, - 'username': self._chan.get_connection().get_extra_info('username') - })() + # Fallback for admin user + if hasattr(conn, '_cfg') and username == conn._cfg.adminuser[0]: + return conn._cfg.adminuser[2] - # Call the same file_management function used by FTP - file_management.on_file_received(mock_handler, newpath) - except Exception as e: - logger.error(f"Error processing SFTP uploaded file {newpath}: {e}", exc_info=True) + return None + + def open(self, path, pflags, attrs): + """Track files being opened for writing.""" + result = super().open(path, pflags, attrs) + + # If file is opened for writing (pflags contains FXF_WRITE) + if pflags & 0x02: # FXF_WRITE flag + real_path = self.map_path(path) + # Convert bytes to str if necessary + if isinstance(real_path, bytes): + real_path = real_path.decode('utf-8') + self._open_files[result] = real_path + logger.debug(f"File opened for writing: {real_path}") return result + async def close(self, file_obj): + """Process file after it's closed.""" + # Call parent close first (this doesn't return anything useful) + result = super().close(file_obj) + + # Check if this file was tracked + if file_obj in self._open_files: + filepath = self._open_files.pop(file_obj) + + # Process CSV files + if filepath.lower().endswith('.csv'): + try: + logger.info(f"CSV file closed after upload via SFTP: {filepath}") + + # Get username + username = self._chan.get_connection().get_extra_info('username') + + # Create a mock handler object with required attributes + mock_handler = type('obj', (object,), { + 'cfg': self.cfg, + 'username': username + })() + + # Call the file processing function + from utils.connect import file_management + await file_management.on_file_received_async(mock_handler, filepath) + except Exception as e: + logger.error(f"Error processing SFTP file on close: {e}", exc_info=True) + + return result + + async def exit(self): + """Handle session close.""" + await super().exit() + + # Note: File processing is handled in close() method, not here + # This avoids double-processing when both close and rename are called + async def start_sftp_server(cfg, host='0.0.0.0', port=22): """ diff --git a/vm2/src/utils/authorizers/database_authorizer.py b/vm2/src/utils/authorizers/database_authorizer.py index 080acd3..75f4f9f 100644 --- a/vm2/src/utils/authorizers/database_authorizer.py +++ b/vm2/src/utils/authorizers/database_authorizer.py @@ -113,8 +113,23 @@ class DatabaseAuthorizer(DummyAuthorizer): # Temporarily add user to in-memory table for this session # This allows pyftpdlib to work correctly for the duration of the session - if username not in self.user_table: - self.add_user(ftpuser, stored_hash, virtpath, perm) + # We add/update directly to avoid issues with add_user() checking if user exists + if username in self.user_table: + # User already exists, just update credentials + self.user_table[username]['pwd'] = stored_hash + self.user_table[username]['home'] = virtpath + self.user_table[username]['perm'] = perm + self.user_table[username]['operms'] = {} + else: + # User doesn't exist, add to table directly with all required fields + self.user_table[username] = { + 'pwd': stored_hash, + 'home': virtpath, + 'perm': perm, + 'operms': {}, # Optional per-directory permissions + 'msg_login': '230 Login successful.', + 'msg_quit': '221 Goodbye.' + } logger.info(f"Successful login for user: {username}") diff --git a/vm2/src/utils/servers/sftp_server.py b/vm2/src/utils/servers/sftp_server.py index 002c3f0..8743559 100644 --- a/vm2/src/utils/servers/sftp_server.py +++ b/vm2/src/utils/servers/sftp_server.py @@ -36,12 +36,14 @@ class ASESSHServer(asyncssh.SSHServer): def __init__(self, cfg): """Initialize SSH server with configuration.""" self.cfg = cfg + self.user_home_dirs = {} # Store user home directories after authentication super().__init__() def connection_made(self, conn): """Called when connection is established.""" # Store config in connection for later use conn._cfg = self.cfg + conn._ssh_server = self # Store reference to server for accessing user_home_dirs logger.info(f"SSH connection from {conn.get_extra_info('peername')[0]}") def connection_lost(self, exc): @@ -62,7 +64,9 @@ class ASESSHServer(asyncssh.SSHServer): # Check if user is admin if username == self.cfg.adminuser[0]: if self.cfg.adminuser[1] == password_hash: - logger.info(f"Admin user '{username}' authenticated successfully") + # Store admin home directory + self.user_home_dirs[username] = self.cfg.adminuser[2] + logger.info(f"Admin user '{username}' authenticated successfully (home: {self.cfg.adminuser[2]})") return True else: logger.warning(f"Failed admin login attempt for user: {username}") @@ -106,7 +110,10 @@ class ASESSHServer(asyncssh.SSHServer): logger.error(f"Failed to create directory for user {username}: {e}") return False - logger.info(f"Successful SFTP login for user: {username}") + # Store the user's home directory for chroot + self.user_home_dirs[username] = virtpath + + logger.info(f"Successful SFTP login for user: {username} (home: {virtpath})") return True except Exception as e: @@ -127,39 +134,79 @@ class SFTPFileHandler(asyncssh.SFTPServer): """Extended SFTP server with file upload handling.""" def __init__(self, chan): - super().__init__(chan) + super().__init__(chan, chroot=self._get_user_home(chan)) self.cfg = chan.get_connection()._cfg + self._open_files = {} # Track open files for processing - async def close(self): - """Handle session close.""" - await super().close() + @staticmethod + def _get_user_home(chan): + """Get the home directory for the authenticated user.""" + conn = chan.get_connection() + username = conn.get_extra_info('username') + ssh_server = getattr(conn, '_ssh_server', None) - # Override file operations to add custom handling - async def rename(self, oldpath, newpath): - """ - Handle file rename/move - called when upload completes. - This is where we trigger the CSV processing like in FTP. - """ - result = await super().rename(oldpath, newpath) + if ssh_server and username in ssh_server.user_home_dirs: + return ssh_server.user_home_dirs[username] - # Check if it's a CSV file that was uploaded - if newpath.lower().endswith('.csv'): - try: - # Trigger file processing (same as FTP on_file_received) - logger.info(f"CSV file uploaded via SFTP: {newpath}") - # Create a mock handler object with required attributes - mock_handler = type('obj', (object,), { - 'cfg': self.cfg, - 'username': self._chan.get_connection().get_extra_info('username') - })() + # Fallback for admin user + if hasattr(conn, '_cfg') and username == conn._cfg.adminuser[0]: + return conn._cfg.adminuser[2] - # Call the same file_management function used by FTP - file_management.on_file_received(mock_handler, newpath) - except Exception as e: - logger.error(f"Error processing SFTP uploaded file {newpath}: {e}", exc_info=True) + return None + + def open(self, path, pflags, attrs): + """Track files being opened for writing.""" + result = super().open(path, pflags, attrs) + + # If file is opened for writing (pflags contains FXF_WRITE) + if pflags & 0x02: # FXF_WRITE flag + real_path = self.map_path(path) + # Convert bytes to str if necessary + if isinstance(real_path, bytes): + real_path = real_path.decode('utf-8') + self._open_files[result] = real_path + logger.debug(f"File opened for writing: {real_path}") return result + async def close(self, file_obj): + """Process file after it's closed.""" + # Call parent close first (this doesn't return anything useful) + result = super().close(file_obj) + + # Check if this file was tracked + if file_obj in self._open_files: + filepath = self._open_files.pop(file_obj) + + # Process CSV files + if filepath.lower().endswith('.csv'): + try: + logger.info(f"CSV file closed after upload via SFTP: {filepath}") + + # Get username + username = self._chan.get_connection().get_extra_info('username') + + # Create a mock handler object with required attributes + mock_handler = type('obj', (object,), { + 'cfg': self.cfg, + 'username': username + })() + + # Call the file processing function + from utils.connect import file_management + await file_management.on_file_received_async(mock_handler, filepath) + except Exception as e: + logger.error(f"Error processing SFTP file on close: {e}", exc_info=True) + + return result + + async def exit(self): + """Handle session close.""" + await super().exit() + + # Note: File processing is handled in close() method, not here + # This avoids double-processing when both close and rename are called + async def start_sftp_server(cfg, host='0.0.0.0', port=22): """