Update app.py

app.py

@@ -6,153 +6,176 @@ import subprocess
 import json
 import re
 from bs4 import BeautifulSoup
+from urllib.parse import urlparse
+import tempfile
+from pyvis.network import Network
+
+THEME = "gradio/soft"
+
+# --- Shared HTTP session for speed & headers ---
+session = requests.Session()
+session.headers.update({
+    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64)"
+})
 
-# --- VirusTotal helper functions ---
 def scan_url_vt(url, api_key):
+    """Return True if VirusTotal sees no malicious hits."""
     headers = {"x-apikey": api_key}
-    resp = requests.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 until complete
     while True:
         time.sleep(5)
-        status_resp = requests.get(
-            f"https://www.virustotal.com/api/v3/analyses/{analysis_id}", headers=headers
-        )
-        status_resp.raise_for_status()
-        attr = status_resp.json()["data"]["attributes"]
+        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":
-            stats = attr.get("stats", {})
-            return stats.get("malicious", 0) == 0
+            return attr.get("stats", {}).get("malicious", 0) == 0
 
-# --- FFprobe metadata extraction ---
 def extract_ffprobe_metadata(url_or_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", url_or_path
     ]
     out = subprocess.check_output(cmd)
     return json.loads(out)
 
-# --- Scrape basic page metadata (title + og: tags) ---
 def fetch_page_metadata(url):
+    """Grab <title>, og: and twitter: meta tags from any page."""
     try:
-        resp = requests.get(url, timeout=5)
-        resp.raise_for_status()
-        html = resp.text
-        soup = BeautifulSoup(html, "html.parser")
-        meta = {"url": url, "title": soup.title.string if soup.title else None}
-        # grab OpenGraph tags
+        r = session.get(url, timeout=5)
+        r.raise_for_status()
+        soup = BeautifulSoup(r.text, "html.parser")
+        data = {
+            "url": url,
+            "title": soup.title.string if soup.title else None
+        }
         for tag in soup.find_all("meta"):
             prop = tag.get("property") or tag.get("name")
             if prop and prop.startswith(("og:", "twitter:")):
-                meta[prop] = tag.get("content")
-        return meta
+                data[prop] = tag.get("content")
+        return data
     except Exception as e:
         return {"url": url, "error": str(e)}
 
-# --- Core search & scan logic ---
 def fetch_clean_videos(keywords, api_key, scan_enabled):
-    query = " OR ".join([f"{kw.strip().replace(' ', '+')}" for kw in keywords.split(",")])
-    ia_query = f"mediatype:(movies) AND ({query})"
+    """Search IA, filter for common video formats, optional 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]
 
-    clean_urls = []
-    for res in results:
-        identifier = res["identifier"]
-        item = get_item(identifier)
+    clean = []
+    for r in results:
+        ident = r["identifier"]
+        item = get_item(ident)
         for f in item.files:
             fmt = f.get("format", "").lower()
-            if fmt.startswith(("mpeg","mp4","avi","mov","webm","m4v")):
-                url = f"https://archive.org/download/{identifier}/{f['name']}"
+            if fmt.startswith(("mp4", "avi", "mov", "webm", "m4v")):
+                url = f"https://archive.org/download/{ident}/{f['name']}"
                 if scan_enabled and api_key:
                     try:
-                        is_clean = scan_url_vt(url, api_key)
-                    except Exception:
+                        if not scan_url_vt(url, api_key):
+                            continue
+                    except:
                         continue
-                else:
-                    is_clean = True
-                if is_clean:
-                    clean_urls.append(url)
-    return clean_urls
-
-# --- Gradio UI setup ---
-with gr.Blocks() as demo:
-    gr.Markdown("# 📼 IA Drone‑Strike Explorer  \nEnable VT scan, FFprobe & Origin Tracing")
+                clean.append(url)
+    return clean
+
+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)
+        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()
+
+with gr.Blocks(theme=THEME) as demo:
+    gr.Markdown("## 📼 IA Drone-Strike Explorer")
     with gr.Row():
-        kw_input     = gr.Textbox(label="Search keywords", value="drone strike, military uav")
-        vt_key_input = gr.Textbox(label="VirusTotal API Key", type="password")
-    scan_toggle    = gr.Checkbox(label="Enable VT scan", value=True)
-    ffprobe_toggle = gr.Checkbox(label="Enable FFprobe metadata", value=False)
-    run_btn        = gr.Button("Search & Scan")
-
-    url_dropdown   = gr.Dropdown(label="Clean Video URLs", choices=[], interactive=True)
-    video_player   = gr.Video(label="Video Player")
-    ia_meta_json   = gr.JSON(label="► Raw IA Metadata")
-    ffprobe_json   = gr.JSON(label="► FFprobe Metadata")
-    origins_json   = gr.JSON(label="► Source‑Origin Metadata")
+        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")
+
+    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():
+        with gr.TabItem("IA Metadata"):
+            ia_json = gr.JSON()
+        with gr.TabItem("FFprobe"):
+            ff_json = gr.JSON()
+        with gr.TabItem("Origins"):
+            graph_html
+            orig_json = gr.JSON()
 
     def search_and_populate(keywords, api_key, scan_enabled):
         urls = fetch_clean_videos(keywords, api_key, scan_enabled)
         return gr.update(choices=urls, value=urls[0] if urls else None)
 
-    def update_all(selected_url, ff_on, api_key):
-        # no selection guard
-        if not selected_url:
-            return None, {}, {}, []
+    def update_all(url_sel, ff_on, api_key):
+        if not url_sel:
+            return None, {}, {}, "<p>No origin graph.</p>", []
 
-        # 1) IA metadata + file list
-        parts = selected_url.split("/")
-        identifier = parts[4] if len(parts) > 4 else None
-        raw_ia = {"identifier": identifier, "metadata": {}, "files": []}
-        if identifier:
-            try:
-                item = get_item(identifier)
-                raw_ia["metadata"] = item.metadata
-                raw_ia["files"] = [
-                    {
-                        "name": f.get("name"),
-                        "format": f.get("format"),
-                        "size": f.get("size"),
-                        "md5": f.get("md5"),
-                        **{k: v for k,v in f.items() if k not in ("name","format","size","md5")}
-                    }
-                    for f in item.files
-                ]
-            except Exception:
-                raw_ia["error"] = "could not fetch IA metadata"
-
-        # 2) FFprobe metadata if toggled
-        ff_md = {}
+        # --- IA metadata for just this file ---
+        parts = url_sel.split("/")
+        ident = parts[4]
+        fn = parts[-1]
+        ia_data = {"identifier": ident, "file": {}}
+        try:
+            item = get_item(ident)
+            ia_data["file"] = next(
+                {**{k: item.metadata.get(k)}, **f}
+                for f in item.files if f["name"] == fn
+            )
+        except Exception as e:
+            ia_data["error"] = str(e)
+
+        # --- FFprobe if requested ---
+        ff_data = {}
         if ff_on:
             try:
-                ff_md = extract_ffprobe_metadata(selected_url)
+                ff_data = extract_ffprobe_metadata(url_sel)
             except Exception as e:
-                ff_md = {"error": str(e)}
+                ff_data = {"error": str(e)}
 
-        # 3) Origin tracing: scrape each URL in description
-        origins = []
-        desc = raw_ia["metadata"].get("description", "")
-        urls_found = re.findall(r'https?://[^\s"<]+', desc)
-        for url in urls_found:
-            meta = fetch_page_metadata(url)
-            origins.append(meta)
-            # stop at first “real” origin (you can remove this break to collect all)
-            break
+        # --- Origin tracing: all links in description ---
+        desc = item.metadata.get("description", "") or ""
+        found = re.findall(r"https?://[^\s\"'<]+", desc)
+        origins = [fetch_page_metadata(u) for u in found]
 
-        return selected_url, raw_ia, ff_md, origins
+        # --- Build the interactive graph HTML ---
+        graph = build_origin_graph(origins) if origins else "<p>No origins.</p>"
+        return url_sel, ia_data, ff_data, graph, origins
 
-    run_btn.click(
+    btn.click(
         fn=search_and_populate,
-        inputs=[kw_input, vt_key_input, scan_toggle],
-        outputs=[url_dropdown]
+        inputs=[kw, vt_key, scan],
+        outputs=[dropdown]
     )
-    url_dropdown.change(
+    dropdown.change(
         fn=update_all,
-        inputs=[url_dropdown, ffprobe_toggle, vt_key_input],
-        outputs=[video_player, ia_meta_json, ffprobe_json, origins_json]
+        inputs=[dropdown, ff, vt_key],
+        outputs=[video_player, ia_json, ff_json, graph_html, orig_json]
     )
 
 if __name__ == "__main__":