app.py

import os
import re
import json
import asyncio
import ipaddress
import urllib.parse
from typing import List, Dict, Any, Optional

import httpx
from bs4 import BeautifulSoup
from pydantic import BaseModel, Field, field_validator
from openai import AsyncOpenAI
import gradio as gr

# Selenium imports
from selenium import webdriver
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC

# =====================================================================
# 1. CONFIGURATION & CREDENTIAL MANAGEMENT
# =====================================================================

DEFAULT_API_KEY = os.environ.get("DEEPSEEK_API_KEY", "")
DEFAULT_BASE_URL = os.environ.get("DEEPSEEK_BASE_URL", "https://api.deepseek.com")
DEFAULT_MODEL = os.environ.get("DEEPSEEK_MODEL", "deepseek-chat")
DEFAULT_MAX_STEPS = int(os.environ.get("MAX_STEPS", "15"))

DEFAULT_SYSTEM_PROMPT = os.environ.get(
    "SYSTEM_PROMPT",
    """You are an automated System Administration and Infrastructure Diagnostics Agent.
You assist system administrators and network engineers in auditing network assets, inspecting configurations, executing terminal diagnostics, and mapping infrastructure surfaces.

Operational Guidelines:
1. Systematically investigate target infrastructure using the provided diagnostic, terminal, and browser automation tools.
2. For JavaScript-rendered pages or interactive web workflows, use `headless_browser_action` to navigate, click selectors, fill inputs, or extract dynamic DOM structures.
3. If tool outputs indicate further investigation is necessary, call subsequent tools autonomously.
4. Present final summaries structured with: Target Profile, Technical Diagnostics, Configuration Observations, and Hardening Recommendations.
5. Maintain a neutral, engineering-focused reporting format."""
)

# =====================================================================
# 2. STRICT VALIDATION & INPUT SANITIZATION LAYER
# =====================================================================

DOMAIN_REGEX = re.compile(
    r"^(?:[a-zA-Z0-9]"
    r"(?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?\.)+"
    r"[a-zA-Z]{2,63}$"
)

def validate_target_host(target: str) -> str:
    """Validate whether input is a strict IPv4, IPv6, or FQDN string."""
    target = target.strip()
    try:
        ipaddress.ip_address(target)
        return target
    except ValueError:
        pass

    if DOMAIN_REGEX.match(target):
        return target

    raise ValueError(f"Invalid target format: '{target}'. Must be a valid IP address or FQDN.")

def validate_http_url(url: str) -> str:
    """Validate HTTP/HTTPS URL scheme and network location."""
    url = url.strip()
    parsed = urllib.parse.urlparse(url)
    if parsed.scheme not in ("http", "https"):
        raise ValueError(f"Invalid URL scheme '{parsed.scheme}'. Only http and https are permitted.")
    if not parsed.netloc:
        raise ValueError("URL must contain a valid domain/host network location.")
    return url


# Pydantic Schemas
class SubdomainEnumInput(BaseModel):
    domain: str = Field(..., description="Target domain to inspect (e.g., example.com)")

    @field_validator("domain")
    @classmethod
    def check_domain(cls, v: str) -> str:
        return validate_target_host(v)


class NmapScanInput(BaseModel):
    target: str = Field(..., description="Target hostname or IP address to scan")

    @field_validator("target")
    @classmethod
    def check_target(cls, v: str) -> str:
        return validate_target_host(v)


class RawWebFetchInput(BaseModel):
    url: str = Field(..., description="Full HTTP/HTTPS URL to fetch complete raw HTML from")

    @field_validator("url")
    @classmethod
    def check_url(cls, v: str) -> str:
        return validate_http_url(v)


class BrowserActionInput(BaseModel):
    url: str = Field(..., description="Target HTTP/HTTPS URL to open")
    action: Optional[str] = Field(
        default="get_text",
        description="Browser action: 'get_text' (rendered text), 'get_html' (rendered DOM), 'click', 'fill', or 'eval_js'"
    )
    selector: Optional[str] = Field(default="", description="CSS selector for click/fill actions (e.g. 'button#submit', 'input[name=q]')")
    text_input: Optional[str] = Field(default="", description="Text string to type into input field (for 'fill' action)")
    script: Optional[str] = Field(default="", description="JavaScript code to execute in page context (for 'eval_js' action)")
    wait_seconds: Optional[int] = Field(default=3, description="Time to wait for JavaScript execution/rendering in seconds (1-15)")

    @field_validator("url")
    @classmethod
    def check_url(cls, v: str) -> str:
        return validate_http_url(v)


class SiteMirrorInput(BaseModel):
    url: str = Field(..., description="Target root URL to mirror recursively")
    max_depth: Optional[int] = Field(default=2, description="Maximum link crawl depth (1-5)")

    @field_validator("url")
    @classmethod
    def check_url(cls, v: str) -> str:
        return validate_http_url(v)


class PacketCaptureInput(BaseModel):
    pcap_path: str = Field(..., description="Path to local capture file (.pcap, .pcapng, .cap)")
    display_filter: Optional[str] = Field(default="", description="Wireshark display filter string")

    @field_validator("pcap_path")
    @classmethod
    def sanitize_path(cls, v: str) -> str:
        clean = os.path.basename(v.strip())
        if not clean.endswith((".pcap", ".pcapng", ".cap")):
            raise ValueError("Target file must end with .pcap, .pcapng, or .cap")
        return clean

    @field_validator("display_filter")
    @classmethod
    def sanitize_filter(cls, v: Optional[str]) -> str:
        if not v:
            return ""
        if not re.match(r"^[a-zA-Z0-9._\s=!<>&|()\"'-]+$", v):
            raise ValueError("Display filter contains unsupported characters.")
        return v


class BashCommandInput(BaseModel):
    command: str = Field(..., description="Shell command to execute in the container environment")


class PythonCodeInput(BaseModel):
    code: str = Field(..., description="Python 3 script code to execute in the container")


class CurlRequestInput(BaseModel):
    url: str = Field(..., description="Target URL to query with curl")
    method: Optional[str] = Field(default="GET", description="HTTP Method (GET, POST, HEAD, OPTIONS)")
    headers: Optional[List[str]] = Field(default=[], description="List of headers in 'Header: Value' format")
    data: Optional[str] = Field(default="", description="Request body payload")
    include_headers: Optional[bool] = Field(default=True, description="Whether to include response headers (-i)")


class DnsLookupInput(BaseModel):
    domain: str = Field(..., description="Domain name or host to query")
    record_type: Optional[str] = Field(default="ANY", description="DNS record type (A, AAAA, MX, NS, TXT, CNAME, ANY)")


class SSLCheckInput(BaseModel):
    host: str = Field(..., description="Target host to inspect SSL/TLS certificate for")
    port: Optional[int] = Field(default=443, description="Port number")

    @field_validator("host")
    @classmethod
    def check_host(cls, v: str) -> str:
        return validate_target_host(v)


# =====================================================================
# 3. ASYNCHRONOUS TOOL IMPLEMENTATIONS
# =====================================================================

def _sync_selenium_execute(
    url: str,
    action: str,
    selector: str,
    text_input: str,
    script: str,
    wait_seconds: int
) -> Dict[str, Any]:
    """Internal synchronous worker for headless Chromium automation."""
    chrome_options = Options()
    chrome_options.add_argument("--headless=new")
    chrome_options.add_argument("--no-sandbox")
    chrome_options.add_argument("--disable-dev-shm-usage")
    chrome_options.add_argument("--disable-gpu")
    chrome_options.add_argument("--window-size=1920,1080")
    chrome_options.add_argument(
        "user-agent=Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36"
    )

    driver_path = os.environ.get("CHROMEDRIVER_PATH", "/usr/bin/chromedriver")
    chrome_bin = os.environ.get("CHROME_BIN", "/usr/bin/chromium")
    if os.path.exists(chrome_bin):
        chrome_options.binary_location = chrome_bin

    service = Service(executable_path=driver_path)
    driver = None

    try:
        driver = webdriver.Chrome(service=service, options=chrome_options)
        driver.set_page_load_timeout(30)
        driver.get(url)

        # Allow dynamic JS execution
        wait_time = min(max(1, wait_seconds), 15)
        WebDriverWait(driver, wait_time).until(
            lambda d: d.execute_script("return document.readyState") == "complete"
        )

        result: Dict[str, Any] = {
            "status": "success",
            "current_url": driver.current_url,
            "title": driver.title
        }

        if action == "click" and selector:
            elem = WebDriverWait(driver, wait_time).until(
                EC.element_to_be_clickable((By.CSS_SELECTOR, selector))
            )
            elem.click()
            WebDriverWait(driver, 2).until(
                lambda d: d.execute_script("return document.readyState") == "complete"
            )
            result["action_performed"] = f"Clicked selector: {selector}"
            result["current_url"] = driver.current_url

        elif action == "fill" and selector:
            elem = WebDriverWait(driver, wait_time).until(
                EC.presence_of_element_located((By.CSS_SELECTOR, selector))
            )
            elem.clear()
            elem.send_keys(text_input)
            result["action_performed"] = f"Filled selector: {selector}"

        elif action == "eval_js" and script:
            eval_output = driver.execute_script(script)
            result["js_eval_result"] = str(eval_output)

        elif action == "get_html":
            result["rendered_dom"] = driver.page_source

        else:  # Default: get_text
            body_text = driver.find_element(By.TAG_NAME, "body").text
            result["rendered_text"] = body_text

        return result
    except Exception as exc:
        return {"status": "error", "message": f"Browser automation error: {str(exc)}"}
    finally:
        if driver:
            try:
                driver.quit()
            except Exception:
                pass


async def run_headless_browser_action(
    url: str,
    action: str = "get_text",
    selector: str = "",
    text_input: str = "",
    script: str = "",
    wait_seconds: int = 3
) -> Dict[str, Any]:
    """Asynchronous wrapper that executes headless browser workflows in a worker thread."""
    try:
        validated = BrowserActionInput(
            url=url,
            action=action,
            selector=selector,
            text_input=text_input,
            script=script,
            wait_seconds=wait_seconds
        )
        return await asyncio.to_thread(
            _sync_selenium_execute,
            validated.url,
            validated.action,
            validated.selector,
            validated.text_input,
            validated.script,
            validated.wait_seconds
        )
    except Exception as exc:
        return {"status": "error", "message": f"Browser tool validation failure: {str(exc)}"}


async def run_bash_executor(command: str) -> Dict[str, Any]:
    """Executes arbitrary shell commands in the persistent container workspace."""
    try:
        validated = BashCommandInput(command=command)
        process = await asyncio.create_subprocess_shell(
            validated.command,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            cwd="/app"
        )
        stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=90.0)

        return {
            "status": "success" if process.returncode == 0 else "error",
            "exit_code": process.returncode,
            "stdout": stdout.decode("utf-8", errors="replace"),
            "stderr": stderr.decode("utf-8", errors="replace")
        }
    except asyncio.TimeoutError:
        return {"status": "error", "message": "Command execution timed out (90s limit)."}
    except Exception as exc:
        return {"status": "error", "message": f"Execution error: {str(exc)}"}


async def run_python_interpreter(code: str) -> Dict[str, Any]:
    """Executes Python code scripts asynchronously in the runtime container."""
    try:
        validated = PythonCodeInput(code=code)
        process = await asyncio.create_subprocess_exec(
            "python3", "-c", validated.code,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            cwd="/app"
        )
        stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=45.0)

        return {
            "status": "success" if process.returncode == 0 else "error",
            "exit_code": process.returncode,
            "stdout": stdout.decode("utf-8", errors="replace"),
            "stderr": stderr.decode("utf-8", errors="replace")
        }
    except asyncio.TimeoutError:
        return {"status": "error", "message": "Python execution timed out (45s limit)."}
    except Exception as exc:
        return {"status": "error", "message": f"Python execution error: {str(exc)}"}


async def run_raw_web_fetcher(url: str) -> Dict[str, Any]:
    """Fetches the complete, unmodified raw HTML content and response headers from a URL."""
    try:
        validated = RawWebFetchInput(url=url)
        headers = {
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36"
        }

        async with httpx.AsyncClient(timeout=30.0, follow_redirects=True) as client:
            resp = await client.get(validated.url, headers=headers)
            return {
                "status": "success",
                "effective_url": str(resp.url),
                "http_status": resp.status_code,
                "response_headers": dict(resp.headers),
                "cookies": list(resp.cookies.keys()),
                "raw_html": resp.text
            }
    except Exception as exc:
        return {"status": "error", "message": f"Raw web fetch error: {str(exc)}"}


async def run_site_mirror_tool(url: str, max_depth: int = 2) -> Dict[str, Any]:
    """Recursively downloads website assets to a local directory using wget for offline analysis."""
    try:
        validated = SiteMirrorInput(url=url, max_depth=min(max(1, max_depth), 5))
        parsed = urllib.parse.urlparse(validated.url)
        safe_domain = re.sub(r"[^a-zA-Z0-9.-]", "_", parsed.netloc)
        output_dir = os.path.join("/app", f"mirror_{safe_domain}")
        os.makedirs(output_dir, exist_ok=True)

        cmd = [
            "wget",
            "--mirror",
            "--convert-links",
            "--adjust-extension",
            "--page-requisites",
            "--no-parent",
            f"--level={validated.max_depth}",
            "-P", output_dir,
            validated.url
        ]

        process = await asyncio.create_subprocess_exec(
            *cmd,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE
        )
        stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=120.0)

        file_count = sum(len(files) for _, _, files in os.walk(output_dir))

        return {
            "status": "success" if process.returncode == 0 else "warning",
            "exit_code": process.returncode,
            "output_directory": output_dir,
            "downloaded_files_count": file_count,
            "stderr_summary": stderr.decode("utf-8", errors="replace")[-500:]
        }
    except asyncio.TimeoutError:
        return {"status": "error", "message": "Site mirror process timed out (120s limit)."}
    except Exception as exc:
        return {"status": "error", "message": f"Site mirroring error: {str(exc)}"}


async def run_curl_requester(
    url: str,
    method: str = "GET",
    headers: Optional[List[str]] = None,
    data: Optional[str] = "",
    include_headers: bool = True
) -> Dict[str, Any]:
    """Performs HTTP/HTTPS network operations using system curl."""
    try:
        validated = CurlRequestInput(
            url=url,
            method=method,
            headers=headers or [],
            data=data or "",
            include_headers=include_headers
        )
        
        cmd = ["curl", "-s", "-S", "-X", validated.method.upper()]
        if validated.include_headers:
            cmd.append("-i")
        if validated.headers:
            for h in validated.headers:
                cmd.extend(["-H", h])
        if validated.data:
            cmd.extend(["--data", validated.data])
        cmd.append(validated.url)

        process = await asyncio.create_subprocess_exec(
            *cmd,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE
        )
        stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=30.0)

        return {
            "status": "success" if process.returncode == 0 else "error",
            "exit_code": process.returncode,
            "stdout": stdout.decode("utf-8", errors="replace"),
            "stderr": stderr.decode("utf-8", errors="replace")
        }
    except Exception as exc:
        return {"status": "error", "message": f"Curl execution error: {str(exc)}"}


async def run_dns_lookup(domain: str, record_type: str = "ANY") -> Dict[str, Any]:
    """Queries DNS records using dig."""
    try:
        validated = DnsLookupInput(domain=domain, record_type=record_type)
        cmd = ["dig", "+noall", "+answer", "+comments", validated.domain, validated.record_type.upper()]

        process = await asyncio.create_subprocess_exec(
            *cmd,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE
        )
        stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=15.0)

        return {
            "status": "success" if process.returncode == 0 else "error",
            "exit_code": process.returncode,
            "stdout": stdout.decode("utf-8", errors="replace"),
            "stderr": stderr.decode("utf-8", errors="replace")
        }
    except Exception as exc:
        return {"status": "error", "message": f"DNS query error: {str(exc)}"}


async def run_ssl_checker(host: str, port: int = 443) -> Dict[str, Any]:
    """Inspects TLS/SSL certificate details and expiration dates using openssl."""
    try:
        validated = SSLCheckInput(host=host, port=port)
        cmd = f"openssl s_client -connect {validated.host}:{validated.port} -servername {validated.host} < /dev/null 2>/dev/null | openssl x509 -text -noout"

        process = await asyncio.create_subprocess_shell(
            cmd,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE
        )
        stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=15.0)

        return {
            "status": "success" if process.returncode == 0 else "error",
            "exit_code": process.returncode,
            "certificate_info": stdout.decode("utf-8", errors="replace"),
            "stderr": stderr.decode("utf-8", errors="replace")
        }
    except Exception as exc:
        return {"status": "error", "message": f"SSL inspection error: {str(exc)}"}


async def run_subdomain_enumerator(domain: str) -> Dict[str, Any]:
    """Query certificate transparency logs (crt.sh) for domain asset mapping."""
    try:
        validated = SubdomainEnumInput(domain=domain)
        crt_url = f"https://crt.sh/?q=%25.{validated.domain}&output=json"

        async with httpx.AsyncClient(timeout=15.0) as client:
            response = await client.get(crt_url)
            if response.status_code != 200:
                return {
                    "status": "error",
                    "message": f"crt.sh query returned status code {response.status_code}",
                    "subdomains": []
                }
            
            data = response.json()
            subdomains = set()
            for entry in data:
                name_value = entry.get("name_value", "")
                for sub in name_value.split("\n"):
                    sub = sub.strip().lower()
                    if "*" not in sub and sub.endswith(validated.domain):
                        subdomains.add(sub)

            return {
                "status": "success",
                "target": validated.domain,
                "count": len(subdomains),
                "subdomains": sorted(list(subdomains))[:100]
            }
    except Exception as exc:
        return {"status": "error", "message": f"Subdomain enumeration failed: {str(exc)}", "subdomains": []}


async def run_nmap_port_scanner(target: str) -> Dict[str, Any]:
    """Execute bounded service discovery on top 100 ports using subprocess execution."""
    try:
        validated = NmapScanInput(target=target)
        cmd = ["nmap", "-sV", "-T4", "--top-ports", "100", validated.target]

        process = await asyncio.create_subprocess_exec(
            *cmd,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE
        )
        stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=45.0)

        return {
            "status": "success" if process.returncode == 0 else "warning",
            "exit_code": process.returncode,
            "stdout": stdout.decode("utf-8", errors="replace"),
            "stderr": stderr.decode("utf-8", errors="replace")
        }
    except asyncio.TimeoutError:
        return {"status": "error", "message": "Nmap scan execution timed out (45s limit)."}
    except Exception as exc:
        return {"status": "error", "message": f"Nmap execution error: {str(exc)}"}


async def run_packet_capture_analyzer(pcap_path: str, display_filter: str = "") -> Dict[str, Any]:
    """Parse local packet captures using tshark with display filters."""
    try:
        validated = PacketCaptureInput(pcap_path=pcap_path, display_filter=display_filter)
        if not os.path.exists(validated.pcap_path):
            return {"status": "error", "message": f"PCAP file '{validated.pcap_path}' not found in runtime directory."}

        cmd = ["tshark", "-r", validated.pcap_path, "-c", "50"]
        if validated.display_filter:
            cmd.extend(["-Y", validated.display_filter])

        process = await asyncio.create_subprocess_exec(
            *cmd,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE
        )
        stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=20.0)

        return {
            "status": "success" if process.returncode == 0 else "warning",
            "exit_code": process.returncode,
            "packet_lines": stdout.decode("utf-8", errors="replace").splitlines()[:50],
            "stderr": stderr.decode("utf-8", errors="replace")
        }
    except Exception as exc:
        return {"status": "error", "message": f"Tshark analysis error: {str(exc)}"}


# =====================================================================
# 4. TOOL REGISTRY & SCHEMA DEFINITIONS
# =====================================================================

DIAGNOSTIC_TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "headless_browser_action",
            "description": "Controls a headless Chromium browser using Selenium to navigate JavaScript single-page apps (SPAs), click elements, type into inputs, execute custom JS, or extract dynamically rendered text and DOM without vision.",
            "parameters": {
                "type": "object",
                "properties": {
                    "url": {"type": "string", "description": "Target HTTP/HTTPS URL."},
                    "action": {
                        "type": "string",
                        "description": "Action: 'get_text' (rendered text), 'get_html' (rendered DOM), 'click' (click selector), 'fill' (type text into selector), 'eval_js' (execute script).",
                        "default": "get_text"
                    },
                    "selector": {"type": "string", "description": "CSS selector for 'click' or 'fill' actions."},
                    "text_input": {"type": "string", "description": "Text to input when action is 'fill'."},
                    "script": {"type": "string", "description": "JavaScript to run when action is 'eval_js'."},
                    "wait_seconds": {"type": "integer", "description": "Seconds to wait for JS rendering (1-15).", "default": 3}
                },
                "required": ["url"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "bash_executor",
            "description": "Executes shell commands directly inside the persistent Linux container workspace (/app).",
            "parameters": {
                "type": "object",
                "properties": {
                    "command": {"type": "string", "description": "The bash shell command to run."}
                },
                "required": ["command"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "python_interpreter",
            "description": "Runs Python 3 code in the container runtime environment.",
            "parameters": {
                "type": "object",
                "properties": {
                    "code": {"type": "string", "description": "Python code snippet to execute."}
                },
                "required": ["code"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "raw_web_fetcher",
            "description": "Fetches the full, complete raw HTML content, response headers, and cookies of a web page.",
            "parameters": {
                "type": "object",
                "properties": {
                    "url": {"type": "string", "description": "Target HTTP/HTTPS URL."}
                },
                "required": ["url"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "site_mirror_tool",
            "description": "Recursively clones/mirrors an entire website or web path locally to disk for inspection.",
            "parameters": {
                "type": "object",
                "properties": {
                    "url": {"type": "string", "description": "Root URL of the website to mirror."},
                    "max_depth": {"type": "integer", "description": "Recursion depth (default: 2, max: 5).", "default": 2}
                },
                "required": ["url"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "curl_requester",
            "description": "Issues custom HTTP/HTTPS network requests using curl.",
            "parameters": {
                "type": "object",
                "properties": {
                    "url": {"type": "string", "description": "Target URL."},
                    "method": {"type": "string", "description": "HTTP Method.", "default": "GET"},
                    "headers": {"type": "array", "items": {"type": "string"}, "description": "List of 'Header: Value' strings."},
                    "data": {"type": "string", "description": "Request payload."},
                    "include_headers": {"type": "boolean", "description": "Include headers in output.", "default": True}
                },
                "required": ["url"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "dns_lookup",
            "description": "Queries DNS records using dig for domain discovery and verification.",
            "parameters": {
                "type": "object",
                "properties": {
                    "domain": {"type": "string", "description": "Domain to look up."},
                    "record_type": {"type": "string", "description": "DNS record type.", "default": "ANY"}
                },
                "required": ["domain"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "ssl_checker",
            "description": "Extracts and parses the SSL/TLS certificate chain and expiration data for a given host.",
            "parameters": {
                "type": "object",
                "properties": {
                    "host": {"type": "string", "description": "Hostname or IP."},
                    "port": {"type": "integer", "description": "Port number (default 443).", "default": 443}
                },
                "required": ["host"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "subdomain_enumerator",
            "description": "Performs DNS asset discovery via certificate transparency records.",
            "parameters": {
                "type": "object",
                "properties": {
                    "domain": {"type": "string", "description": "Domain to audit (e.g. 'example.com')."}
                },
                "required": ["domain"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "nmap_port_scanner",
            "description": "Performs service and port detection on the top 100 ports of a target host.",
            "parameters": {
                "type": "object",
                "properties": {
                    "target": {"type": "string", "description": "IPv4, IPv6, or FQDN to scan."}
                },
                "required": ["target"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "packet_capture_analyzer",
            "description": "Reads and filters network packet captures from a PCAP file using tshark.",
            "parameters": {
                "type": "object",
                "properties": {
                    "pcap_path": {"type": "string", "description": "Path to the PCAP file."},
                    "display_filter": {"type": "string", "description": "Wireshark display filter."}
                },
                "required": ["pcap_path"]
            }
        }
    }
]

TOOL_HANDLERS = {
    "headless_browser_action": run_headless_browser_action,
    "bash_executor": run_bash_executor,
    "python_interpreter": run_python_interpreter,
    "raw_web_fetcher": run_raw_web_fetcher,
    "site_mirror_tool": run_site_mirror_tool,
    "curl_requester": run_curl_requester,
    "dns_lookup": run_dns_lookup,
    "ssl_checker": run_ssl_checker,
    "subdomain_enumerator": run_subdomain_enumerator,
    "nmap_port_scanner": run_nmap_port_scanner,
    "packet_capture_analyzer": run_packet_capture_analyzer
}


# =====================================================================
# 5. RECURSIVE AGENTIC REACTION LOOP
# =====================================================================

async def execute_tool(tool_name: str, arguments_json: str) -> str:
    """Safely parse arguments and route to designated tool wrapper."""
    if tool_name not in TOOL_HANDLERS:
        return json.dumps({"error": f"Tool '{tool_name}' not recognized."})

    try:
        kwargs = json.loads(arguments_json)
    except json.JSONDecodeError:
        return json.dumps({"error": f"Invalid JSON arguments supplied for '{tool_name}'."})

    handler = TOOL_HANDLERS[tool_name]
    try:
        result = await handler(**kwargs)
        return json.dumps(result, indent=2)
    except Exception as exc:
        return json.dumps({"error": f"Execution error in '{tool_name}': {str(exc)}"})


async def run_agent_loop(
    user_message: str,
    chat_history: List[Dict[str, str]],
    api_key: str,
    base_url: str,
    model_name: str,
    system_prompt: str,
    max_steps: int = DEFAULT_MAX_STEPS
):
    """
    Executes the asynchronous multi-turn tool calling loop.
    Streams execution progress, reasoning traces, and the final synthesis.
    """
    if not api_key:
        yield chat_history + [{"role": "assistant", "content": "API Key is required to initialize the agent client."}], "Configuration Error: Missing API Key."
        return

    client = AsyncOpenAI(api_key=api_key, base_url=base_url)

    messages = [{"role": "system", "content": system_prompt}]
    for msg in chat_history:
        messages.append({"role": msg["role"], "content": msg["content"]})
    messages.append({"role": "user", "content": user_message})

    logs = [f"**[INITIALIZE]** Starting session (Max Steps: {max_steps}) for query: `{user_message}`\n"]
    current_chat = chat_history + [{"role": "user", "content": user_message}]

    step = 0
    while step < max_steps:
        step += 1
        logs.append(f"\n**[ITERATION {step}/{max_steps}] Querying Model...**")
        yield current_chat, "\n".join(logs)

        try:
            response = await client.chat.completions.create(
                model=model_name,
                messages=messages,
                tools=DIAGNOSTIC_TOOLS,
                tool_choice="auto",
                temperature=0.2
            )
        except Exception as api_err:
            err_text = f"API Request Failed: {str(api_err)}"
            logs.append(f"❌ `{err_text}`")
            current_chat.append({"role": "assistant", "content": err_text})
            yield current_chat, "\n".join(logs)
            return

        choice = response.choices[0]
        message = choice.message
        reasoning = getattr(message, "reasoning_content", None)

        if reasoning:
            logs.append(f"\n🧠 **Analytical Reasoning:**\n```text\n{reasoning}\n```")
            yield current_chat, "\n".join(logs)

        messages.append(message)

        if message.tool_calls:
            for tool_call in message.tool_calls:
                fn_name = tool_call.function.name
                fn_args = tool_call.function.arguments

                logs.append(f"⚙️ **Invoking:** `{fn_name}`\nArguments:\n```json\n{fn_args}\n```")
                yield current_chat, "\n".join(logs)

                tool_output = await execute_tool(fn_name, fn_args)

                output_preview = tool_output if len(tool_output) <= 600 else f"{tool_output[:600]}... [truncated for UI preview]"
                logs.append(f"📊 **Result from `{fn_name}`:**\n```json\n{output_preview}\n```")
                yield current_chat, "\n".join(logs)

                # Send full output to LLM context
                messages.append({
                    "role": "tool",
                    "tool_call_id": tool_call.id,
                    "name": fn_name,
                    "content": tool_output
                })
        else:
            final_response = message.content or "Diagnostic review complete."
            logs.append("\n✅ **Audit Completed.**")
            current_chat.append({"role": "assistant", "content": final_response})
            yield current_chat, "\n".join(logs)
            return

    fallback_msg = f"Diagnostic workflow reached the maximum iteration limit of {max_steps} steps."
    current_chat.append({"role": "assistant", "content": fallback_msg})
    yield current_chat, "\n".join(logs)


# =====================================================================
# 6. GRADIO USER INTERFACE
# =====================================================================

CUSTOM_CSS = """
body, .gradio-container {
    background-color: #0b0f19 !important;
    color: #f1f5f9 !important;
    font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif;
}
.gr-box, .gr-panel, .gr-form {
    background-color: #111827 !important;
    border-color: #1f2937 !important;
}
#terminal-output textarea {
    background-color: #030712 !important;
    color: #38bdf8 !important;
    font-family: ui-monospace, SFMono-Regular, Menlo, Monaco, Consolas, monospace !important;
    font-size: 0.82rem !important;
}
"""

def create_ui():
    theme = gr.themes.Soft(
        primary_hue="cyan",
        secondary_hue="slate",
        neutral_hue="slate"
    )

    with gr.Blocks(title="Network Diagnostics Agent") as demo:
        gr.Markdown(
            """
            # 🌐 Infrastructure Diagnostics & Asset Audit Platform
            ### Asynchronous ReAct Diagnostic Engine with Containerized Tooling & Headless Browser
            """
        )

        with gr.Accordion("⚙️ Endpoint & System Settings", open=False):
            with gr.Row():
                api_key_input = gr.Textbox(
                    label="API Key",
                    type="password",
                    value=DEFAULT_API_KEY,
                    placeholder="Enter API Key"
                )
                base_url_input = gr.Textbox(
                    label="Base URL",
                    value=DEFAULT_BASE_URL
                )
                model_name_input = gr.Textbox(
                    label="Model Identifier",
                    value=DEFAULT_MODEL
                )
            with gr.Row():
                max_steps_input = gr.Slider(
                    label="Max Agent Iteration Steps",
                    minimum=3,
                    maximum=30,
                    value=DEFAULT_MAX_STEPS,
                    step=1
                )
            system_prompt_input = gr.Textbox(
                label="System Prompt",
                value=DEFAULT_SYSTEM_PROMPT,
                lines=4
            )

        with gr.Row():
            with gr.Column(scale=5):
                chatbot = gr.Chatbot(
                    label="Diagnostic Workflow Session",
                    height=520
                )
                with gr.Row():
                    user_input = gr.Textbox(
                        label="Audit Command / Target Specification",
                        placeholder="E.g., Open https://example.com in headless browser, extract rendered text, or run a port scan",
                        lines=2,
                        scale=4
                    )
                    submit_btn = gr.Button("Run Audit", variant="primary", scale=1)

                clear_btn = gr.Button("Clear Session")

            with gr.Column(scale=4):
                terminal_logs = gr.Markdown(
                    value="*Execution logs, tool output traces, and model reasoning will stream here...*",
                    elem_id="terminal-output"
                )

        async def on_submit(user_msg, chat_hist, key, url, model, steps, sys_prompt):
            if not user_msg.strip():
                yield chat_hist, "*No input provided.*"
                return

            chat_hist = chat_hist or []
            async for updated_chat, updated_logs in run_agent_loop(
                user_message=user_msg,
                chat_history=chat_hist,
                api_key=key,
                base_url=url,
                model_name=model,
                system_prompt=sys_prompt,
                max_steps=int(steps)
            ):
                yield updated_chat, updated_logs

        inputs_list = [
            user_input, chatbot, api_key_input, base_url_input,
            model_name_input, max_steps_input, system_prompt_input
        ]

        submit_btn.click(
            fn=on_submit,
            inputs=inputs_list,
            outputs=[chatbot, terminal_logs]
        ).then(
            fn=lambda: "",
            inputs=None,
            outputs=[user_input]
        )

        user_input.submit(
            fn=on_submit,
            inputs=inputs_list,
            outputs=[chatbot, terminal_logs]
        ).then(
            fn=lambda: "",
            inputs=None,
            outputs=[user_input]
        )

        clear_btn.click(
            fn=lambda: ([], "*Session reset.*"),
            inputs=None,
            outputs=[chatbot, terminal_logs]
        )

    return demo, theme


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