app.py

import gradio as gr
import os
import subprocess
import paramiko
import socket
from dotenv import load_dotenv
from ddgs import DDGS
from langchain_core.messages import HumanMessage, ToolMessage, SystemMessage, AIMessage
from langchain_openai import ChatOpenAI
from langchain_core.tools import tool
from typing import Annotated
from typing_extensions import TypedDict
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode, tools_condition
from error_analyzer import analyze_error
import logging
import re
import json

# Load environment variables from .env file (for local development)
load_dotenv()

# Configure logging
logging.basicConfig(level=logging.INFO, filename="agent.log", format="%(asctime)s - %(levelname)s - %(message)s")

# SSH details from environment variables (for Hugging Face Spaces)
# These will be overridden by user input if provided
DEFAULT_SSH_DETAILS = {
    "linux": {
        "hostname": os.getenv("LINUX_SSH_HOST", ""),
        "port": int(os.getenv("LINUX_SSH_PORT", "22")),
        "username": os.getenv("LINUX_SSH_USER", ""),
        "password": os.getenv("LINUX_SSH_PASS", "")
    },
    "windows": {
        "hostname": os.getenv("WINDOWS_SSH_HOST", ""),
        "port": int(os.getenv("WINDOWS_SSH_PORT", "22")),
        "username": os.getenv("WINDOWS_SSH_USER", ""),
        "password": os.getenv("WINDOWS_SSH_PASS", "")
    }
}

# Global variable to store user-provided SSH details
user_ssh_details = {
    "linux": DEFAULT_SSH_DETAILS["linux"].copy(),
    "windows": DEFAULT_SSH_DETAILS["windows"].copy()
}

# System prompt
SYSTEM_PROMPT = """
You are Mr. Robot, an expert terminal assistant specializing in pentesting, capture-the-flag (CTF) challenges, and command-line operations on both Linux and Windows (PowerShell) environments. Your primary goal is to assist users by intelligently interpreting natural language inputs and generating and executing accurate, secure, and efficient commands tailored to their intent on the specified system (Linux or Windows). Follow these guidelines:

1. **Role and Expertise**:
   - You are connected to a machine via SSH and can run commands on it using the remote_execute tool, which supports both Kali Linux and Windows Server PowerShell environments.
   - Act as a knowledgeable pentesting and CTF expert, proficient in Linux tools (e.g., nmap, metasploit, netcat) and Windows PowerShell commands (e.g., Get-NetIPAddress, Invoke-WebRequest).
   - Dynamically generate and execute commands based on user intent, considering the target system (Linux or Windows) and task requirements (e.g., network scanning, file manipulation).
   - Prioritize safe and ethical practices, warning users about potentially dangerous commands (e.g., rm, del, chmod, network scans) and ensuring they have permission.
   - Always start with web_search for any query involving recent techniques, tools, installations, vulnerabilities, or CTF solutions. Craft precise, keyword-rich queries (e.g., 'latest nmap stealth scanning flags for pentesting site:github.com' instead of 'nmap commands'). Use search operators like site:, filetype:, or "exact phrase" for relevance. If results are poor, chain another web_search with refined queries.
   - For installations related tasks, first prioritise knowing relevant system info and compatible versions, use web_search for installation guides and documentations and bug resolving.

2. **Response Style**:
   - Use a professional yet approachable tone, like a seasoned cybersecurity mentor.
   - Break down complex tasks into clear, actionable steps with brief explanations.
   - Respond concisely unless detailed explanations are requested.
   - After web_search, summarize key snippets/links and directly integrate them into your reasoning or commands (e.g., 'Based on search result [snippet], use nmap -sS'). If the web_search tool returns JSON results, parse the JSON, extract relevant titles, links, and snippets, and summarize them clearly in your response.

3. **Command Generation**:
   - Interpret natural language inputs to generate appropriate commands for the target system (e.g., "find ip of machine" might generate `ip addr show` for Linux or `Get-NetIPAddress` for Windows).
   - For ambiguous requests, use chat_with_human to clarify intent.
   - Detect long-running commands (e.g., python3 -m http.server, Start-Process, nc -l) and:
     - Apply a short timeout (e.g., 5 seconds) to capture initial output.
     - Inform the user the process is running and provide instructions (e.g., `ps aux | grep <command>` for Linux, `Get-Process` for Windows, or stop with `kill <pid>`/`Stop-Process`).
     - Suggest background execution (e.g., appending `&` for Linux or `Start-Process` for Windows) if appropriate.

4. **Error Awareness and Auto-Debugging**:
   - Analyze command outputs for errors (e.g., "input device is not a TTY", "Permission denied", "Host seems down" for Linux; "Access is denied", "Command not found" for Windows).
   - Use reasoning to suggest context-aware fixes (e.g., add `sudo` for Linux permission issues, `Run as Administrator` for Windows, adjust nmap flags for unreachable hosts).
   - If an error is unclear, escalate to chat_with_human.
   - Log errors in the state to improve future responses.
   - After hitting errors, gather system-specific context (e.g., directories, files, versions, libraries) using remote_execute.

5. **Context Retention**:
   - Maintain conversation context using message history for coherent multi-turn interactions.
   - Track the user's current task (e.g., CTF challenge, network scan) and tailor commands to the target system.

6. **Tool Usage**:
   - Use remote_execute to execute commands on the specified system (Kali Linux or Windows Server) via SSH.
   - Use error_analyzer to analyze command outputs only if you think you can't resolve error.
   - Use chat_with_human for human confirmation before executing any risky commands on server, for clarifications, or casual interactions.
   - Use web_search to fetch the latest articles, commands, or information from the web to assist in tasks, especially for up-to-date pentesting techniques or CTF solutions. Analyze and apply results relevantly.
   - Avoid TTY issues in Docker commands by generating non-interactive commands.
   - After remote_execute, call error_analyzer if the output suggests an issue or error otherwise proceed.

7. **Constraints**:
   - Do not execute commands without confirmation from human if they could modify the system (e.g., rm, del, chmod, Set-ItemProperty).
   - Warn users about legal and ethical implications of network scanning (e.g., nmap, Test-NetConnection) and confirm permission.
   - If unsure about command safety or user intent, use chat_with_human.
   - Ensure all commands are runnable without further interaction, using techniques like `cat <<EOF` for Linux or `Out-File` for Windows for input.

8. **Additional Tool Usage**:
   - Call `respond_directly` for direct answers, summaries, or explanations without needing tools.
   - Call `end_task` when the task is resolved to finalize.
   - Call `chat_with_human` for user doubts or casual talk, to seek human input anytime you want.

Respond intelligently, generating commands that are accurate, safe, and tailored to the user's needs and target system. Empower the user with knowledge while maintaining security and ethical standards.
"""

# LLM initialization
llm = ChatOpenAI(
    api_key=os.getenv("OPENAI_API_KEY", ""),
    base_url=os.getenv("OPENAI_BASE_URL", "https://api.groq.com/openai/v1"),
    model=os.getenv("OPENAI_MODEL", "meta-llama/llama-4-maverick-17b-128e-instruct"),
)

# SSH connection functions
def connected_kali(command, timeout=30):
    """Execute a command on a remote Kali Linux machine via SSH."""
    ssh_client = paramiko.SSHClient()
    ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())

    try:
        ssh_details = user_ssh_details.get("linux", DEFAULT_SSH_DETAILS["linux"])
        
        # Check if SSH details are configured
        if not ssh_details.get("hostname") or not ssh_details.get("username"):
            return "[Error]: Linux SSH details not configured. Please configure SSH settings first."
        
        logging.info(f"Executing Linux command: {command}")
        ssh_client.connect(
            hostname=ssh_details["hostname"],
            port=ssh_details["port"],
            username=ssh_details["username"],
            password=ssh_details["password"],
            timeout=10
        )
        stdin, stdout, stderr = ssh_client.exec_command(command, get_pty=False, timeout=timeout)
        output = stdout.read().decode()
        error = stderr.read().decode()

        logging.info(f"Output: {output}")
        logging.info(f"Error: {error}")

        if output:
            return output
        if error:
            return f"[Error]: {error}"
        return "[No output]"

    except socket.timeout:
        return f"[Timeout]: Command '{command}' timed out after {timeout} seconds. If this is a server (e.g., python3 -m http.server), it may be running. Check with `ps aux | grep {command.split()[0]}` or stop with `kill <pid>`."
    except Exception as e:
        return f"[Error]: Failed to execute '{command}': {str(e)}"
    finally:
        ssh_client.close()

def connected_windows(command, timeout=30):
    """Execute a PowerShell command on a remote Windows Server via SSH."""
    ssh_client = paramiko.SSHClient()
    ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())

    try:
        ssh_details = user_ssh_details.get("windows", DEFAULT_SSH_DETAILS["windows"])
        
        # Check if SSH details are configured
        if not ssh_details.get("hostname") or not ssh_details.get("username"):
            return "[Error]: Windows SSH details not configured. Please configure SSH settings first."
        
        logging.info(f"Executing Windows PowerShell command: {command}")
        ssh_client.connect(
            hostname=ssh_details["hostname"],
            port=ssh_details["port"],
            username=ssh_details["username"],
            password=ssh_details["password"],
            timeout=10
        )
        ps_command = f"powershell -Command \"{command}\""
        stdin, stdout, stderr = ssh_client.exec_command(ps_command, get_pty=False, timeout=timeout)
        output = stdout.read().decode()
        error = stderr.read().decode()

        logging.info(f"Output: {output}")
        logging.info(f"Error: {error}")

        if output:
            return output
        if error:
            return f"[Error]: {error}"
        return "[No output]"

    except socket.timeout:
        return f"[Timeout]: Command '{command}' timed out after {timeout} seconds. If this is a server process, it may be running. Check with `Get-Process` or stop with `Stop-Process -Id <pid>`."
    except Exception as e:
        return f"[Error]: Failed to execute '{command}': {str(e)}"
    finally:
        ssh_client.close()

# Tools
@tool
def remote_execute(query: str, system: str = "linux") -> str:
    """Execute a command on a remote system (Kali Linux or Windows Server) via SSH."""
    logging.info(f"Executing command on {system}: {query}")
    try:
        timeout = 30 if "nmap" in query.lower() or "Test-NetConnection" in query.lower() else 5 if any(cmd in query.lower() for cmd in ["python3 -m http.server", "nc -l", "apache2", "nginx", "flask run", "node server.js", "Start-Process"]) else 30
        if system.lower() == "linux":
            output = connected_kali(query, timeout=timeout)
        elif system.lower() == "windows":
            output = connected_windows(query, timeout=timeout)
        else:
            return f"[Error]: Invalid system specified: {system}. Use 'linux' or 'windows'."
        return output
    except Exception as e:
        return f"[Error]: Failed to execute '{query}' on {system}: {str(e)}. Please check the command or use chat_with_human for clarification."

@tool
def say_hello(name: str) -> str:
    """Greet a user by name."""
    logging.info(f"Greeting tool called: {name}")
    return f"Hello, {name} from Own_The_Net"

@tool
def error_analyzer(command: str, output: str) -> str:
    """Analyze command output for errors and suggest fixes using LLM reasoning."""
    logging.info(f"Analyzing output for errors: {output[:30]}...")
    return analyze_error(command, output, llm_with_tools, SYSTEM_PROMPT)

@tool
def respond_directly(reasoning: str) -> str:
    """Use this tool when no further execution or analysis is needed."""
    logging.info(f"Direct response tool called: {reasoning}")
    return reasoning

@tool
def end_task(summary: str = "Task completed.") -> str:
    """Call this tool to signal that the task is complete."""
    logging.info(f"End task tool called: {summary}")
    return f"[END_TASK]: {summary} - Interaction finalized."

@tool
def chat_with_human(query: str, context: str = "") -> str:
    """Use this for casual conversation, user doubts, clarifications, or risky command confirmations."""
    logging.info(f"Chat with human tool called: {query}")
    return f"[HUMAN_INPUT_REQUIRED]: {query}\nContext: {context or 'No additional context.'}"

@tool
def web_search(query: str) -> str:
    """Search the web for the latest articles and commands using DuckDuckGo."""
    logging.info(f"Performing web search for: {query}")
    try:
        results = []
        with DDGS() as ddgs:
            results = list(ddgs.text(query, region='us-en', max_results=5))
        if not results:
            broadened_query = re.sub(r'\b\d{4}\b', '', query).strip() + " OR tutorial OR guide"
            logging.info(f"No results for '{query}'. Retrying with broadened query: '{broadened_query}'")
            results = list(ddgs.text(broadened_query, region='us-en', max_results=5))
        if results:
            formatted_results = [
                {"title": r['title'], "link": r['href'], "snippet": r['body']}
                for r in results
            ]
            logging.info(f"Web search results: {formatted_results}")
            return json.dumps(formatted_results, indent=2)
        else:
            logging.warning(f"No results found even after broadening query for '{query}'")
            return json.dumps({"warning": "No search results found. Try refining the query."})
    except Exception as e:
        logging.error(f"Web search failed: {str(e)}")
        return json.dumps({"error": f"Failed to perform web search for '{query}': {str(e)}. Try a different query."})

# State
class State(TypedDict):
    messages: Annotated[list, add_messages]
    command_output: str
    human_reviewed: bool
    error_detected: bool
    target_system: str

# Build graph
def build_graph():
    graph_builder = StateGraph(State)

    def chatbot(state: State):
        messages = [SystemMessage(content=SYSTEM_PROMPT)] + state["messages"]
        logging.info(f"Invoking LLM with messages: {[msg.type for msg in messages]}")
        message = llm_with_tools.invoke(messages)
        if isinstance(message, AIMessage):
            logging.info(f"LLM reasoning: {message.content}")
            if hasattr(message, 'tool_calls') and message.tool_calls:
                for tc in message.tool_calls:
                    logging.info(f"Agent tool call: {tc['name']} with args: {tc['args']}")
        elif message.content:
            logging.info(f"Agent response: {message.content[:100]}...")
        else:
            logging.warning("Agent returned empty response or invalid tool call")
            message.content = "[Error]: Agent failed to provide a valid response or tool call."
        return {
            "messages": [message],
            "command_output": state.get("command_output", ""),
            "target_system": state.get("target_system", "linux")
        }

    def tools_to_next(state: State):
        if not state["messages"] or not isinstance(state["messages"][-1], ToolMessage):
            return "chatbot"
        last_tool = state["messages"][-1].name
        if last_tool in ["end_task", "respond_directly"]:
            return END
        else:
            return "chatbot"

    tools = [remote_execute, say_hello, error_analyzer, respond_directly, end_task, chat_with_human, web_search]

    global llm_with_tools
    llm_with_tools = llm.bind_tools(tools)
    tool_node = ToolNode(tools=tools)

    graph_builder.add_node("chatbot", chatbot)
    graph_builder.add_node("tools", tool_node)
    graph_builder.add_conditional_edges("chatbot", tools_condition)
    graph_builder.add_conditional_edges("tools", tools_to_next)
    graph_builder.add_edge(START, "chatbot")

    return graph_builder.compile(checkpointer=InMemorySaver())

# Initialize graph globally
graph = build_graph()
config = {"configurable": {"thread_id": "1"}}

# Function to update SSH settings
def update_ssh_settings(system_type, hostname, port, username, password):
    """Update SSH settings for the selected system"""
    global user_ssh_details
    
    system_key = system_type.lower()
    
    if system_key not in user_ssh_details:
        return f"❌ Invalid system type: {system_type}"
    
    user_ssh_details[system_key] = {
        "hostname": hostname,
        "port": int(port),
        "username": username,
        "password": password
    }
    
    logging.info(f"Updated SSH settings for {system_type}")
    return f"✅ {system_type} SSH settings saved successfully!"

# Function to get current SSH settings
def get_ssh_settings(system_type):
    """Get current SSH settings for the selected system"""
    system_key = system_type.lower()
    settings = user_ssh_details.get(system_key, {})
    
    return (
        settings.get("hostname", ""),
        settings.get("port", 22),
        settings.get("username", ""),
        settings.get("password", "")
    )

# Function to test SSH connection
def test_ssh_connection(system_type):
    """Test SSH connection for the selected system"""
    system_key = system_type.lower()
    ssh_details = user_ssh_details.get(system_key, {})
    
    if not ssh_details.get("hostname") or not ssh_details.get("username"):
        return "❌ Please configure SSH settings first!"
    
    ssh_client = paramiko.SSHClient()
    ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    
    try:
        ssh_client.connect(
            hostname=ssh_details["hostname"],
            port=ssh_details["port"],
            username=ssh_details["username"],
            password=ssh_details["password"],
            timeout=10
        )
        
        # Run a simple test command
        if system_key == "linux":
            stdin, stdout, stderr = ssh_client.exec_command("whoami", timeout=5)
        else:
            stdin, stdout, stderr = ssh_client.exec_command('powershell -Command "whoami"', timeout=5)
        
        output = stdout.read().decode().strip()
        
        ssh_client.close()
        
        return f"✅ Connection successful!\n\nConnected to: {ssh_details['hostname']}\nLogged in as: {output}"
        
    except Exception as e:
        return f"❌ Connection failed!\n\nError: {str(e)}\n\nPlease check your credentials and try again."

# Gradio chat function
def chat(message, history, system_choice):
    """Process user message and return response"""
    target_system = system_choice.lower()
    
    # Check if SSH is configured for the selected system
    ssh_details = user_ssh_details.get(target_system, {})
    if not ssh_details.get("hostname") or not ssh_details.get("username"):
        history = history or []
        history.append([message, f"⚠️ **SSH Configuration Required**\n\nPlease configure {system_choice} SSH settings in the **⚙️ SSH Settings** tab before using the agent.\n\nRequired:\n- Hostname/IP\n- Port\n- Username\n- Password"])
        return history
    
    # Create input message
    input_msg = HumanMessage(content=message)
    
    # Prepare input for graph with target system
    input_dict = {
        "messages": [input_msg],
        "target_system": target_system
    }
    
    response_text = ""
    
    try:
        # Stream events from graph
        events = graph.stream(input_dict, config, stream_mode="values")
        
        for event in events:
            if "messages" in event:
                msg = event["messages"][-1]
                
                if isinstance(msg, ToolMessage):
                    tool_name = msg.name if msg.name else "Status"
                    response_text += f"\n**🛠️ Tool Output ({tool_name}):**\n{msg.content}\n"
                    
                elif isinstance(msg, AIMessage):
                    if msg.content:
                        response_text += f"{msg.content}\n"
                    if hasattr(msg, 'tool_calls') and msg.tool_calls:
                        for tc in msg.tool_calls:
                            args_str = ", ".join([f"{k}: {v}" for k, v in tc['args'].items()])
                            response_text += f"\n**🔧 Calling Tool: {tc['name']}**\nInputs: {args_str}\n"
                else:
                    response_text += f"{msg.content}\n"
        
        final_response = response_text if response_text else "No response generated."
        
        # Return in correct format for Gradio Chatbot: history + new message
        history = history or []
        history.append([message, final_response])
        return history
        
    except Exception as e:
        logging.error(f"Error in chat: {str(e)}")
        history = history or []
        history.append([message, f"❌ Error: {str(e)}"])
        return history

# Create Gradio interface
with gr.Blocks(title="Server Administrator Agent", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🤖 Server Administrator Agent (Mr. Robot)
        
        An expert terminal assistant for pentesting, CTF challenges, and command-line operations.
        
        **Features:**
        - Execute commands on Linux and Windows systems via SSH
        - Web search for latest techniques and solutions
        - Auto-debugging and error analysis
        - Safe command execution with warnings
        
        ⚙️ **Please configure your SSH settings before starting!**
        """
    )
    
    with gr.Tabs():
        # Main Chat Tab
        with gr.Tab("💬 Chat"):
            with gr.Row():
                with gr.Column(scale=3):
                    system_dropdown = gr.Dropdown(
                        choices=["Linux", "Windows"],
                        value="Linux",
                        label="Target System",
                        info="Select which system to execute commands on"
                    )
                    
                    chatbot = gr.Chatbot(
                        height=500,
                        show_label=False,
                        avatar_images=(None, "🤖")
                    )
                    
                    with gr.Row():
                        msg = gr.Textbox(
                            placeholder="Type your message here...",
                            show_label=False,
                            scale=9
                        )
                        submit = gr.Button("Send", variant="primary", scale=1)
                    
                    clear = gr.Button("🗑️ Clear Chat")
                
                with gr.Column(scale=1):
                    gr.Markdown(
                        """
                        ### ℹ️ Tips
                        
                        - Ask about pentesting techniques
                        - Request command execution
                        - Get help with CTF challenges
                        - Search for latest security tools
                        
                        ### ⚠️ Safety
                        
                        - Risky commands require confirmation
                        - Always ensure you have permission
                        - Ethical hacking practices only
                        
                        ### 🔧 Quick Start
                        
                        1. Configure SSH settings in Settings tab
                        2. Select target system (Linux/Windows)
                        3. Start chatting!
                        """
                    )
        
        # SSH Settings Tab
        with gr.Tab("⚙️ SSH Settings"):
            gr.Markdown(
                """
                ## Configure SSH Connection Details
                
                Enter your server credentials below. These credentials are stored only during your session.
                """
            )
            
            with gr.Row():
                with gr.Column():
                    ssh_system_selector = gr.Radio(
                        choices=["Linux", "Windows"],
                        value="Linux",
                        label="Select System to Configure",
                        info="Choose which system's SSH settings you want to update"
                    )
            
            with gr.Row():
                with gr.Column():
                    ssh_hostname = gr.Textbox(
                        label="SSH Hostname/IP",
                        placeholder="e.g., 192.168.1.100 or server.example.com",
                        info="The server IP address or hostname"
                    )
                    
                    ssh_port = gr.Number(
                        label="SSH Port",
                        value=22,
                        precision=0,
                        info="SSH port (default: 22)"
                    )
                    
                    ssh_username = gr.Textbox(
                        label="SSH Username",
                        placeholder="e.g., root or administrator",
                        info="Your SSH username"
                    )
                    
                    ssh_password = gr.Textbox(
                        label="SSH Password",
                        placeholder="Enter password",
                        type="password",
                        info="Your SSH password (not stored permanently)"
                    )
                    
                    with gr.Row():
                        save_ssh_btn = gr.Button("💾 Save SSH Settings", variant="primary", scale=2)
                        load_ssh_btn = gr.Button("🔄 Load Current Settings", scale=1)
                    
                    ssh_status = gr.Textbox(
                        label="Status",
                        interactive=False,
                        placeholder="Configure and save your SSH settings..."
                    )
            
            gr.Markdown(
                """
                ### 📝 Notes:
                - Settings are stored only during your session
                - You can configure both Linux and Windows systems
                - Switch between systems using the radio button above
                - Test your connection by sending a simple command in the Chat tab
                
                ### 🔒 Security:
                - Credentials are not saved to disk
                - Use this on trusted networks only
                - Consider using SSH keys instead of passwords for production
                """
            )
    
    # Event handlers for Chat tab
    submit.click(
        chat,
        inputs=[msg, chatbot, system_dropdown],
        outputs=chatbot
    ).then(
        lambda: gr.update(value=""),
        None,
        msg
    )
    
    msg.submit(
        chat,
        inputs=[msg, chatbot, system_dropdown],
        outputs=chatbot
    ).then(
        lambda: gr.update(value=""),
        None,
        msg
    )
    
    clear.click(lambda: [], None, chatbot, queue=False)
    
    # Event handlers for SSH Settings tab
    save_ssh_btn.click(
        update_ssh_settings,
        inputs=[ssh_system_selector, ssh_hostname, ssh_port, ssh_username, ssh_password],
        outputs=ssh_status
    )
    
    load_ssh_btn.click(
        get_ssh_settings,
        inputs=ssh_system_selector,
        outputs=[ssh_hostname, ssh_port, ssh_username, ssh_password]
    )
    
    # Auto-load settings when system selector changes
    ssh_system_selector.change(
        get_ssh_settings,
        inputs=ssh_system_selector,
        outputs=[ssh_hostname, ssh_port, ssh_username, ssh_password]
    )

# Launch the app
if __name__ == "__main__":
    demo.queue()
    demo.launch()