Initial commit
Build and Push Multi-Arch Docker Image / build (push) Has been cancelled
Build and Push Multi-Arch Docker Image / build (push) Has been cancelled
This commit is contained in:
dido committed
commit
8ce01cb00a
34 files changed
+8346
No files matched your search
@@ -0,0 +1,214 @@
|
||||
"""
|
||||
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}")
|
||||
Reference in new issue
Block a user