295 lines
11 KiB
Python
295 lines
11 KiB
Python
# agent_server.py
|
|
"""
|
|
AgentServer — The Living Organism Host
|
|
December 25, 2025
|
|
|
|
Preliminary but runnable implementation.
|
|
|
|
This is the body: one process, one secure WebSocket endpoint,
|
|
hosting many concurrent AgentService organs sharing a single
|
|
tamper-proof MessageBus from xml-pipeline.
|
|
|
|
Features in this version:
|
|
- Mandatory WSS (TLS)
|
|
- First-message TOTP authentication
|
|
- Per-user capability control via config/users.yaml
|
|
- Personalized <catalog/> on connect
|
|
- Ed25519 identity generation helper
|
|
- Boot-time agent registration
|
|
- Hooks for future signed privileged commands
|
|
|
|
XML wins.
|
|
"""
|
|
|
|
import asyncio
|
|
import os
|
|
import ssl
|
|
import time
|
|
from typing import Optional, Dict, Any
|
|
|
|
import pyotp
|
|
import yaml
|
|
from cryptography.hazmat.primitives.asymmetric import ed25519
|
|
from cryptography.hazmat.primitives import serialization
|
|
from websockets.server import serve, WebSocketServerProtocol
|
|
|
|
from xml_pipeline import MessageBus
|
|
from xml_pipeline.service import AgentService
|
|
from xml_pipeline.message import repair_and_canonicalize, XmlTamperError
|
|
|
|
|
|
class AgentServer:
|
|
"""
|
|
The body of the organism.
|
|
One instance = one living, multi-personality swarm.
|
|
"""
|
|
|
|
# Default identity location — can be overridden if needed
|
|
IDENTITY_DIR = os.path.expanduser("~/.agent_server")
|
|
PRIVATE_KEY_PATH = os.path.join(IDENTITY_DIR, "identity.ed25519")
|
|
PUBLIC_KEY_PATH = os.path.join(IDENTITY_DIR, "identity.ed25519.pub")
|
|
|
|
def __init__(
|
|
self,
|
|
host: str = "0.0.0.0",
|
|
port: int = 8765,
|
|
ssl_context: Optional[ssl.SSLContext] = None,
|
|
users_config_path: str = "config/users.yaml",
|
|
identity_pubkey_path: Optional[str] = None,
|
|
):
|
|
self.host = host
|
|
self.port = port
|
|
self.ssl_context = ssl_context # None = ws:// (dev only), set for wss://
|
|
self.bus = MessageBus()
|
|
|
|
# Load per-user TOTP secrets + allowed root tags
|
|
self.users_config: Dict[str, Dict[str, Any]] = self._load_users_config(users_config_path)
|
|
|
|
# Load organism public key for future privileged command verification
|
|
self.pubkey: Optional[bytes] = None
|
|
pubkey_path = identity_pubkey_path or self.PUBLIC_KEY_PATH
|
|
if os.path.exists(pubkey_path):
|
|
self.pubkey = self._load_pubkey(pubkey_path)
|
|
|
|
# Built-in platform listeners will be added here in future versions
|
|
|
|
@staticmethod
|
|
def _load_users_config(path: str) -> Dict[str, Dict[str, Any]]:
|
|
"""Load users.yaml → {user_id: {totp_secret: ..., allowed_root_tags: [...]}}"""
|
|
if not os.path.exists(path):
|
|
raise FileNotFoundError(f"Users config not found: {path}")
|
|
with open(path) as f:
|
|
data = yaml.safe_load(f) or {}
|
|
return data.get("users", {})
|
|
|
|
@staticmethod
|
|
def _load_pubkey(path: str) -> bytes:
|
|
"""Load raw Ed25519 public key bytes"""
|
|
with open(path, "rb") as f:
|
|
content = f.read().strip()
|
|
# Accept either raw bytes or ssh-ed25519 format
|
|
if content.startswith(b"ssh-ed25519 "):
|
|
import base64
|
|
return base64.b64decode(content.split()[1])
|
|
return content
|
|
|
|
def register_agent(
|
|
self,
|
|
agent_class: type[AgentService],
|
|
*,
|
|
system_prompt: str,
|
|
max_concurrent: int = 10,
|
|
session_timeout: float = 1800.0,
|
|
version: str = "1.0",
|
|
public: bool = True,
|
|
) -> None:
|
|
"""Register a permanent agent at boot time."""
|
|
# Wrapper to store public flag for catalog building
|
|
self.bus.register_agent(
|
|
agent_class=agent_class,
|
|
system_prompt=system_prompt,
|
|
max_concurrent=max_concurrent,
|
|
session_timeout=session_timeout,
|
|
version=version,
|
|
metadata={"public": public},
|
|
)
|
|
|
|
async def _handle_client(self, websocket: WebSocketServerProtocol):
|
|
"""Per-connection handler: authenticate → send catalog → pump messages"""
|
|
context = {
|
|
"authenticated": False,
|
|
"user": None,
|
|
"allowed_tags": set(),
|
|
"bad_message_count": 0,
|
|
"last_bad_time": 0.0,
|
|
}
|
|
|
|
try:
|
|
# 1. Authentication — first message must be <authenticate totp="..."/>
|
|
first_raw = await asyncio.wait_for(websocket.recv(), timeout=15.0)
|
|
auth_msg = repair_and_canonicalize(first_raw)
|
|
|
|
if auth_msg.getroot().tag != "authenticate":
|
|
await websocket.close(code=1008, reason="First message must be <authenticate/>")
|
|
return
|
|
|
|
totp_code = auth_msg.getroot().get("totp")
|
|
if not totp_code:
|
|
await websocket.close(code=1008, reason="Missing TOTP code")
|
|
return
|
|
|
|
user_id = self._authenticate_totp(totp_code)
|
|
if not user_id:
|
|
await websocket.close(code=1008, reason="Invalid TOTP")
|
|
return
|
|
|
|
user_cfg = self.users_config[user_id]
|
|
allowed_tags = set(user_cfg.get("allowed_root_tags", []))
|
|
if "*" in allowed_tags:
|
|
# Wildcard = all current + future tags
|
|
allowed_tags = None # Special sentinel
|
|
|
|
context.update({
|
|
"authenticated": True,
|
|
"user": user_id,
|
|
"allowed_tags": allowed_tags,
|
|
})
|
|
|
|
# 2. Send personalized catalog
|
|
catalog_xml = self.bus.build_catalog_for_user(allowed_tags)
|
|
await websocket.send(catalog_xml)
|
|
|
|
# 3. Message pump
|
|
async def inbound():
|
|
async for raw in websocket:
|
|
try:
|
|
yield repair_and_canonicalize(raw)
|
|
except XmlTamperError:
|
|
await websocket.close(code=1008, reason="Invalid/tampered XML")
|
|
raise
|
|
|
|
async def outbound(message: bytes):
|
|
await websocket.send(message)
|
|
|
|
await self.bus.run(
|
|
inbound=inbound(),
|
|
outbound=outbound,
|
|
context=context, # For ACL checks in listeners
|
|
)
|
|
|
|
except asyncio.TimeoutError:
|
|
await websocket.close(code=1008, reason="Authentication timeout")
|
|
except Exception as e:
|
|
print(f"Client error ({websocket.remote_address}): {e}")
|
|
|
|
def _authenticate_totp(self, code: str) -> Optional[str]:
|
|
"""Validate TOTP and return user identifier if successful"""
|
|
for user_id, cfg in self.users_config.items():
|
|
totp = pyotp.TOTP(cfg["totp_secret"])
|
|
if totp.verify(code, valid_window=1): # 30s tolerance
|
|
return user_id
|
|
return None
|
|
|
|
async def start(self):
|
|
"""Start the organism — runs forever"""
|
|
scheme = "wss" if self.ssl_context else "ws"
|
|
print(f"AgentServer starting on {scheme}://{self.host}:{self.port}")
|
|
print("Organism awakening...")
|
|
|
|
async with serve(self._handle_client, self.host, self.port, ssl=self.ssl_context):
|
|
await asyncio.Future() # Run forever
|
|
|
|
@classmethod
|
|
def generate_identity(cls, force: bool = False) -> None:
|
|
"""
|
|
Generate the organism's permanent Ed25519 identity.
|
|
Run once on first deployment.
|
|
"""
|
|
os.makedirs(cls.IDENTITY_DIR, exist_ok=True)
|
|
|
|
if os.path.exists(cls.PRIVATE_KEY_PATH) and not force:
|
|
print("Identity already exists:")
|
|
print(f" Private key: {cls.PRIVATE_KEY_PATH}")
|
|
print(f" Public key : {cls.PUBLIC_KEY_PATH}")
|
|
print("Use --force to regenerate (will overwrite!).")
|
|
return
|
|
|
|
print("Generating organism Ed25519 identity...")
|
|
|
|
private_key = ed25519.Ed25519PrivateKey.generate()
|
|
public_key = private_key.public_key()
|
|
|
|
# Private key — PEM PKCS8 unencrypted (rely on file permissions)
|
|
private_pem = private_key.private_bytes(
|
|
encoding=serialization.Encoding.PEM,
|
|
format=serialization.PrivateFormat.PKCS8,
|
|
encryption_algorithm=serialization.NoEncryption(),
|
|
)
|
|
|
|
# Public key — raw bytes + ssh-ed25519 format for readability
|
|
public_raw = public_key.public_bytes(
|
|
encoding=serialization.Encoding.Raw,
|
|
format=serialization.PublicFormat.Raw,
|
|
)
|
|
public_ssh = f"ssh-ed25519 {public_raw.hex()} organism@{os.uname().nodename}"
|
|
|
|
# Write with secure permissions
|
|
with open(cls.PRIVATE_KEY_PATH, "wb") as f:
|
|
os.fchmod(f.fileno(), 0o600)
|
|
f.write(private_pem)
|
|
|
|
with open(cls.PUBLIC_KEY_PATH, "w") as f:
|
|
f.write(public_ssh + "\n")
|
|
|
|
print("Organism identity created!")
|
|
print(f"Private key (KEEP SAFE): {cls.PRIVATE_KEY_PATH}")
|
|
print(f"Public key : {cls.PUBLIC_KEY_PATH}")
|
|
print("\nBackup the private key offline. Lose it → lose structural control forever.")
|
|
|
|
|
|
# ————————————————————————
|
|
# Example CLI entrypoint
|
|
# ————————————————————————
|
|
if __name__ == "__main__":
|
|
import argparse
|
|
|
|
parser = argparse.ArgumentParser(description="AgentServer — the living organism")
|
|
subparsers = parser.add_subparsers(dest="command", required=True)
|
|
|
|
# Run the server
|
|
run_p = subparsers.add_parser("run", help="Start the organism")
|
|
run_p.add_argument("--host", default="0.0.0.0")
|
|
run_p.add_argument("--port", type=int, default=8765)
|
|
run_p.add_argument("--cert", help="Path to TLS fullchain.pem")
|
|
run_p.add_argument("--key", help="Path to TLS privkey.pem")
|
|
run_p.add_argument("--users-config", default="config/users.yaml")
|
|
|
|
# Generate identity
|
|
gen_p = subparsers.add_parser("generate-identity", help="Create cryptographic identity")
|
|
gen_p.add_argument("--force", action="store_true")
|
|
|
|
args = parser.parse_args()
|
|
|
|
if args.command == "generate-identity":
|
|
AgentServer.generate_identity(force=args.force)
|
|
|
|
else: # run
|
|
ssl_ctx = None
|
|
if args.cert and args.key:
|
|
ssl_ctx = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
|
|
ssl_ctx.load_cert_chain(args.cert, args.key)
|
|
|
|
server = AgentServer(
|
|
host=args.host,
|
|
port=args.port,
|
|
ssl_context=ssl_ctx,
|
|
users_config_path=args.users_config,
|
|
)
|
|
|
|
# Example boot-time listeners (uncomment and customise)
|
|
# from your_agents import CodingAgent, ResearchAgent, GrokAgent
|
|
# server.register_agent(CodingAgent, system_prompt="You are an elite Python engineer...", max_concurrent=20)
|
|
# server.register_agent(ResearchAgent, system_prompt="You are a thorough researcher...", max_concurrent=10)
|
|
# server.register_agent(GrokAgent, system_prompt="You are Grok, built by xAI...", max_concurrent=15)
|
|
|
|
asyncio.run(server.start())
|