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()