215 lines
8.3 KiB
Python
215 lines
8.3 KiB
Python
"""
|
|
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}")
|