0796e45412
- SystemPipeline: Entry point for console/webhook/API messages - TextInput/TextOutput: Generic primitives for human text I/O - Server: WebSocket "send" command routes through SystemPipeline - Console: @target message now injects into pipeline Flow: Console → WebSocket → SystemPipeline → XML envelope → pump Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
276 lines
8.5 KiB
Python
276 lines
8.5 KiB
Python
"""
|
|
Console client that connects to the agent server.
|
|
|
|
Provides SSH-style login with username/password authentication.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import getpass
|
|
import json
|
|
import sys
|
|
from typing import Optional
|
|
|
|
try:
|
|
import aiohttp
|
|
AIOHTTP_AVAILABLE = True
|
|
except ImportError:
|
|
AIOHTTP_AVAILABLE = False
|
|
|
|
try:
|
|
from prompt_toolkit import PromptSession
|
|
from prompt_toolkit.history import InMemoryHistory
|
|
from prompt_toolkit.styles import Style
|
|
PROMPT_TOOLKIT_AVAILABLE = True
|
|
except ImportError:
|
|
PROMPT_TOOLKIT_AVAILABLE = False
|
|
|
|
|
|
DEFAULT_HOST = "127.0.0.1"
|
|
DEFAULT_PORT = 8765
|
|
MAX_LOGIN_ATTEMPTS = 3
|
|
|
|
|
|
class ConsoleClient:
|
|
"""
|
|
Text-based console client for the agent server.
|
|
|
|
Usage:
|
|
client = ConsoleClient()
|
|
asyncio.run(client.run())
|
|
"""
|
|
|
|
def __init__(self, host: str = DEFAULT_HOST, port: int = DEFAULT_PORT):
|
|
self.host = host
|
|
self.port = port
|
|
self.base_url = f"http://{host}:{port}"
|
|
self.ws_url = f"ws://{host}:{port}/ws"
|
|
self.token: Optional[str] = None
|
|
self.username: Optional[str] = None
|
|
self.session: Optional[aiohttp.ClientSession] = None
|
|
self.ws: Optional[aiohttp.ClientWebSocketResponse] = None
|
|
self.running = False
|
|
|
|
async def login(self) -> bool:
|
|
"""
|
|
Perform SSH-style login.
|
|
|
|
Returns:
|
|
True if login successful, False otherwise
|
|
"""
|
|
print(f"Connecting to {self.host}:{self.port}...")
|
|
|
|
for attempt in range(1, MAX_LOGIN_ATTEMPTS + 1):
|
|
try:
|
|
username = input("Username: ")
|
|
password = getpass.getpass("Password: ")
|
|
except (EOFError, KeyboardInterrupt):
|
|
print("\nLogin cancelled.")
|
|
return False
|
|
|
|
if not username or not password:
|
|
print("Username and password required.")
|
|
continue
|
|
|
|
try:
|
|
async with aiohttp.ClientSession() as session:
|
|
async with session.post(
|
|
f"{self.base_url}/auth/login",
|
|
json={"username": username, "password": password},
|
|
) as resp:
|
|
data = await resp.json()
|
|
|
|
if resp.status == 200:
|
|
self.token = data["token"]
|
|
self.username = username
|
|
print(f"Welcome, {username}!")
|
|
return True
|
|
else:
|
|
error = data.get("error", "Authentication failed")
|
|
remaining = MAX_LOGIN_ATTEMPTS - attempt
|
|
if remaining > 0:
|
|
print(f"{error}. {remaining} attempt(s) remaining.")
|
|
else:
|
|
print(f"{error}. No attempts remaining.")
|
|
except aiohttp.ClientError as e:
|
|
print(f"Connection error: {e}")
|
|
return False
|
|
|
|
return False
|
|
|
|
async def connect_ws(self) -> bool:
|
|
"""Connect to WebSocket after authentication."""
|
|
if not self.token:
|
|
return False
|
|
|
|
try:
|
|
self.session = aiohttp.ClientSession(
|
|
headers={"Authorization": f"Bearer {self.token}"}
|
|
)
|
|
self.ws = await self.session.ws_connect(self.ws_url)
|
|
|
|
# Wait for connected message
|
|
msg = await self.ws.receive_json()
|
|
if msg.get("type") == "connected":
|
|
return True
|
|
return False
|
|
except Exception as e:
|
|
print(f"WebSocket connection failed: {e}")
|
|
return False
|
|
|
|
async def send_command(self, cmd: str) -> Optional[dict]:
|
|
"""Send a command via WebSocket and get response."""
|
|
if not self.ws:
|
|
return None
|
|
|
|
await self.ws.send_json(cmd)
|
|
return await self.ws.receive_json()
|
|
|
|
def print_help(self):
|
|
"""Print available commands."""
|
|
print("""
|
|
Available commands:
|
|
/help - Show this help
|
|
/status - Show server status
|
|
/listeners - List available targets
|
|
/targets - Alias for /listeners
|
|
/quit - Disconnect and exit
|
|
|
|
Send messages:
|
|
@target message - Send message to a target listener
|
|
Example: @greeter Hello there!
|
|
""")
|
|
|
|
async def handle_command(self, line: str) -> bool:
|
|
"""
|
|
Handle a command line.
|
|
|
|
Returns:
|
|
False if should quit, True otherwise
|
|
"""
|
|
line = line.strip()
|
|
if not line:
|
|
return True
|
|
|
|
if line == "/help":
|
|
self.print_help()
|
|
elif line == "/quit" or line == "/exit":
|
|
return False
|
|
elif line == "/status":
|
|
resp = await self.send_command({"type": "status"})
|
|
if resp:
|
|
threads = resp.get("threads", 0)
|
|
print(f"Active threads: {threads}")
|
|
elif line == "/listeners" or line == "/targets":
|
|
resp = await self.send_command({"type": "listeners"})
|
|
if resp:
|
|
listeners = resp.get("listeners", [])
|
|
if listeners:
|
|
print("Available targets:")
|
|
for name in listeners:
|
|
print(f" - {name}")
|
|
else:
|
|
print("No targets available (pipeline not running)")
|
|
elif line.startswith("/"):
|
|
print(f"Unknown command: {line}")
|
|
elif line.startswith("@"):
|
|
# Send message to target: @target message
|
|
resp = await self.send_command({"type": "send", "raw": line})
|
|
if resp:
|
|
if resp.get("type") == "sent":
|
|
thread_id = resp.get("thread_id", "")[:8]
|
|
target = resp.get("target", "unknown")
|
|
print(f"Sent to {target} (thread: {thread_id}...)")
|
|
elif resp.get("type") == "error":
|
|
print(f"Error: {resp.get('error')}")
|
|
else:
|
|
print("Use @target message to send. Example: @greeter Hello!")
|
|
print("Type /listeners to see available targets.")
|
|
|
|
return True
|
|
|
|
async def run(self):
|
|
"""Main client loop."""
|
|
if not AIOHTTP_AVAILABLE:
|
|
print("Error: aiohttp not installed")
|
|
sys.exit(1)
|
|
|
|
# Login
|
|
if not await self.login():
|
|
print("Authentication failed.")
|
|
sys.exit(1)
|
|
|
|
# Connect WebSocket
|
|
if not await self.connect_ws():
|
|
print("Failed to connect to server.")
|
|
sys.exit(1)
|
|
|
|
print("Connected. Type /help for commands, /quit to exit.")
|
|
|
|
self.running = True
|
|
|
|
try:
|
|
if PROMPT_TOOLKIT_AVAILABLE:
|
|
await self._run_prompt_toolkit()
|
|
else:
|
|
await self._run_simple()
|
|
finally:
|
|
await self.cleanup()
|
|
|
|
async def _run_prompt_toolkit(self):
|
|
"""Run with prompt_toolkit for better UX."""
|
|
style = Style.from_dict({
|
|
"prompt": "ansicyan bold",
|
|
})
|
|
|
|
session = PromptSession(
|
|
history=InMemoryHistory(),
|
|
style=style,
|
|
)
|
|
|
|
while self.running:
|
|
try:
|
|
line = await asyncio.get_event_loop().run_in_executor(
|
|
None,
|
|
lambda: session.prompt(f"{self.username}> ")
|
|
)
|
|
if not await self.handle_command(line):
|
|
break
|
|
except (EOFError, KeyboardInterrupt):
|
|
break
|
|
|
|
async def _run_simple(self):
|
|
"""Run with simple input (fallback)."""
|
|
while self.running:
|
|
try:
|
|
line = input(f"{self.username}> ")
|
|
if not await self.handle_command(line):
|
|
break
|
|
except (EOFError, KeyboardInterrupt):
|
|
break
|
|
|
|
async def cleanup(self):
|
|
"""Clean up connections."""
|
|
if self.ws:
|
|
await self.ws.close()
|
|
if self.session:
|
|
await self.session.close()
|
|
print("Disconnected.")
|
|
|
|
|
|
def main():
|
|
"""Entry point."""
|
|
import argparse
|
|
|
|
parser = argparse.ArgumentParser(description="XML Pipeline Console")
|
|
parser.add_argument("--host", default=DEFAULT_HOST, help="Server host")
|
|
parser.add_argument("--port", type=int, default=DEFAULT_PORT, help="Server port")
|
|
args = parser.parse_args()
|
|
|
|
client = ConsoleClient(host=args.host, port=args.port)
|
|
asyncio.run(client.run())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|