Update app.py

app.py

@@ -1,450 +1,172 @@
-import os, re, json, urllib.parse
-from typing import Dict, Any, List, Optional, Tuple
-
 import gradio as gr
+import orjson
 import requests
-
-APP_TITLE = "🔎 Query Builder (Educational)"
-APP_DESC = """
-Type what you want (e.g., *"find RDP in Germany for Org Contoso"*) and/or use the structured fields.
-This builds **query strings only** for training: **no live scans, no API calls to data sources**.
-
-**Platforms:** Shodan, ZoomEye, Google (dorks)
-**LLM parsing:** Optional — choose a model from the dropdown. Token box appears only when needed.
-"""
-
-# ----------------------------
-# Heuristic NL parsing (no LLM)
-# ----------------------------
-
-COUNTRY_NAME_TO_ISO2 = {
-    "united states": "US", "usa": "US", "us": "US", "america": "US",
-    "germany": "DE", "deutschland": "DE",
-    "japan": "JP",
-    "brazil": "BR",
-    "india": "IN",            # NOTE: removed bare "in" to avoid preposition collision
-    "united kingdom": "GB", "uk": "GB", "britain": "GB",
-    "france": "FR",
-    "canada": "CA",
-    "australia": "AU",
-    "netherlands": "NL",
-    "spain": "ES",
-    "italy": "IT",
-    "sweden": "SE",
-    "switzerland": "CH",
-}
-
-US_STATE_ABBREV = {
-    "AL","AK","AZ","AR","CA","CO","CT","DE","FL","GA","HI","ID","IL","IN","IA",
-    "KS","KY","LA","ME","MD","MA","MI","MN","MS","MO","MT","NE","NV","NH","NJ",
-    "NM","NY","NC","ND","OH","OK","OR","PA","RI","SC","SD","TN","TX","UT","VT",
-    "VA","WA","WV","WI","WY","DC"
-}
-US_STATE_NAMES = {
-    "alabama","alaska","arizona","arkansas","california","colorado","connecticut","delaware",
-    "florida","georgia","hawaii","idaho","illinois","indiana","iowa","kansas","kentucky",
-    "louisiana","maine","maryland","massachusetts","michigan","minnesota","mississippi",
-    "missouri","montana","nebraska","nevada","new hampshire","new jersey","new mexico",
-    "new york","north carolina","north dakota","ohio","oklahoma","oregon","pennsylvania",
-    "rhode island","south carolina","south dakota","tennessee","texas","utah","vermont",
-    "virginia","washington","west virginia","wisconsin","wyoming","district of columbia","dc"
-}
-
-# tokens we don't want to leak into queries
-STOPWORDS = {
-    "find","show","in","the","and","or","for","of","with",
-    "zip","zipcode","code","ks","us"  # cut noise like "zip code 67401", "KS"
-}
-
-SERVICE_HINTS = {
-    "rdp": (3389, "Remote Desktop"),
-    "ssh": (22, "OpenSSH"),
-    "ftp": (21, "FTP"),
-    "mongodb": (27017, "MongoDB"),
-    "mysql": (3306, "MySQL"),
-    "postgres": (5432, "PostgreSQL"),
-    "jenkins": (8080, "Jenkins"),
-    "elasticsearch": (9200, "Elasticsearch"),
-    "kibana": (5601, "Kibana"),
-    "redis": (6379, "Redis"),
-    "rabbitmq": (15672, "RabbitMQ"),
-    "rdp windows": (3389, "Windows RDP"),
-}
-
-def guess_country_iso2(text: str) -> Optional[str]:
-    t = text.lower()
-
-    # explicit "country:XX"
-    m = re.search(r"\bcountry\s*:\s*([A-Za-z]{2})\b", t)
-    if m:
-        return m.group(1).upper()
-
-    # country names (no "in" shortcut!)
-    for name, iso2 in COUNTRY_NAME_TO_ISO2.items():
-        if re.search(rf"\b{name}\b", t):
-            return iso2
-
-    # US states imply US
-    if re.search(r"\b(" + "|".join(US_STATE_NAMES) + r")\b", t):
-        return "US"
-    if re.search(r"\b(" + "|".join(US_STATE_ABBREV) + r")\b", t):
-        return "US"
-    return None
-
-def guess_port(text: str) -> Optional[int]:
-    m = re.search(r"\bport\s*[:=]?\s*(\d{1,5})\b", text, flags=re.I)
-    return int(m.group(1)) if m else None
-
-def guess_service(text: str) -> Tuple[Optional[str], Optional[int]]:
-    t = text.lower()
-    for k, (p, _) in SERVICE_HINTS.items():
-        if k in t:
-            return (k, p)
-    return (None, None)
-
-def guess_org(text: str) -> Optional[str]:
-    m = re.search(r"\borg(?:anization)?\s*[:=]?\s*([A-Za-z0-9 _\-.&]+)", text, flags=re.I)
-    return m.group(1).strip() if m else None
-
-def guess_asn(text: str) -> Optional[str]:
-    m = re.search(r"\bAS(\d{1,10})\b", text, flags=re.I)
-    return f"AS{m.group(1)}" if m else None
-
-def guess_city(text: str) -> Optional[str]:
-    m = re.search(r"\bcity\s*[:=]?\s*([A-Za-z \-'.]+)\b", text, flags=re.I)
-    return m.group(1).strip() if m else None
-
-def guess_os(text: str) -> Optional[str]:
-    m = re.search(r"\b(windows\s?\d+|ubuntu|debian|centos|linux|macos|ios|android)\b", text, flags=re.I)
-    return m.group(1) if m else None
-
-def guess_product(text: str) -> Optional[str]:
-    for k, (_, prod) in SERVICE_HINTS.items():
-        if k in text.lower():
-            return prod
-    m = re.search(r"\bproduct\s*[:=]?\s*([A-Za-z0-9 _\-.]+)", text, flags=re.I)
-    return m.group(1).strip() if m else None
-
-def guess_title(text: str) -> Optional[str]:
-    m = re.search(r'(?:title|intitle)\s*[:=]?\s*"([^"]+)"', text, flags=re.I)
-    return m.group(1) if m else None
-
-def extract_zip(text: str) -> Optional[str]:
-    """Return a 5-digit US ZIP if present; used only for display/JSON (not added to queries)."""
-    m = re.search(r"\b(?:zip(?:\s*code)?|zipcode)?\s*[:=]?\s*(\d{5})(?:-\d{4})?\b", text, flags=re.I)
-    return m.group(1) if m else None
-
-def nl_to_structured_heuristic(nl: str) -> Dict[str, Any]:
-    iso2 = guess_country_iso2(nl) or ""
-    port = guess_port(nl)
-    _, svc_port = guess_service(nl)
-    org = guess_org(nl) or ""
-    asn = guess_asn(nl) or ""
-    city = guess_city(nl) or ""
-    os_name = guess_os(nl) or ""
-    product = guess_product(nl) or ""
-    title = guess_title(nl) or ""
-    postal = extract_zip(nl) or ""
-
-    if port is None and svc_port:
-        port = svc_port
-
-    # harvest terms but drop stopwords and very short noise
-    raw_terms = re.findall(r"[A-Za-z0-9_.:/+-]{3,}", nl)
-    terms = []
-    for w in raw_terms:
-        lw = w.lower()
-        if lw in STOPWORDS:
-            continue
-        terms.append(w)
-    terms = list(dict.fromkeys(terms))[:12]
-
-    return {
-        "terms": terms,
-        "country_code": iso2,
-        "city": city,
-        "port": port,
-        "product": product,
-        "os": os_name,
-        "org": org,
-        "asn": asn,
-        "title": title,
-        "postal": postal,      # kept for transparency; not injected into queries
-        "filetype": "",
-        "site": "",
-        "after": "",
-        "before": "",
-    }
-
-# ----------------------------
-# LLM parsing via HF Inference API (model chosen in UI)
-# ----------------------------
-
-MODEL_CHOICES = {
-    "No LLM (fastest, heuristic only)": {"id": "", "requires_token": False},
-    "LLM 1 (token required) — TinyLlama-1.1B-Chat": {"id": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", "requires_token": True},
-    "LLM 2 (token required) — Phi-3-mini-4k": {"id": "microsoft/Phi-3-mini-4k-instruct", "requires_token": True},
-    "LLM 3 (token required) — Qwen2.5-7B-Instruct": {"id": "Qwen/Qwen2.5-7B-Instruct", "requires_token": True},
-}
-
-PROMPT = """You are a JSON-only extractor. Convert the user's intent into fields for building search queries.
-
-Return a single JSON object with keys:
-terms (list of up to 10 short tokens),
-country_code (ISO-2 or empty),
-city (string),
-port (int or null),
-product (string),
-os (string),
-org (string),
-asn (like "AS15169" or empty),
-title (string),
-postal (string),
-filetype (string),
-site (string),
-after (YYYY-MM-DD or empty),
-before (YYYY-MM-DD or empty).
-
-User: {text}
-JSON:"""
-
-def call_hf_inference(model_id: str, token: str, text: str, timeout: int = 25) -> Optional[Dict[str, Any]]:
+import time
+from pathlib import Path
+from typing import Dict, List, Optional
+from loguru import logger
+from registry import Registry
+from config import DATA_DIR, HF_TOKEN, SHODAN_API_KEY, ZOOMEYE_API_KEY, ENABLE_EXECUTE, RATE_LIMIT_PER_MIN, OWNER_PIN, LLM_MODE, ALLOWED_MODELS, THEME, COUNTRIES
+
+DATA_DIR.mkdir(exist_ok=True)
+logger.add(DATA_DIR / "app.log", rotation="1 MB", retention="7 days")
+registry = Registry()
+session_counts: Dict[str, List[float]] = {}
+
+def rate_limit(session_id: str) -> bool:
+    now = time.time()
+    if session_id not in session_counts:
+        session_counts[session_id] = []
+    session_counts[session_id] = [t for t in session_counts[session_id] if now - t < 60]
+    if len(session_counts[session_id]) >= RATE_LIMIT_PER_MIN:
+        return False
+    session_counts[session_id].append(now)
+    return True
+
+def build_query(query: str, platform: str, session_id: str) -> str:
+    if not rate_limit(session_id):
+        return "Rate limit exceeded. Please try again later."
+    parsed = registry.parse_query(query)
+    intent = parsed["intents"][0] if parsed["intents"] else None
+    service = parsed["services"][0] if parsed["services"] else None
+    base_query = []
+    if intent and intent in registry.store.intents:
+        templates = registry.store.intents[intent].platform_templates.get(platform, [])
+        defaults = registry.store.intents[intent].defaults
+        for template in templates:
+            formatted = template.format(**defaults)
+            base_query.append(formatted)
+    elif service:
+        for s in registry.store.services:
+            if s["name"] == service:
+                base_query.append(f'port:{s["port"]}' if s.get("port") else f'product:"{s["name"]}"')
+    for key, value in parsed["filters"].items():
+        if key in registry.store.operators:
+            base_query.append(f'{key}:"{value}"')
+    sanitized_query = registry.sanitize_operators(" ".join(base_query + parsed["operators"]))
+    logger.info(f"Built {platform} query: {sanitized_query}")
+    return sanitized_query
+
+def exec_shodan(query: str, session_id: str) -> str:
+    if not ENABLE_EXECUTE or not SHODAN_API_KEY:
+        return "Shodan execution is disabled or API key is missing."
+    if not rate_limit(session_id):
+        return "Rate limit exceeded. Please try again later."
     try:
-        headers = {"Authorization": f"Bearer {token}"} if token else {}
-        payload = {"inputs": PROMPT.format(text=text),
-                   "parameters": {"max_new_tokens": 256, "temperature": 0.2, "return_full_text": False}}
-        resp = requests.post(f"https://api-inference.huggingface.co/models/{model_id}",
-                             headers=headers, json=payload, timeout=timeout)
-        resp.raise_for_status()
-        data = resp.json()
-        gen = (data[0]["generated_text"] if isinstance(data, list) and data and "generated_text" in data[0]
-               else data.get("generated_text") if isinstance(data, dict) else json.dumps(data))
-        m = re.search(r"\{.*\}", gen, flags=re.S)
-        return json.loads(m.group(0)) if m else None
-    except Exception:
-        return None
-
-def nl_to_structured(nl: str, model_label: str, token: str) -> Dict[str, Any]:
-    baseline = nl_to_structured_heuristic(nl)
-    model_info = MODEL_CHOICES.get(model_label, {"id": "", "requires_token": False})
-    model_id = model_info["id"]
-    if not model_id:
-        return baseline
-    if model_info["requires_token"] and not token:
-        return baseline
-    parsed = call_hf_inference(model_id, token, nl)
-    if isinstance(parsed, dict):
-        merged = baseline.copy()
-        for k in merged.keys():
-            if k in parsed and parsed[k] not in (None, "", []):
-                merged[k] = parsed[k]
-        return merged
-    return baseline
-
-# ----------------------------
-# Query builders
-# ----------------------------
-
-def build_shodan(params: Dict[str, Any]) -> str:
-    parts = []
-    terms = [t for t in params["terms"] if t.lower() not in STOPWORDS]
-    parts += terms
-    if params.get("country_code"): parts.append(f"country:{params['country_code']}")
-    if params.get("city"): parts.append(f"city:{json.dumps(params['city'])}")
-    if params.get("port"): parts.append(f"port:{params['port']}")
-    if params.get("product"): parts.append(f'product:{json.dumps(params["product"])}')
-    if params.get("os"): parts.append(f'os:{json.dumps(params["os"])}')
-    if params.get("org"): parts.append(f'org:{json.dumps(params["org"])}')
-    if params.get("asn"): parts.append(f"asn:{params['asn']}")
-    if params.get("title"): parts.append(f'http.title:{json.dumps(params["title"])}')
-    # Deliberately *not* inserting ZIP/postal to avoid unsupported filters
-    return " ".join(parts).strip()
-
-def build_zoomeye(params: Dict[str, Any]) -> str:
-    parts = []
-    terms = [t for t in params["terms"] if t.lower() not in STOPWORDS]
-    parts += terms
-    if params.get("product"): parts.append(f'app:{json.dumps(params["product"])}')
-    if params.get("port"): parts.append(f"port:{params['port']}")
-    if params.get("os"): parts.append(f'os:{json.dumps(params["os"])}')
-    if params.get("country_code"): parts.append(f"country:{params['country_code']}")
-    if params.get("city"): parts.append(f'city:{json.dumps(params["city"])}')
-    if params.get("org"): parts.append(f'org:{json.dumps(params["org"])}')
-    if params.get("asn"): parts.append(f"asn:{params['asn']}")
-    return " ".join(parts).strip()
-
-def build_google(params: Dict[str, Any]) -> str:
-    parts = []
-    if params.get("title"): parts.append(f'intitle:{json.dumps(params["title"])}')
-    if params.get("site"): parts.append(f"site:{params['site']}")
-    if params.get("filetype"): parts.append(f"filetype:{params['filetype']}")
-    if params.get("after"): parts.append(f"after:{params['after']}")
-    if params.get("before"): parts.append(f"before:{params['before']}")
-    extras = []
-    if params.get("product"): extras.append(params["product"])
-    if params.get("os"): extras.append(params["os"])
-    # include ZIP as plain term for Google only, if present
-    if params.get("postal"): extras.append(params["postal"])
-    for t in params["terms"][:8]:
-        extras.append(t)
-    parts += [json.dumps(x) if re.search(r"\s", x) else x for x in extras if x]
-    return " ".join(parts).strip()
-
-def to_click_link(base: str, q: str) -> str:
-    return f"{base}{urllib.parse.quote_plus(q)}" if q else ""
-
-# ----------------------------
-# Gradio callbacks
-# ----------------------------
-
-def build_queries(nl: str,
-                  country_code: str, city: str, port: str,
-                  product: str, os_name: str, org: str, asn: str,
-                  title: str, site: str, filetype: str, after: str, before: str,
-                  enable_shodan: bool, enable_zoomeye: bool, enable_google: bool,
-                  model_label: str, hf_token: str) -> Dict[str, str]:
-
-    parsed = nl_to_structured(nl or "", model_label, hf_token or "")
-
-    # UI overrides (explicit beats inferred)
-    if country_code: parsed["country_code"] = country_code
-    if city: parsed["city"] = city
-    if port.strip().isdigit(): parsed["port"] = int(port.strip())
-    if product: parsed["product"] = product
-    if os_name: parsed["os"] = os_name
-    if org: parsed["org"] = org
-    if asn: parsed["asn"] = asn.upper() if not asn.upper().startswith("AS") else asn.upper()
-    if title: parsed["title"] = title
-    if site: parsed["site"] = site
-    if filetype: parsed["filetype"] = filetype
-    if after: parsed["after"] = after
-    if before: parsed["before"] = before
-
-    out = {}
-    if enable_shodan:
-        shodan_q = build_shodan(parsed)
-        out["Shodan Query"] = shodan_q
-        out["Open in Shodan"] = to_click_link("https://www.shodan.io/search?query=", shodan_q)
-    if enable_zoomeye:
-        zq = build_zoomeye(parsed)
-        out["ZoomEye Query"] = zq
-        out["Open in ZoomEye"] = to_click_link("https://www.zoomeye.org/searchResult?q=", zq)
-    if enable_google:
-        gq = build_google(parsed)
-        out["Google Dork"] = gq
-        out["Open in Google"] = to_click_link("https://www.google.com/search?q=", gq)
-
-    out["Parsed (JSON)"] = json.dumps(parsed, indent=2)
-    return out
-
-def on_model_change(selected_label: str):
-    """Show token box only when a model that requires it is chosen. Clear token when not needed."""
-    requires = MODEL_CHOICES.get(selected_label, {"requires_token": False})["requires_token"]
-    if requires:
-        return gr.update(visible=True), gr.update(value="")
-    else:
-        return gr.update(visible=False), gr.update(value="")
-
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown(f"# {APP_TITLE}")
-    gr.Markdown(APP_DESC)
-
-    with gr.Row():
-        with gr.Column(scale=3):
-            nl = gr.Textbox(label="Natural language", placeholder='e.g., "Find RDP in Germany for Org Contoso"', lines=3)
-
-            with gr.Accordion("Structured inputs (optional overrides)", open=False):
-                country = gr.Dropdown(
-                    label="Country (ISO-2)",
-                    choices=["", "US","DE","JP","BR","IN","GB","FR","CA","AU","NL","ES","IT","SE","CH"],
-                    value="",
-                    allow_custom_value=False
-                )
-                city = gr.Textbox(label="City")
-                port = gr.Textbox(label="Port (number)")
-                product = gr.Textbox(label="Product/App (e.g., Jenkins, MongoDB)")
-                os_name = gr.Textbox(label="OS (e.g., Windows 10, Ubuntu)")
-                org = gr.Textbox(label="Org")
-                asn = gr.Textbox(label="ASN (e.g., AS15169)")
-                title = gr.Textbox(label='Page Title contains (exact match)')
-                with gr.Row():
-                    site = gr.Textbox(label="site: (Google only)")
-                    filetype = gr.Textbox(label="filetype: (Google only)")
-                with gr.Row():
-                    after = gr.Textbox(label="after: YYYY-MM-DD (Google)")
-                    before = gr.Textbox(label="before: YYYY-MM-DD (Google)")
-
-            gr.Markdown("### LLM (optional)")
-            model_label = gr.Dropdown(
-                label="Model",
-                choices=list(MODEL_CHOICES.keys()),
-                value="No LLM (fastest, heuristic only)",
-                info="Choose an LLM to improve NL parsing. Token required for LLM options.",
-                allow_custom_value=False
-            )
-            hf_token = gr.Textbox(
-                label="Hugging Face API Token",
-                type="password",
-                placeholder="hf_...",
-                visible=False
-            )
-
-            with gr.Row():
-                enable_shodan = gr.Checkbox(value=True, label="Build Shodan query")
-                enable_zoomeye = gr.Checkbox(value=True, label="Build ZoomEye query")
-                enable_google = gr.Checkbox(value=True, label="Build Google dork")
-
-            build_btn = gr.Button("Build Queries")
-
-        with gr.Column(scale=4):
-            # OUTPUTS ARE NOW READ-ONLY (non-editable) BUT COPYABLE
-            shodan_out = gr.Textbox(label="Shodan Query", show_copy_button=True, interactive=False, lines=2)
-            shodan_link = gr.Markdown()
-            zoomeye_out = gr.Textbox(label="ZoomEye Query", show_copy_button=True, interactive=False, lines=2)
-            zoomeye_link = gr.Markdown()
-            google_out = gr.Textbox(label="Google Dork", show_copy_button=True, interactive=False, lines=2)
-            google_link = gr.Markdown()
-            parsed_json = gr.Code(label="Parsed (JSON)", interactive=False)
-
-    model_label.change(on_model_change, inputs=[model_label], outputs=[hf_token, hf_token])
-
-    def _on_build(nl, country, city, port, product, os_name, org, asn, title, site, filetype, after, before,
-                  en_shodan, en_zoomeye, en_google, model_label, hf_token):
-        d = build_queries(
-            nl, country, city, port, product, os_name, org, asn, title, site, filetype, after, before,
-            en_shodan, en_zoomeye, en_google, model_label, hf_token
+        response = requests.get(
+            f"https://api.shodan.io/shodan/host/search?key={SHODAN_API_KEY}&query={query}",
+            timeout=10
         )
-        return (
-            d.get("Shodan Query",""), f"[Open in Shodan]({d['Open in Shodan']})" if "Open in Shodan" in d else "",
-            d.get("ZoomEye Query",""), f"[Open in ZoomEye]({d['Open in ZoomEye']})" if "Open in ZoomEye" in d else "",
-            d.get("Google Dork",""), f"[Open in Google]({d['Open in Google']})" if "Open in Google" in d else "",
-            d.get("Parsed (JSON)","{}")
+        response.raise_for_status()
+        data = response.json()
+        logger.info(f"Shodan query executed: {query}")
+        return orjson.dumps(data, option=orjson.OPT_INDENT_2).decode()
+    except Exception as e:
+        logger.error(f"Shodan query failed: {e}")
+        return f"Error: {e}"
+
+def exec_zoomeye(query: str, session_id: str) -> str:
+    if not ENABLE_EXECUTE or not ZOOMEYE_API_KEY:
+        return "ZoomEye execution is disabled or API key is missing."
+    if not rate_limit(session_id):
+        return "Rate limit exceeded. Please try again later."
+    try:
+        response = requests.get(
+            f"https://api.zoomeye.org/host/search?query={query}&key={ZOOMEYE_API_KEY}",
+            timeout=10
         )
-
+        response.raise_for_status()
+        data = response.json()
+        logger.info(f"ZoomEye query executed: {query}")
+        return orjson.dumps(data, option=orjson.OPT_INDENT_2).decode()
+    except Exception as e:
+        logger.error(f"ZoomEye query failed: {e}")
+        return f"Error: {e}"
+
+def query_llm(query: str, model: str) -> str:
+    if model == "No LLM" or not HF_TOKEN:
+        return "LLM is disabled or HF token is missing."
+    try:
+        response = requests.post(
+            f"https://api-inference.huggingface.co/models/{model}",
+            headers={"Authorization": f"Bearer {HF_TOKEN}"},
+            json={"inputs": query},
+            timeout=10
+        )
+        response.raise_for_status()
+        data = response.json()
+        logger.info(f"LLM query executed: {query} with model {model}")
+        return data[0]["generated_text"] if isinstance(data, list) else str(data)
+    except Exception as e:
+        logger.error(f"LLM query failed: {e}")
+        return f"Error: {e}"
+
+def process_batch(queries: str, platform: str, session_id: str) -> List[Dict[str, str]]:
+    results = []
+    for query in queries.split("\n"):
+        query = query.strip()
+        if query:
+            built_query = build_query(query, platform, session_id)
+            results.append({"query": query, "built_query": built_query})
+    return results
+
+def ingest_data(url: str, pin: str) -> str:
+    if pin != OWNER_PIN:
+        return "Invalid owner PIN."
+    registry.ingest_from_url(url)
+    return f"Data ingested from {url}"
+
+with gr.Blocks(theme=THEME, css=".gradio-container { background-color: black; color: crimson; }") as app:
+    gr.Markdown("# DorkForge v3.2")
+    with gr.Row():
+        with gr.Column():
+            query_input = gr.Textbox(label="Enter Query", placeholder="e.g., Grafana in Germany port:3000")
+            platform = gr.Dropdown(choices=["google", "shodan", "zoomeye"], label="Platform", value="google")
+            build_btn = gr.Button("Build Query")
+            execute_btn = gr.Button("Execute Query")
+            query_output = gr.Textbox(label="Built Query", interactive=False)
+            result_output = gr.JSON(label="Query Result")
+        with gr.Column():
+            batch_input = gr.Textbox(label="Batch Queries (one per line)", lines=5)
+            batch_platform = gr.Dropdown(choices=["google", "shodan", "zoomeye"], label="Batch Platform", value="google")
+            batch_btn = gr.Button("Process Batch")
+            batch_output = gr.Dataframe(label="Batch Results")
+    with gr.Row():
+        with gr.Column():
+            llm_model = gr.Dropdown(choices=ALLOWED_MODELS, label="LLM Model", value=LLM_MODE)
+            llm_btn = gr.Button("Query LLM")
+            llm_output = gr.Textbox(label="LLM Response", interactive=False)
+    with gr.Row():
+        with gr.Accordion("Owner Controls", open=False):
+            ingest_url = gr.Textbox(label="Ingest URL", placeholder="e.g., https://example.com/data.json")
+            owner_pin = gr.Textbox(label="Owner PIN", type="password")
+            ingest_btn = gr.Button("Ingest Data")
+            ingest_output = gr.Textbox(label="Ingest Result", interactive=False)
+    
     build_btn.click(
-        _on_build,
-        inputs=[nl, country, city, port, product, os_name, org, asn, title, site, filetype, after, before,
-                enable_shodan, enable_zoomeye, enable_google, model_label, hf_token],
-        outputs=[shodan_out, shodan_link, zoomeye_out, zoomeye_link, google_out, google_link, parsed_json]
+        fn=build_query,
+        inputs=[query_input, platform, gr.State(value="session_id")],
+        outputs=query_output
     )
-
-    nl.change(
-        _on_build,
-        inputs=[nl, country, city, port, product, os_name, org, asn, title, site, filetype, after, before,
-                enable_shodan, enable_zoomeye, enable_google, model_label, hf_token],
-        outputs=[shodan_out, shodan_link, zoomeye_out, zoomeye_link, google_out, google_link, parsed_json]
+    execute_btn.click(
+        fn=lambda q, p, s: exec_shodan(q, s) if p == "shodan" else exec_zoomeye(q, s),
+        inputs=[query_output, platform, gr.State(value="session_id")],
+        outputs=result_output
+    )
+    batch_btn.click(
+        fn=process_batch,
+        inputs=[batch_input, batch_platform, gr.State(value="session_id")],
+        outputs=batch_output
+    )
+    llm_btn.click(
+        fn=query_llm,
+        inputs=[query_input, llm_model],
+        outputs=llm_output
+    )
+    ingest_btn.click(
+        fn=ingest_data,
+        inputs=[ingest_url, owner_pin],
+        outputs=ingest_output
     )
 
-    gr.Markdown("""
----
-**Notes & Docs**
-
-- Shodan filters: `country:DE`, `port:3389`, `org:"Contoso"`, `asn:AS15169`, `http.title:"XYZ"`.
-- ZoomEye fields: `app:"Jenkins"`, `port:8080`, `os:"Windows"`, `country:DE`, `city:"Berlin"`.
-- Google dorks: `site:`, `filetype:`, `inurl:`, `intitle:`, `after:`, `before:`.
-
-**Safety:** This Space creates strings only. No live scanning. Respect platform ToS and laws.
-""")
-
-if __name__ == "__main__":
-    demo.launch()
+app.launch()