124 lines
3.7 KiB
Python
124 lines
3.7 KiB
Python
"""
|
|
Thin client for talking to the OffSec Shield Node (offsec-agents MCP backend).
|
|
|
|
Usage:
|
|
|
|
from scripts.offsec_node_client import OffsecNodeClient
|
|
|
|
client = OffsecNodeClient() # uses OFFSEC_NODE_URL env
|
|
agents = await client.command("agents list")
|
|
status = await client.command("tem status")
|
|
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import os
|
|
from dataclasses import dataclass
|
|
from typing import Any, Dict, Optional
|
|
|
|
import aiohttp
|
|
|
|
|
|
DEFAULT_OFFSEC_NODE_URL = "http://shield-vm:8081"
|
|
|
|
|
|
@dataclass
|
|
class OffsecNodeError(Exception):
|
|
message: str
|
|
status: Optional[int] = None
|
|
details: Optional[Dict[str, Any]] = None
|
|
|
|
def __str__(self) -> str:
|
|
base = self.message
|
|
if self.status is not None:
|
|
base += f" (status={self.status})"
|
|
if self.details:
|
|
base += f" details={self.details}"
|
|
return base
|
|
|
|
|
|
@dataclass
|
|
class OffsecNodeClient:
|
|
base_url: str = DEFAULT_OFFSEC_NODE_URL
|
|
timeout_seconds: int = 10
|
|
|
|
@classmethod
|
|
def from_env(cls) -> "OffsecNodeClient":
|
|
url = os.getenv("OFFSEC_NODE_URL", DEFAULT_OFFSEC_NODE_URL)
|
|
return cls(base_url=url)
|
|
|
|
async def health(self) -> Dict[str, Any]:
|
|
url = f"{self.base_url.rstrip('/')}/health"
|
|
async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=self.timeout_seconds)) as session:
|
|
async with session.get(url) as resp:
|
|
if resp.status != 200:
|
|
raise OffsecNodeError("Shield node health check failed", status=resp.status)
|
|
return await resp.json()
|
|
|
|
async def command(
|
|
self,
|
|
command: str,
|
|
session_id: str = "vaultmesh-client",
|
|
user: str = "vaultmesh",
|
|
) -> Dict[str, Any]:
|
|
"""
|
|
Send a command to the Shield Node MCP backend.
|
|
|
|
Example commands:
|
|
"agents list"
|
|
"status"
|
|
"shield status"
|
|
"proof latest"
|
|
"agent spawn recon"
|
|
"""
|
|
url = f"{self.base_url.rstrip('/')}/mcp/command"
|
|
payload: Dict[str, Any] = {
|
|
"session_id": session_id,
|
|
"user": user,
|
|
"command": command,
|
|
}
|
|
|
|
async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=self.timeout_seconds)) as session:
|
|
async with session.post(url, json=payload) as resp:
|
|
text = await resp.text()
|
|
if resp.status != 200:
|
|
# Try to parse JSON details if present
|
|
try:
|
|
data = json.loads(text)
|
|
except json.JSONDecodeError:
|
|
data = None
|
|
raise OffsecNodeError(
|
|
"Shield node command failed",
|
|
status=resp.status,
|
|
details=data or {"raw": text},
|
|
)
|
|
try:
|
|
return json.loads(text)
|
|
except json.JSONDecodeError as exc:
|
|
raise OffsecNodeError("Invalid JSON from Shield node", details={"raw": text}) from exc
|
|
|
|
|
|
# Optional: CLI entrypoint for quick manual tests
|
|
async def _demo() -> None:
|
|
client = OffsecNodeClient.from_env()
|
|
print(f"[offsec-node] base_url={client.base_url}")
|
|
try:
|
|
health = await client.health()
|
|
print("Health:", json.dumps(health, indent=2))
|
|
except OffsecNodeError as e:
|
|
print("Health check failed:", e)
|
|
return
|
|
|
|
try:
|
|
agents = await client.command("agents list")
|
|
print("Agents:", json.dumps(agents, indent=2))
|
|
except OffsecNodeError as e:
|
|
print("Command failed:", e)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(_demo())
|