comportamento sftp come ftp
This commit is contained in:
@@ -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):
|
||||
"""
|
||||
|
||||
Reference in New Issue
Block a user