app.py

import gradio as gr
import orjson
import requests
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:
        response = requests.get(
            f"https://api.shodan.io/shodan/host/search?key={SHODAN_API_KEY}&query={query}",
            timeout=10
        )
        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(
        fn=build_query,
        inputs=[query_input, platform, gr.State(value="session_id")],
        outputs=query_output
    )
    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
    )

app.launch()