Replace app.py with complete Fury broker

app.py

@@ -1,18 +1,30 @@
 import os
 import uuid
 from enum import Enum
+from html import escape
 from typing import Optional
 
-from fastapi import FastAPI, Header, HTTPException, Form
+from fastapi import FastAPI, Form, Header, HTTPException
 from fastapi.responses import HTMLResponse, RedirectResponse
 from pydantic import BaseModel
 
 
+# ---------------------------------------------------------------------
+# Configuration
+# ---------------------------------------------------------------------
+
 BROKER_TOKEN = os.environ.get("BROKER_TOKEN")
 
 if not BROKER_TOKEN:
-    raise RuntimeError("Missing BROKER_TOKEN environment variable")
+    raise RuntimeError(
+        "Missing BROKER_TOKEN. Add it in Hugging Face Space "
+        "Settings → Variables and secrets → New secret."
+    )
+
 
+# ---------------------------------------------------------------------
+# Data model
+# ---------------------------------------------------------------------
 
 class JobStatus(str, Enum):
     queued = "queued"
@@ -28,69 +40,198 @@ class Job(BaseModel):
     result: Optional[str] = None
 
 
-app = FastAPI()
+# ---------------------------------------------------------------------
+# FastAPI app
+# ---------------------------------------------------------------------
 
+app = FastAPI(title="Fury Broker")
+
+# In-memory job storage.
+# Note: jobs disappear if the Space restarts.
 jobs: dict[str, Job] = {}
 
 
+# Only commands listed here can be executed by Fury.
+# The browser/Space user submits the logical command name.
+# The Fury worker receives the shell command.
 ALLOWED_COMMANDS = {
     "hostname": "hostname",
     "whoami": "whoami",
     "pwd": "pwd",
     "disk": "df -h",
     "date": "date",
+    "list_home": "ls -lah ~ | head -50",
 }
 
 
-def verify_token(authorization: Optional[str]) -> None:
-    expected = f"Bearer {BROKER_TOKEN}"
-    if authorization != expected:
+# ---------------------------------------------------------------------
+# Security
+# ---------------------------------------------------------------------
+
+def verify_broker_token(x_broker_token: Optional[str]) -> None:
+    """
+    Broker-level authentication.
+
+    Important:
+    - Hugging Face private Spaces use:
+        Authorization: Bearer HF_TOKEN
+
+    - This app uses:
+        X-Broker-Token: BROKER_TOKEN
+
+    This avoids conflict between Hugging Face authentication and
+    the app's own authentication.
+    """
+    if x_broker_token != BROKER_TOKEN:
         raise HTTPException(status_code=401, detail="Unauthorized")
 
 
+# ---------------------------------------------------------------------
+# Health and diagnostics
+# ---------------------------------------------------------------------
+
+@app.get("/health")
+def health():
+    return {
+        "status": "ok",
+        "service": "fury-broker",
+        "jobs_count": len(jobs),
+    }
+
+
+@app.get("/api/commands")
+def list_commands(x_broker_token: Optional[str] = Header(default=None)):
+    verify_broker_token(x_broker_token)
+
+    return {
+        "commands": [
+            {
+                "name": name,
+                "shell_command": shell_command,
+            }
+            for name, shell_command in ALLOWED_COMMANDS.items()
+        ]
+    }
+
+
+# ---------------------------------------------------------------------
+# Browser UI
+# ---------------------------------------------------------------------
+
 @app.get("/", response_class=HTMLResponse)
 def home():
     rows = ""
 
     for job in reversed(list(jobs.values())):
-        result = job.result or ""
+        safe_id = escape(job.id)
+        safe_command = escape(job.command)
+        safe_status = escape(job.status.value)
+        safe_result = escape(job.result or "")
+
         rows += f"""
         <tr>
-            <td>{job.id}</td>
-            <td>{job.command}</td>
-            <td>{job.status}</td>
-            <td><pre>{result}</pre></td>
+            <td><code>{safe_id}</code></td>
+            <td>{safe_command}</td>
+            <td>{safe_status}</td>
+            <td><pre>{safe_result}</pre></td>
         </tr>
         """
 
     options = "\n".join(
-        f'<option value="{name}">{name} -> {cmd}</option>'
+        f'<option value="{escape(name)}">{escape(name)} → {escape(cmd)}</option>'
         for name, cmd in ALLOWED_COMMANDS.items()
     )
 
     return f"""
+    <!doctype html>
     <html>
         <head>
             <title>Fury Broker</title>
             <style>
-                body {{ font-family: sans-serif; margin: 40px; }}
-                table {{ border-collapse: collapse; width: 100%; }}
-                th, td {{ border: 1px solid #ccc; padding: 8px; vertical-align: top; }}
-                pre {{ white-space: pre-wrap; max-width: 700px; }}
+                body {{
+                    font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif;
+                    margin: 40px;
+                    max-width: 1200px;
+                    line-height: 1.45;
+                }}
+                h1 {{
+                    margin-bottom: 0.2rem;
+                }}
+                .subtitle {{
+                    color: #555;
+                    margin-bottom: 2rem;
+                }}
+                form {{
+                    margin: 1.5rem 0;
+                    padding: 1rem;
+                    border: 1px solid #ddd;
+                    border-radius: 8px;
+                    background: #fafafa;
+                }}
+                select {{
+                    padding: 6px;
+                    min-width: 260px;
+                }}
+                button {{
+                    padding: 6px 12px;
+                    cursor: pointer;
+                }}
+                table {{
+                    border-collapse: collapse;
+                    width: 100%;
+                    margin-top: 20px;
+                }}
+                th, td {{
+                    border: 1px solid #ccc;
+                    padding: 8px;
+                    vertical-align: top;
+                }}
+                th {{
+                    background: #f5f5f5;
+                    text-align: left;
+                }}
+                pre {{
+                    white-space: pre-wrap;
+                    max-width: 800px;
+                    max-height: 400px;
+                    overflow: auto;
+                    margin: 0;
+                }}
+                code {{
+                    font-size: 0.9em;
+                }}
+                .note {{
+                    color: #666;
+                    font-size: 0.95rem;
+                }}
             </style>
         </head>
         <body>
             <h1>Fury Broker</h1>
+            <div class="subtitle">
+                Hugging Face Space broker for running approved commands on Fury.
+            </div>
+
+            <p>
+                This Space does not SSH into Fury. Instead, the Fury worker polls this Space
+                for queued jobs, executes only predefined commands, and posts the result back.
+            </p>
 
             <form method="post" action="/submit">
-                <label>Command:</label>
-                <select name="command">
+                <label for="command"><strong>Command:</strong></label>
+                <select id="command" name="command">
                     {options}
                 </select>
                 <button type="submit">Submit job</button>
             </form>
 
+            <p class="note">
+                Refresh the page after submitting a job. The worker should pick it up within
+                a few seconds if it is running on Fury.
+            </p>
+
             <h2>Jobs</h2>
+
             <table>
                 <tr>
                     <th>ID</th>
@@ -107,10 +248,18 @@ def home():
 
 @app.post("/submit")
 def submit_from_ui(command: str = Form(...)):
+    """
+    Browser form endpoint.
+
+    This endpoint is intentionally not protected by X-Broker-Token because
+    the Space itself is private. If you later make the Space public, protect
+    this endpoint too.
+    """
     if command not in ALLOWED_COMMANDS:
         raise HTTPException(status_code=400, detail="Command is not allowed")
 
     job_id = str(uuid.uuid4())
+
     jobs[job_id] = Job(
         id=job_id,
         command=command,
@@ -120,33 +269,89 @@ def submit_from_ui(command: str = Form(...)):
     return RedirectResponse("/", status_code=303)
 
 
+# ---------------------------------------------------------------------
+# API: submit job
+# ---------------------------------------------------------------------
+
 @app.post("/api/jobs")
 def submit_job(
+    command: str = Form(...),
+    x_broker_token: Optional[str] = Header(default=None),
+):
+    """
+    Submit a job via API.
+
+    Example:
+
+        curl -X POST \\
+          -H "Authorization: Bearer $HF_TOKEN" \\
+          -H "X-Broker-Token: $BROKER_TOKEN" \\
+          -F "command=hostname" \\
+          https://rasa2-fury.hf.space/api/jobs
+    """
+    verify_broker_token(x_broker_token)
+
+    if command not in ALLOWED_COMMANDS:
+        raise HTTPException(status_code=400, detail="Command is not allowed")
+
+    job_id = str(uuid.uuid4())
+
+    job = Job(
+        id=job_id,
+        command=command,
+        status=JobStatus.queued,
+    )
+
+    jobs[job_id] = job
+
+    return job
+
+
+# Also allow query-param style:
+# POST /api/jobs/query?command=hostname
+@app.post("/api/jobs/query")
+def submit_job_query(
     command: str,
-    authorization: Optional[str] = Header(default=None),
+    x_broker_token: Optional[str] = Header(default=None),
 ):
-    verify_token(authorization)
+    verify_broker_token(x_broker_token)
 
     if command not in ALLOWED_COMMANDS:
         raise HTTPException(status_code=400, detail="Command is not allowed")
 
     job_id = str(uuid.uuid4())
+
     job = Job(
         id=job_id,
         command=command,
         status=JobStatus.queued,
     )
+
     jobs[job_id] = job
+
     return job
 
 
+# ---------------------------------------------------------------------
+# API: worker polling
+# ---------------------------------------------------------------------
+
 @app.get("/api/next-job")
-def get_next_job(authorization: Optional[str] = Header(default=None)):
-    verify_token(authorization)
+def get_next_job(x_broker_token: Optional[str] = Header(default=None)):
+    """
+    Called by the Fury worker.
+
+    The worker asks:
+        "Do you have a job for me?"
+
+    If a queued job exists, this endpoint marks it as running and returns it.
+    """
+    verify_broker_token(x_broker_token)
 
     for job in jobs.values():
         if job.status == JobStatus.queued:
             job.status = JobStatus.running
+
             return {
                 "id": job.id,
                 "command": job.command,
@@ -156,14 +361,21 @@ def get_next_job(authorization: Optional[str] = Header(default=None)):
     return {"id": None}
 
 
+# ---------------------------------------------------------------------
+# API: worker posts result
+# ---------------------------------------------------------------------
+
 @app.post("/api/jobs/{job_id}/result")
 def post_result(
     job_id: str,
-    result: str,
-    success: bool,
-    authorization: Optional[str] = Header(default=None),
+    result: str = Form(...),
+    success: bool = Form(...),
+    x_broker_token: Optional[str] = Header(default=None),
 ):
-    verify_token(authorization)
+    """
+    Called by the Fury worker after it executes a command.
+    """
+    verify_broker_token(x_broker_token)
 
     if job_id not in jobs:
         raise HTTPException(status_code=404, detail="Job not found")
@@ -173,3 +385,45 @@ def post_result(
     job.status = JobStatus.done if success else JobStatus.failed
 
     return job
+
+
+# ---------------------------------------------------------------------
+# API: inspect jobs
+# ---------------------------------------------------------------------
+
+@app.get("/api/jobs")
+def list_jobs(x_broker_token: Optional[str] = Header(default=None)):
+    verify_broker_token(x_broker_token)
+
+    return {
+        "jobs": list(jobs.values())
+    }
+
+
+@app.get("/api/jobs/{job_id}")
+def get_job(
+    job_id: str,
+    x_broker_token: Optional[str] = Header(default=None),
+):
+    verify_broker_token(x_broker_token)
+
+    if job_id not in jobs:
+        raise HTTPException(status_code=404, detail="Job not found")
+
+    return jobs[job_id]
+
+
+# ---------------------------------------------------------------------
+# API: clear jobs
+# ---------------------------------------------------------------------
+
+@app.post("/api/clear")
+def clear_jobs(x_broker_token: Optional[str] = Header(default=None)):
+    verify_broker_token(x_broker_token)
+
+    jobs.clear()
+
+    return {
+        "status": "cleared",
+        "jobs_count": len(jobs),
+    }