Update app.py

app.py

@@ -20,6 +20,19 @@ 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_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 and terminal tools.
+2. When necessary, use the bash executor, curl requester, or Python interpreter to inspect live responses, process outputs, or run diagnostic commands.
+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
 # =====================================================================
@@ -55,6 +68,7 @@ def validate_http_url(url: str) -> str:
     return url
 
 
+# Pydantic Schemas for Tools
 class SubdomainEnumInput(BaseModel):
     domain: str = Field(..., description="Target domain to inspect (e.g., example.com)")
 
@@ -104,10 +118,146 @@ class PacketCaptureInput(BaseModel):
         return v
 
 
+class BashCommandInput(BaseModel):
+    command: str = Field(..., description="Shell command to execute in the persistent 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)")
+
+
 # =====================================================================
-# 3. ASYNCHRONOUS DIAGNOSTIC TOOL WRAPPERS
+# 3. ASYNCHRONOUS TOOL IMPLEMENTATIONS
 # =====================================================================
 
+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=60.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 (60s 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=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 asyncio.TimeoutError:
+        return {"status": "error", "message": "Python script execution timed out (30s limit)."}
+    except Exception as exc:
+        return {"status": "error", "message": f"Python execution 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_subdomain_enumerator(domain: str) -> Dict[str, Any]:
     """Query certificate transparency logs (crt.sh) for domain asset mapping."""
     try:
@@ -140,11 +290,7 @@ async def run_subdomain_enumerator(domain: str) -> Dict[str, Any]:
                 "subdomains": sorted(list(subdomains))[:100]
             }
     except Exception as exc:
-        return {
-            "status": "error",
-            "message": f"Subdomain enumeration failed: {str(exc)}",
-            "subdomains": []
-        }
+        return {"status": "error", "message": f"Subdomain enumeration failed: {str(exc)}", "subdomains": []}
 
 
 async def run_nmap_port_scanner(target: str) -> Dict[str, Any]:
@@ -169,15 +315,9 @@ async def run_nmap_port_scanner(target: str) -> Dict[str, Any]:
             "stderr": stderr.decode("utf-8", errors="replace")
         }
     except asyncio.TimeoutError:
-        return {
-            "status": "error",
-            "message": "Nmap scan execution timed out (45s limit)."
-        }
+        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)}"
-        }
+        return {"status": "error", "message": f"Nmap execution error: {str(exc)}"}
 
 
 async def run_web_scraper_and_analyzer(url: str) -> Dict[str, Any]:
@@ -186,9 +326,7 @@ async def run_web_scraper_and_analyzer(url: str) -> Dict[str, Any]:
         validated = WebScraperInput(url=url)
         target_url = validated.url
 
-        headers = {
-            "User-Agent": "Diagnostic-Audit-Agent/1.0 (Compliance; Non-Intrusive)"
-        }
+        headers = {"User-Agent": "Diagnostic-Audit-Agent/1.0 (Compliance; Non-Intrusive)"}
 
         async with httpx.AsyncClient(timeout=15.0, follow_redirects=True) as client:
             resp = await client.get(target_url, headers=headers)
@@ -223,10 +361,7 @@ async def run_web_scraper_and_analyzer(url: str) -> Dict[str, Any]:
                 "external_scripts": scripts
             }
     except Exception as exc:
-        return {
-            "status": "error",
-            "message": f"Web analyzer error: {str(exc)}"
-        }
+        return {"status": "error", "message": f"Web analyzer error: {str(exc)}"}
 
 
 async def run_packet_capture_analyzer(pcap_path: str, display_filter: str = "") -> Dict[str, Any]:
@@ -237,10 +372,7 @@ async def run_packet_capture_analyzer(pcap_path: str, display_filter: str = "")
         filt = validated.display_filter
 
         if not os.path.exists(target_path):
-            return {
-                "status": "error",
-                "message": f"PCAP file '{target_path}' not found in runtime directory."
-            }
+            return {"status": "error", "message": f"PCAP file '{target_path}' not found in runtime directory."}
 
         cmd = ["tshark", "-r", target_path, "-c", "50"]
         if filt:
@@ -260,10 +392,7 @@ async def run_packet_capture_analyzer(pcap_path: str, display_filter: str = "")
             "stderr": stderr.decode("utf-8", errors="replace")
         }
     except Exception as exc:
-        return {
-            "status": "error",
-            "message": f"Tshark analysis error: {str(exc)}"
-        }
+        return {"status": "error", "message": f"Tshark analysis error: {str(exc)}"}
 
 
 # =====================================================================
@@ -271,6 +400,98 @@ async def run_packet_capture_analyzer(pcap_path: str, display_filter: str = "")
 # =====================================================================
 
 DIAGNOSTIC_TOOLS = [
+    {
+        "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": "curl_requester",
+            "description": "Issues HTTP/HTTPS network requests using curl with custom methods, headers, and payloads.",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "url": {
+                        "type": "string",
+                        "description": "Target URL."
+                    },
+                    "method": {
+                        "type": "string",
+                        "description": "HTTP Method (GET, POST, HEAD, OPTIONS).",
+                        "default": "GET"
+                    },
+                    "headers": {
+                        "type": "array",
+                        "items": {"type": "string"},
+                        "description": "List of headers in 'Header: Value' format."
+                    },
+                    "data": {
+                        "type": "string",
+                        "description": "Request body data."
+                    },
+                    "include_headers": {
+                        "type": "boolean",
+                        "description": "Include response 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 or host to look up."
+                    },
+                    "record_type": {
+                        "type": "string",
+                        "description": "Record type: A, AAAA, MX, NS, TXT, CNAME, ANY.",
+                        "default": "ANY"
+                    }
+                },
+                "required": ["domain"]
+            }
+        }
+    },
     {
         "type": "function",
         "function": {
@@ -346,6 +567,10 @@ DIAGNOSTIC_TOOLS = [
 ]
 
 TOOL_HANDLERS = {
+    "bash_executor": run_bash_executor,
+    "python_interpreter": run_python_interpreter,
+    "curl_requester": run_curl_requester,
+    "dns_lookup": run_dns_lookup,
     "subdomain_enumerator": run_subdomain_enumerator,
     "nmap_port_scanner": run_nmap_port_scanner,
     "web_scraper_and_analyzer": run_web_scraper_and_analyzer,
@@ -357,18 +582,8 @@ TOOL_HANDLERS = {
 # 5. RECURSIVE AGENTIC REACTION LOOP
 # =====================================================================
 
-SYSTEM_PROMPT = """You are an automated System Administration and Infrastructure Diagnostics Agent.
-You assist system administrators and network engineers in auditing network assets, inspecting security configurations, and mapping infrastructure surfaces.
-
-Operational Guidelines:
-1. Systematically investigate target infrastructure using the provided diagnostic tools.
-2. If tool outputs indicate further investigation is necessary, call subsequent tools autonomously.
-3. Present final summaries structured with: Target Profile, Technical Diagnostics, Configuration Observations, and Remediation/Hardening Recommendations.
-4. Maintain a neutral, engineering-focused reporting format."""
-
-
 async def execute_tool(tool_name: str, arguments_json: str) -> str:
-    """Safely parse arguments and route to designated diagnostic wrapper."""
+    """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."})
 
@@ -391,10 +606,11 @@ async def run_agent_loop(
     api_key: str,
     base_url: str,
     model_name: str,
-    max_steps: int = 5
+    system_prompt: str,
+    max_steps: int = 6
 ):
     """
-    Executes the asynchronous multi-turn tool calling loop against the OpenAI-compatible API.
+    Executes the asynchronous multi-turn tool calling loop.
     Streams execution progress, reasoning traces, and the final synthesis.
     """
     if not api_key:
@@ -403,7 +619,7 @@ async def run_agent_loop(
 
     client = AsyncOpenAI(api_key=api_key, base_url=base_url)
 
-    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
+    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})
@@ -511,7 +727,7 @@ def create_ui():
             """
         )
 
-        with gr.Accordion("⚙️ Endpoint & API Settings", open=False):
+        with gr.Accordion("⚙️ Endpoint & System Settings", open=False):
             with gr.Row():
                 api_key_input = gr.Textbox(
                     label="API Key",
@@ -527,10 +743,14 @@ def create_ui():
                     label="Model Identifier",
                     value=DEFAULT_MODEL
                 )
+            system_prompt_input = gr.Textbox(
+                label="System Prompt (Loaded from SYSTEM_PROMPT env by default)",
+                value=DEFAULT_SYSTEM_PROMPT,
+                lines=4
+            )
 
         with gr.Row():
             with gr.Column(scale=5):
-                # Cleaned keyword arguments for universal Gradio 4/5/6 compatibility
                 chatbot = gr.Chatbot(
                     label="Diagnostic Workflow Session",
                     height=520
@@ -538,7 +758,7 @@ def create_ui():
                 with gr.Row():
                     user_input = gr.Textbox(
                         label="Audit Command / Target Specification",
-                        placeholder="E.g., Audit security headers and map open services for scanme.nmap.org",
+                        placeholder="E.g., Query DNS records, check HTTP response headers, or run a diagnostic script for scanme.nmap.org",
                         lines=2,
                         scale=4
                     )
@@ -552,7 +772,7 @@ def create_ui():
                     elem_id="terminal-output"
                 )
 
-        async def on_submit(user_msg, chat_hist, key, url, model):
+        async def on_submit(user_msg, chat_hist, key, url, model, sys_prompt):
             if not user_msg.strip():
                 yield chat_hist, "*No input provided.*"
                 return
@@ -563,14 +783,14 @@ def create_ui():
                 chat_history=chat_hist,
                 api_key=key,
                 base_url=url,
-                model_name=model
+                model_name=model,
+                system_prompt=sys_prompt
             ):
                 yield updated_chat, updated_logs
 
-        # Submit handlers (Button click + Enter key press)
         submit_btn.click(
             fn=on_submit,
-            inputs=[user_input, chatbot, api_key_input, base_url_input, model_name_input],
+            inputs=[user_input, chatbot, api_key_input, base_url_input, model_name_input, system_prompt_input],
             outputs=[chatbot, terminal_logs]
         ).then(
             fn=lambda: "",
@@ -580,7 +800,7 @@ def create_ui():
 
         user_input.submit(
             fn=on_submit,
-            inputs=[user_input, chatbot, api_key_input, base_url_input, model_name_input],
+            inputs=[user_input, chatbot, api_key_input, base_url_input, model_name_input, system_prompt_input],
             outputs=[chatbot, terminal_logs]
         ).then(
             fn=lambda: "",