"""
🛡️ AI Security Scanner
Downloads tools on first launch, then runs normally.
"""

import gradio as gr
import subprocess
import os
import json
import re
import urllib.request
import zipfile
import tarfile
import stat
import shutil
from datetime import datetime
from pathlib import Path

# ═══════════════════════════════════════════════════
# TOOL SETUP - Downloads pre-built binaries on first run
# ═══════════════════════════════════════════════════

BIN_DIR = os.path.join(os.path.expanduser("~"), "bin")
RESULTS_DIR = "/tmp/scan-results"
WORDLIST_PATH = "/tmp/wordlist.txt"
os.makedirs(BIN_DIR, exist_ok=True)
os.makedirs(RESULTS_DIR, exist_ok=True)
os.environ["PATH"] = BIN_DIR + ":" + os.environ.get("PATH", "")

TOOLS = {
    "subfinder": {
        "url": "https://github.com/projectdiscovery/subfinder/releases/download/v2.6.7/subfinder_2.6.7_linux_amd64.zip",
        "type": "zip",
        "binary": "subfinder"
    },
    "httpx": {
        "url": "https://github.com/projectdiscovery/httpx/releases/download/v1.6.9/httpx_1.6.9_linux_amd64.zip",
        "type": "zip",
        "binary": "httpx"
    },
    "nuclei": {
        "url": "https://github.com/projectdiscovery/nuclei/releases/download/v3.3.7/nuclei_3.3.7_linux_amd64.zip",
        "type": "zip",
        "binary": "nuclei"
    },
    "gau": {
        "url": "https://github.com/lc/gau/releases/download/v2.2.4/gau_2.2.4_linux_amd64.tar.gz",
        "type": "tar",
        "binary": "gau"
    },
    "dalfox": {
        "url": "https://github.com/hahwul/dalfox/releases/download/v2.9.3/dalfox_2.9.3_linux_amd64.tar.gz",
        "type": "tar",
        "binary": "dalfox"
    },
    "ffuf": {
        "url": "https://github.com/ffuf/ffuf/releases/download/v2.1.0/ffuf_2.1.0_linux_amd64.tar.gz",
        "type": "tar",
        "binary": "ffuf"
    },
}

WORDLIST_CONTENT = """admin
api
.env
.git
.gitignore
backup
config
console
dashboard
db
debug
dev
docs
download
error
files
graphql
health
help
images
info
internal
login
logout
logs
metrics
monitor
panel
phpinfo
private
readme
register
reset
robots.txt
rss
search
security
server-status
settings
setup
signin
signup
sitemap.xml
staging
static
status
swagger
system
test
tmp
token
upload
uploads
user
users
version
wp-admin
wp-content
wp-login.php
xmlrpc.php
"""


def download_tool(name, info):
    """Download and extract a single tool."""
    binary_path = os.path.join(BIN_DIR, info["binary"])
    if os.path.exists(binary_path):
        return True  # Already downloaded

    tmp = f"/tmp/dl_{name}"
    os.makedirs(tmp, exist_ok=True)

    try:
        print(f"⬇️  Downloading {name}...")
        if info["type"] == "zip":
            dl_path = f"{tmp}/{name}.zip"
            urllib.request.urlretrieve(info["url"], dl_path)
            with zipfile.ZipFile(dl_path, 'r') as z:
                z.extractall(tmp)
        else:
            dl_path = f"{tmp}/{name}.tar.gz"
            urllib.request.urlretrieve(info["url"], dl_path)
            with tarfile.open(dl_path, 'r:gz') as t:
                t.extractall(tmp)

        # Find and move binary
        for root, dirs, files in os.walk(tmp):
            for f in files:
                if f == info["binary"]:
                    src = os.path.join(root, f)
                    shutil.copy2(src, binary_path)
                    os.chmod(binary_path, os.stat(binary_path).st_mode | stat.S_IEXEC | stat.S_IXGRP | stat.S_IXOTH)
                    print(f"✅ {name} installed")
                    return True

        print(f"⚠️  Binary not found for {name}")
        return False
    except Exception as e:
        print(f"❌ Failed to install {name}: {e}")
        return False
    finally:
        shutil.rmtree(tmp, ignore_errors=True)


def setup_tools():
    """Download all tools on startup."""
    print("=" * 50)
    print("🛡️  Setting up security tools...")
    print("=" * 50)

    status = {}
    for name, info in TOOLS.items():
        status[name] = download_tool(name, info)

    # Create wordlist
    if not os.path.exists(WORDLIST_PATH):
        with open(WORDLIST_PATH, "w") as f:
            f.write(WORDLIST_CONTENT.strip())
        print("✅ Wordlist created")

    installed = sum(1 for v in status.values() if v)
    print(f"\n✅ {installed}/{len(TOOLS)} tools ready")
    print("=" * 50)
    return status


# ═══════════════════════════════════════════════════
# SCANNER FUNCTIONS
# ═══════════════════════════════════════════════════

def clean(target):
    t = target.strip().lower()
    t = re.sub(r'^https?://', '', t).rstrip('/')
    if not re.match(r'^[a-zA-Z0-9][a-zA-Z0-9\-\.]*\.[a-zA-Z]{2,}$', t):
        return None
    return t


def run(cmd, timeout=300):
    try:
        p = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
        return p.stdout, p.stderr, p.returncode
    except subprocess.TimeoutExpired:
        return "", "⏰ Timeout", 1
    except FileNotFoundError:
        return "", "❌ Tool not installed", 1
    except Exception as e:
        return "", str(e), 1


def hdr(tool, target):
    return f"{'='*55}\n🔧 {tool}\n🎯 {target}\n🕐 {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n{'='*55}\n"


def do_subfinder(target):
    t = clean(target)
    if not t: return "❌ Invalid domain format. Use: example.com"
    out = f"{RESULTS_DIR}/subfinder_{t}.txt"
    if os.path.exists(out): os.remove(out)
    stdout, stderr, _ = run(["subfinder", "-d", t, "-silent", "-o", out], 120)
    res = hdr("Subfinder – Subdomain Discovery", t)
    if os.path.exists(out):
        data = open(out).read().strip()
        n = len(data.splitlines()) if data else 0
        res += f"\n✅ Found {n} subdomains\n\n" + (data if data else "No subdomains found.")
    else:
        res += f"\n{stdout}\n{stderr}"
    return res


def do_httpx(target):
    t = clean(target)
    if not t: return "❌ Invalid domain format. Use: example.com"
    sf = f"{RESULTS_DIR}/subfinder_{t}.txt"
    if os.path.exists(sf):
        cmd = ["httpx", "-l", sf, "-silent", "-status-code", "-title", "-tech-detect"]
    else:
        cmd = ["httpx", "-u", t, "-silent", "-status-code", "-title", "-tech-detect"]
    stdout, stderr, _ = run(cmd, 180)
    res = hdr("HTTPX – HTTP Probe & Tech Detect", t)
    res += f"\n{stdout}" if stdout.strip() else "\nNo live hosts found."
    return res


def do_gau(target):
    t = clean(target)
    if not t: return "❌ Invalid domain format. Use: example.com"
    out = f"{RESULTS_DIR}/gau_{t}.txt"
    if os.path.exists(out): os.remove(out)
    run(["gau", t, "--threads", "5", "--o", out], 120)
    res = hdr("GAU – URL Discovery", t)
    if os.path.exists(out):
        data = open(out).read().strip()
        lines = data.splitlines() if data else []
        res += f"\n✅ Found {len(lines)} URLs\n\n"
        res += "\n".join(lines[:150])
        if len(lines) > 150: res += f"\n\n... +{len(lines)-150} more"
    else:
        res += "\nNo URLs found."
    return res


def do_nuclei(target, severity):
    t = clean(target)
    if not t: return "❌ Invalid domain format. Use: example.com"
    out = f"{RESULTS_DIR}/nuclei_{t}.json"
    if os.path.exists(out): os.remove(out)
    cmd = ["nuclei", "-u", f"https://{t}", "-severity", severity,
           "-silent", "-jsonl", "-o", out, "-rate-limit", "50",
           "-no-update-templates"]
    stdout, stderr, _ = run(cmd, 600)
    res = hdr("Nuclei – Vulnerability Scan", t)
    res += f"📊 Severity: {severity}\n"
    if os.path.exists(out):
        data = open(out).read().strip()
        if data:
            findings = data.splitlines()
            res += f"\n🚨 {len(findings)} vulnerabilities found\n\n"
            for line in findings:
                try:
                    j = json.loads(line)
                    s = j.get("info", {}).get("severity", "?").upper()
                    name = j.get("info", {}).get("name", "N/A")
                    url = j.get("matched-at", "N/A")
                    tid = j.get("template-id", "")
                    icon = {"CRITICAL": "🔴", "HIGH": "🟠", "MEDIUM": "🟡", "LOW": "🟢"}.get(s, "⚪")
                    res += f"{icon} [{s}] {name}\n   Template: {tid}\n   URL: {url}\n\n"
                except:
                    res += line + "\n"
        else:
            res += "\n✅ No vulnerabilities found at this severity."
    else:
        res += f"\n{stdout}\n{stderr}"
    return res


def do_dalfox(target):
    t = clean(target)
    if not t: return "❌ Invalid domain format. Use: example.com"
    out = f"{RESULTS_DIR}/dalfox_{t}.txt"
    if os.path.exists(out): os.remove(out)
    gau_f = f"{RESULTS_DIR}/gau_{t}.txt"
    if os.path.exists(gau_f):
        urls = [l.strip() for l in open(gau_f) if "?" in l and "=" in l]
        if urls:
            inp = f"{RESULTS_DIR}/dalfox_in_{t}.txt"
            open(inp, "w").write("\n".join(urls[:30]))
            cmd = ["dalfox", "file", inp, "-o", out, "--silence", "--worker", "5", "--timeout", "10"]
        else:
            cmd = ["dalfox", "url", f"https://{t}", "-o", out, "--silence", "--worker", "5"]
    else:
        cmd = ["dalfox", "url", f"https://{t}", "-o", out, "--silence", "--worker", "5"]
    stdout, stderr, _ = run(cmd, 300)
    res = hdr("DalFox – XSS Scanner", t)
    if os.path.exists(out):
        data = open(out).read().strip()
        if data:
            res += f"\n🚨 {len(data.splitlines())} potential XSS found\n\n{data}"
        else:
            res += "\n✅ No XSS found."
    else:
        res += f"\n{stdout}\n{stderr}"
    return res


def do_ffuf(target):
    t = clean(target)
    if not t: return "❌ Invalid domain format. Use: example.com"
    out = f"{RESULTS_DIR}/ffuf_{t}.json"
    if os.path.exists(out): os.remove(out)
    cmd = ["ffuf", "-u", f"https://{t}/FUZZ", "-w", WORDLIST_PATH,
           "-mc", "200,201,301,302,403", "-t", "20", "-timeout", "10",
           "-o", out, "-of", "json", "-s"]
    stdout, stderr, _ = run(cmd, 180)
    res = hdr("FFUF – Directory Fuzzer", t)
    if os.path.exists(out):
        try:
            items = json.load(open(out)).get("results", [])
            res += f"\n✅ Found {len(items)} paths\n\n"
            for r in items:
                s = r.get("status", "?")
                u = r.get("url", "?")
                icon = {200: "🟢", 301: "🔵", 302: "🔵", 403: "🟡"}.get(s, "⚪")
                res += f"{icon} [{s}] {u}\n"
        except:
            res += stdout
    else:
        res += f"\n{stdout}\n{stderr}"
    return res


def do_full(target, severity):
    t = clean(target)
    if not t: return ("❌ Invalid domain format. Use: example.com",) * 6
    return (do_subfinder(t), do_httpx(t), do_gau(t),
            do_nuclei(t, severity), do_dalfox(t), do_ffuf(t))


def check_tools():
    """Check which tools are installed."""
    lines = []
    for name in TOOLS:
        path = os.path.join(BIN_DIR, name)
        if os.path.exists(path):
            lines.append(f"✅ {name} – installed")
        else:
            lines.append(f"❌ {name} – NOT installed")
    return "\n".join(lines)


# ═══════════════════════════════════════════════════
# INSTALL TOOLS ON STARTUP
# ═══════════════════════════════════════════════════
TOOL_STATUS = setup_tools()


# ═══════════════════════════════════════════════════
# GRADIO UI
# ═══════════════════════════════════════════════════

CSS = """
.gradio-container { max-width: 1200px !important; }
.hdr { text-align:center; padding:1.5rem; background:linear-gradient(135deg,#0a0e17,#1a1a2e,#16213e);
       border-radius:12px; margin-bottom:1rem; border-bottom:2px solid #06d6a0; }
.hdr h1 { font-size:2rem; background:linear-gradient(90deg,#06d6a0,#00b4d8);
           -webkit-background-clip:text; -webkit-text-fill-color:transparent; margin:0; }
.hdr p { color:#94a3b8; margin-top:0.3rem; font-size:0.9rem; }
.warn { background:rgba(239,68,68,0.08); border:1px solid #ef4444; border-radius:8px;
        padding:0.8rem; text-align:center; color:#f59e0b; font-weight:600; margin-bottom:1rem; }
footer { display:none !important; }
"""

with gr.Blocks(title="Security Scanner", css=CSS) as app:

    gr.HTML("""
    <div class="hdr"><h1>🛡️ AI SECURITY SCANNER</h1>
    <p>subfinder · httpx · gau · nuclei · dalfox · ffuf</p></div>
    <div class="warn">⚠️ For authorized penetration testing ONLY. Scanning without permission is illegal.</div>
    """)

    with gr.Tabs():

        # ── Full Scan ────────────────────────────────────
        with gr.Tab("🚀 Full Scan"):
            gr.Markdown("### Run all tools sequentially on the target")
            with gr.Row():
                t0 = gr.Textbox(label="🎯 Target Domain", placeholder="example.com", scale=3)
                s0 = gr.Dropdown(
                    ["critical", "high", "medium", "low", "critical,high", "critical,high,medium"],
                    value="critical,high,medium", label="Nuclei Severity", scale=1)
                b0 = gr.Button("🚀 Start Full Scan", variant="primary", scale=1)
            with gr.Accordion("🔍 Subfinder", open=False):
                o1 = gr.Textbox(label="Subdomains", lines=8)
            with gr.Accordion("🌐 HTTPX", open=False):
                o2 = gr.Textbox(label="Live Hosts", lines=8)
            with gr.Accordion("📦 GAU", open=False):
                o3 = gr.Textbox(label="URLs", lines=8)
            with gr.Accordion("☢️ Nuclei", open=True):
                o4 = gr.Textbox(label="Vulnerabilities", lines=8)
            with gr.Accordion("⚡ DalFox", open=True):
                o5 = gr.Textbox(label="XSS Results", lines=8)
            with gr.Accordion("🔓 FFUF", open=True):
                o6 = gr.Textbox(label="Paths Found", lines=8)
            b0.click(do_full, [t0, s0], [o1, o2, o3, o4, o5, o6])

        # ── Individual Tools ─────────────────────────────
        with gr.Tab("🔍 Subfinder"):
            with gr.Row():
                t1 = gr.Textbox(label="Domain", placeholder="example.com", scale=3)
                b1 = gr.Button("🔍 Discover Subdomains", variant="primary", scale=1)
            r1 = gr.Textbox(label="Results", lines=15)
            b1.click(do_subfinder, t1, r1)

        with gr.Tab("🌐 HTTPX"):
            with gr.Row():
                t2 = gr.Textbox(label="Domain", placeholder="example.com", scale=3)
                b2 = gr.Button("🌐 Probe", variant="primary", scale=1)
            r2 = gr.Textbox(label="Results", lines=15)
            b2.click(do_httpx, t2, r2)

        with gr.Tab("📦 GAU"):
            with gr.Row():
                t3 = gr.Textbox(label="Domain", placeholder="example.com", scale=3)
                b3 = gr.Button("📦 Fetch URLs", variant="primary", scale=1)
            r3 = gr.Textbox(label="Results", lines=15)
            b3.click(do_gau, t3, r3)

        with gr.Tab("☢️ Nuclei"):
            with gr.Row():
                t4 = gr.Textbox(label="Domain", placeholder="example.com", scale=2)
                s4 = gr.Dropdown(
                    ["critical", "high", "medium", "low", "critical,high", "critical,high,medium"],
                    value="critical,high", label="Severity", scale=1)
                b4 = gr.Button("☢️ Scan", variant="primary", scale=1)
            r4 = gr.Textbox(label="Results", lines=15)
            b4.click(do_nuclei, [t4, s4], r4)

        with gr.Tab("⚡ DalFox"):
            with gr.Row():
                t5 = gr.Textbox(label="Domain", placeholder="example.com", scale=3)
                b5 = gr.Button("⚡ Scan XSS", variant="primary", scale=1)
            r5 = gr.Textbox(label="Results", lines=15)
            b5.click(do_dalfox, t5, r5)

        with gr.Tab("🔓 FFUF"):
            with gr.Row():
                t6 = gr.Textbox(label="Domain", placeholder="example.com", scale=3)
                b6 = gr.Button("🔓 Fuzz Dirs", variant="primary", scale=1)
            r6 = gr.Textbox(label="Results", lines=15)
            b6.click(do_ffuf, t6, r6)

        with gr.Tab("⚙️ Tool Status"):
            gr.Markdown("### Check installed tools")
            ts_btn = gr.Button("🔄 Check Tools", variant="secondary")
            ts_out = gr.Textbox(label="Status", lines=8)
            ts_btn.click(check_tools, [], ts_out)

    gr.Markdown("<center style='color:#64748b;font-size:0.85rem;'>🛡️ AI Security Scanner v3.0 | Authorized Use Only</center>")


if __name__ == "__main__":
    app.launch(server_name="0.0.0.0", server_port=7860)