Update app.py

app.py

@@ -1,39 +1,72 @@
 import gradio as gr
 from internetarchive import search_items, get_item
-import requests, time, subprocess, json, re, tempfile
+import requests
+import time
+import subprocess
+import json
+import re
 from bs4 import BeautifulSoup
 from urllib.parse import urlparse
+import tempfile
 from pyvis.network import Network
+from functools import lru_cache
+from concurrent.futures import ThreadPoolExecutor
 
 THEME = "gradio/soft"
+
+# --- Shared HTTP session for speed & consistent headers ---
 session = requests.Session()
-session.headers.update({"User-Agent": "Mozilla/5.0"})
+session.headers.update({
+    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64)"
+})
+
+@lru_cache(maxsize=64)
+def get_cached_item(identifier):
+    """Cache IA get_item() calls to avoid duplicate downloads."""
+    return get_item(identifier)
 
 def scan_url_vt(url, api_key):
+    """Return True if VirusTotal sees no malicious hits."""
     headers = {"x-apikey": api_key}
-    resp = session.post("https://www.virustotal.com/api/v3/urls", headers=headers, data={"url": url})
+    resp = session.post(
+        "https://www.virustotal.com/api/v3/urls",
+        headers=headers,
+        data={"url": url}
+    )
     resp.raise_for_status()
     analysis_id = resp.json()["data"]["id"]
+
+    # Poll every second until complete
     while True:
-        time.sleep(5)
-        st = session.get(f"https://www.virustotal.com/api/v3/analyses/{analysis_id}", headers=headers)
+        time.sleep(1)
+        st = session.get(
+            f"https://www.virustotal.com/api/v3/analyses/{analysis_id}",
+            headers=headers
+        )
         st.raise_for_status()
         attr = st.json()["data"]["attributes"]
         if attr.get("status") == "completed":
             return attr.get("stats", {}).get("malicious", 0) == 0
 
-def extract_ffprobe_metadata(url_or_path):
-    out = subprocess.check_output([
+def extract_ffprobe_metadata(path):
+    """Run ffprobe and parse its JSON output."""
+    cmd = [
         "ffprobe", "-v", "error", "-print_format", "json",
-        "-show_format", "-show_streams", url_or_path
-    ])
+        "-show_format", "-show_streams", path
+    ]
+    out = subprocess.check_output(cmd)
     return json.loads(out)
 
 def fetch_page_metadata(url):
+    """Grab <title>, og: and twitter: meta tags from any page."""
     try:
-        r = session.get(url, timeout=5); r.raise_for_status()
+        r = session.get(url, timeout=5)
+        r.raise_for_status()
         soup = BeautifulSoup(r.text, "html.parser")
-        data = {"url": url, "title": getattr(soup.title, "string", None)}
+        data = {
+            "url": url,
+            "title": getattr(soup.title, "string", None)
+        }
         for tag in soup.find_all("meta"):
             prop = tag.get("property") or tag.get("name")
             if prop and prop.startswith(("og:", "twitter:")):
@@ -43,56 +76,63 @@ def fetch_page_metadata(url):
         return {"url": url, "error": str(e)}
 
 def fetch_clean_videos(keywords, api_key, scan_enabled):
+    """Search IA, filter by file extension, optional parallel VT scan."""
     terms = [kw.strip().replace(" ", "+") for kw in keywords.split(",")]
     ia_query = f"mediatype:(movies) AND ({' OR '.join(terms)})"
-    results = list(search_items(ia_query))[:50]
+    results = list(search_items(ia_query))[:10]  # top 10 to stay speedy
 
-    clean = []
+    # collect candidate video URLs
+    candidates = []
     for r in results:
         ident = r["identifier"]
-        item = get_item(ident)
+        item = get_cached_item(ident)
         for f in item.files:
             name = f.get("name", "")
             ext  = name.lower().split(".")[-1]
             if ext in ("mp4", "mkv", "avi", "mov", "webm", "m4v"):
-                url = f"https://archive.org/download/{ident}/{name}"
-                if scan_enabled and api_key:
-                    try:
-                        if not scan_url_vt(url, api_key):
-                            continue
-                    except:
-                        continue
-                clean.append(url)
-    return clean
+                candidates.append(f"https://archive.org/download/{ident}/{name}")
+
+    # optionally scan in parallel
+    if scan_enabled and api_key:
+        with ThreadPoolExecutor(max_workers=5) as exe:
+            safe_urls = exe.map(lambda u: u if scan_url_vt(u, api_key) else None, candidates)
+        return [u for u in safe_urls if u]
+    return candidates
 
 def get_favicon_url(page_url):
+    """Assume https://domain/favicon.ico exists."""
     dom = urlparse(page_url).netloc
     return f"https://{dom}/favicon.ico"
 
 def build_origin_graph(origins):
+    """Render a PyVis network; return its full HTML."""
     net = Network(height="400px", width="100%", directed=True)
     for i, m in enumerate(origins):
-        fav = get_favicon_url(m["url"])
-        label = urlparse(m["url"]).netloc
-        title = json.dumps(m, indent=2)
-        net.add_node(i, label=label, title=title, shape="image", image=fav)
+        net.add_node(
+            i,
+            label=urlparse(m["url"]).netloc,
+            title=json.dumps(m, indent=2),
+            shape="image",
+            image=get_favicon_url(m["url"])
+        )
         if i > 0:
             net.add_edge(i - 1, i)
-    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".html").name
-    net.save_graph(tmp)
-    return open(tmp, encoding="utf-8").read()
+    tmpfile = tempfile.NamedTemporaryFile(delete=False, suffix=".html").name
+    net.save_graph(tmpfile)
+    return open(tmpfile, encoding="utf-8").read()
 
 with gr.Blocks(theme=THEME) as demo:
-    gr.Markdown("## 📼 IA Drone‑Strike Explorer")
+    gr.Markdown("## 📼 IA Drone-Strike Explorer")
+
     with gr.Row():
         kw     = gr.Textbox(label="Search keywords", value="drone strike, military uav")
         vt_key = gr.Textbox(label="VirusTotal API Key", type="password")
-    scan   = gr.Checkbox(label="Enable VT scan", value=True)
-    ff     = gr.Checkbox(label="Enable FFprobe metadata", value=False)
-    btn    = gr.Button("Search & Scan")
+    scan = gr.Checkbox(label="Enable VT scan", value=True)
+    ff   = gr.Checkbox(label="Enable FFprobe metadata", value=False)
+    btn  = gr.Button("Search & Scan")
 
-    dropdown    = gr.Dropdown(label="Clean Video URLs", choices=[])
-    graph_html  = gr.HTML("<p>No origin graph yet.</p>")
+    dropdown     = gr.Dropdown(label="Clean Video URLs", choices=[])
+    graph_html   = gr.HTML("<p>No origin graph yet.</p>")
     video_player = gr.Video()
 
     with gr.Tabs():
@@ -111,34 +151,55 @@ with gr.Blocks(theme=THEME) as demo:
     def update_all(url_sel, ff_on, api_key):
         if not url_sel:
             return None, {}, {}, "<p>No origin graph.</p>", []
-        parts = url_sel.split("/")
-        ident = parts[4]; fn = parts[-1]
-        ia_data = {"identifier": ident, "file": {}}
+
+        parts    = url_sel.split("/")
+        ident, fn = parts[4], parts[-1]
+        ia_data  = {"identifier": ident, "file": {}}
+
+        # — Fix metadata merge
         try:
-            item = get_item(ident)
-            ia_data["file"] = next(
-                {**{k: item.metadata.get(k)}, **f}
-                for f in item.files if f["name"] == fn
-            )
+            item = get_cached_item(ident)
+            file_rec = next((f for f in item.files if f["name"] == fn), None)
+            if file_rec:
+                ia_data["file"] = {**item.metadata, **file_rec}
+            else:
+                ia_data["file"] = {}
         except Exception as e:
             ia_data["error"] = str(e)
 
+        # — FFprobe
         ff_data = {}
         if ff_on:
-            try:    ff_data = extract_ffprobe_metadata(url_sel)
-            except Exception as e: ff_data = {"error": str(e)}
+            try:
+                ff_data = extract_ffprobe_metadata(url_sel)
+            except Exception as e:
+                ff_data = {"error": str(e)}
 
-        desc = item.metadata.get("description", "") or ""
+        # — Origin tracing
+        desc  = (item.metadata.get("description") or "")
         found = re.findall(r"https?://[^\s\"'<]+", desc)
-        origins = [fetch_page_metadata(u) for u in found]
-        graph   = build_origin_graph(origins) if origins else "<p>No origins.</p>"
+        if found:
+            with ThreadPoolExecutor(max_workers=5) as exe:
+                origins = list(exe.map(fetch_page_metadata, found))
+            graph   = build_origin_graph(origins)
+        else:
+            origins, graph = [], "<p>No origins.</p>"
 
         return url_sel, ia_data, ff_data, graph, origins
 
-    btn.click(fn=search_and_populate, inputs=[kw, vt_key, scan], outputs=[dropdown])
-    dropdown.change(fn=update_all,
-                    inputs=[dropdown, ff, vt_key],
-                    outputs=[video_player, ia_json, ff_json, graph_html, orig_json])
+    btn.click(
+        fn=search_and_populate,
+        inputs=[kw, vt_key, scan],
+        outputs=[dropdown]
+    )
+    dropdown.change(
+        fn=update_all,
+        inputs=[dropdown, ff, vt_key],
+        outputs=[video_player, ia_json, ff_json, graph_html, orig_json]
+    )
+
+    # <<< enable async queue & spinner >>>
+    demo.queue()
 
 if __name__ == "__main__":
     demo.launch()