Update app.py

app.py

@@ -16,12 +16,15 @@ import gradio as gr
 from huggingface_hub import InferenceClient
 
 # ---------------------------------------------------------------------
-# Task Registry (auto-loads your MCP tasks)
+# Task Registry (auto-loads MCP tasks dynamically)
 # ---------------------------------------------------------------------
 
 TASK_REGISTRY: Dict[str, Any] = {}
 
 def _register_tasks():
+    """
+    Import tasks.* modules dynamically and pull their run() functions.
+    """
     def _try(name, module):
         try:
             m = __import__(f"tasks.{module}", fromlist=["*"])
@@ -29,7 +32,7 @@ def _register_tasks():
             if callable(fn):
                 TASK_REGISTRY[name] = fn
         except Exception:
-            pass
+            pass  # In HF Spaces, missing tasks won't break startup.
 
     _try("lookup_ip", "lookup_ip")
     _try("lookup_domain", "lookup_domain")
@@ -42,8 +45,9 @@ def _register_tasks():
 
 _register_tasks()
 
+
 # ---------------------------------------------------------------------
-# Core Task Execution
+# Task Execution + Normalization
 # ---------------------------------------------------------------------
 
 def call_task(name: str, payload: Dict[str, Any]):
@@ -61,24 +65,28 @@ def call_task(name: str, payload: Dict[str, Any]):
 
 
 def normalize_result(res: Dict[str, Any]):
-    """Formats UI fields cleanly."""
+    """Ensures consistent UI formatting."""
     pretty = json.dumps(res, indent=2, default=str)
     summary = res.get("summary", "")
     markdown = res.get("markdown") or res.get("report") or ""
+
     if not markdown and summary:
         markdown = f"## Summary\n\n{summary}"
 
+    safe_json = lambda v: json.dumps(v, indent=2, default=str) if v else ""
+
     return {
         "summary": summary,
         "markdown": markdown,
         "json": pretty,
-        "mitre": json.dumps(res.get("mitre", ""), indent=2, default=str) if res.get("mitre") else "",
-        "stix": json.dumps(res.get("stix", ""), indent=2, default=str) if res.get("stix") else "",
-        "sarif": json.dumps(res.get("sarif", ""), indent=2, default=str) if res.get("sarif") else "",
+        "mitre": safe_json(res.get("mitre")),
+        "stix": safe_json(res.get("stix")),
+        "sarif": safe_json(res.get("sarif")),
     }
 
+
 # ---------------------------------------------------------------------
-# ANALYST COPILOT (LLM)
+# Analyst Copilot LLM
 # ---------------------------------------------------------------------
 
 def respond(
@@ -91,14 +99,21 @@ def respond(
     top_p,
     max_tokens,
 ):
-    """Streaming response from WhiteRabbit Neo or Cybertron."""
-    client = InferenceClient(model=model_name, token=hf_token.token)
+    """
+    Streaming LLM output using WhiteRabbit Neo or Cybertron.
+    hf_token is a raw string entered by user.
+    """
+    client = InferenceClient(
+        model=model_name,
+        token=hf_token,  # Direct string, no OAuthToken object
+    )
 
     msgs = [{"role": "system", "content": system_prompt}]
     msgs.extend(history)
     msgs.append({"role": "user", "content": message})
 
-    buf = ""
+    buffer = ""
+
     for chunk in client.chat_completion(
         messages=msgs,
         max_tokens=max_tokens,
@@ -108,12 +123,12 @@ def respond(
     ):
         delta = chunk.choices[0].delta.content
         if delta:
-            buf += delta
-            yield buf
+            buffer += delta
+            yield buffer
 
 
 def inject_osint(history, osint_obj):
-    """Inject raw JSON results into the chat context."""
+    """Inject OSINT result JSON into copilot context."""
     pretty = json.dumps(osint_obj, indent=2, default=str)
     history.append({
         "role": "system",
@@ -121,6 +136,7 @@ def inject_osint(history, osint_obj):
     })
     return history
 
+
 # ---------------------------------------------------------------------
 # OSINT Dashboard Callbacks
 # ---------------------------------------------------------------------
@@ -130,27 +146,32 @@ def ui_lookup_ip(ip, enrich, mitre):
     norm = normalize_result(raw)
     return norm["summary"], norm["markdown"], norm["json"], norm["mitre"], norm["stix"], raw
 
+
 def ui_lookup_domain(domain, enrich, mitre):
     raw = call_task("lookup_domain", {"domain": domain, "enrich": enrich, "map_mitre": mitre})
     norm = normalize_result(raw)
     return norm["summary"], norm["markdown"], norm["json"], norm["mitre"], norm["stix"], raw
 
+
 def ui_lookup_hash(h, ht, enrich, mitre):
     raw = call_task("lookup_hash", {"hash": h, "hash_type": ht, "enrich": enrich, "map_mitre": mitre})
     norm = normalize_result(raw)
     return norm["summary"], norm["markdown"], norm["json"], norm["mitre"], norm["stix"], raw
 
+
 def ui_correlate_iocs(iocs):
     lst = [x.strip() for x in iocs.splitlines() if x.strip()]
     raw = call_task("correlate_iocs", {"iocs": lst})
     norm = normalize_result(raw)
     return norm["summary"], norm["markdown"], norm["json"], norm["mitre"], raw
 
+
 def ui_quickscan(target):
     raw = call_task("quickscan", {"target": target})
     norm = normalize_result(raw)
     return norm["summary"], norm["markdown"], norm["json"], raw
 
+
 # ---------------------------------------------------------------------
 # MCP Bridge
 # ---------------------------------------------------------------------
@@ -160,17 +181,19 @@ def ui_bridge(tool, args_json):
         payload = json.loads(args_json)
     except Exception as e:
         return json.dumps({"error": str(e)}, indent=2), "", {}
+
     raw = call_task(tool, payload)
     norm = normalize_result(raw)
     return norm["json"], norm["markdown"], raw
 
+
 # ---------------------------------------------------------------------
-# BUILD UI
+# UI Layout
 # ---------------------------------------------------------------------
 
 def build_interface():
     with gr.Blocks(title="Parrot OSINT MCP Console") as demo:
-        gr.Markdown("# Parrot OSINT MCP Console\nA multi-mode intelligence workstation.")
+        gr.Markdown("# 🦜 Parrot OSINT MCP Console\nMulti-mode Intelligence Workstation.")
 
         osint_state = gr.State({})
 
@@ -178,43 +201,55 @@ def build_interface():
         # OSINT Dashboard
         # -------------------------
         with gr.Tab("OSINT Dashboard"):
-            with gr.Tab("IP"):
-                ip = gr.Textbox(label="IP Address")
-                enrich = gr.Checkbox(value=True, label="Enrich")
-                mitre = gr.Checkbox(value=True, label="MITRE Map")
+
+            # IP Lookup
+            with gr.Tab("IP Lookup"):
+                ip = gr.Textbox(label="IP Address", placeholder="8.8.8.8")
+                enrich = gr.Checkbox(value=True, label="Enrich data")
+                mitre = gr.Checkbox(value=True, label="MITRE ATT&CK Mapping")
                 run = gr.Button("Run IP Lookup")
 
-                s = gr.Textbox(label="Summary")
-                md = gr.Markdown()
-                js = gr.Code(language="json")
-                mt = gr.Code(language="json")
-                st = gr.Code(language="json")
+                out_s = gr.Textbox(label="Summary")
+                out_md = gr.Markdown()
+                out_json = gr.Code(language="json")
+                out_mitre = gr.Code(language="json")
+                out_stix = gr.Code(language="json")
 
-                run.click(ui_lookup_ip, [ip, enrich, mitre], [s, md, js, mt, st, osint_state])
+                run.click(
+                    ui_lookup_ip,
+                    [ip, enrich, mitre],
+                    [out_s, out_md, out_json, out_mitre, out_stix, osint_state]
+                )
 
         # -------------------------
         # MCP Bridge
         # -------------------------
         with gr.Tab("MCP Bridge"):
-            tool = gr.Dropdown(sorted(TASK_REGISTRY.keys()))
-            args = gr.Code(language="json")
-            btn = gr.Button("Call Tool")
-            out_js = gr.Code(language="json")
-            out_md = gr.Markdown()
+            tool = gr.Dropdown(sorted(TASK_REGISTRY.keys()), label="Tool")
+            args = gr.Code(language="json", label="Args JSON")
+            btn = gr.Button("Run Tool")
 
-            btn.click(ui_bridge, [tool, args], [out_js, out_md, osint_state])
+            out_bridge_json = gr.Code(language="json")
+            out_bridge_md = gr.Markdown()
+
+            btn.click(
+                ui_bridge,
+                [tool, args],
+                [out_bridge_json, out_bridge_md, osint_state]
+            )
 
         # -------------------------
         # Analyst Copilot
         # -------------------------
         with gr.Tab("Analyst Copilot"):
-            gr.Markdown("### WhiteRabbit Neo + Cybertron TI Assistant")
+            gr.Markdown("### WhiteRabbit Neo + Cybertron Threat Intelligence Assistant")
 
             system_prompt = gr.Textbox(
                 label="System Prompt",
                 value=(
                     "You are a threat intelligence analyst. "
-                    "You classify TTPs, map MITRE ATT&CK, and provide investigation guidance."
+                    "You classify TTPs, extract indicators, map MITRE ATT&CK, "
+                    "and provide investigation guidance."
                 ),
             )
 
@@ -229,8 +264,12 @@ def build_interface():
                 value="berkeley-nest/WhiteRabbitNeo-8B",
             )
 
-            gr.Markdown("### Provide your HuggingFace API Token for the LLM:")
-            hf_token = gr.OAuthToken()
+            gr.Markdown("### HuggingFace API Token (required for LLM inference)")
+            hf_token = gr.Textbox(
+                label="HF Token",
+                type="password",
+                placeholder="hf_xxx...",
+            )
 
             chatbot = gr.ChatInterface(
                 respond,
@@ -245,17 +284,18 @@ def build_interface():
                 ],
             )
 
-            inject_btn = gr.Button("Inject Last OSINT Result into Chat")
+            inject_btn = gr.Button("Inject Last OSINT Result into Copilot")
             inject_btn.click(
                 inject_osint,
-                inputs=[chatbot._chatbot_state, osint_state],
-                outputs=[chatbot._chatbot_state],
+                [chatbot._chatbot_state, osint_state],
+                [chatbot._chatbot_state],
             )
 
     return demo
 
+
 # ---------------------------------------------------------------------
-# MAIN
+# MAIN ENTRY
 # ---------------------------------------------------------------------
 
 if __name__ == "__main__":