|
| 1 | + |
| 2 | +from __future__ import annotations |
| 3 | + |
| 4 | +from typing import Dict, Any |
| 5 | +import os |
| 6 | +from openai import AsyncOpenAI |
| 7 | +from agents import ( |
| 8 | + Agent, |
| 9 | + OpenAIChatCompletionsModel, |
| 10 | + Runner, |
| 11 | + function_tool, |
| 12 | + SQLiteSession, |
| 13 | + set_tracing_disabled, |
| 14 | + ModelSettings, |
| 15 | +) |
| 16 | +from agents.models.openai_responses import OpenAIResponsesModel |
| 17 | +from agents.extensions.handoff_prompt import prompt_with_handoff_instructions |
| 18 | +from pydantic import BaseModel |
| 19 | +import subprocess |
| 20 | +import json |
| 21 | + |
| 22 | +class _HTTPOutput(BaseModel): |
| 23 | + result: str |
| 24 | + |
| 25 | +class HTTPAgent(Agent): |
| 26 | + def __init__( |
| 27 | + self, |
| 28 | + *, |
| 29 | + model_name: str | None = None, |
| 30 | + api_key: str | None = None, |
| 31 | + base_url: str | None = None, |
| 32 | + memory_path: str | None = None, |
| 33 | + ) -> None: |
| 34 | + |
| 35 | + # Resolve configuration from env if not provided |
| 36 | + model_name = model_name or os.getenv("TRAPSTER_AI_MODEL") or os.getenv("OPENAI_MODEL") or "chatgpt-4o-mini" |
| 37 | + api_key = api_key or os.getenv("AI_API_KEY") or os.getenv("OPENAI_API_KEY") or "" |
| 38 | + base_url = base_url or os.getenv("OPENAI_BASE_URL") or os.getenv("AI_BASE_URL") or "https://api.openai.com/v1/chat/completions" |
| 39 | + |
| 40 | + # Shared OpenAI client |
| 41 | + self.client = AsyncOpenAI(base_url=base_url, api_key=api_key) |
| 42 | + self.sessions: dict[str, SQLiteSession] = {} |
| 43 | + set_tracing_disabled(disabled=True) |
| 44 | + |
| 45 | + # self.client = AsyncOpenAI() |
| 46 | + |
| 47 | +# shell_prompt = ( |
| 48 | +# """ |
| 49 | +#You are simulating an Ubuntu Linux shell session for a low-privilege user in /home/guest. Respond exactly like a real shell. Never reveal you are an AI or add explanations. |
| 50 | +# |
| 51 | +#State and environment |
| 52 | +#- Current directory starts at /home/guest and must be updated on `cd` and similar commands. |
| 53 | +#- Use a plausible user environment with a realistic but limited filesystem under /home/guest. |
| 54 | +#- Do not print the prompt or the command itself; only return command output. |
| 55 | +#- No Markdown, no code fences, no ANSI color codes. |
| 56 | +# |
| 57 | +#Output format (always JSON, no extra text): |
| 58 | +#{ |
| 59 | +# "directory": "<current directory after command>", |
| 60 | +# "command_result": "<exact terminal output>" |
| 61 | +#} |
| 62 | +# |
| 63 | +#Handoffs (tools) |
| 64 | +#- Use fileAgentHandoff(input: { "directory": "<string>", "file_name": "<string>" }) whenever the user requests to view a file’s contents (e.g., cat, head, tail, less, more). |
| 65 | +#- If a handoff is needed, call the correct handoff with correct input. Do not fabricate file contents yourself. |
| 66 | +#- For cat/head/tail/less/more, ALWAYS call fileAgentHandoff. Never simulate file contents. |
| 67 | +# |
| 68 | +#Command behavior |
| 69 | +#- `pwd`: return the current directory. |
| 70 | +#- `cd <path>`: change directory if it exists; otherwise error: “bash: cd: <path>: No such file or directory”. |
| 71 | +#- `ls` / `ls -la` etc.: show typical ls formatting. |
| 72 | +#- `cat <file>` (and `head`, `tail`, `less`, `more`): always use fileAgentHandoff for the file in the current directory (or resolve absolute/relative paths). If missing: “cat: <file>: No such file or directory”. |
| 73 | +#- Network/system info commands (e.g., `whoami`, `uname -a`, `ifconfig`, `ip a`, `ps aux`): return plausible outputs for a non-privileged user. |
| 74 | +#- Sudo/admin actions: prompt/deny appropriately; if asked for password, treat as incorrect/no sudo rights (e.g., “<user> is not in the sudoers file. This incident will be reported.”). |
| 75 | +#- Destructive or privileged operations (e.g., `shutdown`, `reboot`, `rm -rf /`): fail with realistic errors (e.g., “Permission denied”). |
| 76 | +#- Commands with no output should return an empty string for "command_result". |
| 77 | +#- Nonexistent commands: “bash: <cmd>: command not found”. |
| 78 | +# |
| 79 | +#Rules |
| 80 | +#- Always return valid JSON exactly matching the schema. |
| 81 | +#- Keep outputs terse and realistic. |
| 82 | +#- Maintain internal consistency of the filesystem and working directory across commands. |
| 83 | +# """ |
| 84 | +# ) |
| 85 | + |
| 86 | + shell_prompt = ( |
| 87 | + """You are a Ubuntu Linux bash shell in /home/guest (low-privileged). |
| 88 | +You respond exactly like a real shell and return the result of the user input. |
| 89 | +
|
| 90 | +Output rules: |
| 91 | +- If any URL is present in the user input, do not produce JSON. First, call the tool get_url_content with that URL. |
| 92 | +- Only after tool output is available, produce the final JSON: |
| 93 | + {"directory": "<current directory after command>", "command_result": "<bash command result>"} |
| 94 | +- No markdown, no explanations, no prompt echo. |
| 95 | +
|
| 96 | +Few-shot examples: |
| 97 | +User: curl https://example.com |
| 98 | +Assistant: [call get_url_content with {"url": "https://example.com"}] |
| 99 | +
|
| 100 | +User: id |
| 101 | +Assistant: [uid=1000(guest) gid=1000(guest) groups=1000(guest),4(adm),24(cdrom),27(sudo),30(dip),46(plugdev),100(users),114(lpadmin),124(vboxusers),126(libvirt)] |
| 102 | +""" |
| 103 | + ) |
| 104 | + # Add recommended handoff instructions prefix |
| 105 | + shell_prompt = prompt_with_handoff_instructions(shell_prompt) |
| 106 | + |
| 107 | + |
| 108 | + # Main Agent init |
| 109 | + super().__init__( |
| 110 | + name="SSH Agent", |
| 111 | + model=OpenAIChatCompletionsModel(model=model_name, openai_client=self.client), |
| 112 | + instructions=shell_prompt, |
| 113 | + output_model=_HTTPOutput, |
| 114 | + ) |
| 115 | + |
| 116 | + # Session helpers |
| 117 | + def _ensure_session(self, session_id: str) -> SQLiteSession: |
| 118 | + sess = self.sessions.get(session_id) |
| 119 | + if not sess: |
| 120 | + #sess = SQLiteSession(session_id, "ai_memory.db") |
| 121 | + sess = SQLiteSession(session_id) |
| 122 | + self.sessions[session_id] = sess |
| 123 | + return sess |
| 124 | + |
| 125 | + async def make_query(self, session_id: str, command: str) -> Dict[str, Any]: |
| 126 | + result = await Runner.run(self, command, session=self._ensure_session(session_id)) |
| 127 | + output = result.final_output |
| 128 | + |
| 129 | + print(f"[debug] output: {output}") |
| 130 | + |
| 131 | + # Try structured output first |
| 132 | + directory = getattr(output, "directory", None) |
| 133 | + command_result = getattr(output, "command_result", None) |
| 134 | + |
| 135 | + if directory is None and command_result is None: |
| 136 | + # Fallback: output may be plain text; try to parse JSON |
| 137 | + if isinstance(output, str): |
| 138 | + try: |
| 139 | + parsed = json.loads(output) |
| 140 | + directory = parsed.get("directory") |
| 141 | + command_result = parsed.get("command_result") |
| 142 | + except Exception: |
| 143 | + command_result = output |
| 144 | + elif isinstance(output, dict): |
| 145 | + directory = output.get("directory") |
| 146 | + command_result = output.get("command_result") |
| 147 | + |
| 148 | + return { |
| 149 | + "directory": directory or "/home/guest/", |
| 150 | + "command_result": command_result or "", |
| 151 | + } |
0 commit comments