670 lines
25 KiB
Python
670 lines
25 KiB
Python
"""
|
|
VM Bench — Proxmox Image Conversion Frontend
|
|
|
|
FastAPI web application serving as the user GUI.
|
|
Runs in the vm-bench LXC on port 5000.
|
|
Communicates with the backend at http://10.2.0.2:9000.
|
|
|
|
Routes:
|
|
GET / New session form
|
|
POST /session/start Upload + analyse source image
|
|
POST /session/confirm Submit conversion job
|
|
GET /session/status/{id} Poll job status (JSON)
|
|
POST /session/cleanup/{id} Clean up / reuse staging files
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from logging.handlers import RotatingFileHandler
|
|
import os
|
|
import shutil
|
|
import subprocess
|
|
import threading
|
|
import time
|
|
import uuid
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
import requests as http_requests
|
|
|
|
from fastapi import FastAPI, Form, Request, UploadFile, File
|
|
from fastapi.responses import HTMLResponse, JSONResponse, StreamingResponse
|
|
from fastapi.staticfiles import StaticFiles
|
|
|
|
from jinja2 import Environment, FileSystemLoader
|
|
|
|
from api_client import ApiClient, ApiError, BACKEND_URL
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# App setup
|
|
# ---------------------------------------------------------------------------
|
|
app = FastAPI(title="VM Bench — Frontend", version="1.0.0")
|
|
|
|
# ── Request logging middleware (logs before handler, catches silent failures) ──
|
|
@app.middleware("http")
|
|
async def log_requests(request: Request, call_next):
|
|
cl = request.headers.get("content-length", "?")
|
|
logger.info("→ %s %s (body %s bytes)", request.method, request.url.path, cl)
|
|
try:
|
|
response = await call_next(request)
|
|
# Prevent browser caching (especially important for JS-heavy pages)
|
|
response.headers["Cache-Control"] = "no-cache, no-store, must-revalidate"
|
|
response.headers["Pragma"] = "no-cache"
|
|
response.headers["Expires"] = "0"
|
|
return response
|
|
except Exception:
|
|
logger.exception("← %s %s FAILED", request.method, request.url.path)
|
|
raise
|
|
|
|
BASE = Path(__file__).parent
|
|
app.mount("/static", StaticFiles(directory=str(BASE / "static")), name="static")
|
|
|
|
_jinja = Environment(loader=FileSystemLoader(str(BASE / "templates")), autoescape=True)
|
|
STAGING_ROOT = Path("/mnt/converter/tmp")
|
|
|
|
def _staging(session_id: str = "") -> Path:
|
|
"""Session-aware staging directory. Returns /mnt/converter/tmp/{guid}/ or /mnt/converter/tmp/"""
|
|
return STAGING_ROOT / session_id if session_id else STAGING_ROOT
|
|
|
|
def _get_session_id(request: Request) -> str:
|
|
"""Extract session ID from cookie or query param."""
|
|
sid = request.cookies.get("vm_bench_sid", "")
|
|
if not sid:
|
|
sid = request.query_params.get("session_id", "")
|
|
return sid
|
|
|
|
api = ApiClient()
|
|
|
|
logger = logging.getLogger("vm-bench")
|
|
logger.setLevel(logging.INFO)
|
|
|
|
# Console handler → systemd journal via stdout
|
|
_ch = logging.StreamHandler()
|
|
_ch.setFormatter(logging.Formatter("%(asctime)s %(levelname)s %(message)s"))
|
|
logger.addHandler(_ch)
|
|
|
|
# File handler → shared log directory on /mnt/converter
|
|
LOG_DIR = Path("/mnt/converter/logs")
|
|
LOG_DIR.mkdir(parents=True, exist_ok=True)
|
|
_fh = RotatingFileHandler(
|
|
LOG_DIR / "vm-bench.log", maxBytes=10 * 1024 * 1024, backupCount=5,
|
|
)
|
|
_fh.setFormatter(logging.Formatter("%(asctime)s %(levelname)s %(name)s %(message)s"))
|
|
logger.addHandler(_fh)
|
|
|
|
logger.info("Frontend starting — log file: %s", LOG_DIR / "vm-bench.log")
|
|
|
|
# Config
|
|
VM_ID_MIN = 21000
|
|
VM_ID_MAX = 21100
|
|
DEFAULT_STORAGE = "local-lvm"
|
|
MIN_FREE_DISK_GB = 2 # keep 2 GB headroom
|
|
DOWNLOAD_TIMEOUT = 14400 # 4 hours absolute max (safety net)
|
|
SPEED_CHECK_AFTER = 30 # wait N seconds before judging speed
|
|
MAX_ETA_SECONDS = 3600 # kill download if ETA > 1 hour
|
|
UPLOAD_PROGRESS_INTERVAL = 1024**3 # log every 1 GiB during upload
|
|
|
|
# Track active background downloads for progress polling
|
|
_active_downloads: dict[str, dict] = {}
|
|
|
|
def _validate_vmid(vmid: int) -> Optional[str]:
|
|
if not (VM_ID_MIN <= vmid <= VM_ID_MAX):
|
|
return f"VM ID must be between {VM_ID_MIN} and {VM_ID_MAX}."
|
|
return None
|
|
|
|
def _check_disk_space(path: Path, needed_gb: int) -> Optional[str]:
|
|
"""Return an error message if <path> has less than <needed_gb> + headroom free."""
|
|
usage = shutil.disk_usage(path.parent if path.is_file() or not path.exists() else path)
|
|
free_gb = usage.free / (1024**3)
|
|
required = needed_gb + MIN_FREE_DISK_GB
|
|
if free_gb < required:
|
|
return (
|
|
f"Insufficient disk space on {path.parent}: "
|
|
f"{free_gb:.1f} GB free, {required:.1f} GB needed "
|
|
f"({needed_gb} GB file + {MIN_FREE_DISK_GB} GB headroom). "
|
|
f"Free up space or use a smaller file."
|
|
)
|
|
return None
|
|
|
|
def render(name: str, status: int = 200, **ctx) -> HTMLResponse:
|
|
tpl = _jinja.get_template(name)
|
|
return HTMLResponse(tpl.render(**ctx), status_code=status)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Routes
|
|
# ---------------------------------------------------------------------------
|
|
|
|
@app.get("/", response_class=HTMLResponse)
|
|
async def index(request: Request, vmid: Optional[int] = None, source: Optional[str] = None):
|
|
"""Landing page — new session form. Optionally pre-fills vmid + source for reuse."""
|
|
return render("index.html", request=request, backend_url=BACKEND_URL,
|
|
prefill_vmid=vmid or "", prefill_source=source or "",
|
|
prefill_vmname="")
|
|
|
|
|
|
@app.post("/session/upload")
|
|
def session_upload(
|
|
request: Request,
|
|
vmid: int = Form(...),
|
|
vm_name: str = Form(""),
|
|
source_type: str = Form("upload"),
|
|
source_file: Optional[UploadFile] = File(None),
|
|
source_url: Optional[str] = Form(None),
|
|
session_id: str = Form(""),
|
|
):
|
|
"""Phase 1 — acquire the source file. Returns JSON so the frontend can
|
|
show progress, then call /session/analyze separately."""
|
|
sdir = _staging(session_id) / "in"
|
|
sdir.mkdir(parents=True, exist_ok=True)
|
|
err = _validate_vmid(vmid)
|
|
if err:
|
|
return JSONResponse({"phase": "error", "error": err}, status_code=400)
|
|
|
|
vm_name = vm_name.strip()
|
|
|
|
# ── Upload ──────────────────────────────────────────────────────
|
|
if source_type == "upload":
|
|
if not source_file or not source_file.filename:
|
|
return JSONResponse({"phase": "error", "error": "No file uploaded."}, status_code=400)
|
|
|
|
filename = source_file.filename
|
|
dest = sdir / filename
|
|
|
|
content_length = request.headers.get("content-length")
|
|
if content_length:
|
|
estimated_gb = int(content_length) / (1024**3)
|
|
err = _check_disk_space(dest, estimated_gb)
|
|
if err:
|
|
return JSONResponse({"phase": "error", "error": err}, status_code=400)
|
|
|
|
try:
|
|
logger.info("Receiving upload: %s (%s bytes)", filename, content_length or "unknown")
|
|
written = 0
|
|
with dest.open("wb") as f:
|
|
while True:
|
|
chunk = source_file.file.read(8 * 1024 * 1024)
|
|
if not chunk:
|
|
break
|
|
f.write(chunk)
|
|
written += len(chunk)
|
|
if written % UPLOAD_PROGRESS_INTERVAL < len(chunk):
|
|
logger.info("Upload progress: %s — %.1f GiB", filename, written / (1024**3))
|
|
file_size_gb = round(written / (1024**3), 1)
|
|
logger.info("Upload complete: %s (%.1f GiB)", filename, file_size_gb)
|
|
except OSError as exc:
|
|
if dest.exists():
|
|
dest.unlink(missing_ok=True)
|
|
return JSONResponse({"phase": "error", "error": f"Upload failed (disk full?): {exc}"}, status_code=500)
|
|
except Exception as exc:
|
|
if dest.exists():
|
|
dest.unlink(missing_ok=True)
|
|
return JSONResponse({"phase": "error", "error": f"Upload failed: {exc}"}, status_code=500)
|
|
|
|
return JSONResponse({
|
|
"phase": "staged",
|
|
"filename": filename,
|
|
"vmid": vmid,
|
|
"vm_name": vm_name,
|
|
"file_size_gb": file_size_gb,
|
|
})
|
|
|
|
else:
|
|
# ── Download ────────────────────────────────────────────────
|
|
url = (source_url or "").strip()
|
|
if not url:
|
|
return JSONResponse({"phase": "error", "error": "No URL provided."}, status_code=400)
|
|
|
|
filename = Path(url).name or f"download_{uuid.uuid4().hex[:8]}"
|
|
dest = sdir / filename
|
|
|
|
_active_downloads.pop(filename, None)
|
|
|
|
usage = shutil.disk_usage(_staging(session_id))
|
|
free_gb = usage.free / (1024**3)
|
|
if free_gb < 50:
|
|
logger.warning("Low disk: %.1f GB free — download may fail", free_gb)
|
|
|
|
logger.info("Starting background download: %s → %s", url, dest)
|
|
try:
|
|
proc = subprocess.Popen(
|
|
["aria2c", "-x8", "-s8", "-d", str(dest.parent), "-o", dest.name, url],
|
|
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
|
|
)
|
|
except Exception as exc:
|
|
return JSONResponse({"phase": "error", "error": f"Failed to start download: {exc}"}, status_code=500)
|
|
|
|
_active_downloads[filename] = {
|
|
"proc": proc, "dest": dest, "vmid": vmid, "vm_name": vm_name,
|
|
"start_time": time.time(), "content_length": 0, "_last_logged_bytes": 0,
|
|
}
|
|
|
|
content_length = 0
|
|
def _fetch_cl():
|
|
nonlocal content_length
|
|
try:
|
|
hr = http_requests.head(url, timeout=5, allow_redirects=True)
|
|
cl = hr.headers.get("Content-Length")
|
|
if cl:
|
|
content_length = int(cl)
|
|
if filename in _active_downloads:
|
|
_active_downloads[filename]["content_length"] = int(cl)
|
|
logger.info("Download size: %.1f GiB", int(cl) / (1024**3))
|
|
except Exception:
|
|
pass
|
|
threading.Thread(target=_fetch_cl, daemon=True).start()
|
|
|
|
return JSONResponse({
|
|
"phase": "downloading",
|
|
"filename": filename, "vmid": vmid, "vm_name": vm_name,
|
|
"content_length_gb": round(content_length / (1024**3), 1) if content_length else None,
|
|
})
|
|
|
|
|
|
@app.post("/session/upload-raw")
|
|
async def upload_raw(request: Request):
|
|
"""Raw streaming upload — bypasses multipart parsing for large files."""
|
|
filename = request.headers.get("X-Filename", "upload.bin")
|
|
vmid_str = request.headers.get("X-VMID", "")
|
|
vm_name = request.headers.get("X-VM-Name", "")
|
|
session_id = request.headers.get("X-Session-ID", "")
|
|
|
|
try:
|
|
vmid = int(vmid_str)
|
|
except ValueError:
|
|
return JSONResponse({"phase": "error", "error": "Invalid VM ID."}, status_code=400)
|
|
|
|
err = _validate_vmid(vmid)
|
|
if err:
|
|
return JSONResponse({"phase": "error", "error": err}, status_code=400)
|
|
|
|
sdir = _staging(session_id) / "in"
|
|
sdir.mkdir(parents=True, exist_ok=True)
|
|
dest = sdir / filename
|
|
content_length = request.headers.get("content-length")
|
|
if content_length:
|
|
estimated_gb = int(content_length) / (1024**3)
|
|
err = _check_disk_space(dest, estimated_gb)
|
|
if err:
|
|
return JSONResponse({"phase": "error", "error": err}, status_code=400)
|
|
|
|
logger.info("Raw upload: %s (%s bytes)", filename, content_length or "unknown")
|
|
written = 0
|
|
try:
|
|
with dest.open("wb") as f:
|
|
async for chunk in request.stream():
|
|
f.write(chunk)
|
|
written += len(chunk)
|
|
if written % UPLOAD_PROGRESS_INTERVAL < len(chunk):
|
|
logger.info("Upload progress: %s — %.1f GiB", filename, written / (1024**3))
|
|
file_size_gb = round(written / (1024**3), 1)
|
|
logger.info("Raw upload complete: %s (%.1f GiB)", filename, file_size_gb)
|
|
except Exception as exc:
|
|
if dest.exists():
|
|
dest.unlink(missing_ok=True)
|
|
logger.exception("Raw upload failed: %s", filename)
|
|
return JSONResponse({"phase": "error", "error": f"Upload failed: {exc}"}, status_code=500)
|
|
|
|
return JSONResponse({
|
|
"phase": "staged",
|
|
"filename": filename,
|
|
"vmid": vmid,
|
|
"vm_name": vm_name,
|
|
"file_size_gb": file_size_gb,
|
|
})
|
|
|
|
|
|
@app.get("/session/progress/{filename}")
|
|
async def session_progress(filename: str):
|
|
"""Poll download progress — returns current file size and phase."""
|
|
info = _active_downloads.get(filename)
|
|
if not info:
|
|
# Check if file exists on disk (download already completed in a
|
|
# previous session, or it was an upload)
|
|
dest = _staging("") / "in" / filename
|
|
if dest.exists():
|
|
return JSONResponse({
|
|
"phase": "complete",
|
|
"file_size_bytes": dest.stat().st_size,
|
|
"file_size_gb": round(dest.stat().st_size / (1024**3), 1),
|
|
})
|
|
return JSONResponse({"phase": "unknown", "error": "No active download for this file."}, status_code=404)
|
|
|
|
proc = info["proc"]
|
|
dest = info["dest"]
|
|
|
|
# Current bytes on disk
|
|
current_bytes = dest.stat().st_size if dest.exists() else 0
|
|
|
|
# Check if process still running
|
|
poll = proc.poll()
|
|
if poll is not None:
|
|
# Process exited
|
|
_active_downloads.pop(filename, None)
|
|
if poll != 0:
|
|
logger.error("Download failed: %s (wget exited with code %d)", filename, poll)
|
|
if dest.exists():
|
|
dest.unlink(missing_ok=True)
|
|
return JSONResponse({
|
|
"phase": "error",
|
|
"error": f"Download failed (wget exited with code {poll}).",
|
|
"file_size_bytes": current_bytes,
|
|
})
|
|
# Success
|
|
final_bytes = dest.stat().st_size
|
|
logger.info("Download complete: %s (%.1f GiB)", filename, final_bytes / (1024**3))
|
|
return JSONResponse({
|
|
"phase": "complete",
|
|
"file_size_bytes": final_bytes,
|
|
"file_size_gb": round(final_bytes / (1024**3), 1),
|
|
})
|
|
|
|
# Still downloading — check speed and estimate ETA
|
|
elapsed = max(time.time() - info.get("start_time", 0), 1)
|
|
speed_bps = current_bytes / elapsed
|
|
speed_mbps = round(speed_bps / 1_000_000, 1)
|
|
|
|
content_length = info.get("content_length", 0)
|
|
eta_str = None
|
|
too_slow = False
|
|
if content_length > 0 and elapsed > SPEED_CHECK_AFTER:
|
|
remaining = content_length - current_bytes
|
|
eta = remaining / max(speed_bps, 1) # seconds
|
|
if eta > MAX_ETA_SECONDS:
|
|
too_slow = True
|
|
eta_hours = round(eta / 3600, 1)
|
|
# Kill the download
|
|
try:
|
|
proc.kill()
|
|
except Exception:
|
|
pass
|
|
_active_downloads.pop(filename, None)
|
|
if dest.exists():
|
|
dest.unlink(missing_ok=True)
|
|
logger.warning(
|
|
"Download %s killed: ETA %.1f h at %.1f MB/s (threshold %d h)",
|
|
filename, eta_hours, speed_mbps, MAX_ETA_SECONDS // 3600,
|
|
)
|
|
return JSONResponse({
|
|
"phase": "too_slow",
|
|
"speed_mbps": speed_mbps,
|
|
"eta_hours": eta_hours,
|
|
"content_length_gb": round(content_length / (1024**3), 1),
|
|
"downloaded_gb": round(current_bytes / (1024**3), 1),
|
|
"message": (
|
|
f"Download would take ~{eta_hours} hours at {speed_mbps} MB/s "
|
|
f"(file is {round(content_length / (1024**3), 1)} GiB). "
|
|
f"Consider downloading to your computer manually, then use File Upload."
|
|
),
|
|
})
|
|
eta_str = f"~{round(eta / 60)} min remaining"
|
|
|
|
if too_slow:
|
|
# Already handled above; this line is unreachable but kept for clarity
|
|
pass
|
|
|
|
# Log progress periodically (every ~1 GiB)
|
|
last_logged = info.get("_last_logged_bytes", 0)
|
|
if current_bytes - last_logged >= 1024**3:
|
|
info["_last_logged_bytes"] = current_bytes
|
|
logger.info("Downloading: %s — %.1f GiB (%.1f MB/s)%s",
|
|
filename, current_bytes / (1024**3), speed_mbps,
|
|
f" ETA {eta_str}" if eta_str else "")
|
|
|
|
return JSONResponse({
|
|
"phase": "downloading",
|
|
"file_size_bytes": current_bytes,
|
|
"file_size_gb": round(current_bytes / (1024**3), 1),
|
|
"speed_mbps": speed_mbps,
|
|
"eta": eta_str,
|
|
"content_length_gb": round(content_length / (1024**3), 1) if content_length else None,
|
|
})
|
|
|
|
|
|
@app.post("/session/analyze", response_class=HTMLResponse)
|
|
async def session_analyze(
|
|
request: Request,
|
|
vmid: int = Form(...),
|
|
filename: str = Form(...),
|
|
vm_name: str = Form(""),
|
|
session_id: str = Form(""),
|
|
):
|
|
"""Phase 2 — call the backend /analyze endpoint and render the result."""
|
|
err = _validate_vmid(vmid)
|
|
if err:
|
|
return render("_analysis.html", request=request, error=err)
|
|
|
|
# Prepend session subdirectory to the filename for the backend
|
|
# Files are in tmp/{guid}/in/, so the backend needs {guid}/in/filename
|
|
backend_filename = f"{session_id}/in/{filename}" if session_id else filename
|
|
|
|
try:
|
|
analysis = api.analyze(vmid=vmid, filename=backend_filename)
|
|
except ApiError as exc:
|
|
return render("_analysis.html", request=request,
|
|
error=f"Backend analysis failed: {exc.detail}")
|
|
|
|
vm_name = vm_name.strip()
|
|
if not vm_name:
|
|
vm_name = (analysis.get("os_type") or "vm") + f"-{vmid}"
|
|
|
|
return render("_analysis.html", request=request,
|
|
vmid=vmid, source_filename=filename,
|
|
vm_name=vm_name, analysis=analysis)
|
|
|
|
|
|
@app.post("/session/confirm", response_class=HTMLResponse)
|
|
async def confirm_session(
|
|
request: Request,
|
|
vmid: int = Form(...),
|
|
source_filename: str = Form(...),
|
|
disk_format: str = Form(...),
|
|
vm_name: str = Form(...),
|
|
cpu_cores: int = Form(2),
|
|
ram_mb: int = Form(4096),
|
|
target_storage: str = Form("local-lvm"),
|
|
target_disk_size_gb: Optional[int] = Form(None),
|
|
auto_detect_boot: str = Form("true"),
|
|
boot_type: str = Form("uefi"),
|
|
session_id: str = Form(""),
|
|
):
|
|
"""Build the job payload and submit to the backend."""
|
|
# Validate VM ID range
|
|
err = _validate_vmid(vmid)
|
|
if err:
|
|
return render("_analysis.html", request=request, error=err)
|
|
|
|
payload = {
|
|
"vmid": vmid,
|
|
"vm_name": vm_name,
|
|
"cpu_cores": cpu_cores,
|
|
"ram_mb": ram_mb,
|
|
"target_storage": target_storage,
|
|
"auto_detect_boot": auto_detect_boot == "true",
|
|
"boot_type": boot_type,
|
|
"session_id": session_id,
|
|
"boot_disk": {
|
|
"disk_type": "image_file",
|
|
"source_filename": source_filename,
|
|
"format": disk_format,
|
|
},
|
|
}
|
|
if target_disk_size_gb is not None:
|
|
payload["target_disk_size_gb"] = target_disk_size_gb
|
|
|
|
try:
|
|
job = api.create_job(payload)
|
|
except ApiError as exc:
|
|
return render("_analysis.html", request=request,
|
|
error=f"Job submission failed: {exc.detail}")
|
|
|
|
return render("polling.html", request=request,
|
|
job_id=job["job_id"], vmid=vmid,
|
|
vm_name=vm_name, source_filename=source_filename)
|
|
|
|
|
|
@app.get("/session/status/{job_id}")
|
|
async def session_status(job_id: str):
|
|
"""Poll backend for job status (returns JSON for AJAX polling)."""
|
|
try:
|
|
return api.get_job(job_id)
|
|
except ApiError as exc:
|
|
return JSONResponse({"error": exc.detail, "status": "error"}, status_code=502)
|
|
|
|
|
|
@app.post("/session/cleanup/{job_id}")
|
|
async def session_cleanup(job_id: str, request: Request):
|
|
"""Forward cleanup request to backend."""
|
|
try:
|
|
body = await request.json()
|
|
delete = body.get("delete_staging_files", False)
|
|
session_id = body.get("session_id", "")
|
|
return api.cleanup_job(job_id, delete, session_id)
|
|
except ApiError as exc:
|
|
return JSONResponse({"error": exc.detail}, status_code=502)
|
|
|
|
|
|
@app.post("/session/clone")
|
|
async def session_clone(request: Request):
|
|
"""Proxy clone request to backend."""
|
|
try:
|
|
body = await request.json()
|
|
result = api.clone_vm(
|
|
source_vmid=body["source_vmid"],
|
|
target_vmid=body["target_vmid"],
|
|
target_name=body["target_name"],
|
|
)
|
|
return result
|
|
except ApiError as exc:
|
|
return JSONResponse({"error": exc.detail}, status_code=502)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Run
|
|
# ---------------------------------------------------------------------------
|
|
# Track active SCP pulls for progress polling
|
|
_active_scp: dict[str, dict] = {}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# SCP Pull — large file transfer from remote servers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
@app.get("/scp", response_class=HTMLResponse)
|
|
async def scp_page(request: Request):
|
|
"""SCP pull page — fetch files directly from remote servers."""
|
|
return render("scp.html", request=request, backend_url=BACKEND_URL,
|
|
session_id=uuid.uuid4().hex)
|
|
|
|
|
|
@app.post("/scp/start")
|
|
def scp_start(
|
|
request: Request,
|
|
scp_host: str = Form(...),
|
|
scp_port: int = Form(22),
|
|
scp_user: str = Form(...),
|
|
scp_pass: str = Form(...),
|
|
scp_path: str = Form(...),
|
|
session_id: str = Form(""),
|
|
):
|
|
"""Start an SCP pull in the background. Returns JSON with phase + filename."""
|
|
remote = f"{scp_user}@{scp_host}:{scp_path}"
|
|
filename = Path(scp_path).name
|
|
if not filename:
|
|
return JSONResponse({"phase": "error", "error": "Invalid remote path."}, status_code=400)
|
|
|
|
sdir = _staging(session_id) / "in"
|
|
sdir.mkdir(parents=True, exist_ok=True)
|
|
dest = sdir / filename
|
|
|
|
# Check disk space
|
|
usage = shutil.disk_usage(STAGING)
|
|
free_gb = usage.free / (1024**3)
|
|
if free_gb < 10:
|
|
logger.warning("Low disk: %.1f GB free — SCP pull may fail", free_gb)
|
|
|
|
env = os.environ.copy()
|
|
env["SSHPASS"] = scp_pass
|
|
|
|
cmd = [
|
|
"sshpass", "-e",
|
|
"scp",
|
|
"-o", "StrictHostKeyChecking=no",
|
|
"-o", "ConnectTimeout=10",
|
|
"-P", str(scp_port),
|
|
remote, str(dest),
|
|
]
|
|
|
|
logger.info("Starting SCP pull: %s → %s", remote, dest)
|
|
try:
|
|
proc = subprocess.Popen(cmd, env=env, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
|
|
except Exception as exc:
|
|
return JSONResponse({"phase": "error", "error": f"Failed to start SCP: {exc}"}, status_code=500)
|
|
|
|
key = f"{session_id}/{filename}"
|
|
_active_scp[key] = {
|
|
"proc": proc,
|
|
"dest": dest,
|
|
"start_time": time.time(),
|
|
}
|
|
|
|
return JSONResponse({
|
|
"phase": "pulling",
|
|
"filename": filename,
|
|
"session_id": session_id,
|
|
})
|
|
|
|
|
|
@app.get("/scp/progress/{session_id}/{filename}")
|
|
def scp_progress(session_id: str, filename: str):
|
|
"""Poll SCP progress by checking output file size."""
|
|
key = f"{session_id}/{filename}"
|
|
info = _active_scp.get(key)
|
|
if not info:
|
|
dest = _staging(session_id) / "in" / filename
|
|
if dest.exists():
|
|
return JSONResponse({
|
|
"phase": "complete",
|
|
"file_size_bytes": dest.stat().st_size,
|
|
"file_size_gb": round(dest.stat().st_size / (1024**3), 1),
|
|
})
|
|
return JSONResponse({"phase": "unknown", "error": "No active SCP pull."}, status_code=404)
|
|
|
|
proc = info["proc"]
|
|
dest = info["dest"]
|
|
current_bytes = dest.stat().st_size if dest.exists() else 0
|
|
|
|
poll = proc.poll()
|
|
if poll is not None:
|
|
_active_scp.pop(key, None)
|
|
if poll != 0:
|
|
if dest.exists():
|
|
dest.unlink(missing_ok=True)
|
|
return JSONResponse({
|
|
"phase": "error",
|
|
"error": f"SCP pull failed (exit code {poll}). Check credentials and remote path.",
|
|
"file_size_bytes": current_bytes,
|
|
})
|
|
final_bytes = dest.stat().st_size
|
|
logger.info("SCP pull complete: %s (%.1f GiB)", filename, final_bytes / (1024**3))
|
|
return JSONResponse({
|
|
"phase": "complete",
|
|
"file_size_bytes": final_bytes,
|
|
"file_size_gb": round(final_bytes / (1024**3), 1),
|
|
})
|
|
|
|
elapsed = max(time.time() - info.get("start_time", 0), 1)
|
|
speed_mbps = round(current_bytes / elapsed / 1_000_000, 1)
|
|
|
|
return JSONResponse({
|
|
"phase": "pulling",
|
|
"file_size_bytes": current_bytes,
|
|
"file_size_gb": round(current_bytes / (1024**3), 1),
|
|
"speed_mbps": speed_mbps,
|
|
})
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import uvicorn
|
|
uvicorn.run(app, host="0.0.0.0", port=5000)
|