""" SFTP Handler module - WebSocket SFTP operations. """ import asyncio import base64 import io import json import stat from typing import Optional import paramiko import websockets from websockets.server import WebSocketServerProtocol from .auth import verify_token from .config import HOST, MAX_UPLOAD_SIZE, SFTP_WS_PORT, SSH_TIMEOUT, logger from .database import creds_manager from .utils import get_client_ip async def handle_sftp_client(ws: WebSocketServerProtocol): """Handle SFTP file operations via WebSocket with progress tracking.""" client_ip = get_client_ip(ws) # Verify JWT token first username = await verify_token(ws) if not username: logger.warning(f"SFTP authentication failed for {client_ip}") await ws.close(1008, "Authentication required") return logger.info(f"SFTP client {username} connected from {client_ip}") ssh = None sftp = None current_file = None upload_state = {'bytes_received': 0, 'total_size': 0, 'remote_path': ''} try: async for message in ws: try: data = json.loads(message) msg_type = data.get('type') if msg_type == 'init': # Initialize SFTP connection conn_id = data.get('conn_id') connections = creds_manager.get_connections(username) conn_data = next((c for c in connections if c['id'] == conn_id), None) if not conn_data: await ws.send(json.dumps({'type': 'error', 'message': 'Connection not found'})) continue # Connect SSH ssh = paramiko.SSHClient() ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy()) connect_kwargs = { "hostname": conn_data['host'], "port": conn_data.get('port', 22), "username": conn_data['username'], "timeout": SSH_TIMEOUT, "look_for_keys": False, "allow_agent": False } if conn_data.get('key'): for key_type in [paramiko.RSAKey, paramiko.Ed25519Key, paramiko.ECDSAKey]: try: connect_kwargs["pkey"] = key_type.from_private_key(io.StringIO(conn_data['key'])) break except: continue elif conn_data.get('password'): connect_kwargs["password"] = conn_data['password'] await asyncio.get_event_loop().run_in_executor(None, lambda: ssh.connect(**connect_kwargs)) sftp = ssh.open_sftp() await ws.send(json.dumps({'type': 'ready'})) logger.info(f"SFTP ready for {conn_data['host']}") elif msg_type == 'upload_start': # Start file upload if not sftp: await ws.send(json.dumps({'type': 'error', 'message': 'SFTP not connected'})) continue upload_state['remote_path'] = data.get('remote_path') upload_state['total_size'] = data.get('total_size', 0) upload_state['bytes_received'] = 0 # Check file size limit if upload_state['total_size'] > MAX_UPLOAD_SIZE: await ws.send(json.dumps({ 'type': 'error', 'message': f'File too large. Max size: {MAX_UPLOAD_SIZE / 1024 / 1024:.1f}MB' })) continue # Open file for writing current_file = sftp.file(upload_state['remote_path'], 'wb') await ws.send(json.dumps({'type': 'upload_ready'})) elif msg_type == 'upload_chunk': # Receive file chunk if not current_file: await ws.send(json.dumps({'type': 'error', 'message': 'No upload in progress'})) continue chunk_data = base64.b64decode(data.get('data', '')) current_file.write(chunk_data) upload_state['bytes_received'] += len(chunk_data) # Send progress progress = int((upload_state['bytes_received'] / upload_state['total_size']) * 100) if upload_state['total_size'] > 0 else 0 await ws.send(json.dumps({ 'type': 'progress', 'bytes_received': upload_state['bytes_received'], 'total_size': upload_state['total_size'], 'progress': progress })) elif msg_type == 'upload_complete': # Finish upload if current_file: current_file.close() current_file = None await ws.send(json.dumps({ 'type': 'upload_done', 'success': True, 'bytes_received': upload_state['bytes_received'] })) logger.info(f"Upload complete: {upload_state['remote_path']} ({upload_state['bytes_received']} bytes)") elif msg_type == 'list': # List directory if not sftp: await ws.send(json.dumps({'type': 'error', 'message': 'SFTP not connected'})) continue path = data.get('path', '.') entries = sftp.listdir_attr(path) files = [] for entry in entries: files.append({ 'name': entry.filename, 'size': entry.st_size, 'is_dir': stat.S_ISDIR(entry.st_mode), 'mtime': entry.st_mtime, 'mode': oct(entry.st_mode)[-3:] }) files.sort(key=lambda x: (not x['is_dir'], x['name'].lower())) await ws.send(json.dumps({'type': 'file_list', 'files': files, 'path': path})) elif msg_type == 'mkdir': if not sftp: await ws.send(json.dumps({'type': 'error', 'message': 'SFTP not connected'})) continue sftp.mkdir(data.get('path')) await ws.send(json.dumps({'type': 'success', 'action': 'mkdir'})) elif msg_type == 'rm': if not sftp: await ws.send(json.dumps({'type': 'error', 'message': 'SFTP not connected'})) continue path = data.get('path') is_dir = data.get('is_dir', False) if is_dir: sftp.rmdir(path) else: sftp.remove(path) await ws.send(json.dumps({'type': 'success', 'action': 'rm'})) elif msg_type == 'mv': if not sftp: await ws.send(json.dumps({'type': 'error', 'message': 'SFTP not connected'})) continue sftp.rename(data.get('src'), data.get('dst')) await ws.send(json.dumps({'type': 'success', 'action': 'mv'})) except Exception as e: logger.error(f"SFTP error for {client_ip}: {e}") try: await ws.send(json.dumps({'type': 'error', 'message': str(e)})) except: pass except websockets.exceptions.ConnectionClosed: pass finally: if current_file: try: current_file.close() except: pass if sftp: try: sftp.close() except: pass if ssh: try: ssh.close() except: pass logger.info(f"SFTP client disconnected: {client_ip}")