# app.py - HF Spaces Compatible with demo.launch() pattern
# Gradio 6.0: theme in launch(), not Blocks()
# Secure token handling: UI input -> server memory -> API calls (never in templates)
# Auto-detects C2 server URL, randomizes deployment names, no emojis in UI
# FIXED: Docker SDK support, create_repo parameter fix, queue debugging

import os
import time
import threading
import json
import base64
import secrets
import string
import requests
from datetime import datetime
from urllib.parse import urlparse
import gradio as gr

# Try to import huggingface_hub (optional, for HF API calls)
try:
    from huggingface_hub import HfApi, create_repo
    HF_HUB_AVAILABLE = True
except ImportError:
    HF_HUB_AVAILABLE = False

# ==================================================
# STATE
# ==================================================
class AppState:
    def __init__(self):
        self.commands = []
        self.reports = []
        self.lock = threading.Lock()
        # Runtime token storage (NEVER written to templates)
        self.user_tokens = {
            "hf_token": None,
            "github_token": None
        }
        # Auto-detected C2 server URL
        self.server_url = None
        self._detect_server_url()
    
    def _detect_server_url(self):
        """Auto-detect the C2 server URL from environment or runtime context"""
        # Priority 1: Explicit environment variable
        if os.getenv("C2_SERVER_URL"):
            self.server_url = os.getenv("C2_SERVER_URL").rstrip("/")
            return
        
        # Priority 2: HF Spaces environment (normalize to lowercase for consistency)
        space_host = os.getenv("SPACE_HOST")
        if space_host:
            # Normalize to lowercase to match HF's URL handling
            self.server_url = f"https://{space_host.lower()}".rstrip("/")
            return
        
        # Priority 3: Render/Fly.io environment
        render_url = os.getenv("RENDER_EXTERNAL_URL")
        if render_url:
            self.server_url = render_url.rstrip("/")
            return
        
        # Priority 4: Default localhost for local development
        port = os.getenv("PORT", "7860")
        self.server_url = f"http://localhost:{port}".rstrip("/")

STATE = AppState()

# ==================================================
# CONFIG
# ==================================================
VALID_TYPES = {
    "fetch_http_flood", "fetch_cache_bypass", "fetch_custom_headers",
    "socket_tcp_probe", "socket_udp_probe",
}
ALLOWED_TARGETS = {"127.0.0.1", "localhost", "0.0.0.0"}

def validate_target(ip: str) -> bool:
    return ip.startswith("127.") or ip in ALLOWED_TARGETS

def generate_random_name(prefix: str = "c2-client", length: int = 8) -> str:
    """Generate a random, URL-safe name for deployments"""
    chars = string.ascii_lowercase + string.digits
    random_suffix = ''.join(secrets.choice(chars) for _ in range(length))
    return f"{prefix}-{random_suffix}"

# ==================================================
# TOKEN MANAGEMENT FUNCTIONS
# ==================================================
def save_user_tokens(hf_token: str, github_token: str) -> str:
    """Store tokens in server memory only - never in templates"""
    with STATE.lock:
        # Validate and store HF token
        if hf_token and hf_token.strip():
            if hf_token.strip().startswith("hf_"):
                STATE.user_tokens["hf_token"] = hf_token.strip()
            else:
                return "Invalid HF token format (must start with 'hf_')"
        
        # Validate and store GitHub token
        if github_token and github_token.strip():
            if github_token.strip().startswith("ghp_") or github_token.strip().startswith("gho_"):
                STATE.user_tokens["github_token"] = github_token.strip()
            else:
                return "Invalid GitHub token format (must start with 'ghp_' or 'gho_')"
    
    return "Tokens saved securely (server memory only)"

def test_token_connections() -> str:
    """Test both tokens without exposing them"""
    results = []
    
    # Test HF
    hf_token = STATE.user_tokens.get("hf_token")
    if hf_token:
        try:
            if HF_HUB_AVAILABLE:
                api = HfApi(token=hf_token)
                user = api.whoami()
                results.append(f"HF: {user.get('name', user.get('id', 'connected'))}")
            else:
                headers = {"Authorization": f"Bearer {hf_token}"}
                res = requests.get("https://huggingface.co/api/whoami-v2", headers=headers, timeout=10)
                if res.status_code == 200:
                    data = res.json()
                    results.append(f"HF: {data.get('name', 'connected')}")
                else:
                    results.append(f"HF: HTTP {res.status_code}")
        except Exception as e:
            results.append(f"HF: {str(e)[:60]}")
    else:
        results.append("HF: Not set")
    
    # Test GitHub
    gh_token = STATE.user_tokens.get("github_token")
    if gh_token:
        try:
            headers = {"Authorization": f"token {gh_token}", "Accept": "application/vnd.github.v3+json"}
            res = requests.get("https://api.github.com/user", headers=headers, timeout=10)
            res.raise_for_status()
            login = res.json().get("login", "unknown")
            results.append(f"GitHub: {login}")
        except Exception as e:
            results.append(f"GitHub: {str(e)[:60]}")
    else:
        results.append("GitHub: Not set")
    
    return " | ".join(results)

def clear_tokens() -> str:
    """Wipe tokens from memory"""
    with STATE.lock:
        STATE.user_tokens["hf_token"] = None
        STATE.user_tokens["github_token"] = None
    return "Tokens cleared from memory"

# ==================================================
# DEPLOYMENT FUNCTIONS (Server uses tokens via API)
# ==================================================
def deploy_to_hf_space(space_sdk: str = "docker") -> str:
    """Create/update HF Space using stored token via API - NO template injection"""
    hf_token = STATE.user_tokens.get("hf_token")
    if not hf_token:
        return "HF token not set. Please enter in Tokens tab first."
    
    try:
        api = HfApi(token=hf_token) if HF_HUB_AVAILABLE else None
        
        # Get current user to build repo_id
        if api:
            user = api.whoami()
            username = user.get('name', user.get('id'))
        else:
            headers = {"Authorization": f"Bearer {hf_token}"}
            res = requests.get("https://huggingface.co/api/whoami-v2", headers=headers, timeout=10)
            res.raise_for_status()
            username = res.json().get('name', 'unknown-user')
        
        # Generate random space name
        space_name = generate_random_name("c2-client")
        repo_id = f"{username}/{space_name}"
        
        # Create repo with correct parameter name: space_sdk (not sdk)
        if api:
            create_repo(
                repo_id=repo_id,
                repo_type="space",
                space_sdk=space_sdk,  # FIXED: parameter is space_sdk, not sdk
                token=hf_token,
                exist_ok=True,
                private=False
            )
        else:
            # Fallback via requests API
            headers = {"Authorization": f"Bearer {hf_token}"}
            requests.post(
                "https://huggingface.co/api/repos/create",
                headers=headers,
                json={"name": space_name, "type": "space", "space_sdk": space_sdk},
                timeout=10
            )
        
        # Upload clean template files (NO tokens embedded)
        template_files = _get_clean_template_files(space_sdk)
        for filename, content in template_files.items():
            if api:
                import tempfile
                with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix=filename) as tmp:
                    tmp.write(content)
                    tmp_path = tmp.name
                try:
                    api.upload_file(
                        path_or_fileobj=tmp_path,
                        path_in_repo=filename,
                        repo_id=repo_id,
                        repo_type="space",
                        token=hf_token
                    )
                finally:
                    os.unlink(tmp_path)
            else:
                encoded = base64.b64encode(content.encode()).decode()
                headers = {"Authorization": f"Bearer {hf_token}"}
                # Use the correct endpoint for spaces
                requests.put(
                    f"https://huggingface.co/api/models/{repo_id}/resolve/main/{filename}",
                    headers=headers,
                    json={"content": encoded, "commit_message": f"Deploy {filename}"},
                    timeout=30
                )
        
        return f"Deployed to: https://huggingface.co/spaces/{repo_id}\nClients will ping: {STATE.server_url}"
    
    except Exception as e:
        return f"Deployment failed: {str(e)}"

def deploy_to_github_pages(branch: str = "gh-pages") -> str:
    """Push to GitHub Pages using stored token via API - NO template injection"""
    github_token = STATE.user_tokens.get("github_token")
    if not github_token:
        return "GitHub token not set. Please enter in Tokens tab first."
    
    try:
        headers = {
            "Authorization": f"token {github_token}",
            "Accept": "application/vnd.github.v3+json"
        }
        
        # Get user login
        user_res = requests.get("https://api.github.com/user", headers=headers, timeout=10)
        user_res.raise_for_status()
        username = user_res.json()["login"]
        
        # Generate random repo name
        repo_name = generate_random_name("c2-client")
        
        # Create repo if needed (ignore 422 if exists)
        requests.post(
            "https://api.github.com/user/repos",
            headers=headers,
            json={"name": repo_name, "private": False, "auto_init": True},
            timeout=10
        )
        
        # Ensure gh-pages branch exists
        try:
            requests.get(
                f"https://api.github.com/repos/{username}/{repo_name}/branches/{branch}",
                headers=headers, timeout=5
            ).raise_for_status()
        except:
            # Create gh-pages branch from main
            main_ref = requests.get(
                f"https://api.github.com/repos/{username}/{repo_name}/git/refs/heads/main",
                headers=headers, timeout=5
            ).json()
            requests.post(
                f"https://api.github.com/repos/{username}/{repo_name}/git/refs",
                headers=headers,
                json={"ref": f"refs/heads/{branch}", "sha": main_ref["object"]["sha"]},
                timeout=10
            )
        
        # Upload clean template files (NO tokens embedded)
        template_files = _get_clean_template_files("docker")  # Same templates for GH
        for filename, content in template_files.items():
            encoded = base64.b64encode(content.encode()).decode()
            
            # Check if file exists to decide PUT vs POST
            try:
                existing = requests.get(
                    f"https://api.github.com/repos/{username}/{repo_name}/contents/{filename}?ref={branch}",
                    headers=headers, timeout=5
                )
                sha = existing.json().get("sha") if existing.status_code == 200 else None
            except:
                sha = None
            
            payload = {
                "message": f"Deploy {filename} via C2",
                "content": encoded,
                "branch": branch
            }
            if sha:
                payload["sha"] = sha
            
            requests.put(
                f"https://api.github.com/repos/{username}/{repo_name}/contents/{filename}",
                headers=headers,
                json=payload,
                timeout=10
            )
        
        return f"Deployed to: https://{username}.github.io/{repo_name}/\nClients will ping: {STATE.server_url}"
    
    except Exception as e:
        return f"GitHub deploy failed: {str(e)}"

def _get_clean_template_files(space_sdk: str = "docker") -> dict:
    """Return template files with SERVER_URL injected, but NO tokens"""
    files = {}
    
    # Clean client.py - inject SERVER_URL only
    client_template = f'''# hf_space_client/client.py - ETHICAL DEMO CLIENT
import requests, socket, time, json, os, random, string
from datetime import datetime

# Configuration
SERVER = os.getenv("SERVER_URL", "{STATE.server_url}")
DEFAULT_ALLOWED_TARGETS = {{"127.0.0.1", "localhost"}}
ALLOWED_TARGETS_HF = set(DEFAULT_ALLOWED_TARGETS)

def validate_target_ip(ip):
    if ip.startswith("127."): return True
    return ip in ALLOWED_TARGETS_HF

def log_client_message(msg):
    print(f"[{{datetime.now().isoformat()}}] HF-Client: {{msg}}")

def exec_fetch_http_flood(ip, port, sim_headers):
    url = f"http://{{ip}}:{{port}}/"
    try:
        t0 = time.time()
        res = requests.get(url, headers=sim_headers, timeout=3)
        return {{"status": f"status_{{res.status_code}}", "latency_ms": round((time.time()-t0)*1000, 2), "bytes": len(res.content), "protocol": "http"}}
    except Exception as e:
        return {{"status": "error", "detail": str(e), "protocol": "http"}}

def exec_fetch_cache_bypass(ip, port, sim_headers):
    key = ''.join(random.choices(string.ascii_lowercase + string.digits, k=12))
    url = f"http://{{ip}}:{{port}}/?_cache_bust={{key}}&_={{int(time.time()*1000)}}"
    try:
        t0 = time.time()
        res = requests.get(url, headers=sim_headers, timeout=3)
        return {{"status": "success", "latency_ms": round((time.time()-t0)*1000, 2), "bytes": len(res.content), "cache_key": key, "protocol": "http"}}
    except Exception as e:
        return {{"status": "error", "detail": str(e), "protocol": "http"}}

def exec_fetch_custom_headers(ip, port, sim_headers):
    url = f"http://{{ip}}:{{port}}/"
    custom_headers = {{ **sim_headers, "X-Custom-Test": "sim_payload", "Accept": "application/json" }}
    try:
        t0 = time.time()
        res = requests.get(url, headers=custom_headers, timeout=3)
        return {{"status": f"status_{{res.status_code}}", "latency_ms": round((time.time()-t0)*1000, 2), "bytes": len(res.content), "protocol": "http"}}
    except Exception as e:
        return {{"status": "error", "detail": str(e), "protocol": "http"}}

def exec_socket_tcp_probe(ip, port):
    try:
        t0 = time.time()
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            s.settimeout(2)
            s.connect((ip, port))
            s.sendall(b"SIM_PKT_TCP_01_HEARTBEAT")
        return {{"status": "tcp_connected", "latency_ms": round((time.time()-t0)*1000, 2), "bytes_sent": 24, "protocol": "tcp"}}
    except Exception as e:
        return {{"status": "tcp_failed", "detail": str(e), "protocol": "tcp"}}

def exec_socket_udp_probe(ip, port):
    try:
        t0 = time.time()
        with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
            s.settimeout(2)
            sent = s.sendto(b"SIM_PKT_UDP_01_HEARTBEAT", (ip, port))
        return {{"status": "udp_sent", "latency_ms": round((time.time()-t0)*1000, 2), "bytes_sent": sent, "protocol": "udp"}}
    except Exception as e:
        return {{"status": "udp_failed", "detail": str(e), "protocol": "udp"}}

EXECUTORS = {{
    "fetch_http_flood": exec_fetch_http_flood,
    "fetch_cache_bypass": exec_fetch_cache_bypass,
    "fetch_custom_headers": exec_fetch_custom_headers,
    "socket_tcp_probe": exec_socket_tcp_probe,
    "socket_udp_probe": exec_socket_udp_probe
}}

def run():
    log_client_message("="*50)
    log_client_message("HUGGING FACE SPACE CLIENT (ETHICAL DEMO)")
    log_client_message("="*50)
    log_client_message(f"Polling Control Server: {{SERVER}}")
    log_client_message("Educational simulation only - no malicious activity")
    log_client_message("="*50)

    while True:
        try:
            cmd_response = requests.get(f"{{SERVER}}/api/poll_command", timeout=5).json()
            if cmd_response.get("status") == "idle":
                time.sleep(2); continue

            cmd = cmd_response
            cmd_type, ip, port = cmd["type"], cmd["ip"], int(cmd["port"])
            sim_headers = cmd.get("simulation_headers", {{"X-Simulation": "true", "User-Agent": "EduSimClient/1.0"}})

            log_client_message(f"Command (ID={{cmd['id']}}): {{cmd_type}} -> {{ip}}:{{port}}")

            if not validate_target_ip(ip):
                report = {{"client": "hf_space", "command_id": cmd["id"], "status": "blocked_by_client_allowlist", "target": f"{{ip}}:{{port}}", "timestamp": time.time()}}
            elif cmd_type in EXECUTORS:
                if cmd_type.startswith("fetch_"):
                    res = EXECUTORS[cmd_type](ip, port, sim_headers)
                else:
                    res = EXECUTORS[cmd_type](ip, port)
                report = {{"client": "hf_space", "command_id": cmd["id"], "type": cmd_type, "target": f"{{ip}}:{{port}}", "requests_sent": 1, "result": res, "timestamp": time.time()}}
            else:
                report = {{"client": "hf_space", "command_id": cmd["id"], "status": "unsupported_command_type", "type": cmd_type, "target": f"{{ip}}:{{port}}", "timestamp": time.time()}}

            requests.post(f"{{SERVER}}/api/report", json=report, timeout=5)
        except Exception as e:
            log_client_message(f"Error: {{e}}")
        time.sleep(2)

if __name__ == "__main__":
    run()
'''
    files["client.py"] = client_template
    
    # Dockerfile template for HF Spaces (Docker SDK)
    if space_sdk == "docker":
        dockerfile_template = f'''FROM python:3.11-slim

WORKDIR /app

COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

COPY client.py .

# Environment variable for server URL (can be overridden at runtime)
ENV SERVER_URL={STATE.server_url}

CMD ["python", "client.py"]
'''
        files["Dockerfile"] = dockerfile_template
    
    # Clean index.html for GitHub Pages
    gh_pages_template = f'''<!DOCTYPE html>
<html lang="en">
<head>
  <meta charset="UTF-8">
  <meta name="viewport" content="width=device-width, initial-scale=1.0">
  <title>GitHub Pages Fetch Sim Client (Ethical Demo)</title>
  <style>
    body {{ font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif; padding: 20px; background: #f8f9fa; color: #333; }}
    pre {{ background: #fff; padding: 15px; border: 1px solid #ddd; border-radius: 8px; max-height: 60vh; overflow-y: auto; white-space: pre-wrap; font-size: 0.9em; }}
    h2 {{ color: #007bff; border-bottom: 2px solid #eee; padding-bottom: 10px; }}
    .warning {{ background: #ffe0e0; border: 1px solid #dc3545; color: #dc3545; padding: 15px; border-radius: 8px; margin-bottom: 25px; text-align: center; }}
    .config-info {{ background: #e2f0fe; border: 1px solid #007bff; color: #007bff; padding: 10px; border-radius: 5px; margin-top: 15px; font-size: 0.9em; }}
  </style>
</head>
<body>
  <div class="warning">
    <h3>ETHICAL DEMONSTRATION ONLY</h3>
    <p>This client demonstrates controlled network testing. NO malicious activity is performed. All requests target pre-approved IPs (loopback only).</p>
  </div>

  <h2>GitHub Pages Fetch Simulation Client</h2>
  <div class="config-info">
    <p><strong>Control Server:</strong> <code id="server-url-display">Loading...</code></p>
    <p><strong>Allowed Targets:</strong> <code>127.x.x.x (loopback only)</code></p>
  </div>
  <pre id="log">Initializing...</pre>

  <script>
    const SERVER = "{STATE.server_url}";
    const SIM_ALLOWED_TARGETS = ["localhost"];
    
    document.getElementById("server-url-display").textContent = SERVER;

    function log(msg) {{
      const el = document.getElementById("log");
      el.textContent += `\\n${{new Date().toLocaleTimeString()}} - ${{msg}}`;
      el.scrollTop = el.scrollHeight;
    }}

    function validateTargetIp(ip) {{
      return ip.startsWith("127.") || SIM_ALLOWED_TARGETS.includes(ip);
    }}

    async function executeFetch(type, ip, port, simHeaders) {{
      if (!validateTargetIp(ip)) {{
        log(`Blocked: IP '${{ip}}' not allowlisted`);
        return {{ status: "blocked_by_client_allowlist", target_ip: ip }};
      }}
      const base = `http://${{ip}}:${{port}}/`;
      const opts = {{ method: "GET", headers: simHeaders, mode: "cors", cache: "no-store" }};
      const t0 = performance.now();
      try {{
        let response;
        if (type === "fetch_http_flood") {{
          log(`HTTP GET to ${{base}}`);
          response = await fetch(base, opts);
          const txt = await response.text();
          return {{ status: `status_${{response.status}}`, latency_ms: performance.now() - t0, bytes: new Blob([txt]).size, protocol: "http" }};
        }} else if (type === "fetch_cache_bypass") {{
          const k = Math.random().toString(36).substring(2, 10);
          const url = `${{base}}?_cache_bust=${{k}}&_=${{Date.now()}}`;
          log(`Cache-bypass GET to ${{url}}`);
          response = await fetch(url, opts);
          const txt = await response.text();
          return {{ status: "success", latency_ms: performance.now() - t0, bytes: new Blob([txt]).size, cache_key: k, protocol: "http" }};
        }} else if (type === "fetch_custom_headers") {{
          opts.headers = {{ ...simHeaders, "X-Custom-Test": "sim_payload", "Accept": "application/json" }};
          log(`Custom headers GET to ${{base}}`);
          response = await fetch(base, opts);
          return {{ status: `status_${{response.status}}`, latency_ms: performance.now() - t0, protocol: "http" }};
        }}
        return {{ status: "unsupported_fetch_type", type: type }};
      }} catch (e) {{
        log(`Fetch error: ${{e.message}}`);
        return {{ status: "fetch_error", detail: e.message }};
      }}
    }}

    async function poll() {{
      try {{
        const cmdResponse = await fetch(`${{SERVER}}/api/poll_command`).then(r => r.json());
        if (cmdResponse.status === "idle") {{ setTimeout(poll, 2000); return; }}
        const cmd = cmdResponse;
        const {{ id, type, ip, port, simulation_headers }} = cmd;
        log(`Command (ID=${{id}}): ${{type}} -> ${{ip}}:${{port}}`);
        if (!type.startsWith("fetch_")) {{
          const report = {{ client: "gh_pages", command_id: id, type, target: `${{ip}}:${{port}}`, requests_sent: 0, result: {{ status: "unsupported_by_browser" }}, timestamp: Date.now()/1000 }};
          await fetch(`${{SERVER}}/api/report`, {{ method: "POST", headers: {{"Content-Type":"application/json"}}, body: JSON.stringify(report) }});
          setTimeout(poll, 2000); return;
        }}
        const result = await executeFetch(type, ip, port, simulation_headers);
        const report = {{ client: "gh_pages", command_id: id, type, target: `${{ip}}:${{port}}`, requests_sent: 1, result, timestamp: Date.now()/1000 }};
        await fetch(`${{SERVER}}/api/report`, {{ method: "POST", headers: {{"Content-Type":"application/json"}}, body: JSON.stringify(report) }});
        log(`Report: Status=${{result.status}}`);
      }} catch (e) {{ log(`Error: ${{e.message}}`); }}
      setTimeout(poll, 2000);
    }}
    log("Client Ready. Polling...");
    poll();
  </script>
</body>
</html>
'''
    files["index.html"] = gh_pages_template
    
    # Requirements file for HF Space
    files["requirements.txt"] = """requests>=2.31.0
"""
    
    # README for deployed space
    files["README.md"] = f"""# C2 Client Simulation (Ethical Demo)

This Space runs a client that polls the C2 server at:
**`{STATE.server_url}`**

## Security Notice
- This is an educational demonstration only
- All network requests are limited to loopback addresses (127.x.x.x)
- No actual malicious activity is performed
- Tokens are managed by the C2 server, not embedded in this Space

## Configuration
The client polls `{STATE.server_url}/api/poll_command` for commands and reports results to `{STATE.server_url}/api/report`.

## Runtime
- SDK: {space_sdk}
- Entry point: client.py
"""
    
    return files

# ==================================================
# UI FUNCTIONS
# ==================================================
def get_dashboard_html():
    with STATE.lock:
        recent = STATE.reports[-10:]
        queued = len(STATE.commands)
    
    rows = []
    for r in reversed(recent):
        ts = datetime.fromtimestamp(r.get("ts", time.time())).strftime("%H:%M:%S")
        client = str(r.get('client', '?'))[:20]
        cmd_type = str(r.get('type', '?'))[:25]
        status = str(r.get('status', 'ok'))[:15]
        color = "#22c55e" if status == "ok" else "#ef4444"
        rows.append(f"""<tr style="border-bottom:1px solid #30363d">
            <td style="padding:8px;font-family:monospace">{client}</td>
            <td style="padding:8px">{cmd_type}</td>
            <td style="padding:8px;color:{color}">{status}</td>
            <td style="padding:8px;color:#8b949e">{ts}</td>
        </tr>""")
    
    rows_html = "".join(rows) if rows else "<tr><td colspan='4' style='padding:20px;text-align:center;color:#8b949e'>No reports yet</td></tr>"
    
    return f"""
    <div style="padding:20px;border-radius:12px;background:#0d1117;color:white;border:1px solid #30363d">
        <h2 style="margin:0 0 15px">System Dashboard</h2>
        <p><b>Queued:</b> {queued} | <b>Reports:</b> {len(recent)}</p>
    </div>
    <div style="margin-top:20px;overflow-x:auto">
        <table style="width:100%;color:#c9d1d9;border-collapse:collapse">
            <thead style="background:#161b22"><tr>
                <th style="padding:12px;text-align:left">Client</th>
                <th style="padding:12px;text-align:left">Type</th>
                <th style="padding:12px;text-align:left">Status</th>
                <th style="padding:12px;text-align:left">Time</th>
            </tr></thead>
            <tbody>{rows_html}</tbody>
        </table>
    </div>"""

def push_command(cmd_type: str, ip: str, port: float) -> str:
    try:
        port_int = int(port)
        if not (1 <= port_int <= 65535): return "Port must be 1-65535"
        if not validate_target(ip.strip()): return f"'{ip}' not allowed"
        if cmd_type not in VALID_TYPES: return "Invalid command type"
        with STATE.lock:
            STATE.commands.append({
                "type": cmd_type, "ip": ip.strip(), "port": port_int,
                "id": int(time.time()*1000), "source": "ui",
                "simulation_headers": {"X-Simulation": "true", "User-Agent": "EduSimClient/1.0"}
            })
        return f"Queued: {cmd_type} -> {ip}:{port_int}"
    except Exception as e: return f"Error: {e}"

# API endpoints for deployed clients to poll/report
def api_poll_command():
    """Endpoint for clients to poll for commands"""
    with STATE.lock:
        if STATE.commands:
            return STATE.commands.pop(0)
        return {"status": "idle"}

def api_report(report: dict):
    """Endpoint for clients to submit reports"""
    with STATE.lock:
        report["ts"] = time.time()
        STATE.reports.append(report)
        # Keep only last 100 reports
        if len(STATE.reports) > 100:
            STATE.reports = STATE.reports[-100:]
    return {"status": "received"}

# ==================================================
# GRADIO UI - Gradio 6.0 Compatible (No Emojis)
# ==================================================
css = """
body { background:#0b0f19; color:#c9d1d9; font-family:system-ui,-apple-system,sans-serif }
.gradio-container { max-width:1000px !important; margin:auto !important; padding:20px !important }
button.primary { background:linear-gradient(135deg,#238636,#2ea043) !important; border:none !important }
input,select,textarea { background:#161b22 !important; border:1px solid #30363d !important; color:#c9d1d9 !important }
.markdown-details summary { cursor: pointer; color: #58a6ff; margin: 10px 0; }
.markdown-details[open] summary { margin-bottom: 15px; }
"""

with gr.Blocks() as demo:
    gr.HTML(f"<style>{css}</style><title>Control Panel</title>")
    gr.Markdown("# Control Panel\n*Network Testing Interface - HF Spaces Compatible*")
    
    with gr.Tabs():
        
        with gr.Tab("Dashboard"):
            dashboard = gr.HTML(value=get_dashboard_html())
            refresh = gr.Button("Refresh", variant="secondary", size="sm")
            refresh.click(fn=get_dashboard_html, outputs=dashboard)
        
        with gr.Tab("Commands"):
            gr.Markdown("### Send Command to Connected Clients")
            with gr.Row():
                with gr.Column(scale=2):
                    cmd_type = gr.Dropdown(choices=sorted(VALID_TYPES), value="fetch_http_flood", label="Command Type")
                    target_ip = gr.Textbox(value="127.0.0.1", label="Target IP", placeholder="127.0.0.1")
                    target_port = gr.Number(value=80, label="Port", minimum=1, maximum=65535, step=1, precision=0)
                with gr.Column(scale=1):
                    gr.Markdown("#### Allowed Targets\n- 127.x.x.x\n- localhost\n- External IPs blocked")
            send = gr.Button("Send Command", variant="primary")
            result = gr.Textbox(label="Result", interactive=False, lines=2)
            send.click(fn=push_command, inputs=[cmd_type, target_ip, target_port], outputs=result)
        
        with gr.Tab("Tokens"):
            gr.Markdown("""
            ### Enter Your Access Tokens
            - Tokens stored only in server memory for this session
            - Tokens NEVER embedded in deployed templates
            - Used by C2 server for API authentication only
            """)
            
            with gr.Row():
                with gr.Column():
                    hf_token_input = gr.Textbox(
                        label="Hugging Face Token", 
                        type="password",
                        placeholder="hf_xxx...",
                        info="Get at: https://huggingface.co/settings/tokens"
                    )
                with gr.Column():
                    gh_token_input = gr.Textbox(
                        label="GitHub Token", 
                        type="password",
                        placeholder="ghp_xxx...",
                        info="Get at: https://github.com/settings/tokens (scopes: repo, gist)"
                    )
            
            with gr.Row():
                save_btn = gr.Button("Save Tokens", variant="primary")
                test_btn = gr.Button("Test Connections", variant="secondary")
                clear_btn = gr.Button("Clear", variant="stop")
            
            token_status = gr.Textbox(label="Status", interactive=False, lines=2)
            
            save_btn.click(fn=save_user_tokens, inputs=[hf_token_input, gh_token_input], outputs=[token_status])
            test_btn.click(fn=test_token_connections, inputs=[], outputs=[token_status])
            clear_btn.click(fn=clear_tokens, inputs=[], outputs=[token_status])
            
            gr.Markdown("""
            <details class="markdown-details">
            <summary>How Token Security Works</summary>
            
            1. Enter tokens - stored in Python server memory only
            2. Deployments use tokens via API calls (huggingface_hub, GitHub API)
            3. Template files contain NO token placeholders - safe to push publicly
            4. Deployed clients ping YOUR C2 server; server authenticates using stored tokens
            5. Tokens lost on Space restart (re-enter if needed)
            
            </details>
            """, elem_classes=["markdown-details"])
        
        with gr.Tab("Deploy"):
            gr.Markdown("### Deploy Clean Templates Using Your Tokens")
            
            # Auto-detected server URL display
            gr.Markdown(f"**C2 Server URL**: `{STATE.server_url}`\n\n_Deployed clients will poll this URL for commands_")
            
            with gr.Tabs():
                with gr.Tab("Hugging Face Spaces"):
                    hf_sdk = gr.Dropdown(choices=["docker"], value="docker", label="SDK Type", interactive=False)
                    hf_deploy_btn = gr.Button("Deploy to HF Spaces", variant="primary")
                    hf_deploy_result = gr.Textbox(label="Result", interactive=False, lines=4)
                    hf_deploy_btn.click(
                        fn=deploy_to_hf_space, 
                        inputs=[hf_sdk], 
                        outputs=[hf_deploy_result]
                    )
                    gr.Markdown("*Space name will be auto-generated (e.g., c2-client-a8x9k2m1)*\n*SDK: docker (Dockerfile template included)*")
                
                with gr.Tab("GitHub Pages"):
                    gh_branch = gr.Textbox(label="Branch", value="gh-pages")
                    gh_deploy_btn = gr.Button("Deploy to GitHub Pages", variant="primary")
                    gh_deploy_result = gr.Textbox(label="Result", interactive=False, lines=4)
                    gh_deploy_btn.click(
                        fn=deploy_to_github_pages, 
                        inputs=[gh_branch], 
                        outputs=[gh_deploy_result]
                    )
                    gr.Markdown("*Repository name will be auto-generated (e.g., c2-client-a8x9k2m1)*")
            
            gr.Markdown("""
            > Important: 
            > - Deployed clients will poll your C2 server at the URL shown above
            > - Templates contain NO tokens - authentication happens server-side
            > - Ensure your C2 server is publicly accessible and implements /api/poll_command and /api/report endpoints
            > - If commands stay queued, check that the deployed client can reach your C2 server URL
            """)
        
        with gr.Tab("Info"):
            space_host = os.getenv("SPACE_HOST", "your-space.hf.space")
            gr.Markdown(f"""
            ### HF Spaces Info
            - Current URL: `https://{space_host}`
            - Gradio Version: 6.0+
            - Security: Server-side target validation (loopback only)
            - Token Storage: Runtime memory only (cleared on restart)
            
            ### API Endpoints for Clients
            - POST /api/poll_command - Clients poll for commands
            - POST /api/report - Clients submit execution reports
            
            ### Template Files Deployed
            - client.py - Python client for HF Spaces/GitHub Pages
            - Dockerfile - For HF Spaces with Docker SDK
            - index.html - JavaScript client for GitHub Pages  
            - requirements.txt - Python dependencies
            - README.md - Documentation
            
            ### Troubleshooting Queued Commands
            1. Verify deployed client can reach: {STATE.server_url}
            2. Check Space/repo logs for connection errors
            3. Ensure C2 server allows CORS if clients are on different origins
            4. Commands expire after 60 seconds if not picked up
            """)

# ==================================================
# ENTRY POINT - With launch() for HF Spaces + Local
# ==================================================
if __name__ == "__main__":
    # Mount API endpoints for client communication (use POST for both)
    demo.app.post("/api/poll_command")(api_poll_command)
    demo.app.post("/api/report")(api_report)
    
    # Configure queue for concurrent requests
    demo.queue(max_size=10)
    
    # Launch with Gradio 6.0 compatible settings
    demo.launch(
        server_name="0.0.0.0",
        server_port=int(os.getenv("PORT", 7860)),
        share=False,
        show_error=True,
        theme=gr.themes.Soft(),
        ssr_mode=False,
        debug=os.getenv("DEBUG", "false").lower() == "true"
    )