ssh_agent.py

#!/usr/bin/env python3
import argparse
import asyncio
import getpass
import json
import os
import re
import sys
import time
from dataclasses import dataclass
from enum import Enum
from typing import Any, Dict, List, Optional

import paramiko

from textual.app import App, ComposeResult
from textual.containers import Horizontal, Vertical
from textual.widgets import Header, Footer, Input, Static, RichLog
from textual.reactive import reactive

# ----------------------------
# Safety / Command policy
# ----------------------------

DANGEROUS_PATTERNS = [
    r"\brm\s+-rf\b",
    r"\brm\s+-r\b",
    r"\brm\s+\-f\b",
    r"\bdd\s+if=",
    r"\bmkfs\.",
    r"\bshutdown\b",
    r"\breboot\b",
    r"\bpoweroff\b",
    r"\binit\s+0\b",
    r"\binit\s+6\b",
    r"\bhalt\b",
    r"\bapt\s+remove\b",
    r"\bapt\s+purge\b",
    r"\bdpkg\s+-r\b",
    r"\buserdel\b",
    r"\bchmod\s+-R\b",
    r"\bchown\s+-R\b",
    r"\btruncate\b",
    r"\b:\s*>\s*/",  # shell truncation
]


def is_dangerous(cmd: str) -> bool:
    c = cmd.strip()
    if not c:
        return False
    for pat in DANGEROUS_PATTERNS:
        if re.search(pat, c, re.IGNORECASE):
            return True
    return False


def normalize_cmd(cmd: str) -> str:
    return cmd.strip()


# ----------------------------
# SSH client wrapper
# ----------------------------

@dataclass
class SSHResult:
    command: str
    exit_status: int
    stdout: str
    stderr: str
    duration_sec: float


class SSHRunner:
    def __init__(self, host: str, username: str, password: Optional[str], port: int = 22, timeout: int = 15):
        self.host = host
        self.username = username
        self.password = password
        self.port = port
        self.timeout = timeout
        self.client: Optional[paramiko.SSHClient] = None

    def connect(self) -> None:
        c = paramiko.SSHClient()
        c.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        c.connect(
            hostname=self.host,
            port=self.port,
            username=self.username,
            password=self.password,
            timeout=self.timeout,
            auth_timeout=self.timeout,
            banner_timeout=self.timeout,
            look_for_keys=True,
            allow_agent=True,
        )
        self.client = c

    def close(self) -> None:
        if self.client:
            self.client.close()
            self.client = None

    def run(self, command: str) -> SSHResult:
        if not self.client:
            raise RuntimeError("SSH not connected")

        cmd = normalize_cmd(command)
        start = time.time()
        stdin, stdout, stderr = self.client.exec_command(cmd)
        out = stdout.read().decode(errors="replace")
        err = stderr.read().decode(errors="replace")
        exit_status = stdout.channel.recv_exit_status()
        dur = time.time() - start
        return SSHResult(cmd, exit_status, out, err, dur)


# ----------------------------
# Tool schema (Gemini-safe)
# ----------------------------

RUN_SSH_TOOL = {
    "type": "function",
    "function": {
        "name": "run_ssh",
        "description": "Run a shell command on the remote host over SSH and return stdout/stderr/exit code.",
        "parameters": {
            "type": "object",
            "properties": {
                "command": {"type": "string", "description": "Shell command to run on the remote host."}
            },
            "required": ["command"],
        },
    },
}


# ----------------------------
# LLM Providers
# ----------------------------

class LLMProviderBase:
    async def chat(self, messages: List[Dict[str, Any]], tools: List[Dict[str, Any]]) -> Dict[str, Any]:
        raise NotImplementedError


class OpenAIProvider(LLMProviderBase):
    def __init__(self, model: str):
        from openai import OpenAI
        self.client = OpenAI()
        self.model = model

    async def chat(self, messages: List[Dict[str, Any]], tools: List[Dict[str, Any]]) -> Dict[str, Any]:
        def _call():
            return self.client.chat.completions.create(
                model=self.model,
                messages=messages,
                tools=tools,
                tool_choice="auto",
            )

        resp = await asyncio.to_thread(_call)
        choice = resp.choices[0].message

        out = {"role": "assistant",
               "content": choice.content or "", "tool_calls": []}
        if choice.tool_calls:
            for tc in choice.tool_calls:
                out["tool_calls"].append({
                    "id": tc.id,
                    "name": tc.function.name,
                    "arguments": tc.function.arguments,
                })
        return out


class GeminiProvider(LLMProviderBase):
    def __init__(self, model: str):
        from google import genai
        api_key = os.environ.get("GEMINI_API_KEY", "")
        if not api_key:
            raise RuntimeError("GEMINI_API_KEY is not set")
        self.client = genai.Client(api_key=api_key)
        self.model = model

    async def chat(self, messages: List[Dict[str, Any]], tools: List[Dict[str, Any]]) -> Dict[str, Any]:
        from google.genai import types

        def to_gemini_contents(msgs: List[Dict[str, Any]]):
            contents = []
            for m in msgs:
                role = m["role"]
                if role == "system":
                    contents.append(types.Content(role="user", parts=[
                        types.Part(text=f"[SYSTEM]\n{m['content']}")
                    ]))
                elif role == "user":
                    contents.append(types.Content(role="user", parts=[
                        types.Part(text=m.get("content", ""))
                    ]))
                elif role == "assistant":
                    txt = m.get("content", "")
                    if txt:
                        contents.append(types.Content(role="model", parts=[
                            types.Part(text=txt)
                        ]))
                elif role == "tool":
                    # Tool results are provided as user text. Gemini sometimes echoes these back;
                    # we suppress that in the UI.
                    contents.append(types.Content(role="user", parts=[
                        types.Part(
                            text=f"[TOOL RESULT]\n{m.get('content','')}")
                    ]))
            return contents

        function_decls = []
        for t in tools:
            if t.get("type") != "function":
                continue
            fn = t["function"]
            function_decls.append(types.FunctionDeclaration(
                name=fn["name"],
                description=fn.get("description", ""),
                parameters=fn.get(
                    "parameters", {"type": "object", "properties": {}}),
            ))

        tool = types.Tool(function_declarations=function_decls)

        def _call():
            return self.client.models.generate_content(
                model=self.model,
                contents=to_gemini_contents(messages),
                config=types.GenerateContentConfig(
                    tools=[tool],
                    temperature=0.2,
                ),
            )

        resp = await asyncio.to_thread(_call)

        out = {"role": "assistant", "content": "", "tool_calls": []}

        try:
            out["content"] = resp.text or ""
        except Exception:
            out["content"] = ""

        # Extract function calls
        try:
            c = resp.candidates[0]
            parts = c.content.parts
            for p in parts:
                if hasattr(p, "function_call") and p.function_call:
                    fc = p.function_call
                    out["tool_calls"].append({
                        "id": f"gemini_{len(out['tool_calls'])+1}",
                        "name": fc.name,
                        "arguments": json.dumps(fc.args or {}),
                    })
        except Exception:
            pass

        return out


# ----------------------------
# Agent core
# ----------------------------

SYSTEM_PROMPT = """You are an SSH troubleshooting agent.
You can run commands on the remote host using the run_ssh tool.
Your job is to diagnose issues, interpret output, and propose fixes.

Rules:
- Prefer safe read-only commands first.
- Be concise and practical.
- When making changes, explain why.
- Avoid destructive commands. If a command might be dangerous, ask for confirmation.
- Do NOT echo tool results back to the user; only interpret them.
"""


class AgentState(str, Enum):
    WAITING_INPUT = "WAITING_INPUT"
    THINKING = "THINKING"
    CONFIRMING = "CONFIRMING"
    RUNNING = "RUNNING"
    RESPONDING = "RESPONDING"
    ERROR = "ERROR"


class Agent:
    def __init__(self, provider: LLMProviderBase, ssh: SSHRunner, app_ref: "SSHAgentApp"):
        self.provider = provider
        self.ssh = ssh
        self.app = app_ref
        self.messages: List[Dict[str, Any]] = [
            {"role": "system", "content": SYSTEM_PROMPT}]

    async def handle_user(self, text: str) -> None:
        self.messages.append({"role": "user", "content": text})
        await self.step_loop()

    async def step_loop(self, max_steps: int = 12) -> None:
        for _ in range(max_steps):
            self.app.set_agent_state(AgentState.THINKING)

            try:
                reply = await self.provider.chat(self.messages, tools=[RUN_SSH_TOOL])
            except Exception as e:
                self.app.set_agent_state(AgentState.ERROR)
                self.app.add_output(f"[LLM ERROR] {e}")
                self.app.set_agent_state(AgentState.WAITING_INPUT)
                return

            content = (reply.get("content") or "").strip()
            tool_calls = reply.get("tool_calls") or []

            # Suppress Gemini echo of tool results
            if content:
                if not content.lstrip().startswith("[TOOL RESULT]"):
                    self.app.add_chat(f"AI: {content}")
                self.messages.append({"role": "assistant", "content": content})

            if not tool_calls:
                self.app.set_agent_state(AgentState.WAITING_INPUT)
                return

            for tc in tool_calls:
                name = tc["name"]
                raw_args = tc["arguments"]

                try:
                    args = json.loads(raw_args) if isinstance(
                        raw_args, str) else raw_args
                except Exception:
                    args = {"command": str(raw_args)}

                if name != "run_ssh":
                    tool_result = f"Unknown tool: {name}"
                    self.messages.append(
                        {"role": "tool", "content": tool_result})
                    self.app.add_output(tool_result)
                    continue

                cmd = normalize_cmd(args.get("command", ""))

                # Safety gating / confirm mode
                if is_dangerous(cmd) or (not self.app.auto_run):
                    self.app.set_agent_state(AgentState.CONFIRMING)
                    self.app.pending_command = cmd
                    self.app.add_chat(
                        f"AI wants to run:\n$ {cmd}\n\nType 'y' to run, 'n' to skip."
                    )
                    return

                await self._run_and_record(cmd)

        self.app.add_chat(
            "AI: Stopping (max steps reached). Type 'go' to continue.")
        self.app.set_agent_state(AgentState.WAITING_INPUT)

    async def _run_and_record(self, cmd: str) -> None:
        self.app.set_agent_state(AgentState.RUNNING)

        self.app.add_output(f"$ {cmd}")
        res: SSHResult = await asyncio.to_thread(self.ssh.run, cmd)

        # Human-friendly output (what you want to see)
        if res.stdout.strip():
            self.app.add_output(res.stdout.rstrip())
        if res.stderr.strip():
            self.app.add_output("[stderr]")
            self.app.add_output(res.stderr.rstrip())

        self.app.add_output(
            f"[exit={res.exit_status}  time={res.duration_sec:.2f}s]")

        tool_payload = {
            "command": res.command,
            "exit_status": res.exit_status,
            "stdout": res.stdout,
            "stderr": res.stderr,
            "duration_sec": res.duration_sec,
        }

        # Feed JSON to model
        self.messages.append(
            {"role": "tool", "content": json.dumps(tool_payload)})

        # Optional debug display
        if getattr(self.app, "show_tool_json", False):
            self.app.add_output("[TOOL RESULT JSON]")
            self.app.add_output(json.dumps(tool_payload, indent=2))

        self.app.set_agent_state(AgentState.RESPONDING)


# ----------------------------
# Textual UI
# ----------------------------

class ModeBar(Static):
    pass


class SSHAgentApp(App):
    CSS = """
    Screen { layout: vertical; }
    #body { height: 1fr; }
    #left { width: 40%; }
    #right { width: 60%; }
    RichLog { border: solid gray; }
    #chatlog { height: 1fr; }
    #outlog { height: 1fr; }
    #inputrow { height: auto; }
    """

    BINDINGS = [
        ("f2", "toggle_mode", "Auto/Confirm"),
        ("f3", "shrink_left", "Shrink left pane"),
        ("f4", "grow_left", "Grow left pane"),
        ("f5", "reset_split", "Reset split"),
        ("f6", "toggle_tool_json", "Tool JSON"),
        ("ctrl+c", "quit", "Quit"),
    ]

    auto_run: bool = reactive(False)
    show_tool_json: bool = reactive(False)

    def __init__(self, ssh_runner: SSHRunner, agent: Optional[Agent]):
        super().__init__()
        self.ssh = ssh_runner
        self.agent = agent
        self.pending_command: Optional[str] = None

        self.chatlog: Optional[RichLog] = None
        self.outlog: Optional[RichLog] = None
        self.input: Optional[Input] = None
        self.modebar: Optional[ModeBar] = None

        self.left_pct: int = 40
        self.right_pct: int = 60

        self.agent_state: str = AgentState.WAITING_INPUT.value

    def compose(self) -> ComposeResult:
        yield Header()
        self.modebar = ModeBar()
        yield self.modebar

        with Horizontal(id="body"):
            with Vertical(id="left"):
                self.chatlog = RichLog(
                    id="chatlog", wrap=True, highlight=True, markup=False)
                yield self.chatlog
                self.input = Input(
                    placeholder="Type message. 'go' to start. 'y'/'n' to confirm.",
                    id="inputrow",
                )
                yield self.input
            with Vertical(id="right"):
                self.outlog = RichLog(
                    id="outlog", wrap=True, highlight=True, markup=False)
                yield self.outlog

        yield Footer()

    def on_mount(self) -> None:
        self.refresh_modebar()
        self.add_chat("Connected. Type 'go' or describe the issue.")
        self.query_one(Input).focus()

    def set_agent_state(self, state: AgentState) -> None:
        self.agent_state = state.value if isinstance(
            state, AgentState) else str(state)
        self.refresh_modebar()

    def refresh_modebar(self) -> None:
        mode = "AUTO-RUN" if self.auto_run else "CONFIRM"
        dbg = "ON" if self.show_tool_json else "OFF"
        if self.modebar:
            self.modebar.update(
                f"Mode: {mode}   |   Status: {self.agent_state}   |   ToolJSON: {dbg}   "
                f"(F2 Auto/Confirm, F3/F4 resize, F5 reset, F6 tool json)"
            )

    def _apply_split(self) -> None:
        left = self.query_one("#left")
        right = self.query_one("#right")
        left.styles.width = f"{self.left_pct}%"
        right.styles.width = f"{self.right_pct}%"

    def action_shrink_left(self) -> None:
        self.left_pct = max(20, self.left_pct - 5)
        self.right_pct = 100 - self.left_pct
        self._apply_split()

    def action_grow_left(self) -> None:
        self.left_pct = min(80, self.left_pct + 5)
        self.right_pct = 100 - self.left_pct
        self._apply_split()

    def action_reset_split(self) -> None:
        self.left_pct = 40
        self.right_pct = 60
        self._apply_split()

    def action_toggle_mode(self) -> None:
        self.auto_run = not self.auto_run
        self.refresh_modebar()
        self.add_chat(f"[Mode changed] auto_run={self.auto_run}")

    def action_toggle_tool_json(self) -> None:
        self.show_tool_json = not self.show_tool_json
        self.refresh_modebar()
        self.add_chat(f"[Tool JSON] show_tool_json={self.show_tool_json}")

    def add_chat(self, text: str) -> None:
        if self.chatlog:
            self.chatlog.write(text)

    def add_output(self, text: str) -> None:
        if self.outlog:
            self.outlog.write(text)

    async def on_input_submitted(self, event: Input.Submitted) -> None:
        text = event.value.strip()
        event.input.value = ""
        if not text:
            return

        # Confirmation flow
        if self.pending_command and self.agent:
            if text.lower() in ("y", "yes"):
                cmd = self.pending_command
                self.pending_command = None
                await self.agent._run_and_record(cmd)
                await self.agent.step_loop()
                self.set_agent_state(AgentState.WAITING_INPUT)
                return
            elif text.lower() in ("n", "no"):
                self.add_chat("Skipped.")
                self.pending_command = None
                self.agent.messages.append(
                    {"role": "tool", "content": json.dumps({"skipped_command": True})})
                await self.agent.step_loop()
                self.set_agent_state(AgentState.WAITING_INPUT)
                return

        self.add_chat(f"You: {text}")
        if self.agent:
            self.set_agent_state(AgentState.THINKING)
            await self.agent.handle_user(text)
            self.set_agent_state(AgentState.WAITING_INPUT)


# ----------------------------
# Main
# ----------------------------

def build_provider(provider_name: str, model: Optional[str]) -> LLMProviderBase:
    provider_name = provider_name.lower().strip()
    if provider_name == "openai":
        m = model or os.environ.get("OPENAI_MODEL", "gpt-5-mini")
        return OpenAIProvider(model=m)
    if provider_name == "gemini":
        m = model or os.environ.get("GEMINI_MODEL", "gemini-2.0-flash")
        return GeminiProvider(model=m)
    raise ValueError(f"Unknown provider: {provider_name}")


def main():
    ap = argparse.ArgumentParser(
        description="SSH Agent (single-file prototype)")
    ap.add_argument("--host", required=True, help="SSH host (IP or hostname)")
    ap.add_argument("--port", type=int, default=22,
                    help="SSH port (default 22)")
    ap.add_argument("--username", required=True, help="SSH username")
    ap.add_argument("--password", default=None,
                    help="SSH password (not recommended; use --ask-password)")
    ap.add_argument("--ask-password", action="store_true",
                    help="Prompt for SSH password")
    ap.add_argument("--provider", choices=["openai", "gemini"],
                    default=os.environ.get("LLM_PROVIDER", "gemini"))
    ap.add_argument("--model", default=None, help="Model name override")
    args = ap.parse_args()

    if args.ask_password:
        pw = getpass.getpass("SSH password: ")
    else:
        pw = args.password or os.environ.get("SSH_PASSWORD")

    ssh = SSHRunner(host=args.host, username=args.username,
                    password=pw, port=args.port)

    try:
        ssh.connect()
    except Exception as e:
        print(f"SSH connect failed: {e}", file=sys.stderr)
        sys.exit(2)

    try:
        provider = build_provider(args.provider, args.model)
    except Exception as e:
        print(f"Provider init failed: {e}", file=sys.stderr)
        ssh.close()
        sys.exit(3)

    app = SSHAgentApp(ssh_runner=ssh, agent=None)
    agent = Agent(provider=provider, ssh=ssh, app_ref=app)
    app.agent = agent

    try:
        app.run()
    finally:
        ssh.close()


if __name__ == "__main__":
    main()