import asyncio import os import json import time import secrets from typing import Dict, List, Optional, Any, Set from fastapi import WebSocket from backend.config import JOBS_FILE class JobManager: def __init__(self): self.active_jobs: Dict[str, Dict[str, Any]] = {} self.subscribers: Dict[str, Set[WebSocket]] = {} self.global_subscribers: Set[WebSocket] = set() self.lock = asyncio.Lock() self.type_locks: Dict[str, asyncio.Lock] = { "backup": asyncio.Lock(), "restore": asyncio.Lock(), "update_pangolin": asyncio.Lock(), "update_system": asyncio.Lock(), } def _load_history(self) -> List[Dict[str, Any]]: if not JOBS_FILE.exists(): return [] try: with open(JOBS_FILE, "r", encoding="utf-8") as f: return json.load(f) except Exception: return [] def _save_history(self, history: List[Dict[str, Any]]) -> None: try: with open(JOBS_FILE, "w", encoding="utf-8") as f: json.dump(history[:50], f, indent=2) except Exception as e: print(f"Error saving jobs history: {e}") def list_jobs(self, limit: int = 25) -> List[Dict[str, Any]]: history = self._load_history() # Merge active jobs at the beginning active_list = list(self.active_jobs.values()) seen_ids = {j["id"] for j in active_list} merged = active_list + [j for j in history if j["id"] not in seen_ids] merged.sort(key=lambda x: x.get("start_time", 0), reverse=True) return merged[:limit] def get_job(self, job_id: str) -> Optional[Dict[str, Any]]: if job_id in self.active_jobs: return self.active_jobs[job_id] history = self._load_history() for j in history: if j["id"] == job_id: return j return None async def register_subscriber(self, job_id: str, ws: WebSocket): if job_id not in self.subscribers: self.subscribers[job_id] = set() self.subscribers[job_id].add(ws) def unregister_subscriber(self, job_id: str, ws: WebSocket): if job_id in self.subscribers: self.subscribers[job_id].discard(ws) if not self.subscribers[job_id]: del self.subscribers[job_id] async def register_global(self, ws: WebSocket): self.global_subscribers.add(ws) def unregister_global(self, ws: WebSocket): self.global_subscribers.discard(ws) async def broadcast(self, job_id: str, message: Dict[str, Any]): # Broadcast to specific job subscribers if job_id in self.subscribers: dead = set() for ws in list(self.subscribers[job_id]): try: await ws.send_json(message) except Exception: dead.add(ws) for ws in dead: self.subscribers[job_id].discard(ws) # Broadcast to global subscribers dead_global = set() for ws in list(self.global_subscribers): try: await ws.send_json(message) except Exception: dead_global.add(ws) for ws in dead_global: self.global_subscribers.discard(ws) async def run_command_job( self, name: str, cmd: List[str] | str, job_type: str = "general", env: Optional[Dict[str, str]] = None, cwd: Optional[str] = None, description: str = "" ) -> Dict[str, Any]: lock = self.type_locks.get(job_type) if lock and lock.locked(): return { "success": False, "error": f"A {job_type} task is already running in the background. Please wait for it to complete." } job_id = f"job_{int(time.time())}_{secrets.token_hex(4)}" job = { "id": job_id, "name": name, "type": job_type, "description": description, "cmd": cmd if isinstance(cmd, str) else " ".join(cmd), "status": "running", "start_time": time.time(), "end_time": None, "duration": None, "exit_code": None, "logs": [] } self.active_jobs[job_id] = job await self.broadcast(job_id, {"type": "job_started", "job": job}) # Run in background asyncio task asyncio.create_task(self._execute_process(job_id, cmd, env, cwd, lock)) return {"success": True, "job_id": job_id, "job": job} async def _execute_process( self, job_id: str, cmd: List[str] | str, env: Optional[Dict[str, str]], cwd: Optional[str], lock: Optional[asyncio.Lock] ): if lock: await lock.acquire() job = self.active_jobs[job_id] run_env = os.environ.copy() run_env["DEBIAN_FRONTEND"] = "noninteractive" run_env["TERM"] = "xterm-256color" run_env["COLUMNS"] = "120" if env: run_env.update(env) try: if isinstance(cmd, list): proc = await asyncio.create_subprocess_exec( *cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.STDOUT, env=run_env, cwd=cwd ) else: proc = await asyncio.create_subprocess_shell( cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.STDOUT, env=run_env, cwd=cwd, executable="/bin/bash" if os.path.exists("/bin/bash") else None ) while True: line = await proc.stdout.readline() if not line: break text = line.decode('utf-8', errors='replace').rstrip('\r\n') timestamp = time.strftime("%H:%M:%S") log_entry = f"[{timestamp}] {text}" job["logs"].append(log_entry) # Keep up to 2000 lines if len(job["logs"]) > 2000: job["logs"] = job["logs"][-2000:] await self.broadcast(job_id, {"type": "log", "job_id": job_id, "data": log_entry}) await proc.wait() exit_code = proc.returncode job["exit_code"] = exit_code job["status"] = "completed" if exit_code == 0 else "failed" except Exception as e: err_msg = f"Execution error: {str(e)}" job["logs"].append(f"[{time.strftime('%H:%M:%S')}] ERROR: {err_msg}") job["status"] = "failed" job["exit_code"] = -1 await self.broadcast(job_id, {"type": "log", "job_id": job_id, "data": err_msg}) finally: job["end_time"] = time.time() job["duration"] = round(job["end_time"] - job["start_time"], 1) await self.broadcast(job_id, {"type": "job_finished", "job": job}) # Save to history file history = self._load_history() history.insert(0, job.copy()) self._save_history(history) # Clean from active_jobs after a short delay so clients can finish reading await asyncio.sleep(5) if job_id in self.active_jobs: del self.active_jobs[job_id] if lock: lock.release() runner = JobManager()