Add UI TOKEN handling

app.py

@@ -9,22 +9,15 @@ from fastapi.responses import HTMLResponse, RedirectResponse
 from pydantic import BaseModel
 
 
-# ---------------------------------------------------------------------
-# Configuration
-# ---------------------------------------------------------------------
-
 BROKER_TOKEN = os.environ.get("BROKER_TOKEN")
+UI_TOKEN = os.environ.get("UI_TOKEN")
 
 if not BROKER_TOKEN:
-    raise RuntimeError(
-        "Missing BROKER_TOKEN. Add it in Hugging Face Space "
-        "Settings → Variables and secrets → New secret."
-    )
+    raise RuntimeError("Missing BROKER_TOKEN")
 
+if not UI_TOKEN:
+    raise RuntimeError("Missing UI_TOKEN")
 
-# ---------------------------------------------------------------------
-# Data model
-# ---------------------------------------------------------------------
 
 class JobStatus(str, Enum):
     queued = "queued"
@@ -40,20 +33,11 @@ class Job(BaseModel):
     result: Optional[str] = None
 
 
-# ---------------------------------------------------------------------
-# 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",
@@ -64,31 +48,15 @@ ALLOWED_COMMANDS = {
 }
 
 
-# ---------------------------------------------------------------------
-# 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
-# ---------------------------------------------------------------------
+def verify_ui_token(ui_token: Optional[str]) -> None:
+    if ui_token != UI_TOKEN:
+        raise HTTPException(status_code=401, detail="Invalid UI token")
+
 
 @app.get("/health")
 def health():
@@ -99,41 +67,17 @@ def health():
     }
 
 
-@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())):
-        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><code>{safe_id}</code></td>
-            <td>{safe_command}</td>
-            <td>{safe_status}</td>
-            <td><pre>{safe_result}</pre></td>
+            <td><code>{escape(job.id)}</code></td>
+            <td>{escape(job.command)}</td>
+            <td>{escape(job.status.value)}</td>
+            <td><pre>{escape(job.result or "")}</pre></td>
         </tr>
         """
 
@@ -154,13 +98,6 @@ def home():
                     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;
@@ -168,13 +105,9 @@ def home():
                     border-radius: 8px;
                     background: #fafafa;
                 }}
-                select {{
+                input, select, button {{
                     padding: 6px;
-                    min-width: 260px;
-                }}
-                button {{
-                    padding: 6px 12px;
-                    cursor: pointer;
+                    margin: 4px;
                 }}
                 table {{
                     border-collapse: collapse;
@@ -197,37 +130,34 @@ def home():
                     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.
+                Submit an approved command from Hugging Face Spaces.
+                The worker running on Fury will pick it up, execute it on Fury,
+                and return the result here.
             </p>
 
             <form method="post" action="/submit">
-                <label for="command"><strong>Command:</strong></label>
-                <select id="command" name="command">
-                    {options}
-                </select>
-                <button type="submit">Submit job</button>
+                <div>
+                    <label><strong>UI token:</strong></label>
+                    <input type="password" name="ui_token" placeholder="Enter UI_TOKEN" required>
+                </div>
+
+                <div>
+                    <label><strong>Command:</strong></label>
+                    <select name="command">
+                        {options}
+                    </select>
+                    <button type="submit">Submit job</button>
+                </div>
             </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>
+                Refresh the page after a few seconds to see updated results.
             </p>
 
             <h2>Jobs</h2>
@@ -247,14 +177,12 @@ def home():
 
 
 @app.post("/submit")
-def submit_from_ui(command: str = Form(...)):
-    """
-    Browser form endpoint.
+def submit_from_ui(
+    command: str = Form(...),
+    ui_token: str = Form(...),
+):
+    verify_ui_token(ui_token)
 
-    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")
 
@@ -269,26 +197,11 @@ 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:
@@ -307,45 +220,8 @@ def submit_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,
-    x_broker_token: Optional[str] = Header(default=None),
-):
-    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(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():
@@ -361,10 +237,6 @@ def get_next_job(x_broker_token: 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,
@@ -372,9 +244,6 @@ def post_result(
     success: bool = Form(...),
     x_broker_token: Optional[str] = Header(default=None),
 ):
-    """
-    Called by the Fury worker after it executes a command.
-    """
     verify_broker_token(x_broker_token)
 
     if job_id not in jobs:
@@ -387,43 +256,19 @@ def post_result(
     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())
-    }
+    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)):
+@app.get("/api/commands")
+def list_commands(x_broker_token: Optional[str] = Header(default=None)):
     verify_broker_token(x_broker_token)
 
-    jobs.clear()
-
     return {
-        "status": "cleared",
-        "jobs_count": len(jobs),
+        "commands": [
+            {"name": name, "shell_command": cmd}
+            for name, cmd in ALLOWED_COMMANDS.items()
+        ]
     }