diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..6abe271 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,53 @@ +# Keep the build context small — ship only what's needed to build the wheel. + +# VCS & editors +.git +.gitignore +.github +.vscode +.idea + +# Python caches +__pycache__ +*.pyc +*.pyo +*.pyd +.pytest_cache +.mypy_cache +.ruff_cache +.coverage +htmlcov +*.egg-info +build +dist + +# Local venvs and env files — never bake secrets into the image +.venv +venv +env +.env +.env.* +!.env.example + +# Runtime state that should NOT go into the image +clients.json +dashboard.db +*.sqlite +*.sqlite3 + +# Local config (user-specific — built image uses /config at runtime) +beaconmcp.yaml + +# Docs / tooling / ops extras that aren't needed at runtime +docs +deploy/install.sh +deploy/beaconmcp.service +.playwright-mcp +.assistant + +# OS cruft +.DS_Store +Thumbs.db + +# Tests aren't shipped in the wheel +tests diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..ece0c9f --- /dev/null +++ b/.env.example @@ -0,0 +1,40 @@ +# BeaconMCP secrets. +# +# Topology lives in beaconmcp.yaml (see beaconmcp.yaml.example). This file +# only holds the values the YAML references as ${VAR}. Every variable below +# is optional — include only the secrets your topology actually uses. + +# --- Proxmox API tokens ---------------------------------------------------- +# One secret per entry under proxmox.nodes[] in beaconmcp.yaml. +PVE1_TOKEN_SECRET=xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx +PVE2_TOKEN_SECRET=xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx + +# --- BMC credentials ------------------------------------------------------- +# One secret per entry under bmc.devices[]. Name the env vars after the +# device id for clarity. +RACK1_ILO_PASSWORD=change-me +RACK2_IPMI_PASSWORD=change-me + +# --- SSH fallback ---------------------------------------------------------- +SSH_PASSWORD=change-me + +# --- Dashboard ------------------------------------------------------------- +# Omit GEMINI_API_KEY to disable the integrated chat (the tokens page still works). +# GEMINI_API_KEY=... + +# Auto-generated by deploy/install.sh on first run. Encrypts client_secret +# at rest for dashboard sessions. Regenerating invalidates every session. +# BEACONMCP_SESSION_KEY= + +# --- Legacy env-var overrides (deprecated, removed in 2.1) ----------------- +# Only used when no beaconmcp.yaml is found. Prefer the YAML file. +# BEACONMCP_CONFIG=/etc/beaconmcp/config.yaml +# BEACONMCP_CLIENTS_FILE=/opt/beaconmcp/clients.json +# BEACONMCP_PORT=8420 +# BEACONMCP_HOST=0.0.0.0 +# BEACONMCP_ALLOWED_HOSTS=mcp.example.com,127.0.0.1:*,localhost:*,[::1]:* +# BEACONMCP_ALLOWED_ORIGINS=https://assistant.ai,https://chat.openai.com,https://gemini.google.com +# BEACONMCP_DASHBOARD_ENABLED=true +# BEACONMCP_DASHBOARD_PUBLIC_URL=https://mcp.example.com +# BEACONMCP_DASHBOARD_LIMIT_5H_USD=2.0 +# BEACONMCP_DASHBOARD_LIMIT_WEEK_USD=10.0 diff --git a/.github/dependabot.yml b/.github/dependabot.yml new file mode 100644 index 0000000..74f7b25 --- /dev/null +++ b/.github/dependabot.yml @@ -0,0 +1,15 @@ +version: 2 +updates: + - package-ecosystem: github-actions + directory: / + schedule: + interval: monthly + commit-message: + prefix: "chore(ci)" + + - package-ecosystem: pip + directory: / + schedule: + interval: monthly + commit-message: + prefix: "chore(deps)" diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..4e605ba --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,127 @@ +name: CI + +on: + # Branch-scoped so a push to a branch with an open PR doesn't run twice. + push: + branches: [main] + pull_request: + +concurrency: + group: ci-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + test: + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + # 3.13 is what the Docker image runs on. + python-version: ["3.11", "3.12", "3.13"] + + steps: + - uses: actions/checkout@v7 + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v7 + with: + python-version: ${{ matrix.python-version }} + cache: pip + + - name: Install dependencies + run: pip install -e ".[wizard,dev]" + + - name: Lint + run: ruff check src/ tests/ + + - name: Test + run: pytest -q + + # beaconmcp.yaml.example is the documented schema. If a config key is renamed + # in the loader without updating the template, this catches it. + config: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v7 + + - uses: actions/setup-python@v7 + with: + python-version: "3.12" + cache: pip + + - name: Install + run: pip install -e . + + - name: Validate the example config + # The template references secrets as ${VAR} and the loader rejects + # unset ones, so feed it throwaway values. + env: + BEACONMCP_SESSION_KEY: MDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDA= + GEMINI_API_KEY: ci-dummy + PVE1_TOKEN_SECRET: ci-dummy + PVE2_TOKEN_SECRET: ci-dummy + RACK1_ILO_PASSWORD: ci-dummy + RACK2_IPMI_PASSWORD: ci-dummy + RACK3_IDRAC_PASSWORD: ci-dummy + VPS2_PW: ci-dummy + run: beaconmcp validate-config --config beaconmcp.yaml.example + + # Templates, static files and the logo ride along implicitly with the package + # directory. A missing one only surfaces when the dashboard renders, so assert + # on the built artifact here. + build: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v7 + + - uses: actions/setup-python@v7 + with: + python-version: "3.12" + cache: pip + + - name: Build sdist + wheel + run: | + pip install build + python -m build + + - name: Assert packaged data files are present + run: | + python - <<'PY' + import pathlib, sys, zipfile + + names = zipfile.ZipFile(next(pathlib.Path("dist").glob("*.whl"))).namelist() + for prefix in ("beaconmcp/assets/logo.webp", + "beaconmcp/dashboard/templates/", + "beaconmcp/dashboard/static/"): + if not any(n.startswith(prefix) for n in names): + sys.exit(f"wheel is missing {prefix}") + PY + + - name: Install the wheel in a clean env and smoke-test the CLI + run: | + python -m venv /tmp/smoke + /tmp/smoke/bin/pip install dist/*.whl + /tmp/smoke/bin/beaconmcp --help + + docker: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v7 + + - uses: docker/setup-buildx-action@v4 + + - name: Build the image (no push) + uses: docker/build-push-action@v7 + with: + context: . + push: false + load: true + tags: beaconmcp:ci + cache-from: type=gha + cache-to: type=gha,mode=max + + - name: Smoke-test the entrypoint + run: docker run --rm beaconmcp:ci --help diff --git a/.gitignore b/.gitignore index b7faf40..30762ca 100644 --- a/.gitignore +++ b/.gitignore @@ -1,207 +1,11 @@ -# Byte-compiled / optimized / DLL files __pycache__/ -*.py[codz] +*.py[cod] *$py.class - -# C extensions -*.so - -# Distribution / packaging -.Python -build/ -develop-eggs/ +*.egg-info/ dist/ -downloads/ -eggs/ +build/ .eggs/ -lib/ -lib64/ -parts/ -sdist/ -var/ -wheels/ -share/python-wheels/ -*.egg-info/ -.installed.cfg *.egg -MANIFEST - -# PyInstaller -# Usually these files are written by a python script from a template -# before PyInstaller builds the exe, so as to inject date/other infos into it. -*.manifest -*.spec - -# Installer logs -pip-log.txt -pip-delete-this-directory.txt - -# Unit test / coverage reports -htmlcov/ -.tox/ -.nox/ -.coverage -.coverage.* -.cache -nosetests.xml -coverage.xml -*.cover -*.py.cover -.hypothesis/ -.pytest_cache/ -cover/ - -# Translations -*.mo -*.pot - -# Django stuff: -*.log -local_settings.py -db.sqlite3 -db.sqlite3-journal - -# Flask stuff: -instance/ -.webassets-cache - -# Scrapy stuff: -.scrapy - -# Sphinx documentation -docs/_build/ - -# PyBuilder -.pybuilder/ -target/ - -# Jupyter Notebook -.ipynb_checkpoints - -# IPython -profile_default/ -ipython_config.py - -# pyenv -# For a library or package, you might want to ignore these files since the code is -# intended to run in multiple environments; otherwise, check them in: -# .python-version - -# pipenv -# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. -# However, in case of collaboration, if having platform-specific dependencies or dependencies -# having no cross-platform support, pipenv may install dependencies that don't work, or not -# install all needed dependencies. -#Pipfile.lock - -# UV -# Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control. -# This is especially recommended for binary packages to ensure reproducibility, and is more -# commonly ignored for libraries. -#uv.lock - -# poetry -# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. -# This is especially recommended for binary packages to ensure reproducibility, and is more -# commonly ignored for libraries. -# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control -#poetry.lock -#poetry.toml - -# pdm -# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. -# pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python. -# https://pdm-project.org/en/latest/usage/project/#working-with-version-control -#pdm.lock -#pdm.toml -.pdm-python -.pdm-build/ - -# pixi -# Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control. -#pixi.lock -# Pixi creates a virtual environment in the .pixi directory, just like venv module creates one -# in the .venv directory. It is recommended not to include this directory in version control. -.pixi - -# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm -__pypackages__/ - -# Celery stuff -celerybeat-schedule -celerybeat.pid - -# SageMath parsed files -*.sage.py - -# Environments .env -.envrc -.venv -env/ +.venv/ venv/ -ENV/ -env.bak/ -venv.bak/ - -# Spyder project settings -.spyderproject -.spyproject - -# Rope project settings -.ropeproject - -# mkdocs documentation -/site - -# mypy -.mypy_cache/ -.dmypy.json -dmypy.json - -# Pyre type checker -.pyre/ - -# pytype static type analyzer -.pytype/ - -# Cython debug symbols -cython_debug/ - -# PyCharm -# JetBrains specific template is maintained in a separate JetBrains.gitignore that can -# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore -# and can be added to the global gitignore or merged into this file. For a more nuclear -# option (not recommended) you can uncomment the following to ignore the entire idea folder. -#.idea/ - -# Abstra -# Abstra is an AI-powered process automation framework. -# Ignore directories containing user credentials, local state, and settings. -# Learn more at https://abstra.io/docs -.abstra/ - -# Visual Studio Code -# Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore -# that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore -# and can be added to the global gitignore or merged into this file. However, if you prefer, -# you could uncomment the following to ignore the entire vscode folder -# .vscode/ - -# Ruff stuff: -.ruff_cache/ - -# PyPI configuration file -.pypirc - -# Cursor -# Cursor is an AI-powered code editor. `.cursorignore` specifies files/directories to -# exclude from AI features like autocomplete and code analysis. Recommended for sensitive data -# refer to https://docs.cursor.com/context/ignore-files -.cursorignore -.cursorindexingignore - -# Marimo -marimo/_static/ -marimo/_lsp/ -__marimo__/ diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 0000000..5e1e477 --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,62 @@ +# BeaconMCP + +Remote MCP server for managing any Proxmox VE cluster together with its BMC-managed hardware (HP iLO, IPMI). Runs over HTTP with OAuth 2.1 + TOTP. Topology is described in a single YAML file; secrets are referenced through `${ENV_VAR}` placeholders. + +## Quick Start + +```bash +pip install -e . +cp beaconmcp.yaml.example beaconmcp.yaml # Describe your infrastructure +cp .env.example .env # Fill in the referenced secrets +beaconmcp validate-config # Dry-run the loader, secrets masked +beaconmcp auth create --name "x" # Create an OAuth client +beaconmcp serve # Start HTTP server on :8420 +``` + +## Project Structure + +``` +src/beaconmcp/ + __main__.py CLI: serve + auth + validate-config + server.py FastMCP server, registers every tool module + config.py YAML loader with ${ENV} resolver + legacy env fallback + auth.py OAuth 2.1 client credentials (ClientStore + TokenStore) + proxmox/ + client.py proxmoxer wrapper (API-token auth, N-node aware) + monitoring.py 6 tools: list_nodes, node_status, list_vms, vm_status, get_logs, get_tasks + vms.py 7 tools: vm_start/stop/restart/create/clone/migrate/config + system.py 5 tools: storage_status, network_config, exec_command (sync+async+get_result) + ssh/ + client.py asyncssh wrapper with configurable VMID->IP template + tools.py 4 tools: ssh_exec_command (sync+async+get_result), ssh_list_sessions + bmc/ + base.py BMCClient Protocol + shared exceptions + stub base class + hp_ilo.py HPILOBackend (python-hpilo, optional SSH jump tunnel) + ipmi.py GenericIPMIBackend (shells out to ipmitool) + idrac.py IDRACStubBackend (TODO) + supermicro.py SupermicroStubBackend (TODO) + registry.py build_registry(config) -> {device_id: BMCClient} + tools.py 8 tools: bmc_list_devices + 7 action tools (device_id param) + dashboard/ Optional web panel: /app/login, /app/chat, /app/tokens +beaconmcp.yaml.example Template describing the full config schema +deploy/ + install.sh One-command install script + beaconmcp.service systemd unit file +``` + +## Configuration + +Two files: + +- **`beaconmcp.yaml`** — topology and feature flags. Resolution order: `--config` flag → `BEACONMCP_CONFIG` env → `./beaconmcp.yaml` → `/etc/beaconmcp/config.yaml`. +- **`.env`** — secrets referenced by the YAML via `${VAR}`. + +Legacy `PVE*_*`, `ILO_*`, `SSH_*` env vars still work when no YAML is found (deprecated; removed in 2.1). + +## Auth + +OAuth 2.1 client credentials + mandatory TOTP. Manage with `beaconmcp auth create/list/revoke`. + +## Design Spec + +See `docs/superpowers/specs/2026-04-16-beaconmcp-design.md`. diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..10c84bc --- /dev/null +++ b/Dockerfile @@ -0,0 +1,52 @@ +FROM python:3.13-slim AS builder + +ENV PIP_NO_CACHE_DIR=1 \ + PIP_DISABLE_PIP_VERSION_CHECK=1 \ + PYTHONDONTWRITEBYTECODE=1 + +RUN apt-get update \ + && apt-get install -y --no-install-recommends build-essential \ + && rm -rf /var/lib/apt/lists/* + +WORKDIR /build +COPY pyproject.toml README.md ./ +COPY src ./src + +RUN pip install --upgrade pip build \ + && python -m build --wheel --outdir /dist + + +FROM python:3.13-slim AS runtime + +ENV PYTHONDONTWRITEBYTECODE=1 \ + PYTHONUNBUFFERED=1 \ + PIP_NO_CACHE_DIR=1 \ + BEACONMCP_CONFIG=/config/beaconmcp.yaml \ + BEACONMCP_DASHBOARD_DB=/state/dashboard.db \ + BEACONMCP_CLIENTS_FILE=/state/clients.json + +RUN apt-get update \ + && apt-get install -y --no-install-recommends ipmitool ca-certificates \ + && rm -rf /var/lib/apt/lists/* + +RUN groupadd --system --gid 10001 beaconmcp \ + && useradd --system --uid 10001 --gid beaconmcp --home /app --shell /usr/sbin/nologin beaconmcp + +COPY --from=builder /dist/*.whl /tmp/ +RUN pip install /tmp/*.whl \ + && rm -f /tmp/*.whl + +RUN mkdir -p /config /state \ + && chown -R beaconmcp:beaconmcp /config /state + +USER beaconmcp +WORKDIR /app + +EXPOSE 8420 + +HEALTHCHECK --interval=30s --timeout=5s --start-period=10s --retries=3 \ + CMD python -c "import urllib.request,sys; urllib.request.urlopen('http://127.0.0.1:8420/health', timeout=3).read(); sys.exit(0)" \ + || exit 1 + +ENTRYPOINT ["beaconmcp"] +CMD ["serve"] diff --git a/LICENSE b/LICENSE index 261eeb9..baf3786 100644 --- a/LICENSE +++ b/LICENSE @@ -1,3 +1,27 @@ +"Commons Clause" License Condition v1.0 + +The Software is provided to you by the Licensor under the License, as +defined below, subject to the following condition. + +Without limiting other conditions in the License, the grant of rights +under the License will not include, and the License does not grant to +you, the right to Sell the Software. + +For purposes of the foregoing, "Sell" means practicing any or all of +the rights granted to you under the License to provide to third parties, +for a fee or other consideration (including without limitation fees for +hosting or consulting/support services related to the Software), a +product or service whose value derives, entirely or substantially, from +the functionality of the Software. Any license notice or attribution +required by the License must also include this Commons Clause License +Condition notice. + +Software: BeaconMCP/TarkaMCP +License: Apache License, Version 2.0 (with Commons Clause) +Licensor: Showdown76py + +------------------------------------------------------------------------------- + Apache License Version 2.0, January 2004 http://www.apache.org/licenses/ @@ -186,7 +210,7 @@ same "printed page" as the copyright notice for easier identification within third-party archives. - Copyright [yyyy] [name of copyright owner] + Copyright 2026 Showdown76py Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. diff --git a/README.md b/README.md new file mode 100644 index 0000000..638c01b --- /dev/null +++ b/README.md @@ -0,0 +1,112 @@ +
+ +# BeaconMCP + +[![Python 3.11+](https://img.shields.io/badge/python-3.11+-3776AB?logo=python&logoColor=white)](https://www.python.org/downloads/) +[![MCP Protocol](https://img.shields.io/badge/MCP-Model_Context_Protocol-5A67D8)](https://modelcontextprotocol.io/) +[![Proxmox VE](https://img.shields.io/badge/Proxmox-VE_8.x-E57000?logo=proxmox&logoColor=white)](https://www.proxmox.com/) +[![License](https://img.shields.io/badge/license-Apache_2.0_%2B_Commons_Clause-red)](LICENSE) + +**One MCP endpoint for a Proxmox cluster, the hardware under it, and your SSH hosts.** + +
+ +--- + +BeaconMCP is a remote MCP server (Streamable HTTP, OAuth 2.1 + TOTP). It exposes a Proxmox VE +cluster, the BMCs of the machines running it (HP iLO, IPMI, Redfish), and arbitrary SSH-reachable +hosts as a single authenticated endpoint. An MCP client can then diagnose a crash, power-cycle a +frozen host, migrate a VM, or run a command inside a guest without you opening four different web +UIs. + +Capabilities are independent: configure a full cluster, three VPS reachable only by SSH, a rack of +IPMI BMCs, or any mix. Tools are registered per capability, so an SSH-only deployment never exposes +`proxmox_*` tools. There is no hard-coded node count. Nodes, BMC devices and SSH hosts are all lists +in one YAML file, with secrets referenced as `${ENV_VAR}`. + +Tested with Assistant (web, mobile, desktop), ChatGPT, Gemini, Mistral, VS Code, Cursor and +OpenCode. + +## Quick start + +Docker is the fastest path. Run it on a Proxmox node, or anywhere on the same LAN. + +```bash +git clone https://github.com/Showdown76py/BeaconMCP.git +cd BeaconMCP +cp beaconmcp.yaml.example beaconmcp.yaml # describe your topology +cp .env.example .env # fill in the ${VAR} secrets +docker compose up -d + +docker compose exec beaconmcp beaconmcp validate-config +docker compose exec beaconmcp beaconmcp auth create --name "Assistant Web" +curl http://localhost:8420/health +``` + +`auth create` prints a client id, a client secret and a TOTP seed as a QR code, **once**. Scan it +into an authenticator app before closing the terminal. + +The server listens on `:8420` in plain HTTP. Put a reverse proxy with TLS in front of it, list the +public hostname under `server.allowed_hosts`, then add the connector to your client. + +There is also a TUI wizard (`beaconmcp init`) that writes the YAML for you, and a bare-metal systemd +install. Both are covered in [Installation](docs/installation.md). + +## Documentation + +| Guide | What's in it | +|-------|--------------| +| [Installation](docs/installation.md) | Requirements, Docker, systemd install, config wizard, reverse proxy, updates | +| [Configuration](docs/configuration.md) | The two config files, every YAML key that matters, where to run the server | +| [Tools](docs/tools.md) | The 46 MCP tools, grouped by module | +| [Updates](docs/updates.md) | The update notice, the self-update tools, and how to turn both off | +| [Client setup](docs/clients.md) | Assistant, ChatGPT, Gemini, Mistral, VS Code, Cursor, OpenCode | +| [Security](docs/security.md) | What to review before approving a tool call, token handling, TOTP hygiene | +| [Dashboard](docs/dashboard.md) | The optional `/app/*` web panel: login, API tokens, Gemini chat | +| [Behind Cloudflare](docs/cloudflare.md) | WAF, Access and caching rules that stop Cloudflare from eating MCP traffic | +| [TOTP automation](docs/totp-automation.md) | Machine-held TOTP for unattended jobs, and why you probably shouldn't | +| [Troubleshooting](docs/troubleshooting.md) | Common errors and what actually fixed them | +| [Tests](docs/tests.md) | Unit tests and the live-cluster integration script | + +## Architecture + +``` +Clients (Assistant, ChatGPT, Gemini, …) + │ + │ HTTPS (reverse proxy / tunnel) + ▼ +┌──────────────────────────────────┐ +│ BeaconMCP (HTTP :8420) │ +│ ├── proxmox/ → Proxmox API │ +│ ├── ssh/ → SSH :22 │ +│ ├── bmc/ → iLO / IPMI │ +│ └── dashboard/ → /app/* │ +└──────────────────────────────────┘ + │ + ▼ +Proxmox nodes (N) · BMC devices (N) · SSH hosts (N) +``` + +Run BeaconMCP on the same local network as the cluster. Every `proxmox.nodes[].host` is then a plain +LAN IP that works for both the Proxmox API (`:8006`) and SSH (`:22`), which is what makes the +`ssh.inherit_proxmox_nodes` shortcut and the iLO SSH-jump tunnel usable. +[Configuration](docs/configuration.md#where-to-run-it) explains what breaks when a node is only +reachable through a public FQDN. + +## Before you point a model at your infrastructure + +BeaconMCP exposes tools that destroy things: `ssh_run`, `proxmox_run`, `bmc_power_off`, +`vm_bulk_action`, `proxmox_vm_create`. A model that misreads a VMID will run the command anyway. + +- Turn off auto-approve in every client. Never accept "always allow this tool". +- Read the `command` argument before approving. If it hit the wrong host, could you recover? +- Treat a `/app/tokens` bearer as root on your nodes for its whole lifetime (30 days by default). + +The full checklist, including what the integrated chat confirms on your behalf, is in +[docs/security.md](docs/security.md). + +## License + +[Apache 2.0 with Commons Clause](LICENSE). Use, fork and modification are free; reselling the +software, including as a hosted service, requires a separate commercial license. The code stays +source-available. diff --git a/beaconmcp.yaml.example b/beaconmcp.yaml.example new file mode 100644 index 0000000..4943986 --- /dev/null +++ b/beaconmcp.yaml.example @@ -0,0 +1,260 @@ +# BeaconMCP configuration file. +# +# Path resolution: --config flag > BEACONMCP_CONFIG env > ./beaconmcp.yaml +# > /etc/beaconmcp/config.yaml. +# +# Secrets use ${VAR_NAME} references resolved against the process environment +# at load time. Put them in a .env file next to this one (loaded by the CLI) +# or export them through your init system. Missing references fail fast with +# the offending YAML path. +# +# Capabilities are independent. Configure only what you have: +# - Proxmox-only: delete `ssh:` and `bmc:` (or leave them empty). +# - SSH-only (VPS fleet, bare-metal, anything reachable by SSH): delete +# `proxmox:` and `bmc:`. +# - BMC-only: delete `proxmox:` and `ssh:`. +# - Any combination: keep what applies. +# At least one capability must be present or the server refuses to start. + +version: 1 + +server: + host: 0.0.0.0 + port: 8420 + + # Host header allowlist for DNS-rebinding protection. Include the public + # FQDN behind the reverse proxy; localhost entries are useful for local dev. + allowed_hosts: + - mcp.example.com + - "127.0.0.1:*" + - "localhost:*" + - "[::1]:*" + # CORS origin allowlist. Browser-based MCP clients send a CORS preflight + # before calling /mcp; their origin MUST be listed here. Desktop / CLI + # clients (Assistant Desktop, Gemini CLI, Cursor, VS Code, Mistral Vibe, + # OpenCode, ...) are NOT browser-based and don't need an entry. + allowed_origins: + - https://assistant.ai + - https://chatgpt.com + - https://chat.openai.com + - https://chat.mistral.ai + - https://www.perplexity.ai + - https://gemini.google.com + + # Direct peers you operate whose forwarded headers are trusted. Governs both + # X-Forwarded-For (auth rate-limit client IP) and X-Forwarded-Host (the host + # advertised in the OAuth issuer and the "paste this MCP URL" strings). When + # empty, both fall back to the request's own peer / Host and forwarded values + # from any peer are ignored. + # IPs and CIDRs are accepted. + # - Keep loopback entries when your reverse proxy runs on the same host. + # - Add `cloudflare` to auto-expand to Cloudflare's published proxy CIDRs. + trusted_proxies: + - "127.0.0.1" + - "::1" + # - cloudflare + + # Persistent storage for OAuth clients and TOTP secrets + clients_file: /opt/beaconmcp/clients.json + session_key: ${BEACONMCP_SESSION_KEY} # optional, generated if omitted + # Enable the OAuth Dynamic Client Registration bootstrap flow used by + # clients that cannot accept a pre-provisioned client_id/secret pair + # (ChatGPT in particular). Each registration is gated by a single-use + # slug minted from the dashboard after you type your TOTP. Off by + # default — turn on only if you need ChatGPT integration. + allow_dynamic_registration: false + + # Staging directory used by `proxmox_upload_file` / `proxmox_download_file` + # to exchange files larger than the QEMU Guest Agent 1MB cap. Files + # placed here (e.g. via SCP, rsync, or the dashboard) become available + # to the upload tool by basename; downloads land here under the + # requested basename. Capped by `transfers_max_mb` (default 500 MB). + transfers_dir: ~/.cache/beaconmcp/transfers + transfers_max_mb: 500 + + # SQLite file persisting *named* API tokens (the ones minted on the + # dashboard's "API tokens" page) across restarts. Defaults to `tokens.db` + # next to `clients_file`. Created owner-only (0600). The + # BEACONMCP_TOKENS_DB env var overrides this. + # tokens_db: /opt/beaconmcp/tokens.db + + # JSON-lines audit log: every MCP tool call, dashboard login, OAuth + # authorize and client revoke. Defaults to /opt/beaconmcp/audit.log; + # set to "-" to disable the file and keep stderr only. Created + # owner-only (0600). The BEACONMCP_AUDIT_LOG env var overrides this. + # audit_log: /opt/beaconmcp/audit.log + + # Lifetime (seconds) of *named* API tokens minted on the dashboard's + # "API tokens" page. They persist across restarts, so the default is + # 30 days (2592000) rather than the 24 h used for internal OAuth/session + # bearers. Set to 0 for tokens that never expire (revoke-only -- make + # sure you trust every client you paste one into). The + # BEACONMCP_NAMED_TOKEN_TTL env var overrides this. + # named_token_ttl: 2592000 + +# -------- Proxmox capability (optional) ------------------------------------ +# Delete this section if you have no Proxmox cluster. Tools starting with +# `proxmox_` won't be registered. +# +# `host:` convention — use the node's **LAN address** whenever BeaconMCP is +# on the same local network as the cluster. It's the one string that works +# cleanly for *both* the Proxmox API (port 8006) and SSH inheritance (port +# 22) — same IP, two services. A `localhost` is fine when BeaconMCP runs +# directly on that node. Avoid public FQDNs with reverse-proxy ports (e.g. +# `pve2.example.com:443`): the port pins the API to HTTPS and breaks the +# `ssh.inherit_proxmox_nodes` shortcut (SSH is on :22, not :443). For a +# fully remote node, declare it explicitly under `ssh.hosts[]` with its +# real SSH address (Tailscale IP, VPN tunnel, bastion…). +proxmox: + verify_ssl: false + # One entry per Proxmox node. N nodes supported; the first one is not special. + # Each node needs an API token (Datacenter > Permissions > API Tokens). + nodes: + - name: pve1 + host: 10.0.0.1 # LAN IP — works for both API and SSH + token_id: "root@pam!beaconmcp" # quotes required: '!' and '@' are YAML-reserved + token_secret: ${PVE1_TOKEN_SECRET} + + - name: pve2 + host: 10.0.0.2 # LAN IP of the second node + token_id: "root@pam!beaconmcp" + token_secret: ${PVE2_TOKEN_SECRET} + +# -------- SSH capability (optional) ---------------------------------------- +# Delete this section if you have no SSH targets. Each `hosts[]` entry is one +# addressable target with its own credentials (password OR key_file, exactly +# one). An `ssh.hosts[]` entry may share a name with a `proxmox.nodes[]` +# entry — the two live in separate tool namespaces (`ssh_*` vs `proxmox_*`) +# so there is no routing ambiguity. +# +# Typical homelab shortcut: set `inherit_proxmox_nodes: true` + `defaults:` +# and every Proxmox node becomes SSH-reachable under its own name with the +# default creds. Explicit `hosts[]` entries still win when they match a node +# by name or address, so you can override per-node. +ssh: + # Optional helper: a numeric identifier passed to ssh_* tools (e.g. a + # Proxmox VMID) is fed into this template, and the resulting IP is matched + # against `hosts[].host` to pick the right credentials. Leave unset to + # disable the shortcut. + vmid_to_ip: "192.168.1.{id}" + + # Default credentials used when `inherit_proxmox_nodes: true` synthesizes + # one ssh.hosts[] entry per Proxmox node. Provide exactly one of + # `password` or `key_file`. + defaults: + user: root + key_file: ~/.ssh/beaconmcp + # password: ${SSH_ROOT_PW} + + # When true, each Proxmox node from `proxmox.nodes[]` that isn't already + # covered by an explicit `hosts[]` entry (matched by name OR address) is + # auto-declared as an SSH target using `defaults`. So `ssh_run(host="pve1")` + # just works without repeating credentials per node. + inherit_proxmox_nodes: true + + # Host-key verification. Unset (default) means "accept any key on first + # contact" -- fine on a trusted LAN, unsafe over the public internet. + # Set `known_hosts` to an OpenSSH-format file to pin keys, or leave it + # unset and flip `strict_host_key_checking: true` to use the system + # `~/.ssh/known_hosts` instead. + # known_hosts: /etc/beaconmcp/known_hosts + # strict_host_key_checking: true + + hosts: + - name: vps1 + host: 198.51.100.10 + port: 22 # optional, defaults to 22 + user: root + key_file: ~/.ssh/id_ed25519 # OR `password:`, not both + # Per-host host-key overrides; unset fields inherit the global + # `ssh.known_hosts` / `ssh.strict_host_key_checking` above. Handy to + # enforce strict checking on one public-internet host while + # trusted-LAN nodes stay permissive. Note: overrides only *tighten* — + # once a global `known_hosts` file is set, a per-host + # `strict_host_key_checking: false` cannot loosen back to + # accept-any-key. + # known_hosts: /etc/beaconmcp/vps1_known_hosts + # strict_host_key_checking: true + + - name: vps2 + host: 198.51.100.11 + user: admin + password: ${VPS2_PW} + + # Example: a Proxmox node BeaconMCP can't reach on the LAN (remote datacenter, + # API behind a Cloudflare tunnel on :443, etc.) — the inherited SSH entry + # would point at the API URL which is wrong for SSH. Declare the node here + # explicitly with its real SSH address and the override wins. + # - name: pve-remote + # host: 100.64.5.3 # Tailscale IP, VPN tunnel, bastion, … + # user: root + # key_file: ~/.ssh/pve_remote + +# -------- BMC capability (optional) ---------------------------------------- +# Delete this section if you have no HP iLO / IPMI / iDRAC / Supermicro +# hardware. `bmc_*` tools won't be registered. +bmc: + devices: + - id: rack1-ilo + type: hp_ilo # hp_ilo | ipmi | idrac (stub) | supermicro (stub) + host: 192.168.10.20 + user: Administrator + password: ${RACK1_ILO_PASSWORD} + # Optional: tunnel the iLO connection through an ssh.hosts[] entry + # (useful when the BMC lives on a private management VLAN only + # reachable from a bastion). References `ssh.hosts[].name`. + jump_host: pve1 + + - id: rack2-bmc + type: ipmi + host: 192.168.10.21 + user: admin + password: ${RACK2_IPMI_PASSWORD} + # no jump_host = direct connection + + - id: rack3-idrac + type: idrac # idrac | supermicro | redfish -> Redfish over HTTPS + host: 192.168.10.22 + user: root + password: ${RACK3_IDRAC_PASSWORD} + # Verify the BMC's TLS certificate on Redfish calls. Defaults to + # false because BMCs ship self-signed certs out of the factory; the + # admin password is sent on every request, so turn this on (after + # installing a cert the host trusts) for any BMC that is not on a + # fully trusted management VLAN. + verify_tls: false + +features: + dashboard: + enabled: true + gemini_api_key: ${GEMINI_API_KEY} # empty disables the chat panel + limits: + per_5h_usd: 2.0 + per_week_usd: 10.0 + public_url: https://mcp.example.com # used to generate MCP URLs in the UI + mcp_mode: local # "local" (default) or "remote" + + # Update notifications. When enabled, the server periodically compares + # this checkout against the upstream default branch and shows a notice + # in the dashboard (signed-in operators only), with instructions matched + # to how BeaconMCP was installed here. Also exposes the + # beaconmcp_check_update / beaconmcp_self_update MCP tools. + updates: + # Set to false on an air-gapped or change-controlled deployment: the + # server then never contacts the git remote at all. + enabled: true + # Set to false to keep the notification but forbid applying it from + # the dashboard or over MCP -- appropriate when updates go through a + # deployment pipeline. Manual instructions are still shown. + allow_self_update: true + +# Free-form infrastructure context exposed as an MCP resource. Edit freely: +# the LLM reads this to understand your topology, naming conventions, and +# operational notes. +infrastructure: + conventions: + vmid_to_ip: "VMIDs map 1:1 to 192.168.1.{VMID}" + naming: "VMs are prefixed by role (web-101, db-102)" + notes: + - "API tokens must be created on each Proxmox node before use." + - "BMC devices are reachable on a private management VLAN." diff --git a/deploy/beaconmcp.service b/deploy/beaconmcp.service new file mode 100644 index 0000000..0b64ff0 --- /dev/null +++ b/deploy/beaconmcp.service @@ -0,0 +1,23 @@ +[Unit] +Description=BeaconMCP - Proxmox MCP Server +After=network.target + +[Service] +Type=simple +User=root +WorkingDirectory=/opt/beaconmcp +EnvironmentFile=/opt/beaconmcp/.env +ExecStart=/opt/beaconmcp/.venv/bin/python -m beaconmcp serve +Restart=on-failure +RestartSec=5 + +# Files the server creates at runtime (clients.json, tokens.db, dashboard.db, +# audit.log) all hold secrets. Default umask 0022 would make them +# world-readable before the explicit chmod calls land; 0077 closes that +# window and covers anything not chmod'd individually. +UMask=0077 +NoNewPrivileges=true +PrivateTmp=true + +[Install] +WantedBy=multi-user.target diff --git a/deploy/install.sh b/deploy/install.sh new file mode 100755 index 0000000..78fe123 --- /dev/null +++ b/deploy/install.sh @@ -0,0 +1,106 @@ +#!/bin/bash +# BeaconMCP - quick install on a Linux host (tested on Debian/Ubuntu) +# Usage: bash deploy/install.sh + +set -e + +INSTALL_DIR="/opt/beaconmcp" +REPO="https://github.com/Showdown76py/BeaconMCP.git" +VENV_DIR="$INSTALL_DIR/.venv" + +echo "=== BeaconMCP - Installation ===" + +# 1. System dependencies +echo "[*] Checking system dependencies..." +apt-get update -qq +apt-get install -y python3 python3-pip python3-venv git + +# Versioned venv package (e.g. python3.11-venv on Debian 12) +PY_VER=$(python3 -c 'import sys; print(f"python{sys.version_info.major}.{sys.version_info.minor}")') +apt-get install -y "${PY_VER}-venv" 2>/dev/null || true + +# 2. Clone or update +if [ -d "$INSTALL_DIR/.git" ]; then + echo "[*] Updating BeaconMCP..." + cd "$INSTALL_DIR" && git pull +else + echo "[*] Cloning BeaconMCP..." + git clone "$REPO" "$INSTALL_DIR" + cd "$INSTALL_DIR" +fi + +# 3. Python virtual environment +if [ ! -x "$VENV_DIR/bin/pip" ]; then + echo "[*] (Re)creating the Python virtualenv..." + rm -rf "$VENV_DIR" + python3 -m venv "$VENV_DIR" +fi + +# 4. Install Python dependencies +echo "[*] Installing Python dependencies..." +"$VENV_DIR/bin/pip" install --upgrade pip --quiet +"$VENV_DIR/bin/pip" install -e . --quiet + +# 5. .env file +if [ ! -f "$INSTALL_DIR/.env" ]; then + echo "[*] Creating .env..." + cp .env.example .env + echo " .env created. Fill it with your Proxmox credentials and BMC secrets." +else + echo "[*] Existing .env preserved." +fi +# The .env holds every secret the server has: Proxmox API tokens, BMC admin +# passwords, SSH passwords, the dashboard session key. `cp` and the shell +# create it with the default umask (0644), which leaves it readable by every +# local user. Make it owner-only, on fresh and pre-existing installs alike. +chmod 600 "$INSTALL_DIR/.env" + +# 5.b beaconmcp.yaml config file +if [ ! -f "$INSTALL_DIR/beaconmcp.yaml" ]; then + echo "[*] Creating beaconmcp.yaml..." + cp beaconmcp.yaml.example beaconmcp.yaml + echo " beaconmcp.yaml created. Edit it to describe your topology." +else + echo "[*] Existing beaconmcp.yaml preserved." +fi + +# 5.c BEACONMCP_SESSION_KEY for the dashboard (encrypts client_secret at rest) +if ! grep -q "^BEACONMCP_SESSION_KEY=" "$INSTALL_DIR/.env"; then + SESSION_KEY=$(openssl rand -base64 32) + echo "" >> "$INSTALL_DIR/.env" + echo "# Auto-generated by install.sh -- DO NOT regenerate or all dashboard sessions invalidate." >> "$INSTALL_DIR/.env" + echo "BEACONMCP_SESSION_KEY=$SESSION_KEY" >> "$INSTALL_DIR/.env" + echo "[*] BEACONMCP_SESSION_KEY generated in .env." +fi + +if ! grep -q "^GEMINI_API_KEY=" "$INSTALL_DIR/.env"; then + echo "[!] GEMINI_API_KEY missing from .env -- the dashboard chat will be disabled." + echo " Add GEMINI_API_KEY=... to /opt/beaconmcp/.env to enable it." +fi + +# 6. beaconmcp wrapper in /usr/local/bin +echo "[*] Installing the 'beaconmcp' wrapper into /usr/local/bin..." +cat > /usr/local/bin/beaconmcp </mcp`) and three +auth paths the dashboard helps you drive: + +- **OAuth 2.1 (pre-registered client)** — Assistant, Codex, Le Chat, Gemini + CLI, Antigravity, OpenCode, Cursor, VS Code. Provision a + `client_id` / `client_secret` pair via `beaconmcp auth create` and + paste them into the client's config. Standard OAuth 2.1 authorization + code + PKCE from there. +- **OAuth + Dynamic Client Registration** — ChatGPT Web / Mobile. Reserved + for clients whose UI won't let you paste credentials. Requires + `server.allow_dynamic_registration: true` in `beaconmcp.yaml` and a + single-use bootstrap slug minted from `/app/connectors`. +- **Static bearer token** — Mistral Vibe, any HTTP-only MCP client that + can't do OAuth. Fallback only. + +> **Security note — always type the TOTP by hand from your phone.** +> The TOTP seed belongs in an authenticator app on a device you physically +> control. Do **not** generate codes programmatically with `oathtool` / +> `pyotp` / a shell alias, and do **not** store the raw seed in a `.env` or +> a secrets manager. Every flow below is designed so you read a 6-digit +> code off your phone. Unattended-service automation is covered separately +> in [totp-automation.md](totp-automation.md). + +> **Redirect + CORS allowlist — both use `server.allowed_origins`.** +> Every HTTPS `redirect_uri` reaching `/oauth/authorize` or +> `/oauth/register/c/` must match an origin in +> `server.allowed_origins` in `beaconmcp.yaml`. The same list is also the +> CORS allowlist for browser preflight to `/mcp`. If a web client fails with +> `invalid_redirect_uri` or CORS errors, add its origin here: +> +> ```yaml +> server: +> allowed_origins: +> - https://claude.ai +> - https://assistant.ai +> - https://chatgpt.com +> - https://chat.mistral.ai +> - https://www.perplexity.ai +> - https://gemini.google.com +> ``` +> +> Non-origin OAuth redirect forms remain built-in for desktop/CLI clients and +> do not need an entry: `vscode://`, `vscode-insiders://`, `cursor://`, and +> loopback callbacks (`http://localhost:*`, `http://127.0.0.1:*`, `http://[::1]:*`). + +The dashboard's [`/app/tokens`](../src/beaconmcp/dashboard/templates/tokens.html) +page presents the same information with copy-pasteable snippets per platform +— this document is the offline reference. + +--- + +## Assistant (web, mobile, desktop) + +Assistant drives the full OAuth 2.1 flow against BeaconMCP, so nothing long-lived is stored on its +side — you type the TOTP into the authorization page whenever a new token is issued. + +1. **Settings → Integrations → Add custom connector.** +2. Fill in: + - **Name:** BeaconMCP + - **Remote MCP server URL:** `https:///mcp` + - **OAuth Client ID** and **OAuth Client Secret** from `beaconmcp auth create`. +3. **Add.** + +On first use, and after each 24-hour token expiry, Assistant redirects to the BeaconMCP +authorization page. Read the current 6-digit code from your authenticator app and type it in. +Assistant never holds the TOTP seed, so a leaked session cannot mint a new token without a fresh +code from your phone. + +For the web app, add its origin to `server.allowed_origins` (see the allowlist note above); +the desktop app uses a loopback callback and needs no entry. + +## ChatGPT (OAuth 2.1) + +ChatGPT's Developer Mode connector only accepts **OAuth with Dynamic Client +Registration (RFC 7591)** — it will not take a pre-provisioned +`client_id` / `client_secret` nor a static bearer header. BeaconMCP supports +this by minting a one-off bootstrap URL from the dashboard: the URL lets +ChatGPT register a derived OAuth client tied to your account. 2FA is +preserved — at authorization time, you still type your own TOTP from your +phone; the derived client has no TOTP seed of its own. + +**One-time setup:** + +1. Enable the feature in `beaconmcp.yaml`: + ```yaml + server: + allow_dynamic_registration: true + ``` + Then restart `beaconmcp serve`. + +**To add ChatGPT (from your phone, no laptop needed):** + +1. In your mobile browser, open `https:///app/connectors`, sign in with your TOTP from your authenticator app. +2. Enter a label (e.g. `ChatGPT iPhone`), type your current TOTP, submit. You get a one-off URL of the shape `https:///mcp/c/`. The URL is **single-use** and expires in 15 min. +3. In the ChatGPT app: **Settings → Connectors → Add custom**. + - **Name:** BeaconMCP + - **URL:** paste the `/mcp/c/` URL. + - **Authentication:** OAuth. +4. ChatGPT fetches the OAuth metadata, POSTs to the slug-gated `/oauth/register/c/` — BeaconMCP consumes the slug atomically and mints a derived client scoped to your account. +5. ChatGPT then redirects you to BeaconMCP's authorization page. Type your TOTP from your phone. Token lifetime: 24 h. +6. From now on, ChatGPT auto-refreshes via the authorization code flow. Every 24 h it re-prompts for your TOTP — no re-registration, no new slug. + +**Revocation:** `https:///app/connectors` lists every active derived client. Revoke one and ChatGPT loses access immediately. Revoking your human account cascades to every derived client automatically. + +**Why not a static bearer?** ChatGPT's connector UI has no "Authorization header" field — only "No authentication" or "OAuth" — and the OAuth path strictly requires DCR. The slug-gated bootstrap is the narrow, audit-friendly way to let it in while keeping your TOTP on your phone. + +--- + +## Perplexity (not supported) + +> ⚠ Perplexity is deprecating MCP. In March 2026, Perplexity's CTO +> announced that the company is moving to direct REST APIs and a +> "Code Mode" execution model, citing OAuth / DCR friction and +> context-window waste from MCP tool schemas. No setup instructions +> here — there is no working integration to document. + +## ChatGPT Codex (OAuth 2.1, terminal/IDE) + +Codex is OpenAI's terminal/IDE MCP client. Unlike the web connector it +lets you pre-register credentials in `config.toml`, so no slug needed. +Codex catches the OAuth redirect on an ephemeral local port. + +1. On the server: `beaconmcp auth create --name "Codex"`. +2. Add BeaconMCP to Codex's `config.toml`: + ```toml + [mcp_servers.beaconmcp] + url = "https:///mcp" + client_id = "beaconmcp_..." + client_secret = "sk_..." + ``` +3. Run `codex mcp login beaconmcp`. Codex binds a loopback listener and + opens your browser on BeaconMCP's authorization page — type your TOTP. + +**Remote dev environments** (Codespaces, SSH container): set +`mcp_oauth_callback_url` in `config.toml` to your ingress URL so the +redirect hits the right host instead of localhost. A matching port can +be pinned via `mcp_oauth_callback_port`. + +--- + +## OpenCode (OAuth 2.1) + +OpenCode accepts a pre-registered `client_id` / `client_secret` in +`opencode.json`, and also supports DCR as a fallback. Tokens live in +`~/.local/share/opencode/mcp-auth.json` and refresh automatically. + +**Recommended — OAuth 2.1 (pre-registered):** + +1. `beaconmcp auth create --name "OpenCode"`. +2. Add to `opencode.json` (or `~/.config/opencode/opencode.json`): + ```json + { + "mcp": { + "beaconmcp": { + "type": "remote", + "url": "https:///mcp", + "enabled": true, + "oauth": { + "clientId": "beaconmcp_...", + "clientSecret": "sk_..." + } + } + } + } + ``` +3. Run `opencode mcp auth beaconmcp`. Type your TOTP in the browser. + +**Alternative — DCR** (requires `allow_dynamic_registration: true`): + +```json +{ + "mcp": { + "beaconmcp": { + "type": "remote", + "url": "https:///mcp/c/", + "enabled": true, + "oauth": true + } + } +} +``` + +**Alternative — Bearer:** + +```json +{ + "mcp": { + "beaconmcp": { + "type": "remote", + "url": "https:///mcp", + "enabled": true, + "oauth": false, + "headers": { "Authorization": "Bearer " } + } + } +} +``` + +--- + +## Gemini + +### Gemini CLI (OAuth 2.1, recommended) + +Gemini CLI accepts a pre-registered `client_id` / `client_secret` in +`settings.json`, so no DCR slug is needed. `/mcp auth beaconmcp` then +opens the browser flow with your TOTP prompt on BeaconMCP's page. + +1. `beaconmcp auth create --name "Gemini CLI"`. +2. Add to `~/.gemini/settings.json`: + ```json + { + "mcpServers": { + "beaconmcp": { + "httpUrl": "https:///mcp", + "oauth": { + "clientId": "beaconmcp_...", + "clientSecret": "sk_..." + } + } + } + } + ``` +3. In the CLI: `/mcp auth beaconmcp`. Type your TOTP in the browser. + +Bearer header is also supported as a fallback: + +```bash +gemini mcp add beaconmcp \ + --url https:///mcp \ + --header "Authorization: Bearer " +``` + +### Gemini Web / Mobile / macOS native app (not supported yet) + +Gemini's consumer web UI (gemini.google.com), the iOS / Android apps, and +the new macOS native app do **not** expose a custom-MCP connector today. +The only Gemini surfaces that can reach BeaconMCP are **Gemini CLI** and +**Antigravity**. + +### Gemini API (google-genai SDK) + +For programmatic Gemini API usage, BeaconMCP is passed as a remote MCP +tool. Obtain the bearer interactively from the dashboard rather than +letting the process derive TOTP codes on its own. + +```python +import os +from google import genai + +token = os.environ["BEACONMCP_TOKEN"] + +client = genai.Client() +response = client.models.generate_content( + model="gemini-2.0-flash", + contents="List the VMs on pve1", + config={ + "tools": [ + { + "mcp_servers": [ + { + "url": "https:///mcp", + "headers": {"Authorization": f"Bearer {token}"}, + } + ] + } + ] + }, +) +``` + +Long-running services should rotate tokens on a schedule (an operator +typing the TOTP) rather than embedding the seed. + +### Google Antigravity + +Antigravity's visual connection manager handles both OAuth 2.1 and +Bearer. OAuth keeps the TOTP prompt on BeaconMCP's side; Bearer is a +quick fallback. + +**Recommended — OAuth 2.1 (pre-registered):** + +1. `beaconmcp auth create --name "Antigravity"`. +2. In Antigravity: *Customizations → Connections → Add MCP server*. Paste + the URL (`https:///mcp`) and the OAuth client credentials. +3. Authorize in the browser popup — your TOTP prompt shows up on + BeaconMCP's page. + +**Alternative — Bearer:** + +Antigravity also reads MCP servers from +`~/.gemini/antigravity/mcp_config.json` (macOS / Linux) or +`%USERPROFILE%\.gemini\antigravity\mcp_config.json` (Windows). Top-level +key is `mcpServers` and the HTTP URL field is **`serverUrl`** (not +`url`): + +```json +{ + "mcpServers": { + "beaconmcp": { + "serverUrl": "https:///mcp", + "headers": { + "Authorization": "Bearer " + } + } + } +} +``` + +If the native HTTP transport misbehaves, fall back to `mcp-remote`: + +```json +{ + "mcpServers": { + "beaconmcp": { + "command": "npx", + "args": [ + "-y", "mcp-remote", + "https:///mcp", + "--header", "Authorization: Bearer " + ] + } + } +} +``` + +--- + +## Mistral + +### Le Chat (OAuth 2.1) + +Le Chat speaks OAuth 2.1 natively. Same flow as Assistant — point it at +the bare `/mcp` URL and it handles the rest. + +1. In Le Chat: *Intelligence → Connecteurs → Ajouter un connecteur → Connecteur MCP personnalisé*. +2. Fill in: + - **Name:** BeaconMCP + - **Description:** (optional) + - **MCP Server URL:** `https:///mcp` +3. Validate. Le Chat discovers the OAuth metadata and redirects you to + BeaconMCP's authorization page — type your TOTP from your phone. + Token lifetime: 24 h; Le Chat refreshes via the authorization code + flow on its own. + +Custom connectors are on Le Chat Pro / Enterprise; the free tier may +hide the panel. + +**CORS:** add `https://chat.mistral.ai` to `server.allowed_origins` +(see the allowlist note at the top of this file). + +### Mistral Vibe + +> ⚠ Unverified — Vibe's bearer support hasn't been tested against a +> live BeaconMCP instance. If it doesn't work out of the box, check the +> latest Vibe docs (the schema has been iterating fast) and report back. + +Vibe reads its config from `./.vibe/config.toml` (per-project) or +`~/.vibe/config.toml` (global). **TOML format**, not JSON: + +```toml +[[mcp_servers]] +name = "beaconmcp" +transport = "http" +url = "https:///mcp" +headers = { "Authorization" = "Bearer " } +``` + +`transport` accepts `"http"`, `"streamable-http"`, or `"stdio"`. Each +server is its own `[[mcp_servers]]` array entry. + +--- + +## VS Code (OAuth 2.1) + +VS Code routes MCP authentication through its native Authentication +Provider system — the same flow used for GitHub / Microsoft Entra +logins. It reads `WWW-Authenticate`, shows a toast to Allow, catches +the redirect on the `vscode://` (or `vscode-insiders://`) OS URI scheme, +and stores the resulting token in your OS keychain. + +**Recommended — OAuth 2.1 (pre-registered):** + +1. `beaconmcp auth create --name "VS Code"`. +2. Add to `.vscode/mcp.json` (or `settings.json → "mcp.servers"`): + ```json + { + "inputs": [ + { "type": "promptString", "id": "beaconmcp-client-id", "description": "client_id" }, + { "type": "promptString", "id": "beaconmcp-client-secret", "description": "client_secret", "password": true } + ], + "servers": { + "beaconmcp": { + "type": "http", + "url": "https:///mcp", + "clientId": "${input:beaconmcp-client-id}", + "clientSecret": "${input:beaconmcp-client-secret}" + } + } + } + ``` +3. VS Code prompts you on first use — TOTP on BeaconMCP's page, OS + keychain stores the bearer afterward. + +**Alternative — Bearer:** + +```json +{ + "servers": { + "beaconmcp": { + "type": "http", + "url": "https:///mcp", + "headers": { "Authorization": "Bearer " } + } + } +} +``` + +Verify with *Command Palette → MCP: List Servers*. + +--- + +## Cursor (OAuth 2.1) + +Cursor is a first-class OAuth 2.1 client since v1.0. It surfaces a +blue *Connect* button in *Settings → Tools & MCP* and catches the +redirect via the `cursor://` OS scheme (or a loopback fallback). + +**Recommended — OAuth 2.1 (pre-registered):** + +1. `beaconmcp auth create --name "Cursor"`. +2. Add to `~/.cursor/mcp.json` (global) or `.cursor/mcp.json` (per + project): + ```json + { + "mcpServers": { + "beaconmcp": { + "url": "https:///mcp", + "clientId": "${env:BEACONMCP_CLIENT_ID}", + "clientSecret": "${env:BEACONMCP_CLIENT_SECRET}" + } + } + } + ``` +3. Export the credentials in your shell. Reload Cursor; click *Connect* + when "Needs authentication" appears. + +**Alternative — Bearer:** + +```json +{ + "mcpServers": { + "beaconmcp": { + "url": "https:///mcp", + "headers": { "Authorization": "Bearer ${env:BEACONMCP_TOKEN}" } + } + } +} +``` + +Cursor expands `${env:VAR}` natively so the bearer can live in your +shell environment rather than in the repo. + +--- + +## Other MCP-over-HTTP clients + +Any client that can send a bearer on `https:///mcp` works the +same way: create a token from `/app/tokens` after typing your TOTP, +configure the client to send `Authorization: Bearer `, revoke from +the same page when you are done. If the client natively speaks OAuth 2.1 +(like Assistant) or OAuth + DCR (like ChatGPT / OpenCode), prefer those flows +— they keep the TOTP prompt at the authorization page instead of relying +on a stored bearer. diff --git a/docs/cloudflare.md b/docs/cloudflare.md new file mode 100644 index 0000000..0c8ace2 --- /dev/null +++ b/docs/cloudflare.md @@ -0,0 +1,176 @@ +# Running BeaconMCP behind Cloudflare + +BeaconMCP works behind `cloudflared` (Cloudflare Tunnel) and the Cloudflare +proxy, but Cloudflare's default security posture **breaks MCP out of the box**. +MCP clients (Claude, ChatGPT, Assistant Desktop) are headless HTTP clients, not +browsers — Cloudflare's bot/WAF/Access defenses treat them as suspicious and +either block the request or strip the `Authorization` header before it reaches +BeaconMCP. + +## Symptom + +The MCP endpoint returns **401 or 403 only via the public Cloudflare URL** while +`curl http://localhost:8420/mcp` (on the box) works fine. If you see a 401 body +containing a `hint` that mentions `cf-ray`, BeaconMCP detected the request came +through Cloudflare with no usable bearer — a Cloudflare rule ate it. + +A second, distinct symptom: a command **works until one of its arguments +contains an attack-shaped string** — a path like `/etc/passwd`, a traversal +`../../`, an `http://host/x.php?...` URL, or a SQL-ish quote — at which point you +get Cloudflare's full-page **"Sorry, you have been blocked"** HTML (with a Ray +ID), *not* a JSON 401/403. That is the OWASP Managed Ruleset, not auth: the +request is dropped at the edge and never reaches BeaconMCP, so there is no +server-side fix — only the WAF skip in §1. Verified against a live deployment: +`ssh_run` with `curl http://evil.example.com/shell.php?x=../../../../etc/passwd` +is blocked, while the same call with a plain URL or no URL is not. + +## Why it happens + +| Cloudflare feature | What it does to MCP | +|--------------------|---------------------| +| Bot Fight Mode / Super Bot Fight Mode / Browser Integrity Check | Challenges or blocks non-browser clients (MCP sends no browser headers) → **403 / challenge HTML** before the app sees the request. | +| WAF Managed Rules (OWASP CRS) | False-positive on a tool **argument** that pattern-matches an attack signature — path traversal (`../../`), `/etc/passwd`, an RFI/LFI URL (`http://host/x.php?...`), SQL-ish quotes. These appear constantly in legitimate infra commands (`cat /etc/passwd`, `grep -r ../`, `curl http://repo/file.tar.gz`), so Cloudflare serves its **"Sorry, you have been blocked"** HTML → **403** and BeaconMCP never sees the request. | +| Cloudflare Access | Sits in front of the hostname and **consumes/strips the `Authorization` header** (it owns that header for its own JWT) → BeaconMCP sees no bearer → **401**. | +| "Cache Everything" / buffering | Buffers or caches the streamable-HTTP / SSE response → the MCP stream hangs or returns stale data. | + +All four are fixed below. Apply them to the **MCP and OAuth paths only** so the +rest of your zone keeps Cloudflare's protection. + +--- + +## 1. WAF custom rule: skip bot/WAF protection for MCP + OAuth paths + +**Dashboard:** *Security → WAF → Custom rules → Create rule* (per-zone). +For Bot Fight Mode skip on the legacy plan, the equivalent is *Security → Bots*; +on Free, prefer creating the skip rule below and disabling **Bot Fight Mode** +zone-wide if you cannot scope it. + +Rule expression (copy-paste into the *Edit expression* box): + +``` +(starts_with(http.request.uri.path, "/mcp")) or (starts_with(http.request.uri.path, "/oauth/")) or (starts_with(http.request.uri.path, "/.well-known/")) +``` + +Action: **Skip**, and check every relevant box: + +- Skip → **All remaining custom rules** +- Skip → **Managed rules** (WAF Managed Ruleset + OWASP) +- Skip → **Super Bot Fight Mode** +- Skip → **Browser Integrity Check** +- (if present) Skip → **Rate limiting rules** + +Place this rule **first** in the custom-rules list so it short-circuits before +anything blocks the request. + +> **What you give up, and what replaces it.** This rule removes Cloudflare's +> protection from your *authentication* endpoints, which is not a decision to +> make blindly. BeaconMCP enforces its own equivalents on those paths: a +> per-IP rate limit on `/oauth/token` and `/authorize`, a 5-strike / 5-minute +> TOTP lockout per client, and mandatory OAuth 2.1 + TOTP on every exchange. +> That is why skipping is safe *here* and only here — keep the expression +> anchored with `starts_with` so it can't widen to the rest of your zone. +> (`contains "/mcp"` would also match `/anything/mcp-foo`.) + +> The paths cover the MCP endpoint (`/mcp`, `/mcp/c/`), the OAuth +> authorization/token/registration endpoints (`/oauth/...`), and RFC 9728 / +> RFC 8414 discovery (`/.well-known/...`). Discovery must be reachable +> unauthenticated or clients can't find your authorization server. + +--- + +## 2. Cloudflare Access (Zero Trust) — exclude /mcp or pass the header through + +If the hostname is protected by **Cloudflare Access**, Access owns the +`Authorization` header and strips the client's bearer. You have two options: + +**A. Bypass Access for the MCP/OAuth paths (simplest).** +*Zero Trust → Access → Applications → your app → Policies* — add a **Bypass** +policy, or scope the application's path so it does **not** cover `/mcp`, +`/oauth/`, or `/.well-known/`. BeaconMCP enforces its own OAuth 2.1 + TOTP on +these paths, so Access in front of them is redundant and harmful. + +**B. Keep Access but pass `Authorization` through.** +If you must keep Access on the path, ensure it does not consume the header: +use an Access **service token** for the automated client and confirm the +`Authorization: Bearer ` header survives to the origin (test +with the diagnostic in §5). In practice, **option A is strongly recommended** — +Access and BeaconMCP both want the `Authorization` header, and only one can win. + +--- + +## 3. No caching / no buffering on /mcp + +Streamable-HTTP and SSE must stream straight through. **Do not** apply +"Cache Everything" to the MCP path. + +**Dashboard:** *Caching → Cache Rules → Create rule*: + +- When incoming requests match: `(http.request.uri.path contains "/mcp")` +- Then: **Bypass cache**. + +If you have a zone-wide "Cache Everything" Page Rule / Cache Rule, add the +bypass rule **above** it. Cloudflare does not buffer responses by default, but +an aggressive cache rule on this path will break the MCP transport. + +--- + +## 4. BeaconMCP-side configuration + +In `beaconmcp.yaml`: + +```yaml +server: + # The public hostname Cloudflare forwards to BeaconMCP. Without it the MCP + # SDK rejects requests with 421 Misdirected Request (DNS-rebinding guard). + allowed_hosts: + - mcp.example.com + # Trust Cloudflare's edge so its forwarded headers are honoured: X-Forwarded- + # For for the auth rate-limit client IP, and X-Forwarded-Host for the host in + # the OAuth issuer and token MCP URLs. The literal "cloudflare" auto-expands + # to Cloudflare's IP ranges. + trusted_proxies: + - cloudflare +``` + +Restart BeaconMCP after editing. (Both also have env equivalents: +`BEACONMCP_ALLOWED_HOSTS`, `BEACONMCP_TRUSTED_PROXIES`.) + +--- + +## 5. Verify + +From your laptop (not the server), with a valid bearer: + +```bash +# Discovery must be reachable unauthenticated (no challenge HTML, HTTP 200): +curl -i https://mcp.example.com/.well-known/oauth-protected-resource + +# An unauthenticated /mcp POST should return a clean JSON 401 from BeaconMCP +# (NOT a Cloudflare challenge / "Just a moment..." HTML page): +curl -i -X POST https://mcp.example.com/mcp + +# With a real bearer it should reach the app (no 403, header survives): +curl -i -X POST https://mcp.example.com/mcp \ + -H "Authorization: Bearer " \ + -H "Content-Type: application/json" \ + -d '{"jsonrpc":"2.0","id":1,"method":"ping"}' +``` + +Checks: + +- A **403** or an HTML "Just a moment…" / "Checking your browser" page on any + of the above → a bot/WAF rule is still blocking. Revisit §1. +- A **401** whose JSON body contains a `hint` about `cf-ray` even though you + *did* send a valid bearer → Access/WAF stripped your `Authorization` header. + Revisit §2. BeaconMCP logs the matching warning to + `journalctl -u beaconmcp` (look for `cf-ray=` + "Cloudflare"). That warning + is throttled to one line per 5 minutes — a public endpoint gets scanned + constantly and every drive-by hit carries a `cf-ray` — so the line reports + how many similar events it stands for rather than one line per request. +- A clean `401 {"error":"unauthorized"}` with no bearer, and a `200` with a + valid bearer → everything is wired correctly. + +## See also + +- `docs/troubleshooting.md` — symptom/fix table. +- `README.md` → *Expose publicly* — `allowed_hosts` / `trusted_proxies`. diff --git a/docs/configuration.md b/docs/configuration.md new file mode 100644 index 0000000..7de278e --- /dev/null +++ b/docs/configuration.md @@ -0,0 +1,81 @@ +# Configuration + +BeaconMCP reads two files at startup. + +**`beaconmcp.yaml`** — topology and feature flags. Path resolution, in order: the `--config` flag, +the `BEACONMCP_CONFIG` env var, `./beaconmcp.yaml`, `/etc/beaconmcp/config.yaml`. The full annotated +schema lives in [`beaconmcp.yaml.example`](../beaconmcp.yaml.example). + +**`.env`** — the secrets the YAML references as `${VAR}`. A missing reference fails the startup +check and names the offending YAML path. + +`beaconmcp validate-config` resolves everything and prints the result with secrets masked, without +starting the server. Run it after every edit. + +Legacy `PVE*_*`, `ILO_*` and `SSH_*` env vars still work when no YAML is found. They are deprecated +and disappear in 2.1. + +## Where to run it + +Put BeaconMCP on the same local network as the cluster: on one of the nodes, in a dedicated LXC or +VM, or in a container with host networking. + +That way every `proxmox.nodes[].host` is a plain LAN IP (`10.0.0.1`, `10.0.0.2`) usable as-is for +the Proxmox API on `:8006` **and** for SSH on `:22`. Two features depend on that being true: the +`ssh.inherit_proxmox_nodes` shortcut, and the SSH-jump tunnel `bmc_*` uses to reach HP iLO devices +on a private management VLAN. + +A public FQDN with a reverse-proxy port (`pve2.example.com:443`) pins the entry to HTTPS and breaks +SSH inheritance, since sshd listens on port 22 of the node and not behind your HTTPS tunnel. For a +genuinely remote node, declare it explicitly under `ssh.hosts[]` with its real SSH address: a +Tailscale IP, a VPN address, a bastion. + +## Server + +| Key | Notes | +|-----|-------| +| `server.allowed_hosts` | DNS-rebinding allowlist. **Must** include the public FQDN behind your reverse proxy, or requests come back `421 Misdirected Request`. | +| `server.allowed_origins` | Web-origin allowlist, used for browser CORS preflights and for OAuth HTTPS redirect URIs. Desktop and CLI callbacks (`vscode://`, `cursor://`, loopback) are handled separately. | +| `server.trusted_proxies` | Direct peers whose forwarded headers are trusted, as IPs or CIDRs. Governs `X-Forwarded-For` (the auth rate-limit client IP) **and** `X-Forwarded-Host` (the host advertised in the OAuth issuer and the token/connector MCP URLs). When empty, both fall back to the request's own peer / `Host` and forwarded values from any peer are ignored. The value `cloudflare` auto-expands to Cloudflare's edge ranges. | +| `server.tokens_db` | SQLite file persisting *named* API tokens (the `/app/tokens` page) across restarts. Created owner-only (0600). Defaults to `tokens.db` next to `clients_file`. Env override: `BEACONMCP_TOKENS_DB`. | +| `server.named_token_ttl` | Lifetime of named API tokens, in seconds. Default `2592000` (30 days); `0` means never expires, revoke-only. Internal OAuth and session bearers keep their fixed 24 h TTL either way. Env override: `BEACONMCP_NAMED_TOKEN_TTL`. | +| `server.audit_log` | JSON-lines audit log covering tool calls, dashboard logins, OAuth authorize and client revokes. Created owner-only (0600). Default `/opt/beaconmcp/audit.log`; `-` keeps stderr only. Env override: `BEACONMCP_AUDIT_LOG`. | +| `server.transfers_max_mb` | Size cap for `proxmox_upload_file` / `proxmox_download_file`. Default 500. | + +## Proxmox + +| Key | Notes | +|-----|-------| +| `proxmox.nodes[]` | One entry per node, each with its own API token. Prefer a LAN IP in `host:` — it's the one string that works for both the API and SSH inheritance. `localhost` is fine when BeaconMCP runs on that node. Reserve FQDN-with-port for nodes you cannot reach on the LAN, and give those an explicit `ssh.hosts[]` entry. | + +## SSH + +| Key | Notes | +|-----|-------| +| `ssh.hosts[]` | One entry per SSH target: VPS, Proxmox node, jump box. Each carries its own `user` plus exactly one of `password` / `key_file`. Names may match `proxmox.nodes[].name`. Per-host `known_hosts` and `strict_host_key_checking` override the global `ssh.*` settings, but only to tighten them. | +| `ssh.defaults` + `ssh.inherit_proxmox_nodes` | Homelab shortcut: set `defaults:` (user + password/key_file), flip `inherit_proxmox_nodes: true`, and every Proxmox node becomes SSH-reachable under its own name without duplication. An explicit `ssh.hosts[]` entry still wins when it matches a node by name or address. | +| `ssh.vmid_to_ip` | Optional template (`"192.168.1.{id}"`) used by `ssh_run` when `host` is a bare VMID. The resolved IP must match an `ssh.hosts[].host` to authenticate. Omit it to disable numeric shortcuts entirely. | + +## BMC + +| Key | Notes | +|-----|-------| +| `bmc.devices[]` | Zero or more BMCs. `type` is `hp_ilo`, `ipmi`, `idrac`, `supermicro` or `redfish`. iDRAC (14G+) and Supermicro (X11+) are served by the Redfish backend. `jump_host` is optional: set it to a `proxmox.nodes[]` name to route the connection over an SSH tunnel, which is how you reach a management VLAN. | + +## Dashboard + +| Key | Notes | +|-----|-------| +| `features.dashboard.limits` | Per-5h and per-week USD caps on the Gemini chat. `0` disables that window. | + +Everything else about the panel — enabling it, the tokens page, cost tracking, the confirmation +modal — is in [dashboard.md](dashboard.md). + +## Updates + +| Key | Notes | +|-----|-------| +| `features.updates.enabled` | Default `true`. Compares this checkout against the upstream default branch and shows a notice to signed-in operators. Set to `false` on an air-gapped or change-controlled box: the server then never contacts the git remote, and the `beaconmcp_*_update` MCP tools are not registered. | +| `features.updates.allow_self_update` | Default `true`. Set to `false` to keep the notification but forbid applying it from the dashboard or over MCP — the right setting when deploys go through a pipeline. Manual instructions are still shown. | + +See [updates.md](updates.md) for the notice, the MCP tools, and what the self-update does. diff --git a/docs/dashboard.md b/docs/dashboard.md new file mode 100644 index 0000000..d27239f --- /dev/null +++ b/docs/dashboard.md @@ -0,0 +1,210 @@ +# BeaconMCP web dashboard + +Optional web panel served by BeaconMCP on the same origin as the MCP endpoint (`https:///app/...`). Three pages: + +- **`/app/login`** — exchanges a client id + client secret + a second factor (TOTP code or passkey) for an MCP bearer, and stores it in a 90-day HttpOnly session cookie. Removes the need to issue `curl` requests from a phone. +- **`/app/chat`** — multi-conversation chat with Gemini 3.6 Flash / 3.5 Flash-Lite (GA) or Gemini 3.1 Pro (preview, Google allowlist required). **Requires `GEMINI_API_KEY`.** +- **`/app/tokens`** — generates named bearers so external MCP clients (Gemini web, ChatGPT, Assistant Desktop) can be wired up without the OAuth dance. **Works without `GEMINI_API_KEY`.** + +## Enabling + +The dashboard is on by default as long as a `BEACONMCP_SESSION_KEY` is set. Two modes: + +| Mode | Condition | Active pages | +|------|-----------|--------------| +| **Full** | `GEMINI_API_KEY` set | `/app/login`, `/app/chat`, `/app/tokens` | +| **Tokens only** | `GEMINI_API_KEY` absent | `/app/login`, `/app/tokens` (chat redirects to tokens) | + +1. *(optional)* Add a Gemini API key to enable the integrated chat: + ```env + GEMINI_API_KEY=... + ``` +2. The session encryption key (`BEACONMCP_SESSION_KEY`) is generated by `install.sh` on first run. For manual deployments: + ```bash + echo "BEACONMCP_SESSION_KEY=$(openssl rand -base64 32)" >> /opt/beaconmcp/.env + ``` +3. Restart: `systemctl restart beaconmcp`. + +At boot, the server logs the active dashboard URL and whether the chat is enabled: + +``` +Dashboard: http://0.0.0.0:8420/app/login (chat: enabled) +``` + +or `(chat: disabled, tokens only)` when the Gemini key is missing, or `disabled` when `BEACONMCP_DASHBOARD_ENABLED=false`. + +## Flow + +1. Navigate to `https:///` on any device — the root path redirects to `/app/login`. +2. Enter client id + client secret, then the current TOTP code — or use a passkey instead, from the + link under the code boxes (see [Passkeys](#passkeys)). Enter validates as soon as the + sixth digit lands. +3. A confirmation screen shows when the MCP access and the signed-in session expire, offers to + register a passkey on this device, and waits for **Finish signing in**. +4. Land on `/app/chat` (or `/app/tokens` in tokens-only mode) with the conversation history restored. +5. Every 24 hours the underlying MCP bearer expires. The dashboard prompts only for a fresh TOTP code on `/app/refresh`; the client id and secret remain encrypted server-side. + +With JavaScript disabled the page falls back to the previous behaviour: a plain form POST that +redirects straight to the landing page, no confirmation screen and no passkeys. + +## Passkeys + +A passkey replaces the 6-digit code, not the client secret — signing in still needs both factors. +Enrol one from the confirmation screen right after a normal sign-in, then on the next visit use +**Use a passkey instead** under the code boxes. + +The same option is on the OAuth consent page (`/oauth/authorize`), which now also confirms before +redirecting: it shows when the granted access expires, offers passkey enrolment, and only mints the +authorization code when you press **Finish signing in**. + +Registered passkeys are listed at the bottom of `/app/tokens` with their label, when they were added +and when they were last used, each with a **Remove** button. + +Caveats — hostname binding, the HTTPS/loopback requirement, and why you should keep your TOTP seed — +are covered in [security.md](security.md#passkeys). + +## Tokens page + +Reached from the sidebar or directly at `/app/tokens`. Intended for users wiring BeaconMCP into an external MCP client rather than using the integrated chat. + +It exposes: + +- The MCP URL to paste into the external client (with a copy button). +- A creation form that requires a **name** (60 characters max, e.g. "Gemini Web", "ChatGPT macOS") plus the current TOTP code. +- The newly generated token, shown **once** in an orange card with a copy button. Reloading the page removes it from view. +- The list of active tokens: name, 12-character prefix, hours until expiry, revoke button. + +Constraints: + +| Setting | Value | +|---------|-------| +| Expiry | **24 h** (inherits `TokenStore.TOKEN_TTL`) | +| Per-client cap | **3 active tokens** | +| TOTP | Re-verified on every creation | +| Revocation | By prefix (6 chars minimum), scoped to the owning client | +| Storage | **In-memory** — `systemctl restart` invalidates every token | + +## Chat — models and thinking + +- **Gemini 3.6 Flash** — GA, available on every AI Studio key. **Used by default** (`gemini-3.6-flash`). +- **Gemini 3.5 Flash-Lite** — GA, the cheap high-throughput option (`gemini-3.5-flash-lite`), 5× cheaper in and out than 3.6 Flash. +- **Gemini 3.1 Pro (preview)** — gated by a Google allowlist. Without allowlist access, the dashboard surfaces a clear message pointing back to the two GA models. +- **Thinking effort** — dropdown with `minimal` / `low` / `medium` / `high`. The Gemini 3 family takes a `thinking_level` enum; the token-budget mapping in `_BUDGET_BY_EFFORT` (and the 128-token floor `gemini-2.5-pro` needed) only applies to models outside it, and is kept for conversations that predate the switch. + +Gemini 2.5 Flash / Pro and `gemini-3-flash-preview` were retired from the picker when 3.6 Flash and 3.5 Flash-Lite went GA (2026-07-21). A conversation sitting on one of them is moved to the closest current model by the schema-6 migration; `messages.model` is left alone, so the transcript keeps naming whichever model actually wrote each reply, and `usage.py` keeps their rates so old turns are never re-priced. +- **Markdown rendering** — the client parses headings (`#`–`######`), ordered and unordered lists, blockquotes, horizontal rules, code fences (with `lang-*` class), inline code, bold/italic/strikethrough, and HTTP(S) links. + +## Mandatory confirmation for dangerous tools + +The model's input is untrusted: a log line, a config file, or a web-search result can carry an injected instruction. So two classes of tool **never** run without manual approval from the chat UI. + +**Arbitrary code execution** — `ssh_run`, `proxmox_run`, and the primitives that become code execution in one hop (`proxmox_write_file`, `proxmox_upload_file`, `proxmox_download_file`, `proxmox_delete_transfer`). Writing `~/.ssh/authorized_keys` or a file under `/etc/cron.d` is exactly as good as a shell, which is why the file tools sit alongside the exec tools. + +**Destructive or irreversible** — `vm_bulk_action`, `proxmox_vm_stop`, `proxmox_vm_restart`, `proxmox_vm_migrate`, `proxmox_snapshot_rollback`, `proxmox_snapshot_delete`, `proxmox_backup_restore`, `bmc_power_off`, `bmc_power_reset`, and `beaconmcp_self_update` with `confirm=True` — which pulls new code, reinstalls dependencies and restarts the service, i.e. replaces the very process enforcing this gate. + +Three call shapes are let through without a modal, because they don't change anything: + +- `ssh_run` / `proxmox_run` carrying only `exec_id=` — that's read-only polling of an already-approved session. +- `proxmox_snapshot_create` / `_rollback` / `_delete` called with `dry_run=True` — they only report what they *would* do. The exemption is limited to those three by name, never inferred from the argument: an undeclared `dry_run` is silently dropped during argument validation, so trusting it would let `ssh_run(command=..., dry_run=True)` past the modal and then run for real. +- `proxmox_vm_config` without `updates`, and `beaconmcp_self_update` without `confirm=True` — the read shape of a read-or-write tool. Reading the argument is sound for these two because both parameters are declared by the tool, so what the gate reads is what the tool acts on; that is precisely what makes it unsound for an undeclared `dry_run`. + +When Gemini fires a gated call: + +1. The tool card switches to an "approval required" state (orange badge, auto-expanded so arguments are visible). +2. Two buttons: **Approve** / **Reject**. +3. The Gemini turn blocks server-side until the decision is made (5-minute timeout). +4. On rejection, Gemini receives a `FunctionResponse {"error": "user_rejected"}` and can revise its reply. + +The allow-list is hard-coded in `src/beaconmcp/dashboard/chat.py` (`_NEEDS_CONFIRMATION`, `_CONFIRM_WHEN_ARG_PRESENT`, `_tool_call_requires_confirmation`). Only the integrated chat enforces this gate; external MCP clients (Assistant Desktop, Gemini CLI, ChatGPT MCP) must enable their own per-call approval mode (see the **Security** section of the root README). + +## Interactive panels (MCP Apps) + +Tools that carry `_meta.ui.resourceUri` — `proxmox_vm_panel`, `proxmox_logs_panel`, `cluster_overview_interactive` — render as a live interface in the chat instead of a block of JSON. The dashboard implements both halves of the [MCP Apps extension](https://modelcontextprotocol.io/extensions/apps/overview): it announces `io.modelcontextprotocol/ui` when it opens its MCP session, and it plays host to the `ui://` document over `postMessage`. + +The extension is declared through `ClientCapabilities.extensions`, a field mcp only types in 2.0. The `<2` pin is not in the way: the model accepts extra fields, so the capability serialises under the name the spec gives it and the server reads the same JSON either way (`dashboard/mcp_bridge.py`). + +### How a panel is isolated + +The document is served by `/app/api/mcp/panel` and framed with `sandbox="allow-scripts"` and **no** `allow-same-origin`. That puts it on an opaque origin: it cannot read the session cookie, cannot read the CSRF token, cannot reach into the parent page. Its response also carries its own `Content-Security-Policy` — `default-src 'none'`, `connect-src 'none'`, `frame-ancestors 'self'` — so it cannot open a socket of its own either. + +What is left is `postMessage` to the parent. Every tool call a panel makes therefore goes through `/app/api/mcp/call`, which is session-authenticated and CSRF-protected, and where the policy below is applied. + +### What a panel may call on its own + +A panel button is a labelled control a human clicked, so the approval modal — which exists because the *model's* input is untrusted — would restate the click rather than check it. Panel calls are therefore not gated. What is *not* granted is a blanket exemption for anything running in a frame: a `ui://` document is HTML the server wrote, and this dashboard is a general MCP host, so a blanket rule would hand every connected server a way around the gate it is documented to be subject to. + +The exemption is a closed list, enforced server-side in `panel_call_allowed()`: + +- `proxmox_vm_start` / `_stop` / `_restart` — one guest per call, visible in the panel, reversible from it. +- `proxmox_vm_config`, but only when every key in `updates` is sizing (`cores`, `sockets`, `memory`, `balloon`, `cpulimit`, `cpuunits`). Exempting the tool itself would exempt `hookscript`, raw QEMU `args` and device passthrough along with it. +- Everything that was never gated in the first place — the read-only tools, including the three panel tools themselves. + +Anything else is refused with `403 confirmation_required`, and the panel shows the reason. It is refused rather than prompted because there is no turn in flight to hang a modal on — and because the panel already has a way through: `ui/message` hands the request to the model, which puts it back under the modal where it belongs. + +### Model context + +A panel that acts on the cluster pushes the fresh state back with `ui/update-model-context`. The page holds the latest update per panel and sends it with the next message, labelled as coming from the panel rather than from the operator. Without it, stopping a VM from the panel would leave the next turn believing it still runs — the button's result goes to the iframe, not into the conversation. + +### Reopening from history + +Only the `ui://` URI is stored with the tool call, never the snapshot behind it. A panel in an older conversation renders as an **Open panel** button; clicking it mounts the frame and refetches. Live figures in a panel that has been sitting in the transcript for a week would be worse than a short spinner. + +## Stored data + +SQLite at `/opt/beaconmcp/dashboard.db` (WAL mode). Five tables: + +- `sessions` — cookie → client id + AES-GCM-encrypted client secret + current bearer. +- `conversations` — title, model, effort, owning client. +- `messages` — user/assistant, content, tool calls JSON, thinking summary. +- `usage_events` — per-turn ledger: client id, tokens (prompt/cached/output), USD cost, timestamp. +- `usage_5h_sessions` — one row per client holding the current 5-hour session (materialized to avoid a `GROUP BY` on every pre-check). + +Everything is scoped by `client_id`; a single client can hold multiple active sessions (phone + laptop). + +## Usage limits and cost + +Each chat turn computes its USD cost from the `usage_metadata` returned by Gemini. Input tokens are billed at the cached-discount rate whenever `cachedContentTokenCount` is non-zero — Gemini 2.5+ applies implicit caching automatically once the prompt crosses 1024 tokens for Flash or 4096 tokens for Pro, with no client-side work. + +Two windows are enforced **per OAuth client**: + +| Window | Semantics | Env variable | Default | +|--------|-----------|--------------|---------| +| **5h** | Anthropic-style session: opens on the first message after ≥5 h of idle time, lasts exactly 5 h, then closes. | `BEACONMCP_DASHBOARD_LIMIT_5H_USD` | `2.0` | +| **Week** | Rolling sum over the last 7 days. | `BEACONMCP_DASHBOARD_LIMIT_WEEK_USD` | `10.0` | + +Setting a variable to `0` disables that window. When a cap is exceeded, the next request is rejected **before** being sent to Gemini, with a message stating when the 5 h window resets. (When `beaconmcp.yaml` is used, the same caps are configured under `features.dashboard.limits`.) + +The chat footer shows a compact `5H XX% · 7D XX%` line updated after every turn via an SSE `usage_update` event. Clicking the bar opens a modal with progress bars, the 5 h reset time, a "rolling 7-day window" label, and a refresh button. + +Rates used (USD per 1 M tokens, aligned with the public Google AI Studio pricing on 2026-07-29). The retired rows are kept because the ledger re-prices stored turns, and dropping a rate would silently re-bill that history at the fallback model's price: + +| Model | Input | Cached | Output | +|-------|-------|--------|--------| +| `gemini-3.6-flash` | $1.50 | $0.15 | $7.50 | +| `gemini-3.5-flash-lite` | $0.30 | $0.03 | $2.50 | +| `gemini-3.1-pro-preview` (≤200k / >200k) | $2.00 / $4.00 | $0.20 / $0.40 | $12.00 / $18.00 | +| `gemini-2.5-flash` *(retired)* | $0.30 | $0.03 | $2.50 | +| `gemini-2.5-pro` *(retired, ≤200k / >200k)* | $1.25 / $2.50 | $0.125 / $0.25 | $10.00 / $15.00 | +| `gemini-3-flash-preview` *(retired)* | $0.50 | $0.05 | $3.00 | + +Constants live in `src/beaconmcp/dashboard/usage.py` — update them when Google adjusts its prices. + +## Internal MCP architecture + +The dashboard keeps its own MCP session (`streamablehttp_client` + `ClientSession`) pointed at the local `/mcp` endpoint (`http://127.0.0.1:8420/mcp`). Tools are converted manually into `FunctionDeclaration` objects and the `function_call` / `function_response` loop is orchestrated server-side (AFC SDK disabled via `AutomaticFunctionCallingConfig(disable=True)`). This works around known `google-genai` bugs with Gemini 2.5 Pro + MCP + streaming + thinking. + +> A remote mode (server-driven `McpServer`) exists but is disabled: requests through Cloudflare Tunnel lose the `Authorization` header, producing systematic 500 INTERNAL responses. Setting `BEACONMCP_DASHBOARD_MCP_MODE=remote` logs a startup warning and returns an actionable error on every chat turn. + +## Resilience to restarts + +`TokenStore` lives in memory. After `systemctl restart beaconmcp`, bearers are invalidated while dashboard sessions (SQLite) persist. The dashboard detects this by calling `TokenStore.validate()` on every sensitive route; when a bearer is gone but the session timestamp is still valid, the user is routed to `/app/refresh` to enter a fresh TOTP code and mint a new bearer. + +Consequence for externally-issued tokens (`/app/tokens`): a service restart forces every Gemini-web / ChatGPT / Assistant-Desktop integration to regenerate its token. If this is operationally annoying, move `TokenStore` to SQLite (not done today). + +## Disabling entirely + +```env +BEACONMCP_DASHBOARD_ENABLED=false +``` + +To keep the tokens page but drop the chat, leave `GEMINI_API_KEY` unset. diff --git a/docs/installation.md b/docs/installation.md new file mode 100644 index 0000000..7a5f743 --- /dev/null +++ b/docs/installation.md @@ -0,0 +1,157 @@ +# Installation + +Two supported paths: Docker, or the bare-metal install script that registers a systemd service. +They expose the same CLI and the same HTTP surface, so pick whichever fits your infra. + +## Requirements + +- Python 3.11+ (bare-metal path only; the image ships its own) +- Proxmox VE 8.x with an API token per node (Datacenter → Permissions → API Tokens) +- `ipmitool` on the BeaconMCP host, if any BMC is declared with `type: ipmi` +- a reachable jump host (usually a Proxmox node) if your HP iLO devices sit on a private + management VLAN +- `GEMINI_API_KEY`, if you want the integrated chat panel + +## Option A — Docker + +Needs Docker Engine 20.10+ with the Compose plugin. Run it on the Proxmox node itself, in an LXC/VM +on the same LAN, or on any box that reaches every node's API and SSH port directly. + +```bash +git clone https://github.com/Showdown76py/BeaconMCP.git +cd BeaconMCP +cp beaconmcp.yaml.example beaconmcp.yaml +cp .env.example .env +docker compose up -d +``` + +The bundled [`docker-compose.yml`](../docker-compose.yml) uses `network_mode: host`, so the +container sits directly on the LAN and LAN IPs in `proxmox.nodes[].host` work for both the Proxmox +API and SSH. State (OAuth clients, dashboard DB, usage history) lives in a named volume +`beaconmcp-state` and survives `docker compose up --build`. + +Run the one-time setup while the container is up: + +```bash +docker compose exec beaconmcp beaconmcp validate-config +docker compose exec beaconmcp beaconmcp auth create --name "Assistant Web" +curl http://localhost:8420/health # {"status":"ok","server":"beaconmcp"} +``` + +**SSH key files.** Host paths such as `~/.ssh/id_ed25519` are resolved inside the container, where +they don't exist. If any `ssh.hosts[]` entry (or `ssh.defaults`) uses `key_file:`, either copy the +keys into the `beaconmcp-state` volume and point at `/state/keys/...`, or uncomment the `~/.ssh` +bind mount in the compose file. + +## Option B — bare metal (systemd) + +SSH to the machine that will host BeaconMCP, usually your primary Proxmox node: + +```bash +git clone https://github.com/Showdown76py/BeaconMCP.git /opt/beaconmcp +cd /opt/beaconmcp +sudo bash deploy/install.sh +``` + +The script creates a `beaconmcp` system user, installs the package in editable mode, registers a +systemd unit, and uses `/opt/beaconmcp` for persistent state. + +Then write the config (below), provision a client, and start the service: + +```bash +sudo systemctl enable --now beaconmcp +curl http://localhost:8420/health +``` + +## Writing the config + +`beaconmcp.yaml` holds the topology, `.env` holds the secrets it references. Two ways to produce +them. + +**Guided.** A terminal wizard walks through each capability (Proxmox nodes, SSH, BMC, server) with a +live YAML preview, and appends `${VAR}` placeholders to `.env` for the secrets you fill in +afterwards. The same command edits an existing file: it parses your YAML back into the forms, so you +can tweak and re-save without losing anything. + +```bash +pip install 'beaconmcp[wizard]' # pulls the optional textual dependency +beaconmcp init # creates OR edits beaconmcp.yaml, extends .env +beaconmcp init --blank # force a fresh draft even if the YAML exists +``` + +Arrow keys browse sections, `enter` opens a form, `ctrl+s` saves without quitting, `q` exits. + +**Manual.** Copy the example and edit it: + +```bash +cp beaconmcp.yaml.example /opt/beaconmcp/beaconmcp.yaml +cp .env.example /opt/beaconmcp/.env +``` + +Either way, check the result before starting the server. `validate-config` resolves every `${VAR}`, +prints the config with secrets masked, and summarises what would be registered: + +```bash +beaconmcp validate-config +``` + +Key-by-key reference: [configuration.md](configuration.md). + +## Provisioning an OAuth client + +```bash +beaconmcp auth create --name "Assistant Web" +``` + +This prints a client id, a client secret and a TOTP seed with an ASCII QR code. **Both secrets are +shown exactly once.** Scan the QR into an authenticator app (Google Authenticator, Authy, 1Password, +Aegis) right away — on a device you physically control. + +Do not park the raw seed next to the client secret, in a `.env` or in a secrets manager: that +collapses the two factors into one. See [security.md#totp](security.md#totp). Unattended services +that genuinely need machine-held codes are covered in +[totp-automation.md](totp-automation.md). + +Create one client per MCP client that should have access. To review or remove them: + +```bash +beaconmcp auth list +beaconmcp auth revoke +``` + +## Exposing it publicly + +Put a reverse proxy in front that terminates TLS and forwards your public hostname to +`http://localhost:8420` (Caddy, nginx, Traefik and Cloudflare Tunnel all work). + +Then declare that hostname under `server.allowed_hosts` in `beaconmcp.yaml`. Without it the MCP SDK +rejects requests with `421 Misdirected Request`, which is its DNS-rebinding protection doing its +job. + +If you proxy through Cloudflare, add `cloudflare` to `server.trusted_proxies` so forwarded client +IPs can be trusted for auth rate limiting. Cloudflare's bot, WAF and Access defaults also block +headless MCP clients or strip the `Authorization` header outright: [cloudflare.md](cloudflare.md) +lists the skip and cache-bypass rules you need. + +Browser-based clients (Assistant Web, ChatGPT, Le Chat, Gemini Web) additionally need their origin +in `server.allowed_origins`, because they send a CORS preflight before reaching `/mcp`, and OAuth +HTTPS `redirect_uri` checks use the same list. + +## Updating + +Pull the new code and restart. + +Docker: + +```bash +cd BeaconMCP +git pull +docker compose up -d --build +``` + +Bare metal — the installer doubles as an updater. It stashes local state, pulls, installs new +dependencies into the virtualenv, and restarts the unit: + +```bash +sudo bash /opt/beaconmcp/deploy/install.sh +``` diff --git a/docs/security.md b/docs/security.md new file mode 100644 index 0000000..ee92b29 --- /dev/null +++ b/docs/security.md @@ -0,0 +1,116 @@ +# Security + +> Never let a model execute shell commands on infrastructure you care about without reading the +> command first. + +BeaconMCP exposes tools that cause irreversible changes: `ssh_run`, `proxmox_run`, `bmc_power_off`, +`proxmox_vm_stop`, `proxmox_vm_create`, `vm_bulk_action`, and more. Models do not reliably grasp +consequences — an errant `rm -rf`, a `systemctl stop` on the wrong unit, a `pct destroy` where `pct +stop` was meant. + +## Reviewing tool calls + +- **Disable auto-approve** on every external MCP client (Assistant Desktop, Gemini CLI, ChatGPT + MCP). Keep per-call approval on and refuse "always allow this tool". +- **Read the `command` argument** before approving any `ssh_run` or `proxmox_run`. The question to + ask: if this ran against the wrong VM or host, could I recover? +- **Prefer read-only tools** for exploration (`cluster_overview`, `cluster_health`, `*_list_*`, + `*_status`, `proxmox_get_logs`). They cannot break anything and are never gated behind a + confirmation. + +The integrated chat at `/app/chat` forces a human confirmation for every code-execution tool +(`ssh_run`, `proxmox_run`, `proxmox_write_file` and the transfer tools) and every destructive one +(`vm_bulk_action`, `proxmox_vm_stop`, snapshot rollback/delete, backup restore, `bmc_power_off`, +`bmc_power_reset`, and `beaconmcp_self_update` with `confirm=True`). Writing a guest file counts as code execution: `~/.ssh/authorized_keys` and +`/etc/cron.d/` are one hop from a shell. Skipping the modal is reserved for calls that cannot +change anything — polling by `exec_id` alone, `dry_run=True` on the snapshot tools that implement +it, and the read shape of `proxmox_vm_config`. The full list is in +[dashboard.md](dashboard.md#mandatory-confirmation-for-dangerous-tools). Read the arguments on the +confirmation card even when you're clicking through fast. No answer within 5 minutes counts as a +refusal. + +The interactive panels are the one place a gated tool runs without that modal, and only for a +closed list: starting, stopping and restarting a single guest, and resizing its CPU or memory. +That is a different question from the one the modal answers — the modal exists because the model's +input is untrusted, and a panel button is a human click on a labelled control. Everything else a +panel asks for is refused outright, including any `proxmox_vm_config` key that is not sizing. The +boundary is enforced server-side, not in the frame: +[dashboard.md](dashboard.md#interactive-panels-mcp-apps). + +## Tokens + +A `/app/tokens` bearer grants arbitrary shell access on your Proxmox nodes for its full lifetime +(`server.named_token_ttl`, 30 days by default). Don't hand one to a client you don't fully control, +and revoke it from `/app/tokens` the moment it leaks. + +`systemctl restart beaconmcp` invalidates dashboard sessions and other internal bearers, but **named +API tokens survive restarts** — they live in `server.tokens_db`. Revoke them individually from +`/app/tokens`, or delete `tokens.db` before restarting to kill all of them at once. + +`security_end_session` lets a client revoke its own bearer at the end of a task, which is a cheap way +to shrink the replay window. + +## TOTP + +Keep the TOTP seed in an authenticator app on a device you physically control: Google Authenticator, +Authy, 1Password, Aegis, a YubiKey with OTP. Type the 6-digit code by hand into the authorization +page or the dashboard. + +Do **not** generate codes programmatically with `oathtool` / `pyotp` / a shell alias, and do not +store the raw seed in a `.env`, in a secrets manager, or anywhere near the client secret. Any of +those collapses two factors into one and removes the entire point of the second one. + +Unattended services (scheduled jobs, CI pipelines) sometimes genuinely need machine-held TOTP. That +case, with its required precautions, is covered separately in +[totp-automation.md](totp-automation.md). Read it end to end before deciding. + +## Passkeys + +A passkey (WebAuthn) replaces the **TOTP factor**, never the client secret. Both login pages keep the +same two-factor shape: + +1. `client_id` + `client_secret` +2. a 6-digit code **or** a passkey assertion + +That ordering is what makes a stolen passkey worthless on its own, and it is also a practical +constraint: the dashboard session encrypts the client secret so it can re-mint MCP bearers later, so +a fully usernameless login could not build a working session anyway. + +Where they work: + +| Page | Sign in with a passkey | Enrol a passkey | +|------|------------------------|-----------------| +| `/app/login` | Link under the 2FA step | On the post-2FA screen | +| `/oauth/authorize` | Link under the 2FA step | On the approval screen | + +Credentials are stored in the dashboard database (`passkeys` table): a credential id, a **public** +key and a signature counter. Nothing secret leaves the authenticator. + +Things worth knowing before you enrol: + +- **Passkeys are bound to the hostname.** The relying-party ID is derived from the request host, so + a credential registered on `beacon.example` will not work on `beacon.internal` or on a raw IP that + differs from the one used at registration. Settle on your public hostname first. +- **A secure context is required.** Browsers only expose the WebAuthn API over HTTPS or on loopback. + On a plain-HTTP LAN deployment the passkey buttons are hidden and TOTP stays the only path in. +- **Keep TOTP working.** Passkeys are an alternative, not a replacement: losing every enrolled device + must not lock you out. The authenticator seed remains the recovery path. +- **Dynamically-registered clients delegate**, exactly like TOTP: a client created through the DCR + bootstrap is authorized by its *owner's* passkeys, so the second factor never leaves the owner. +- Registrations, revocations and passkey sign-ins are recorded in the audit log + (`dashboard.passkey.*`, `auth.passkey.*`, and `login`/`authorize` events tagged `via=passkey`). + +Passkey ceremonies are rate-limited per IP by the same limiter that guards `/app/login`, and a TOTP +lockout also blocks the passkey path for that client. + +Manage enrolled credentials from `/app/tokens`, or drop them all for a client with: + +```sql +DELETE FROM passkeys WHERE client_id = 'beaconmcp_...'; +``` + +## Audit trail + +`server.audit_log` records tool calls, dashboard logins, OAuth authorizations and client revocations +as JSON lines, in a file created owner-only (0600). It's the first place to look after something +unexpected happens. See [configuration.md](configuration.md#server). diff --git a/docs/superpowers/specs/2026-04-16-beaconmcp-design.md b/docs/superpowers/specs/2026-04-16-beaconmcp-design.md new file mode 100644 index 0000000..8f05ec3 --- /dev/null +++ b/docs/superpowers/specs/2026-04-16-beaconmcp-design.md @@ -0,0 +1,263 @@ +# BeaconMCP -- Proxmox Infrastructure MCP Server + +> **Note (post-v1):** the three-tool exec surface described below +> (`proxmox_exec_command` / `_async` / `_get_result` and its SSH twin) +> has been **superseded by the unified `proxmox_run` and `ssh_run` +> tools**. Each unified tool exposes the same three call patterns via +> parameters: sync (default), async start (`wait=False`), and poll +> (`exec_id=…`). The design intent -- auto-detect VM vs CT, in-memory +> session registry, timeout fallback to async -- is preserved; only the +> public tool names changed. See `README.md` § *Available tools* and +> `docs/dashboard.md` for the current contract. + +## Context + +BeaconMCP is an MCP server that gives Assistant direct access to a Proxmox VE infrastructure for diagnostics, VM management, system administration, and hardware management. The motivation: when a server crashes or misbehaves, Assistant should be able to diagnose the issue, check hardware health, and propose/execute resolutions -- rather than the user having to manually SSH, check logs, and relay information back and forth. + +**Infrastructure:** +- **pve1.example.com** -- Proxmox VE node (active), exposed on the internet via HTTPS +- **pve2.example.com** -- Proxmox VE node (currently down) +- **iLO 4** -- HP Integrated Lights-Out, one unit, accessible only from the local network (not publicly exposed) +- **Zyxel USG 210** -- Firewall, deferred from v1 (no REST API available) + +## Architecture + +Single Python MCP server (`beaconmcp`) with modular design, running in **stdio** mode. Three core modules: + +``` +src/beaconmcp/ +├── __init__.py +├── __main__.py # Entry point +├── server.py # FastMCP server, registers all tools +├── config.py # Environment variable loading & validation +├── proxmox/ +│ ├── __init__.py +│ ├── client.py # proxmoxer wrapper, connection management +│ ├── vms.py # VM/CT lifecycle tools +│ ├── monitoring.py # Node & VM monitoring tools +│ └── system.py # Storage, network, command execution tools +├── ilo/ +│ ├── __init__.py +│ └── client.py # python-hpilo wrapper + SSH tunnel management +└── ssh/ + ├── __init__.py + └── client.py # asyncssh wrapper, session management +``` + +**Dependencies:** +- `mcp` -- MCP Python SDK (FastMCP) +- `proxmoxer` + `requests` -- Proxmox VE API client +- `python-hpilo` -- HP iLO 4 management (synchronous library, runs in asyncio executor) +- `asyncssh` -- Async SSH connections +- `python-dotenv` -- Environment variable loading + +## MCP Tools + +### Module Proxmox -- Monitoring & Diagnostic + +| Tool | Description | Key Parameters | +|------|-------------|----------------| +| `proxmox_list_nodes` | List cluster nodes with status (online/offline/unknown) | -- | +| `proxmox_node_status` | Detailed node status: CPU, RAM, disk, uptime, kernel version, PVE version | `node` | +| `proxmox_list_vms` | List all VMs/CTs with status, resource usage | `node` (optional, all nodes if omitted) | +| `proxmox_vm_status` | Detailed VM/CT status: CPU, RAM, disk I/O, network I/O, uptime | `node`, `vmid` | +| `proxmox_get_logs` | Retrieve system logs (syslog, tasks, journal) | `node`, `source` (syslog/tasks), `limit` | +| `proxmox_get_tasks` | List recent Proxmox tasks (migrations, backups, etc.) | `node` (optional), `limit` | + +### Module Proxmox -- VM/CT Management + +| Tool | Description | Key Parameters | +|------|-------------|----------------| +| `proxmox_vm_start` | Start a VM or CT | `node`, `vmid` | +| `proxmox_vm_stop` | Stop a VM or CT (clean shutdown or force) | `node`, `vmid`, `force` | +| `proxmox_vm_restart` | Restart a VM or CT | `node`, `vmid` | +| `proxmox_vm_create` | Create a new VM or CT | `node`, `config` (dict) | +| `proxmox_vm_clone` | Clone an existing VM/CT | `node`, `vmid`, `newid`, `name` | +| `proxmox_vm_migrate` | Migrate a VM/CT to another node | `node`, `vmid`, `target_node` | +| `proxmox_vm_config` | Read or modify VM/CT configuration | `node`, `vmid`, `updates` (optional) | + +### Module Proxmox -- System Administration + +| Tool | Description | Key Parameters | +|------|-------------|----------------| +| `proxmox_storage_status` | Storage status across the cluster | `node` (optional) | +| `proxmox_network_config` | Network configuration of a node | `node` | +| `proxmox_run` | Execute a command inside a QEMU VM (QEMU Guest Agent). Sync by default; pass `wait=False` to start async (returns `exec_id`), or `exec_id=…` to poll an existing session. LXC exec is not exposed by the Proxmox API; use `ssh_run` + `pct exec` on the host node. | `node`, `vmid`, `command`, `timeout` (default 60s), `wait`, `exec_id` | + +**Command execution design:** +- The tool auto-detects whether the target is a VM or CT. VMs execute via QEMU Guest Agent; CTs return an actionable error that points to `ssh_run` + `pct exec -- ` on the node. +- `proxmox_run` (default, `wait=True`) blocks until the command completes or timeout is reached. Returns `{"status": "ok", "stdout": "...", "stderr": "...", "exit_code": N, "duration_s": ...}`. On timeout it auto-switches to async and returns `{"status": "running", "exec_id": "..."}`. +- `proxmox_run(..., wait=False)` returns immediately with `{"status": "running", "exec_id": "..."}`. Internally uses QEMU Guest Agent's native async exec for VMs (start -> PID -> poll). +- `proxmox_run(exec_id="...")` polls an existing session and returns `{"status": "running|ok|failed|timeout", "stdout": "...", "stderr": "...", "exit_code": N}`. +- Async exec state is held in-memory in the server process. A dict of `{exec_id: {pid, node, vmid, type, status, output}}`. + +### Module iLO + +| Tool | Description | Key Parameters | +|------|-------------|----------------| +| `ilo_server_info` | Server model, serial number, firmware versions (iLO, BIOS) | -- | +| `ilo_health_status` | Full health: temperatures, fans, power supplies, disks, memory | -- | +| `ilo_power_status` | Current power state of the server | -- | +| `ilo_power_on` | Power on the physical server | -- | +| `ilo_power_off` | Power off the physical server (use when server is unresponsive) | `force` (default false) | +| `ilo_power_reset` | Hard reset the physical server | -- | +| `ilo_get_event_log` | iLO event log (hardware errors, reboots, etc.) | `limit` | + +**iLO access via SSH tunnel:** +Since iLO is only accessible from the local network, the module establishes an SSH tunnel through pve1: +1. asyncssh opens a tunnel: `localhost:dynamic_port -> pve1 -> ilo_local_ip:443` +2. python-hpilo connects to `localhost:dynamic_port` +3. Tunnel is created on-demand and reused for subsequent calls +4. If pve1 is unreachable, iLO tools return an error explaining the dependency + +### Module SSH + +| Tool | Description | Key Parameters | +|------|-------------|----------------| +| `ssh_run` | Execute a command on any host via SSH. Sync by default; `wait=False` starts async and returns `exec_id`; `exec_id=…` polls an existing session. | `host`, `command`, `timeout` (default 60s), `wait`, `exec_id` | +| `ssh_list_sessions` | List active async command sessions with their status | -- | + +SSH uses password authentication. The `host` parameter accepts: +- A Proxmox node name (`pve1`, `pve2`) -- resolved to the configured host from env vars +- A VMID (e.g., `101`) -- resolved to IP via the infrastructure.yaml convention (192.168.1.{VMID}) +- A direct IP or hostname (e.g., `192.168.1.50`) + +## Configuration + +All configuration via environment variables, loaded from `.env` file by `python-dotenv`: + +```env +# Proxmox nodes -- API tokens (to be created on the Proxmox nodes) +PVE1_HOST=pve1.example.com +PVE1_TOKEN_ID=root@pam!beaconmcp +PVE1_TOKEN_SECRET=xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx + +PVE2_HOST=pve2.example.com +PVE2_TOKEN_ID=root@pam!beaconmcp +PVE2_TOKEN_SECRET=xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx + +# iLO -- single unit, local network only +ILO_HOST=192.168.x.x +ILO_USER=Administrator +ILO_PASSWORD=xxxxx +ILO_JUMP_HOST=pve1 # Proxmox node used as SSH jump host + +# SSH credentials (fallback access) +SSH_USER=root +SSH_PASSWORD=xxxxx + +# Options +PVE_VERIFY_SSL=false # Set to true if using valid SSL certificates +``` + +**Startup validation:** +- PVE1 credentials are required (server won't start without them) +- PVE2 credentials are optional (graceful degradation if missing or node is down) +- iLO credentials are optional (iLO tools disabled if not configured) +- SSH credentials are optional (SSH tools disabled if not configured) + +## Error Handling + +- **Node unreachable:** Tools return a clear error message indicating which node is unreachable, rather than raising exceptions. Assistant can then suggest remediation (check iLO, try SSH, etc.). +- **Authentication failures:** Logged and returned as structured errors with guidance (check token, check password, etc.). +- **Command timeouts:** Async commands that exceed timeout are marked as `timeout` status. Partial output is preserved. +- **iLO tunnel failure:** If pve1 (jump host) is unreachable, iLO tools return an error explaining that iLO is only accessible through pve1. + +## Assistant Code Integration + +Add to `~/.assistant/settings.json` or project `.assistant/settings.json`: + +```json +{ + "mcpServers": { + "beaconmcp": { + "command": "python", + "args": ["-m", "beaconmcp"], + "cwd": "/path/to/BeaconMCP/src", + "env": { + "PVE1_HOST": "pve1.example.com", + "PVE1_TOKEN_ID": "root@pam!beaconmcp", + "PVE1_TOKEN_SECRET": "..." + } + } + } +} +``` + +Or use a `.env` file in the project directory and configure only the command. + +## Verification Plan + +1. **Unit:** Test each module's client wrapper independently with mocked API responses +2. **Integration:** Test against pve1 with real API token: + - List nodes, check node status + - List VMs, start/stop a test VM + - Execute a simple command via QEMU Guest Agent (`echo hello`) + - Run an async command and poll for result +3. **iLO:** Test tunnel creation + health check against the real iLO +4. **SSH:** Test direct SSH command execution on pve1 +5. **End-to-end:** Start the MCP server, use it from Assistant Code to diagnose a real scenario (e.g., "why is pve2 down?") + +## MCP Resources & Prompts + +### Infrastructure Context Resource + +An `infrastructure.yaml` file at the project root provides contextual information about the infrastructure. The MCP server exposes it as a resource so Assistant can read it automatically. + +```yaml +# infrastructure.yaml +conventions: + vmid_to_ip: "CT VMID corresponds to local IP 192.168.1.{VMID}" + naming: "VMs are prefixed by their role (e.g., web-101, db-102)" + +nodes: + pve1: + host: pve1.example.com + role: "Primary node" + local_network: "192.168.1.0/24" + pve2: + host: pve2.example.com + role: "Secondary node" + notes: "Currently down" + +ilo: + host: "192.168.x.x" + access: "Local network only, via SSH tunnel through pve1" + +firewall: + model: "Zyxel USG 210" + notes: "No API, managed via web GUI" + +notes: + - "iLO is accessible only through pve1 as SSH jump host" + - "Zyxel USG 210 is the network gateway" + - "API tokens must be created on each Proxmox node before use" +``` + +The server exposes this as `beaconmcp://infrastructure` -- a readable resource that provides Assistant with the full infrastructure context. + +### MCP Prompt: Infrastructure Overview + +The server registers an MCP prompt `beaconmcp-context` that injects a concise infrastructure summary into the conversation. This follows prompt engineering best practices (from `docs/prompt-engineering-guide.md`): +- Role definition: "You are managing a Proxmox VE infrastructure" +- Context: node topology, naming conventions, access constraints +- Positive instructions: what to check first, how to diagnose + +### Tool Description Quality + +All MCP tool descriptions follow best practices: +- **Self-sufficient**: each description is understandable without external context +- **Namespaced**: `proxmox_*`, `ilo_*`, `ssh_*` prefixes +- **When to use / when not to use**: each tool specifies its use case and alternatives +- **Actionable errors**: error messages include what went wrong and what to try next +- **Semantic parameter names**: `vmid` not `id`, `target_node` not `dest` + +Reference: `/docs/prompt-engineering-guide.md` -- sections 4.1 through 4.6. + +## Out of Scope (v1) + +- Zyxel USG 210 firewall integration (no API available) +- Proxmox built-in firewall management (can be added later) +- Backup management (can be added later via Proxmox Backup Server API) +- User/permission management on Proxmox +- Automated alerting/monitoring (this is a tool for Assistant, not a monitoring stack) diff --git a/docs/superpowers/specs/2026-04-17-beaconmcp-dashboard-design.md b/docs/superpowers/specs/2026-04-17-beaconmcp-dashboard-design.md new file mode 100644 index 0000000..5a35ee0 --- /dev/null +++ b/docs/superpowers/specs/2026-04-17-beaconmcp-dashboard-design.md @@ -0,0 +1,525 @@ +# BeaconMCP Dashboard — Design + +**Date:** 2026-04-17 +**Status:** Draft (awaiting review) +**Scope:** Web dashboard adjacent to the existing MCP server, served on the same origin (`mcp.example.com`). + +## 1. Goals + +Add a web-accessible companion dashboard to BeaconMCP providing two surfaces: + +- **Login panel** — collect `client_id` + `client_secret` + TOTP, exchange for an MCP bearer, and persist a browser session for 90 days so users (especially on mobile) no longer need to run `curl + oathtool` to obtain a token. +- **Chat panel** — multi-conversation chat interface powered by Gemini 3 Flash / 3.1 Pro, invoking BeaconMCP tools through the Gemini SDK's native MCP integration. Usable from phone or desktop. + +### Success criteria + +1. From a phone browser, an operator can log in, open a chat, ask "status de pve2", and see tool-call results without typing in a shell. +2. Session survives closing the browser and coming back within 90 days; only TOTP is re-prompted every 24h. +3. UI feels smooth (streaming responses, no full-page reloads, < 100 ms perceived click latency). +4. The dashboard is optional — if `GEMINI_API_KEY` is absent, the rest of BeaconMCP runs unchanged. + +### Non-goals (explicit) + +- No multi-user model — one `client_id` = one operator. Conversations are scoped to the client that owns them. +- No rate-limiting beyond the existing TOTP lockout (5 failures / 5 min). +- No conversation sharing, export, or public links. +- No voice / STT. +- No push notifications. +- No syntax-highlighting in code blocks (v1). +- No full-text search over history (v1). + +## 2. Decisions (from brainstorming) + +| Decision | Choice | +|---|---| +| Gemini API key location | Server-side env var `GEMINI_API_KEY`. Server proxies all Gemini calls. | +| Tool execution path | Approach **A** — Gemini SDK invoked with an `McpServer` tool pointing to `https://mcp.example.com/mcp` + user bearer in headers. Google's backend calls BeaconMCP directly. | +| Session durability | 90-day HttpOnly server-side session cookie. TOTP re-prompt every 24 h to refresh MCP bearer. | +| Conversation model | Multiple conversations with sidebar, server-persisted in SQLite. | +| Thinking control | Gemini 3 effort presets: `minimal` / `low` / `medium` / `high`. Per-conversation setting. | +| Tool-call visualization | Inline collapsible cards showing `name · duration · status`; expand for args + result preview. | +| Frontend stack | Vanilla HTML + CSS custom properties + ES modules. Jinja2 server-rendered templates. Marked + DOMPurify vendored for markdown. No build step. | +| Chat history storage | SQLite at `/opt/beaconmcp/dashboard.db` (WAL, foreign keys). Single DB for sessions + conversations + messages. | +| Theme | Light / dark via `prefers-color-scheme`, slate-neutral palette + Proxmox orange `#e57000` accent. | +| Model selector | Default `gemini-3-flash`; dropdown to switch to `gemini-3.1-pro`. Choice persisted per conversation. | +| Iconography | Inline monochrome SVG icons (plus, chevron, arrow, ellipsis, status check/warn/spinner). **No Unicode emojis anywhere.** The "no AI slop" rule bans emojis 🎉✨🤖 etc., not vector icons. | + +## 3. Module layout + +``` +src/beaconmcp/ + dashboard/ NEW MODULE + __init__.py register_dashboard_routes(app, ...) + app.py Starlette routes: login, refresh, logout, chat, api/* + session.py SessionStore (SQLite + AES-GCM) and cookie helpers + db.py Connection pooling, migrations (PRAGMA user_version) + chat.py ChatEngine: Gemini SDK wrapper, SSE event generation + csrf.py Double-submit cookie middleware + templates/ + base.html Layout, CSS vars, icon sprite + login.html + totp_refresh.html + chat.html Shell (sidebar + messages + composer) + static/ + app.css + chat.js SSE client, sidebar, markdown render, composer + marked.min.js Vendored 15 kb + dompurify.min.js Vendored 20 kb + icons.svg SVG sprite (6-7 icons) + __main__.py + mount dashboard routes when enabled +``` + +The dashboard is mounted conditionally in `__main__.py::_run_http`, similar to how SSH / iLO modules register themselves. Requires `GEMINI_API_KEY` set. Can be disabled via `BEACONMCP_DASHBOARD_ENABLED=false`. + +## 4. URL routing + +| Path | Method | Auth | Role | +|---|---|---|---| +| `/` | GET | — | 302 to `/app/chat` if session, else `/app/login` | +| `/app/login` | GET | — | Render login form (or single-field refresh if cookie exists) | +| `/app/login` | POST | — | Validate creds + TOTP, create session, set cookie, 302 to `/app/chat` | +| `/app/refresh` | GET | cookie | Render TOTP-only form | +| `/app/refresh` | POST | cookie | Re-issue MCP bearer, update session | +| `/app/logout` | POST | cookie | Revoke bearer, delete session, clear cookie, 302 to `/app/login` | +| `/app/chat` | GET | cookie | Render chat shell HTML | +| `/app/api/conversations` | GET | cookie | JSON list (scoped to `client_id`) | +| `/app/api/conversations` | POST | cookie + CSRF | Create empty conversation | +| `/app/api/conversations/{id}` | GET | cookie | Fetch full conversation + messages | +| `/app/api/conversations/{id}` | PATCH | cookie + CSRF | Rename, change model/effort | +| `/app/api/conversations/{id}` | DELETE | cookie + CSRF | Delete conversation and messages | +| `/app/api/chat/stream` | POST | cookie + CSRF | SSE streaming turn | +| `/app/static/*` | GET | — | Static assets | +| `/mcp`, `/oauth/*`, `/.well-known/*`, `/health` | — | — | **Unchanged** | + +All `/app/*` routes go through a dedicated `DashboardSessionMiddleware`, not the existing bearer middleware that protects `/mcp`. + +## 5. Authentication & sessions + +### Principle + +Cookie holds only an opaque `session_id` (256-bit random). All sensitive material (client_secret, current MCP bearer) lives in SQLite, encrypted at rest. + +### Session table + +```sql +CREATE TABLE sessions ( + session_id TEXT PRIMARY KEY, + client_id TEXT NOT NULL, + client_secret_enc BLOB NOT NULL, + mcp_bearer TEXT, + mcp_bearer_expires_at REAL, + created_at REAL NOT NULL, + last_seen_at REAL NOT NULL, + expires_at REAL NOT NULL, + user_agent TEXT +); +CREATE INDEX idx_sessions_client ON sessions(client_id); +``` + +### Client secret encryption + +- Master key: env var `BEACONMCP_SESSION_KEY` (32 bytes, base64-encoded). Generated by `deploy/install.sh` if absent. +- Algorithm: AES-256-GCM via `cryptography.hazmat.primitives.ciphers.aead.AESGCM`. +- Storage layout: `nonce(12 bytes) || ciphertext || tag`. +- Rotation: not addressed in v1. Compromise recovery = wipe `sessions` table, rotate key, users log in again. + +### Cookie + +``` +Set-Cookie: beaconmcp_session=; + HttpOnly; Secure; SameSite=Strict; + Path=/app; + Max-Age=7776000 +``` + +`Path=/app` prevents the cookie from being sent on `/mcp`, `/oauth/*`, or Google's back-channel calls to the MCP endpoint. + +### Initial login flow (`POST /app/login`) + +``` +1. ClientStore.verify(client_id, client_secret) → 401 if fail +2. totp_locked(client_id) → 429 if locked +3. ClientStore.verify_totp(client_id, totp) → 401 + increment fail count +4. token_store.issue(client_id) → (bearer, 86400) +5. Generate session_id = secrets.token_urlsafe(32) +6. AES-GCM encrypt client_secret with BEACONMCP_SESSION_KEY +7. INSERT INTO sessions (...) +8. Set cookie; 302 → /app/chat +``` + +### 24h bearer refresh flow + +Middleware detects `mcp_bearer_expires_at < now()` when an `/app/api/chat/stream` request hits it: + +``` +1. Return SSE event {type: "session_expired"} +2. Client redirects to /app/refresh +3. /app/refresh GET: render "Code TOTP pour {client_name}" (1 field) +4. /app/refresh POST: verify_totp → token_store.issue → UPDATE sessions +5. 302 → /app/chat; UI auto-retries the last turn +``` + +If the session cookie is missing, unknown, or past `expires_at`: 302 → `/app/login` (full form). + +### Logout + +``` +1. Load session +2. token_store.revoke(session.mcp_bearer) (existing 8s grace) +3. DELETE FROM sessions WHERE session_id=? +4. Clear-Site-Data: "cookies" + Set-Cookie beaconmcp_session=; Max-Age=0 +5. 302 → /app/login +``` + +### Multi-session and admin revocation + +- Multiple sessions per `client_id` (phone + desktop) are allowed. +- `beaconmcp auth revoke ` cascades: deletes the client, deletes all its sessions, revokes all its bearers. +- New CLI subcommand `beaconmcp dashboard sessions [--client-id X]` lists sessions with last-seen timestamps and supports `--kill `. + +### Security + +| Surface | Measure | +|---|---| +| Session cookie | HttpOnly, Secure, SameSite=Strict, Path=/app, Max-Age=7776000 | +| CSRF | Double-submit cookie `beaconmcp_csrf_token` (JS-readable, `SameSite=Strict`, `Path=/app`) + header `X-CSRF-Token` required on POST/PATCH/DELETE | +| Session fixation | Regenerate `session_id` at login | +| Secret at rest | AES-256-GCM with env-derived key | +| Secret in logs | Logging filter redacts `sk_*` and `beaconmcp_*` tokens | +| Login brute-force | Reuses existing 5-failure / 5-minute TOTP lockout | +| Clickjacking | `X-Frame-Options: DENY` on `/app/*` | +| MIME sniffing | `X-Content-Type-Options: nosniff` | +| Referrer | `Referrer-Policy: strict-origin-when-cross-origin` | +| CSP | `default-src 'self'; style-src 'self' 'unsafe-inline'; script-src 'self'; connect-src 'self'; img-src 'self' data:` | +| DNS rebinding | Existing `TransportSecuritySettings` + `BEACONMCP_ALLOWED_HOSTS` | + +## 6. Database schema + +SQLite at `/opt/beaconmcp/dashboard.db` (overridable via `BEACONMCP_DASHBOARD_DB`). Mode WAL, `synchronous=NORMAL`, `foreign_keys=ON`. Versioned via `PRAGMA user_version`. + +```sql +CREATE TABLE sessions (...); -- see §5 + +CREATE TABLE conversations ( + id TEXT PRIMARY KEY, -- UUID v4 + client_id TEXT NOT NULL, + title TEXT, + model TEXT NOT NULL DEFAULT 'gemini-3-flash', + thinking_effort TEXT NOT NULL DEFAULT 'low', -- minimal|low|medium|high + created_at REAL NOT NULL, + updated_at REAL NOT NULL +); +CREATE INDEX idx_conv_client ON conversations(client_id, updated_at DESC); + +CREATE TABLE messages ( + id TEXT PRIMARY KEY, + conversation_id TEXT NOT NULL REFERENCES conversations(id) ON DELETE CASCADE, + role TEXT NOT NULL, -- user|assistant + content TEXT, + tool_calls TEXT, -- JSON array + thinking_summary TEXT, + model TEXT, + effort TEXT, + created_at REAL NOT NULL +); +CREATE INDEX idx_msg_conv ON messages(conversation_id, created_at); +``` + +`tool_calls` is a JSON array: `[{id, name, args, result, status, duration_ms, preview}]`. JSON instead of a child table keeps reads to a single SELECT + parse. + +## 7. Chat flow + +### Turn end-to-end + +``` +POST /app/api/chat/stream { conversation_id, content, model, effort } + ↓ +1. DashboardSessionMiddleware validates cookie + CSRF, loads session +2. If bearer expired → SSE event "session_expired" and close +3. Load conversation history from SQLite (ordered by created_at) +4. INSERT message (role=user) +5. Open SSE response (Content-Type: text/event-stream, no-cache) +6. Call google-genai streaming API with: + - model: the conversation's model + - contents: conversation history + new user message + - thinking config: effort level from the conversation + - tools: MCP server pointed at https://mcp.example.com/mcp + with the session's bearer in the Authorization header + (Exact SDK class names — Tool, McpServer, ThinkingConfig, + GenerateContentConfig — are verified at implementation time against + the installed google-genai version.) +7. For each chunk, emit matching SSE event (see table below) +8. On stream end, INSERT message (role=assistant, content, tool_calls, ...) +9. UPDATE conversations SET updated_at=now +10. If this was the first user turn in the conversation, fire auto-title: + - Secondary short genai call (gemini-3-flash, effort=minimal, no tools) + - UPDATE conversations SET title=? + - SSE event "title_updated" +11. Emit "done" event and close stream +``` + +### SSE event vocabulary + +| Event | Payload | +|---|---| +| `text_delta` | `{"text": "..."}` | +| `thinking_delta` | `{"summary": "..."}` | +| `tool_call` | `{"id": "...", "name": "...", "args": {...}}` | +| `tool_result` | `{"id": "...", "status": "ok"\|"error", "preview": "...", "duration_ms": 234}` | +| `error` | `{"code": "...", "message": "..."}` | +| `session_expired` | `{}` | +| `title_updated` | `{"conversation_id": "...", "title": "..."}` | +| `aborted` | `{}` | +| `done` | `{"message_id": "..."}` | + +Server honors `Request.is_disconnected` — if the client aborts (user pressed "Arrêter"), the loop breaks, a partial assistant message is persisted with a marker, and `"aborted"` is emitted before close. + +### Model and effort + +Per-conversation columns. Default for new chat: `gemini-3-flash` + `low`. Changing model or effort applies to subsequent turns; existing messages are unaffected (each message records the `model` and `effort` used). + +### Auto-title + +Performed after the first *user* message gets its assistant response. Call is isolated (no tools, no streaming), prompt: "Donne un titre de 4 mots maximum, sans emoji, sans ponctuation finale, pour: {user_msg}". Persisted on the conversation. UI updates the sidebar on `title_updated`. + +## 8. UI + +### Login (`/app/login`) + +Centered card layout, styled with the dashboard's light/dark palette (not coupled to the existing `/oauth/authorize` page, which is left untouched). First-time visit shows 3 fields (Client ID, Client Secret, TOTP) + "Rester connecté 90j" checkbox. On subsequent visits where a valid session cookie exists but its bearer is stale, `/app/refresh` renders a single TOTP field with the client name in evidence. + +### Chat desktop layout + +``` +┌──────────────┬──────────────────────────────────────────────┐ +│ BeaconMCP │ │ +│ + Nouveau │ [user bubble, right-aligned] │ +│ ───────── │ │ +│ > pve2 down? │ Gemini 3 Flash · low │ +│ LXC update │ [streaming text …] │ +│ … │ │ +│ │ ▸ proxmox_list_vms · 180 ms · ok │ +│ │ │ +│ ───────── │ ┌──────────────────────────────────────────┐ │ +│ Modèle │ │ Envoyer un message à BeaconMCP… │ │ +│ Flash ▼ │ └──────────────────────────────────────────┘ │ +│ Effort: low │ Flash ▼ · effort ▼ [Envoyer] │ +│ Déconnexion │ │ +└──────────────┴──────────────────────────────────────────────┘ +``` + +- Sidebar 260 px, collapsible on narrow screens. +- Conversation list sorted by `updated_at DESC`, single-line truncation, hover reveals ⋯ menu (rename, delete). +- Active conversation highlighted with `--bg-soft` and a left accent bar. +- Model / effort controls live at the bottom of the sidebar AND in the composer row (both edit the same conversation settings). + +### Chat mobile layout (< 768 px) + +``` +┌─────────────────────────────┐ +│ ☰ pve2 down? ⋯ │ +├─────────────────────────────┤ +│ │ +│ msg user │ +│ │ +│ Gemini 3 Flash · low │ +│ streaming text … │ +│ │ +│ ▸ proxmox_list_vms · 180ms │ +│ │ +├─────────────────────────────┤ +│ [input message ... ] │ +│ Flash ▼ effort ▼ [↑] │ +└─────────────────────────────┘ +``` + +- Hamburger icon opens sidebar as a left drawer with an overlay dim. +- Safe-area insets respected on iOS (`env(safe-area-inset-bottom)`). +- Touch targets ≥ 44 px. + +### Messages zone + +- Centered column, `max-width: 48rem`. +- **User** message: `--user-bubble` background, right-aligned, 80 % max-width, plain text (no markdown parsing). +- **Assistant** message: no bubble, plain flow on canvas, with a small header "`{model} · {effort}`" in `--fg-muted`. Markdown rendered via `marked` + sanitized with `DOMPurify` (strict allowlist, no raw HTML passes through). +- **Tool-call card** — collapsed: + ``` + [chevron-right] proxmox_list_vms · 180 ms · ok + ``` + Expanded: + ``` + [chevron-down] proxmox_list_vms · 180 ms · ok + args: { "node": "pve1" } + result: 12 VMs listées + [vmid=101] web-prod · running + ... + [Copier le résultat] + ``` + Border 1 px, rounded, `--bg-soft` background. Status icon (check / warn / spinner) is an inline SVG, no emoji. + +### Composer + +- Auto-growing ` +
+ +
+ + + + + + + + + + + + + + +{% endblock %} diff --git a/src/beaconmcp/dashboard/templates/connectors.html b/src/beaconmcp/dashboard/templates/connectors.html new file mode 100644 index 0000000..c8811fe --- /dev/null +++ b/src/beaconmcp/dashboard/templates/connectors.html @@ -0,0 +1,136 @@ +{% extends "base.html" %} +{% block title %}OAuth connectors · BeaconMCP{% endblock %} +{% block body_class %}tokens-page{% endblock %} +{% block body %} + + +
+
+
BeaconMCP
+
+ + + Back to API access + +
+ +

OAuth connectors

+

+ Signed in as {{ client_name }}. Clients that use OAuth Dynamic Client + Registration (RFC 7591) — ChatGPT, OpenCode, and similar — can't accept a pre-provisioned + client_id or a bearer header. Mint a single-use bootstrap URL here, paste it + into the client, and complete the 2FA prompt in-app. Each URL is single-use and expires after + {{ slug_ttl_minutes }} min. +

+ + {% if just_created %} +
+ + + Connector URL ready — single-use, dies in {{ slug_ttl_minutes }} min + +
+ {{ just_created.url }} + +
+
+ {% endif %} + +
+
+

New connector

+
+ {% if form_error %} + + {% endif %} +
+ + + + +
+
+ +
+
+

Pending URLs

+
+ {% if pending_slugs %} +
    + {% for s in pending_slugs %} +
  • +
    +
    + {{ s.label }} +
    +
    {{ s.slug[:10] }}… · expires in {{ s.expires_in_minutes }} min
    +
    +
    + + + +
    +
    • + {% endfor %} +
+ {% else %} +

No pending URLs.

+ {% endif %} +
+ +
+
+

Active connectors

+
+ {% if derived_clients %} +
    + {% for c in derived_clients %} +
  • +
    +
    + {{ c.name }} +
    +
    {{ c.client_id }} · registered {{ c.created_at_human }}
    +
    +
    + + + +
    +
    • + {% endfor %} +
+ {% else %} +

No active OAuth connectors.

+ {% endif %} +
+
+ + +{% endblock %} diff --git a/src/beaconmcp/dashboard/templates/login.html b/src/beaconmcp/dashboard/templates/login.html new file mode 100644 index 0000000..a784041 --- /dev/null +++ b/src/beaconmcp/dashboard/templates/login.html @@ -0,0 +1,160 @@ +{% extends "base.html" %} +{% block title %}Sign in · BeaconMCP{% endblock %} +{% block body_class %}auth-page{% endblock %} +{% block body %} +
+
+ BeaconMCP +
+ +
+ + {% if next %}{% endif %} + + {# Step 1 — Client ID + Client Secret #} +
+

Sign in

+

Paste your OAuth client credentials to continue.

+ + {% if banner %} + + {% endif %} + +
+ +
+ +
+
+ +
+ +
+ + +
+
+ +
+ +
+ + + +
+ No account? Ask your administrator +
+
+ + {# Step 2 — TOTP, or a passkey instead #} + +
+ + {# Step 3 — signed in: session lifetime, passkey enrolment, finish #} + +
+ + + +{% endblock %} diff --git a/src/beaconmcp/dashboard/templates/overview.html b/src/beaconmcp/dashboard/templates/overview.html new file mode 100644 index 0000000..ad203a0 --- /dev/null +++ b/src/beaconmcp/dashboard/templates/overview.html @@ -0,0 +1,46 @@ +{% extends "base.html" %} +{% block title %}Status Overview · BeaconMCP{% endblock %} +{% block body_class %}tokens-page{% endblock %} +{% block body %} + + +
+
+
BeaconMCP
+
+ + + Back to chat + +
+ +

Status Overview Hub

+

+ System health, Proxmox cluster status, and hardware metrics overview. +

+ +
+
Overall Health
+
+
Status gathering coming soon
+
+
+ +
+
+

Proxmox Nodes

+
    +
  • (Coming soon)
    • +
+
+
+

Hardware BMCs

+
    +
  • (Coming soon)
    • +
+
+
+
+{% endblock %} diff --git a/src/beaconmcp/dashboard/templates/tokens.html b/src/beaconmcp/dashboard/templates/tokens.html new file mode 100644 index 0000000..1e149be --- /dev/null +++ b/src/beaconmcp/dashboard/templates/tokens.html @@ -0,0 +1,815 @@ +{% extends "base.html" %} +{% block title %}API access · BeaconMCP{% endblock %} +{% block body_class %}tokens-page{% endblock %} +{% block body %} + + +
+ +
+
BeaconMCP
+
+ {% if chat_enabled %} + + + Back to chat + + {% endif %} +
+ +

API access

+

+ Signed in as {{ client_name }}. Pick your client and follow the matching flow — + every flow keeps your TOTP on your phone. +

+ +
+
MCP server URL
+
+ {{ mcp_url }} + +
+
+ +
+ + + Browser clients need their origin allowlisted + +
+

+ A browser calling a different host has to pass a CORS preflight. Web UIs + (Assistant, ChatGPT, Le Chat, + Perplexity, Gemini Web) all do this before they can reach + /mcp. If the origin isn't listed, every request fails silently in the + browser console. +

+

Add the origins you plan to use to beaconmcp.yaml:

+
server:
+  allowed_origins:
+    - https://assistant.ai
+    - https://chatgpt.com
+    - https://chat.mistral.ai
+    - https://www.perplexity.ai
+    - https://gemini.google.com
+

+ Desktop apps and CLI clients (Assistant Desktop, Gemini CLI, Cursor, VS Code, Mistral Vibe, + OpenCode) don't run in a browser and don't need an entry here. +

+
+
+ +
+ + +
+ + +
+ + {# ============================================================ #} + {# VIEW: BY METHOD #} + {# ============================================================ #} +
+ + {% if dcr_enabled %}{% endif %} + + + +
+
+
+
+

OAuth 2.1 with a pre-registered client

+
+

+ The path most clients take: provision client_id / + client_secret here, paste them into the client, and let the client run + authorization code + PKCE against BeaconMCP. TOTP prompt lives on BeaconMCP's + authorization page; tokens last 24 h. +

+
+
+ Assistant + Codex + Le Chat + Gemini CLI + Antigravity + OpenCode + Cursor + VS Code +
+ +
1 · Provision credentials on the server
+

One client per install so you can revoke granularly:

+
beaconmcp auth create --name "Assistant iPhone"
+

The CLI prints a client_id, a client_secret, and a TOTP QR code (scan it immediately — the secret is shown once).

+ +
2 · Paste credentials into the client
+

+ Every client has its own config file or UI panel. Switch to the + By platform view above for the exact snippet per client. The MCP URL + is always {{ mcp_url }}. +

+
+
+ + {% if dcr_enabled %} +
+
+
+
+

OAuth with Dynamic Client Registration

+
+

+ Reserved for clients whose UI won't let you paste client_id / + client_secret — they read BeaconMCP's .well-known metadata + and register themselves. On our side that needs a single-use bootstrap slug to gate + the registration. +

+
+
ChatGPT
+

+ Mint a one-off connector URL, paste it into the client, it auto-registers. The slug is + single-use, expires in 15 min, and the derived client is bound to your account — 2FA + at /oauth/authorize delegates to your TOTP seed, so a leaked client can't + mint a token without your phone. +

+ Open OAuth connectors + + +
+
+ {% endif %} + +
+
+
+
+

Static bearer tokens

+
+

For clients that attach Authorization: Bearer … to every call.

+
+
+ GeminiMistral + VS CodeCursor + + REST +
+ + {% if just_created %} +
+ + + New token: {{ just_created.name }} — copy it now, it won't be shown again + +
+ {{ just_created.token }} + +
+
+ {% endif %} + + {% if form_error %} + + {% endif %} + +
+
+
Create a token {{ count }} / {{ cap }}
+ {% if can_create %} +
+ + + + +
+ {% else %} +

+ You reached the limit of {{ cap }} active tokens. Revoke one before creating another. +

+ {% endif %} +
+ +
+
Active tokens
+ {% if tokens %} +
    + {% for t in tokens %} +
  • +
    +
    + {{ t.name }} +
    +
    {{ t.prefix }}… · {{ t.expires_label }}
    +
    +
    + + + +
    +
    • + {% endfor %} +
+ {% else %} +

No active tokens yet.

+ {% endif %} +
+
+ + {% if passkeys_enabled %} +
+
Passkeys
+ {% if passkeys %} +
+ {% for p in passkeys %} +
+
+
{{ p.label }}
+
+ Added {{ p.created_human }}{% if p.last_used_human %} · last used {{ p.last_used_human }}{% else %} · never used{% endif %} +
+
+
+ + + +
+
+ {% endfor %} +
+ {% else %} +

+ No passkey registered. Add one right after your next sign-in to skip + the 2FA code on this device. +

+ {% endif %} +
+ {% endif %} + +

+ Concrete setup per client is in the By platform tab. +

+
+
+
+ + {# ============================================================ #} + {# VIEW: BY PLATFORM #} + {# ============================================================ #} +
+ + {% if dcr_enabled %}{% endif %} + + + + + + + + + {# ---- Assistant ---- #} +
+
+
+
+

Assistant

+ OAuth 2.1 +
+

+ Assistant uses OAuth with a user-provided client. Same CLI command creates credentials + for every surface — what changes is where you paste them. +

+
+ + + + + +
+
    +
  1. On the server:
    beaconmcp auth create --name "Assistant Web"
    2. +
  2. assistant.ai or iOS/Android app → Settings → Integrations → Add custom connector.
    3. +
  3. URL: {{ mcp_url }}. Paste client_id / client_secret.
    4. +
  4. Type your TOTP on the authorization page when Assistant redirects you.
    5. +
+
+
+

Assistant Desktop loads MCP servers from a local JSON config.

+
// ~/Library/Application Support/Assistant/assistant_desktop_config.json (macOS)
+// %APPDATA%\Assistant\assistant_desktop_config.json (Windows)
+{
+  "mcpServers": {
+    "beaconmcp": {
+      "command": "npx",
+      "args": [
+        "-y", "mcp-remote",
+        "{{ mcp_url }}",
+        "--oauth"
+      ]
+    }
+  }
+}
+

+ mcp-remote is the community OAuth-to-stdio proxy. It triggers the same + authorization code flow as the web app (TOTP prompt in your browser). Pre-provision + the client with beaconmcp auth create first. +

+
+
+
+ + {# ---- ChatGPT ---- #} + {% if dcr_enabled %} +
+
+
+
+

ChatGPT

+ OAuth + DCR +
+
+ + + + + + +
+

+ ChatGPT's Developer Mode connector UI only exposes "No auth" and "OAuth" — no field + for client_id / client_secret. The OAuth path strictly + requires DCR, so you have to mint a single-use bootstrap URL. +

+
    +
  1. Mint a connector URL from OAuth connectors (single-use, 15 min).
    2. +
  2. Settings → Connectors → Developer Mode → Add custom connector.
    3. +
  3. Paste the /mcp/c/<slug> URL. Auth: OAuth.
    4. +
  4. ChatGPT runs DCR, then redirects you — type your TOTP.
    5. +
+ Mint a connector URL + + +
+
+

Same DCR flow as Web/Desktop — ChatGPT's mobile UI has no credentials field either.

+
    +
  1. Mint a connector URL from OAuth connectors.
    2. +
  2. ChatGPT app → profile → Settings → Connectors → Add custom.
    3. +
  3. Paste the /mcp/c/<slug> URL. OAuth. TOTP in the in-app browser.
    4. +
+

Mint a fresh slug per device if you want independent revocation.

+
+
+

OAuth 2.1

+

+ Unlike the web connector, Codex (OpenAI's terminal/IDE MCP client) lets you + pre-register credentials via its config.toml. It catches the OAuth + redirect on an ephemeral local port and opens your default browser for the TOTP + prompt. +

+
    +
  1. Provision a client on the server:
    beaconmcp auth create --name "Codex"
    2. +
  2. Add BeaconMCP to your Codex config: +
    [mcp_servers.beaconmcp]
+url = "{{ mcp_url }}"
+client_id = "beaconmcp_..."
+client_secret = "sk_..."
    3. +
  3. Run codex mcp login beaconmcp. Your browser opens on BeaconMCP's authorization page — type your TOTP.
    4. +
+

+ Remote devbox (Codespaces, SSH container)? Set mcp_oauth_callback_url + in config.toml to your ingress URL so the redirect hits the right host. +

+
+
+
+ {% endif %} + + {# ---- Gemini ---- #} +
+
+
+

Gemini

+
+ + + + + +
+

OAuth 2.1 Bearer

+

+ Gemini CLI accepts a pre-registered client_id / client_secret + in settings.json, so you don't need a DCR slug — just provision a client + via the CLI and paste the credentials. /mcp auth beaconmcp then runs the + browser flow with your TOTP prompt. +

+

Recommended — OAuth 2.1 (pre-registered):

+
    + 
  1. beaconmcp auth create --name "Gemini CLI"
    2. +
  2. Drop the credentials into ~/.gemini/settings.json: +
    {
+  "mcpServers": {
+    "beaconmcp": {
+      "httpUrl": "{{ mcp_url }}",
+      "oauth": {
+        "clientId": "beaconmcp_...",
+        "clientSecret": "sk_..."
+      }
+    }
+  }
+}
    3. +
  3. In the CLI: /mcp auth beaconmcp. Type your TOTP in the browser that opens.
    4. +
+

Alternative — Bearer:

+
gemini mcp add beaconmcp \
+  --url {{ mcp_url }} \
+  --header "Authorization: Bearer <token>"
+
+ +
+

OAuth 2.1 Bearer

+

+ Antigravity's visual connection manager handles both paths. OAuth keeps the TOTP + prompt on BeaconMCP's side; Bearer is a quick fallback when the OAuth flow misbehaves. +

+

Recommended — OAuth 2.1 (pre-registered):

+
    + 
  1. beaconmcp auth create --name "Antigravity"
    2. +
  2. In Antigravity: Customizations → Connections → Add MCP server. Paste the URL and the OAuth client credentials.
    3. +
  3. Authorize in the browser popup — your TOTP prompt shows up on BeaconMCP's page.
    4. +
+

Alternative — Bearer:

+

Edit ~/.gemini/antigravity/mcp_config.json (macOS / Linux) or %USERPROFILE%\.gemini\antigravity\mcp_config.json (Windows):

+
{
+  "mcpServers": {
+    "beaconmcp": {
+      "serverUrl": "{{ mcp_url }}",
+      "headers": {
+        "Authorization": "Bearer <token>"
+      }
+    }
+  }
+}
+

+ Antigravity uses serverUrl (not url). If the native HTTP + transport misbehaves, fall back to npx mcp-remote with --header. +

+
+ +
+

+ Gemini Web (gemini.google.com), the mobile apps, and the macOS native app do + not expose a custom-MCP connector yet. Only Gemini CLI and + Antigravity can reach BeaconMCP today. +

+
+
+
+ + {# ---- Mistral ---- #} +
+
+
+

Mistral

+
+ + + +
+

OAuth 2.1

+

+ Le Chat speaks OAuth 2.1 natively. In the app: + Intelligence → Connecteurs → Ajouter un connecteur → Connecteur MCP personnalisé. +

+
    +
  1. Name: BeaconMCP
    2. +
  2. MCP Server URL: {{ mcp_url }}
    3. +
  3. Validate. Le Chat discovers the OAuth metadata and redirects you to BeaconMCP's authorization page — type your TOTP from your phone. Token lifetime: 24 h.
    4. +
+

+ Custom connectors are on Le Chat Pro / Enterprise plans; on the free tier the panel + may be hidden. Add https://chat.mistral.ai to + server.allowed_origins so the browser CORS preflight succeeds. +

+
+
+

Bearer (unverified)

+

+ Vibe's bearer path hasn't been tested against a live BeaconMCP. If the snippet below + fails, check the latest Vibe docs — the schema has been iterating fast. +

+

Vibe reads ./.vibe/config.toml (per-project) or ~/.vibe/config.toml (global). TOML format, not JSON.

+
[[mcp_servers]]
+name = "beaconmcp"
+transport = "http"
+url = "{{ mcp_url }}"
+headers = { "Authorization" = "Bearer <token>" }
+

+ transport accepts "http", "streamable-http", + or "stdio". +

+
+
+
+ + {# ---- OpenCode ---- #} +
+
+
+
+

OpenCode

+ OAuth 2.1 + Bearer +
+
+

+ OpenCode supports three flavours in opencode.json: pre-registered OAuth + (recommended), auto-registration via DCR, or a plain bearer token. Tokens land in + ~/.local/share/opencode/mcp-auth.json and refresh on their own. +

+ +

Recommended — OAuth 2.1 (pre-registered):

+
    + 
  1. beaconmcp auth create --name "OpenCode"
    2. +
  2. Add to opencode.json: +
    {
+  "mcp": {
+    "beaconmcp": {
+      "type": "remote",
+      "url": "{{ mcp_url }}",
+      "enabled": true,
+      "oauth": {
+        "clientId": "beaconmcp_...",
+        "clientSecret": "sk_..."
+      }
+    }
+  }
+}
    3. +
  3. Run opencode mcp auth beaconmcp. Type your TOTP in the browser.
    4. +
+ + {% if dcr_enabled %} +

Alternative — DCR (no credentials to paste):

+
{
+  "mcp": {
+    "beaconmcp": {
+      "type": "remote",
+      "url": "https://<your-host>/mcp/c/<slug>",
+      "enabled": true,
+      "oauth": true
+    }
+  }
+}
+ Mint a connector URL + + + {% endif %} + +

Alternative — Bearer:

+
{
+  "mcp": {
+    "beaconmcp": {
+      "type": "remote",
+      "url": "{{ mcp_url }}",
+      "enabled": true,
+      "oauth": false,
+      "headers": {
+        "Authorization": "Bearer <token>"
+      }
+    }
+  }
+}
+
+
+ + {# ---- VS Code ---- #} +
+
+
+
+

VS Code

+ OAuth 2.1 + Bearer +
+
+

+ VS Code routes MCP auth through its native Authentication Provider system (same as + GitHub / Microsoft Entra). It reads WWW-Authenticate, shows a toast to + Allow, catches the redirect on vscode:// (or vscode-insiders://), + and stores tokens in the OS keychain. +

+ +

Recommended — OAuth 2.1 (pre-registered):

+
    + 
  1. beaconmcp auth create --name "VS Code"
    2. +
  2. Add to .vscode/mcp.json (or settings.json → "mcp.servers"): +
    {
+  "inputs": [
+    { "type": "promptString", "id": "beaconmcp-client-id", "description": "client_id" },
+    { "type": "promptString", "id": "beaconmcp-client-secret", "description": "client_secret", "password": true }
+  ],
+  "servers": {
+    "beaconmcp": {
+      "type": "http",
+      "url": "{{ mcp_url }}",
+      "clientId": "${input:beaconmcp-client-id}",
+      "clientSecret": "${input:beaconmcp-client-secret}"
+    }
+  }
+}
    3. +
  3. VS Code prompts you on first use — OS keychain stores the tokens after your TOTP.
    4. +
+ + {% if dcr_enabled %} +

Alternative — DCR:

+
{
+  "servers": {
+    "beaconmcp": {
+      "type": "http",
+      "url": "https://<your-host>/mcp/c/<slug>"
+    }
+  }
+}
+ Mint a connector URL + + + {% endif %} + +

Alternative — Bearer:

+
{
+  "servers": {
+    "beaconmcp": {
+      "type": "http",
+      "url": "{{ mcp_url }}",
+      "headers": {
+        "Authorization": "Bearer <token>"
+      }
+    }
+  }
+}
+

+ Command Palette → MCP: List Servers to verify. Revoke access anytime from the + Accounts menu (profile icon, bottom left) — MCP sessions live next to your other + trusted logins. +

+
+
+ + {# ---- Cursor ---- #} +
+
+
+
+

Cursor

+ OAuth 2.1 + Bearer +
+
+

+ Cursor is a first-class OAuth 2.1 client since v1.0. It surfaces a blue + Connect button in Settings → Tools & MCP and catches the redirect + via the cursor:// scheme (or a loopback fallback). PKCE runs in-app. +

+ +

Recommended — OAuth 2.1 (pre-registered):

+
    + 
  1. beaconmcp auth create --name "Cursor"
    2. +
  2. Add to ~/.cursor/mcp.json (global) or .cursor/mcp.json (per project): +
    {
+  "mcpServers": {
+    "beaconmcp": {
+      "url": "{{ mcp_url }}",
+      "clientId": "${env:BEACONMCP_CLIENT_ID}",
+      "clientSecret": "${env:BEACONMCP_CLIENT_SECRET}"
+    }
+  }
+}
    3. +
  3. Export the credentials in your shell (BEACONMCP_CLIENT_ID / BEACONMCP_CLIENT_SECRET). Reload Cursor; click Connect.
    4. +
+ + {% if dcr_enabled %} +

Alternative — DCR:

+
{
+  "mcpServers": {
+    "beaconmcp": {
+      "url": "https://<your-host>/mcp/c/<slug>"
+    }
+  }
+}
+ Mint a connector URL + + + {% endif %} + +

Alternative — Bearer:

+
{
+  "mcpServers": {
+    "beaconmcp": {
+      "url": "{{ mcp_url }}",
+      "headers": {
+        "Authorization": "Bearer ${env:BEACONMCP_TOKEN}"
+      }
+    }
+  }
+}
+

+ Reload the Cursor window after editing; the server shows up under Settings → + Cursor Settings → MCP Servers. Cursor supports ${env:VAR} expansion + so you can keep secrets out of the repo. +

+
+
+ + {# ---- Other ---- #} +
+
+
+
+

Other MCP-HTTP clients

+ Bearer +
+
+

+ Any client that can send Authorization: Bearer … on HTTP POST works. Mint + a token, point the client at {{ mcp_url }}, and attach the bearer on + every request. Named tokens live {{ named_ttl_label }} and + survive restarts — revoke them here any time. +

+
+
+
+
+
+ + +{% endblock %} diff --git a/src/beaconmcp/dashboard/templates/totp_refresh.html b/src/beaconmcp/dashboard/templates/totp_refresh.html new file mode 100644 index 0000000..b552fd0 --- /dev/null +++ b/src/beaconmcp/dashboard/templates/totp_refresh.html @@ -0,0 +1,44 @@ +{% extends "base.html" %} +{% block title %}2FA · BeaconMCP{% endblock %} +{% block body_class %}auth-page{% endblock %} +{% block body %} +
+
+ BeaconMCP +
+ +

Two-factor

+

Session active for {{ client_name }}. Enter the 6-digit code from your authenticator app to renew your MCP token.

+ + {% if banner %} + + {% endif %} + +
+ + {% if next %}{% endif %} + +
+ + + + + + +
+ + + +
+ +
+ + +
+
+ + +{% endblock %} diff --git a/src/beaconmcp/dashboard/templates/usage_cost.html b/src/beaconmcp/dashboard/templates/usage_cost.html new file mode 100644 index 0000000..ff2aeb0 --- /dev/null +++ b/src/beaconmcp/dashboard/templates/usage_cost.html @@ -0,0 +1,50 @@ +{% extends "base.html" %} +{% block title %}Usage & Cost · BeaconMCP{% endblock %} +{% block body_class %}tokens-page{% endblock %} +{% block body %} + + +
+
+
BeaconMCP
+
+ + + Back to chat + +
+ +

Usage & Cost Tracking UI

+

+ Monitor token usage and estimated LLM costs. +

+ +
+
Current Billing Cycle (Estimated)
+
+
Coming soon
+
+
+ +
+

Token Consumption

+ + + + + + + + + + + + + + +
DateModelInput TokensOutput Tokens
(Data pipeline coming soon)
+
+
+{% endblock %} diff --git a/src/beaconmcp/dashboard/usage.py b/src/beaconmcp/dashboard/usage.py new file mode 100644 index 0000000..d83bb65 --- /dev/null +++ b/src/beaconmcp/dashboard/usage.py @@ -0,0 +1,333 @@ +"""Per-client usage accounting and budget enforcement. + +Two moving parts: + +- :class:`UsageMeter` -- stateless pricing calculator. Turns Gemini + ``usage_metadata`` (prompt/cached/output token counts + model) into a + USD cost using the public Google AI Studio rate card. +- :class:`UsageStore` -- SQLite-backed ledger + 5h session tracker. One + row per assistant turn in ``usage_events`` (immutable), plus one + live-session row per ``client_id`` in ``usage_5h_sessions`` that is + reset whenever the 5-hour window expires. + +Windows: + +- **5h session (Anthropic-style)**: a contiguous 5h window that opens on + the first turn after any inactivity of >=5h. While the window is open, + turns accumulate into it. Once the window closes (now - started_at >= + 18000s at check time), the next turn resets the window to 0 and starts + a new one beginning at that moment. +- **Weekly (rolling)**: a trailing 7-day sum over ``usage_events``. + +Both caps are configurable via env vars, read in ``__main__`` and passed +into :class:`Budget`. Cap <= 0 means "unlimited". +""" + +from __future__ import annotations + +import time +import uuid +from dataclasses import dataclass +from typing import Any + +from .db import Database + + +# --------------------------------------------------------------------------- +# Pricing +# --------------------------------------------------------------------------- + +# Rates are in USD per 1M tokens. Pulled from the Google AI Studio pricing +# page on 2026-07-29. Keys with ``_hi`` suffixes apply when the prompt +# token count exceeds :data:`_TIER_THRESHOLD`; only the Pro models have +# a high tier in the public rate card. +# +# The retired entries at the bottom are deliberate. A model leaving +# ``VALID_MODELS`` does not erase the turns it billed, and ``cost_usd`` is +# called again whenever a stored turn is re-priced -- dropping the rates +# would silently re-bill that history at the fallback model's price. +_PRICING: dict[str, dict[str, float]] = { + "gemini-3.6-flash": { + "input": 1.50, "cached": 0.15, "output": 7.50, + }, + "gemini-3.5-flash-lite": { + "input": 0.30, "cached": 0.03, "output": 2.50, + }, + "gemini-3.1-pro-preview": { + "input": 2.00, "cached": 0.20, "output": 12.00, + "input_hi": 4.00, "cached_hi": 0.40, "output_hi": 18.00, + }, + # --- no longer selectable, kept so old turns keep their real price --- + "gemini-2.5-flash": { + "input": 0.30, "cached": 0.03, "output": 2.50, + }, + "gemini-2.5-pro": { + "input": 1.25, "cached": 0.125, "output": 10.00, + "input_hi": 2.50, "cached_hi": 0.25, "output_hi": 15.00, + }, + "gemini-3-flash-preview": { + "input": 0.50, "cached": 0.05, "output": 3.00, + }, +} + +# Prompt tokens above this count trigger the Pro models' "long prompt" +# pricing tier. +_TIER_THRESHOLD = 200_000 + +_FIVE_HOURS_SECS = 5 * 3600 +_SEVEN_DAYS_SECS = 7 * 24 * 3600 + + +# --------------------------------------------------------------------------- +# Meter +# --------------------------------------------------------------------------- + + +class UsageMeter: + """Pure function: (model, token counts) -> cost USD.""" + + @staticmethod + def cost_usd( + model: str, + *, + prompt_tokens: int, + cached_tokens: int, + output_tokens: int, + ) -> float: + rates = _PRICING.get(model) or _PRICING["gemini-3.6-flash"] + use_hi = prompt_tokens > _TIER_THRESHOLD and "input_hi" in rates + in_rate = rates["input_hi"] if use_hi else rates["input"] + out_rate = rates["output_hi"] if use_hi else rates["output"] + ca_rate = rates["cached_hi"] if use_hi else rates["cached"] + + # cached_tokens is the subset of prompt_tokens served from cache; + # bill the remainder at the input rate and cached_tokens at the + # cached rate (which is ~10x cheaper across the board). + billable_input = max(0, prompt_tokens - cached_tokens) + total = ( + billable_input * in_rate + + cached_tokens * ca_rate + + output_tokens * out_rate + ) + return total / 1_000_000 + + +# --------------------------------------------------------------------------- +# Budget config +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class Budget: + """Per-client spending caps (global, identical across clients). + + ``0`` (or any non-positive value) disables the cap on that window. + Units: USD. + """ + + limit_5h_usd: float + limit_week_usd: float + + @property + def has_any_limit(self) -> bool: + return self.limit_5h_usd > 0 or self.limit_week_usd > 0 + + +# --------------------------------------------------------------------------- +# Store +# --------------------------------------------------------------------------- + + +@dataclass +class UsageSnapshot: + """Current usage picture for one client, used by the UI footer.""" + + spent_5h_usd: float + limit_5h_usd: float + session_5h_started_at: float | None # None if session window is empty/expired + session_5h_reset_at: float | None # started_at + 5h, when UI can expect reset + + spent_week_usd: float + limit_week_usd: float + + def to_json(self) -> dict[str, Any]: + return { + "spent_5h_usd": round(self.spent_5h_usd, 6), + "limit_5h_usd": self.limit_5h_usd, + "session_5h_started_at": self.session_5h_started_at, + "session_5h_reset_at": self.session_5h_reset_at, + "spent_week_usd": round(self.spent_week_usd, 6), + "limit_week_usd": self.limit_week_usd, + } + + +@dataclass +class BudgetBlock: + """Returned by ``check_budget`` when a request must be refused.""" + + window: str # "5h" | "week" + spent_usd: float + limit_usd: float + reset_at: float | None # absolute epoch seconds of window reset, if known + + +class UsageStore: + def __init__(self, db: Database, budget: Budget) -> None: + self._db = db + self._budget = budget + + @property + def budget(self) -> Budget: + return self._budget + + # --- write path ------------------------------------------------------- + + def record_turn( + self, + *, + client_id: str, + conversation_id: str | None, + message_id: str | None, + model: str, + prompt_tokens: int, + cached_tokens: int, + output_tokens: int, + cost_usd: float, + now: float | None = None, + ) -> None: + """Append a ledger row and update the 5h session row atomically.""" + ts = now if now is not None else time.time() + conn = self._db.conn() + conn.execute("BEGIN") + try: + conn.execute( + """ + INSERT INTO usage_events (id, client_id, conversation_id, + message_id, ts, model, + prompt_tokens, cached_tokens, + output_tokens, cost_usd) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + str(uuid.uuid4()), client_id, conversation_id, message_id, + ts, model, prompt_tokens, cached_tokens, output_tokens, + cost_usd, + ), + ) + self._apply_to_session(conn, client_id, ts, cost_usd) + conn.execute("COMMIT") + except Exception: + conn.execute("ROLLBACK") + raise + + def _apply_to_session( + self, conn: Any, client_id: str, ts: float, cost_usd: float, + ) -> None: + row = conn.execute( + "SELECT started_at, last_event_at, cost_usd " + " FROM usage_5h_sessions WHERE client_id = ?", + (client_id,), + ).fetchone() + if row is None: + conn.execute( + "INSERT INTO usage_5h_sessions " + " (client_id, started_at, last_event_at, cost_usd) " + " VALUES (?, ?, ?, ?)", + (client_id, ts, ts, cost_usd), + ) + return + started_at = float(row["started_at"]) + if ts - started_at >= _FIVE_HOURS_SECS: + # Previous session expired -- start a fresh one at ``ts``. + conn.execute( + "UPDATE usage_5h_sessions " + " SET started_at = ?, last_event_at = ?, cost_usd = ? " + " WHERE client_id = ?", + (ts, ts, cost_usd, client_id), + ) + else: + conn.execute( + "UPDATE usage_5h_sessions " + " SET last_event_at = ?, cost_usd = cost_usd + ? " + " WHERE client_id = ?", + (ts, cost_usd, client_id), + ) + + # --- read path -------------------------------------------------------- + + def snapshot(self, client_id: str, *, now: float | None = None) -> UsageSnapshot: + """Return current usage for ``client_id``. + + The 5h-window figure reflects the Anthropic-style session: if the + last known session has been dormant for >=5h, we report 0 spent + (the window is closed and the next turn will open a new one). + """ + ts = now if now is not None else time.time() + row = self._db.conn().execute( + "SELECT started_at, cost_usd " + " FROM usage_5h_sessions WHERE client_id = ?", + (client_id,), + ).fetchone() + if row is None: + spent_5h = 0.0 + started_at = None + reset_at = None + else: + started_at = float(row["started_at"]) + if ts - started_at >= _FIVE_HOURS_SECS: + spent_5h = 0.0 + started_at = None + reset_at = None + else: + spent_5h = float(row["cost_usd"]) + reset_at = started_at + _FIVE_HOURS_SECS + + week_row = self._db.conn().execute( + "SELECT COALESCE(SUM(cost_usd), 0) AS total " + " FROM usage_events WHERE client_id = ? AND ts >= ?", + (client_id, ts - _SEVEN_DAYS_SECS), + ).fetchone() + spent_week = float(week_row["total"] if week_row else 0.0) + + return UsageSnapshot( + spent_5h_usd=spent_5h, + limit_5h_usd=self._budget.limit_5h_usd, + session_5h_started_at=started_at, + session_5h_reset_at=reset_at, + spent_week_usd=spent_week, + limit_week_usd=self._budget.limit_week_usd, + ) + + # --- enforcement ------------------------------------------------------ + + def check_budget( + self, client_id: str, *, now: float | None = None, + ) -> BudgetBlock | None: + """Return a :class:`BudgetBlock` if ``client_id`` is over either + cap, or ``None`` if the request may proceed. + + Called before handing a turn to the chat engine. We err on the + side of letting the turn through when both caps are 0 (disabled) + so that users can run an unmetered setup if they choose. + """ + if not self._budget.has_any_limit: + return None + snap = self.snapshot(client_id, now=now) + if snap.limit_5h_usd > 0 and snap.spent_5h_usd >= snap.limit_5h_usd: + return BudgetBlock( + window="5h", + spent_usd=snap.spent_5h_usd, + limit_usd=snap.limit_5h_usd, + reset_at=snap.session_5h_reset_at, + ) + if snap.limit_week_usd > 0 and snap.spent_week_usd >= snap.limit_week_usd: + # Rolling 7 days -> the reset moment isn't a single clock tick, + # so we leave ``reset_at`` unset; the UI formats this as + # "sur 7 jours glissants" rather than an absolute time. + return BudgetBlock( + window="week", + spent_usd=snap.spent_week_usd, + limit_usd=snap.limit_week_usd, + reset_at=None, + ) + return None diff --git a/src/beaconmcp/maintenance/__init__.py b/src/beaconmcp/maintenance/__init__.py new file mode 100644 index 0000000..1d488b2 --- /dev/null +++ b/src/beaconmcp/maintenance/__init__.py @@ -0,0 +1,5 @@ +"""Self-maintenance tools: update checking and applying.""" + +from .tools import register_maintenance_tools + +__all__ = ["register_maintenance_tools"] diff --git a/src/beaconmcp/maintenance/tools.py b/src/beaconmcp/maintenance/tools.py new file mode 100644 index 0000000..ec6364c --- /dev/null +++ b/src/beaconmcp/maintenance/tools.py @@ -0,0 +1,117 @@ +"""MCP tools for keeping the BeaconMCP server itself up to date.""" + +from __future__ import annotations + +from pathlib import Path + +from mcp.server.fastmcp import FastMCP + +from .. import audit, updates +from ..auth import current_client_id +from ..config import UpdatesConfig + + +def register_maintenance_tools( + mcp: FastMCP, + settings: UpdatesConfig | None = None, + *, + config_path: Path | None = None, +) -> None: + """Register ``beaconmcp_check_update`` and ``beaconmcp_self_update``. + + ``settings.enabled`` gates the whole module (no network egress at all); + ``settings.allow_self_update`` keeps the check but refuses to apply. + """ + settings = settings or UpdatesConfig() + if not settings.enabled: + return + + @mcp.tool() + def beaconmcp_check_update() -> dict: + """Check whether a newer BeaconMCP revision is available. + + Reports the running version, how many commits this install is + behind the upstream default branch, the changelog between the two, + and — importantly — any configuration the new revision knows about + that this install has not set yet (new ``.env`` variables, new + ``beaconmcp.yaml`` settings). + + Also returns the exact shell commands that would update *this* + install, which differ between a git checkout, a pip install and a + container. + + Read-only and safe to call at any time: it fetches git objects but + never modifies the working tree. Results are cached for a few hours; + this returns the cached answer when it is still fresh. + """ + info = updates.check_for_update(config_path=config_path) + payload = info.to_json() + payload["self_update_allowed"] = settings.allow_self_update + if info.can_self_update and not settings.allow_self_update: + payload["can_self_update"] = False + payload["blockers"] = [ + *payload.get("blockers", []), + "self-update is disabled by features.updates.allow_self_update", + ] + return payload + + if not settings.allow_self_update: + return + + @mcp.tool() + def beaconmcp_self_update(confirm: bool = False, restart: bool = True) -> dict: + """Update this BeaconMCP server to the latest upstream revision. + + Runs, in order: ``git pull --ff-only`` → reinstall the Python + package and its dependencies → **validate the configuration against + the new code** → schedule a service restart. + + The configuration check is a hard gate. If the new revision cannot + load the operator's config (because a setting was renamed, or a new + one is now required), the checkout is rolled back to exactly where + it started, dependencies are restored, and nothing is restarted. The + return value says so explicitly. + + Requires ``confirm=True``. Call ``beaconmcp_check_update`` first and + show the user what is about to change — including any new config + variables — before asking them to confirm. + + Refuses to run when the checkout has uncommitted changes, or when + this is not a git install; ``beaconmcp_check_update`` reports those + blockers in advance along with manual instructions. + + The restart is deliberately deferred a few seconds so this response + reaches you before the process dies. After that, expect the server + to be briefly unreachable. + """ + if not confirm: + info = updates.check_for_update(config_path=config_path) + return { + "ok": False, + "applied": False, + "reason": "confirmation_required", + "message": ( + "This will pull new code, reinstall dependencies and " + "restart the server. Review the pending changes, then " + "call again with confirm=True." + ), + "pending": info.to_json(), + } + + client_id = current_client_id() + audit.emit("maintenance.self_update.start", client_id=client_id) + result = updates.apply_update(restart=restart, config_path=config_path) + audit.emit( + "maintenance.self_update.finish", + client_id=client_id, + ok=result.ok, + from_ref=result.from_ref, + to_ref=result.to_ref, + rolled_back=result.rolled_back, + ) + # The next check must not serve a stale "update available". + updates.invalidate_cache() + + payload = result.to_json() + payload["applied"] = result.ok + return payload diff --git a/src/beaconmcp/metrics.py b/src/beaconmcp/metrics.py new file mode 100644 index 0000000..ec313b0 --- /dev/null +++ b/src/beaconmcp/metrics.py @@ -0,0 +1,146 @@ +"""Minimal Prometheus text-format metrics for BeaconMCP. + +Deliberately avoids ``prometheus_client`` to keep the dependency tree +small. Two primitive counter types cover everything we need right now: + +* :class:`Counter` -- monotonic integer counter, optionally labelled. +* :class:`Histogram` -- fixed-bucket histogram over milliseconds. + +The :class:`Registry` collects them and renders the Prometheus text +exposition format on demand. Thread-safe via a single registry lock. + +Usage:: + + from beaconmcp.metrics import REGISTRY, tool_calls, tool_latency_ms + + tool_calls.inc(tool="proxmox_run", status="ok") + tool_latency_ms.observe(123.4, tool="proxmox_run") + + text = REGISTRY.render() # served at /metrics +""" + +from __future__ import annotations + +import threading +import time +from contextlib import contextmanager +from typing import Iterator + + +def _labels_key(labels: dict[str, str]) -> tuple[tuple[str, str], ...]: + return tuple(sorted(labels.items())) + + +def _format_labels(labels: tuple[tuple[str, str], ...]) -> str: + if not labels: + return "" + parts = [f'{k}="{str(v).replace(chr(92), chr(92) + chr(92)).replace(chr(34), chr(92) + chr(34))}"' for k, v in labels] + return "{" + ",".join(parts) + "}" + + +class Counter: + def __init__(self, name: str, help_text: str) -> None: + self.name = name + self.help = help_text + self._values: dict[tuple[tuple[str, str], ...], int] = {} + self._lock = threading.Lock() + + def inc(self, amount: int = 1, **labels: str) -> None: + key = _labels_key(labels) + with self._lock: + self._values[key] = self._values.get(key, 0) + amount + + def render(self) -> str: + lines = [f"# HELP {self.name} {self.help}", f"# TYPE {self.name} counter"] + with self._lock: + snapshot = dict(self._values) + for key, value in snapshot.items(): + lines.append(f"{self.name}{_format_labels(key)} {value}") + return "\n".join(lines) + + +class Histogram: + """Fixed-bucket histogram. Buckets are upper bounds in milliseconds.""" + + # Covers <10ms cached calls all the way to 30-second BMC round trips. + DEFAULT_BUCKETS_MS: tuple[float, ...] = ( + 5, 10, 25, 50, 100, 250, 500, 1000, 2500, 5000, 10000, 30000, + ) + + def __init__(self, name: str, help_text: str, buckets_ms: tuple[float, ...] | None = None) -> None: + self.name = name + self.help = help_text + self._buckets: tuple[float, ...] = tuple(buckets_ms or self.DEFAULT_BUCKETS_MS) + # Per-label bucket counts + running sum/count. + self._counts: dict[tuple[tuple[str, str], ...], list[int]] = {} + self._sum: dict[tuple[tuple[str, str], ...], float] = {} + self._total: dict[tuple[tuple[str, str], ...], int] = {} + self._lock = threading.Lock() + + def observe(self, value_ms: float, **labels: str) -> None: + key = _labels_key(labels) + with self._lock: + counts = self._counts.setdefault(key, [0] * len(self._buckets)) + for i, upper in enumerate(self._buckets): + if value_ms <= upper: + counts[i] += 1 + self._sum[key] = self._sum.get(key, 0.0) + value_ms + self._total[key] = self._total.get(key, 0) + 1 + + @contextmanager + def time(self, **labels: str) -> Iterator[None]: + start = time.monotonic() + try: + yield + finally: + self.observe((time.monotonic() - start) * 1000.0, **labels) + + def render(self) -> str: + lines = [f"# HELP {self.name} {self.help}", f"# TYPE {self.name} histogram"] + with self._lock: + counts = {k: list(v) for k, v in self._counts.items()} + sums = dict(self._sum) + totals = dict(self._total) + for key, bucket_counts in counts.items(): + # ``observe`` already increments every bucket whose upper bound + # is >= the value, so each slot holds the cumulative count -- + # render them directly, Prometheus-style. + for i, upper in enumerate(self._buckets): + labels_with_le = tuple(sorted(key + (("le", str(upper)),))) + lines.append(f"{self.name}_bucket{_format_labels(labels_with_le)} {bucket_counts[i]}") + labels_inf = tuple(sorted(key + (("le", "+Inf"),))) + lines.append(f"{self.name}_bucket{_format_labels(labels_inf)} {totals[key]}") + lines.append(f"{self.name}_sum{_format_labels(key)} {sums[key]}") + lines.append(f"{self.name}_count{_format_labels(key)} {totals[key]}") + return "\n".join(lines) + + +class Registry: + def __init__(self) -> None: + self._metrics: list[Counter | Histogram] = [] + + def register(self, metric: Counter | Histogram) -> Counter | Histogram: + self._metrics.append(metric) + return metric + + def render(self) -> str: + parts = [m.render() for m in self._metrics] + return "\n".join(parts) + "\n" + + +# --- Default registry + standard metrics ----------------------------------- + +REGISTRY = Registry() + +tool_calls: Counter = REGISTRY.register( # type: ignore[assignment] + Counter("beaconmcp_tool_calls_total", "Total MCP tool invocations, by tool and status.") +) +tool_latency_ms: Histogram = REGISTRY.register( # type: ignore[assignment] + Histogram("beaconmcp_tool_latency_ms", "Tool call latency in milliseconds, by tool.") +) +auth_events: Counter = REGISTRY.register( # type: ignore[assignment] + Counter("beaconmcp_auth_events_total", "Auth events, by kind (login, token, refresh) and outcome.") +) +http_requests: Counter = REGISTRY.register( # type: ignore[assignment] + Counter("beaconmcp_http_requests_total", "HTTP requests to BeaconMCP endpoints, by path and status.") +) diff --git a/src/beaconmcp/proxmox/__init__.py b/src/beaconmcp/proxmox/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/beaconmcp/proxmox/aggregators.py b/src/beaconmcp/proxmox/aggregators.py new file mode 100644 index 0000000..4a2c1e9 --- /dev/null +++ b/src/beaconmcp/proxmox/aggregators.py @@ -0,0 +1,390 @@ +"""High-level aggregator tools that collapse multi-call workflows into one. + +Rationale +--------- +A typical diagnostic session in an MCP client looks like: +``list_nodes`` -> ``list_vms`` -> ``storage_status`` -> ``node_status`` -> ``get_logs``. +That's five tool calls, five round-trips, and a lot of repeated JSON. + +The helpers in this module return the same information in one call each, at +the cost of slightly larger payloads. Clients keep full access to the +fine-grained tools; these aggregators exist so the LLM can pick a shorter +path when it doesn't yet know what it's looking for. + +All aggregators are careful to: +* gracefully downgrade if a capability is missing (no SSH -> no SSH facts; + no BMC registry -> no hardware facts). +* report errors inline per node/VM rather than failing the whole call, so the + caller can still work with the partial view. +""" + +from __future__ import annotations + +import asyncio +import fnmatch +from concurrent.futures import ThreadPoolExecutor +from typing import Any + +from mcp.server.fastmcp import FastMCP + +from ..config import Config +from ..utils import filter_fields +from .client import ProxmoxClient + + +# --------------------------------------------------------------------------- +# Internal helpers +# --------------------------------------------------------------------------- + +# Cap on concurrent Proxmox/BMC fan-outs. Small homelab clusters (3-10 nodes) +# fit comfortably; bigger clusters still benefit from parallelism without +# hammering the API with hundreds of simultaneous TLS handshakes. +_MAX_PARALLEL = 8 + + +def _parallel_map(fn: Any, items: list[Any]) -> list[Any]: + """Run ``fn(item)`` for each item in parallel threads, preserving order. + + Keeps aggregator fan-out roughly linear in the slowest node rather than + serial sum-of-all-nodes. Errors bubble back as the function's normal + return shape (each helper already returns error dicts inline), so we + don't catch here. + """ + if not items: + return [] + if len(items) == 1: + return [fn(items[0])] + with ThreadPoolExecutor(max_workers=min(_MAX_PARALLEL, len(items))) as ex: + return list(ex.map(fn, items)) + + +def _collect_node_summaries(client: ProxmoxClient) -> list[dict[str, Any]]: + """One row per configured node with key health metrics. + + Mirrors ``proxmox_list_nodes`` but keeps only the fields cluster_overview + actually needs to stay token-efficient. + """ + def _one(node_name: str) -> dict[str, Any]: + data = client.get(node_name, "nodes") + if isinstance(data, dict) and "error" in data: + return {"name": node_name, "status": "unreachable", "error": data["error"]} + if not isinstance(data, list): + return {"name": node_name, "status": "unknown"} + for node in data: + if node.get("node") != node_name: + continue + return { + "name": node_name, + "status": node.get("status", "unknown"), + "cpu": round(node.get("cpu", 0) * 100, 1), + "mem_used_gb": round(node.get("mem", 0) / 1073741824, 1), + "mem_total_gb": round(node.get("maxmem", 0) / 1073741824, 1), + "uptime_h": round(node.get("uptime", 0) / 3600, 1), + } + return {"name": node_name, "status": "unknown"} + + return _parallel_map(_one, list(client.configured_nodes)) + + +def _collect_vm_summaries( + client: ProxmoxClient, target_nodes: list[str] | None = None +) -> tuple[list[dict[str, Any]], int]: + """Flat list of VMs across one or more nodes + total count. + + Returns a flat list (not nested by node) because callers that use this + helper want to filter/count across the whole set; the per-node nesting + shape is already available via ``proxmox_list_vms``. + """ + nodes = list(target_nodes or client.configured_nodes) + + # Each node needs a qemu + lxc fetch. Flatten to (node, vm_type) tuples + # so the whole fan-out runs in parallel instead of 2 * N serial calls. + tasks = [(n, t) for n in nodes for t in ("qemu", "lxc")] + + def _one(task: tuple[str, str]) -> list[dict[str, Any]]: + n, vm_type = task + data = client.get(n, f"nodes/{n}/{vm_type}") + if not isinstance(data, list): + return [] + return [{ + "node": n, + "vmid": vm.get("vmid"), + "name": vm.get("name", ""), + "status": vm.get("status"), + "type": vm_type, + "cpu_pct": round(vm.get("cpu", 0) * 100, 1), + "mem_used_mb": round(vm.get("mem", 0) / 1048576, 0), + } for vm in data] + + rows: list[dict[str, Any]] = [] + for chunk in _parallel_map(_one, tasks): + rows.extend(chunk) + rows.sort(key=lambda v: (v.get("node", ""), v.get("vmid", 0))) + return rows, len(rows) + + +def _collect_storage_summaries(client: ProxmoxClient) -> list[dict[str, Any]]: + def _per_node(n: str) -> list[dict[str, Any]]: + data = client.get(n, f"nodes/{n}/storage") + if isinstance(data, dict) and "error" in data: + return [{"node": n, "error": data["error"]}] + if not isinstance(data, list): + return [] + # Fan out the per-pool status queries within a node too -- 4+ pools + # per node is common (local, zfs, nfs, cephfs). + pools = [s for s in data if s.get("storage")] + + def _pool(s: dict[str, Any]) -> dict[str, Any]: + name = s["storage"] + status = client.get(n, f"nodes/{n}/storage/{name}/status") + used = total = 0 + if isinstance(status, dict) and "error" not in status: + used = status.get("used", 0) + total = status.get("total", 0) + return { + "node": n, + "name": name, + "type": s.get("type"), + "used_gb": round(used / 1073741824, 1), + "total_gb": round(total / 1073741824, 1), + "usage_pct": round(used / total * 100, 1) if total > 0 else 0, + } + + return _parallel_map(_pool, pools) + + rows: list[dict[str, Any]] = [] + for chunk in _parallel_map(_per_node, list(client.configured_nodes)): + rows.extend(chunk) + return rows + + +def _find_vm_location(client: ProxmoxClient, vmid: int) -> tuple[str, str] | None: + """Return (node, vm_type) for a VMID; None if not found.""" + for n in client.configured_nodes: + for vm_type in ("qemu", "lxc"): + data = client.get(n, f"nodes/{n}/{vm_type}/{vmid}/status/current") + if isinstance(data, dict) and "error" in data: + continue + if isinstance(data, dict) and data.get("status"): + return (n, vm_type) + return None + + +async def _bmc_summary(bmc_registry: dict, config: Config, node: str) -> dict[str, Any] | None: + """Return a short BMC status blurb for the given node, if one is mapped. + + Heuristic: a BMC device is "attached" to a node when its ``jump_host`` + matches the node name. That's how HP iLO setups tend to be declared and + is the only mapping the config currently exposes. + """ + if not bmc_registry: + return None + matches = [d for d in config.bmc_devices if d.jump_host == node] + if not matches: + return None + device = matches[0] + client = bmc_registry.get(device.id) + if not client: + return {"device_id": device.id, "error": "BMC device in config but not in registry"} + try: + power = await client.power_status() + health = await client.health() + except Exception as exc: # noqa: BLE001 -- surface anything as a soft error + return {"device_id": device.id, "error": str(exc)} + return { + "device_id": device.id, + "type": device.type, + "power": power.get("power_status") if isinstance(power, dict) else None, + "health_summary": { + k: v + for k, v in (health.items() if isinstance(health, dict) else []) + if k in ("overall", "fans", "temperatures", "power_supplies") + }, + } + + +def _recent_errors(client: ProxmoxClient, node: str, limit: int = 20) -> list[dict[str, Any]]: + """Pull the last ``limit`` failed tasks on ``node``. + + Proxmox exposes task exit status as a string: "OK" for success, anything + else (including "unknown", actual error strings) means not-ok. + """ + data = client.get(node, f"nodes/{node}/tasks", limit=limit) + if not isinstance(data, list): + return [] + errors = [] + for t in data: + status = t.get("status", "") + if status and status != "OK": + errors.append({ + "upid": t.get("upid"), + "type": t.get("type"), + "status": status, + "user": t.get("user"), + "starttime": t.get("starttime"), + "endtime": t.get("endtime"), + }) + return errors + + +# --------------------------------------------------------------------------- +# Registration +# --------------------------------------------------------------------------- + +def register_aggregator_tools( + mcp: FastMCP, + proxmox_client: ProxmoxClient, + config: Config, + bmc_registry: dict | None = None, +) -> None: + """Register the four aggregator tools. + + ``bmc_registry`` is accepted as a plain dict (device_id -> BMCClient) + rather than importing the type, so this module can be registered even + when BMC support is disabled. + """ + + bmc_registry = bmc_registry or {} + + @mcp.tool() + def cluster_overview( + include_storage: bool = True, + fields: list[str] | None = None, + ) -> dict[str, Any]: + """Return cluster state (nodes + VMs + optional storage) in one call. + + Use this as the first diagnostic step -- it replaces + ``proxmox_list_nodes`` + ``proxmox_list_vms`` + ``proxmox_storage_status``. + Set ``include_storage=False`` to skip storage (saves the per-pool status + roundtrip on large clusters). Pass ``fields=[...]`` to trim each entry + to only the keys you need (applied uniformly to nodes/vms/storage). + """ + nodes = _collect_node_summaries(proxmox_client) + vms, total_vms = _collect_vm_summaries(proxmox_client) + out: dict[str, Any] = { + "nodes": filter_fields(nodes, fields), + "vms": filter_fields(vms, fields), + "total_vms": total_vms, + } + if include_storage: + out["storage"] = filter_fields(_collect_storage_summaries(proxmox_client), fields) + return out + + @mcp.tool() + async def cluster_health(node: str = "") -> dict[str, Any]: + """Aggregate health signals for one node (or all nodes): metrics + BMC + recent errors. + + Collapses ``proxmox_node_status`` + ``bmc_health_status`` + + ``proxmox_get_tasks`` into one call. When ``node`` is empty every + configured node is scanned. BMC data is only attached for nodes that + have a BMC device declared with ``jump_host: ``. + """ + target_nodes = [node] if node else list(proxmox_client.configured_nodes) + + async def _one(n: str) -> dict[str, Any]: + # Offload the blocking Proxmox calls to a thread so we can run the + # BMC await + Proxmox fetch in parallel per node, and every node in + # parallel overall via gather. + loop = asyncio.get_running_loop() + status_task = loop.run_in_executor( + None, lambda: proxmox_client.get(n, f"nodes/{n}/status"), + ) + errors_task = loop.run_in_executor( + None, lambda: _recent_errors(proxmox_client, n, limit=20), + ) + bmc_task = _bmc_summary(bmc_registry, config, n) + status, errors, bmc = await asyncio.gather(status_task, errors_task, bmc_task) + if isinstance(status, dict) and "error" in status: + return {"node": n, "error": status["error"]} + entry: dict[str, Any] = { + "node": n, + "cpu_pct": round(status.get("cpu", 0) * 100, 1), + "mem_used_gb": round(status.get("memory", {}).get("used", 0) / 1073741824, 1), + "mem_total_gb": round(status.get("memory", {}).get("total", 0) / 1073741824, 1), + "uptime_h": round(status.get("uptime", 0) / 3600, 1), + "kernel": status.get("kversion"), + "pve_version": status.get("pveversion"), + } + if bmc is not None: + entry["bmc"] = bmc + entry["recent_errors"] = errors + return entry + + results = list(await asyncio.gather(*[_one(n) for n in target_nodes])) + if node: + return results[0] if results else {"error": f"Node {node!r} not configured."} + return {"nodes": results} + + @mcp.tool() + def vm_find(pattern: str, node: str = "") -> dict[str, Any]: + """Find VMs/CTs by name using glob (``web-*``) or substring (``db``). + + Returns a compact hit list so the caller can follow up with + ``proxmox_vm_status`` or ``vm_bulk_action``. Omit ``node`` to search + across every configured node. + """ + target_nodes = [node] if node else None + vms, _ = _collect_vm_summaries(proxmox_client, target_nodes) + pat = pattern.strip() + is_glob = any(ch in pat for ch in "*?[") + hits: list[dict[str, Any]] = [] + for vm in vms: + name = vm.get("name", "") + if is_glob: + if fnmatch.fnmatchcase(name, pat): + hits.append(vm) + elif pat.lower() in name.lower(): + hits.append(vm) + return {"pattern": pat, "total": len(hits), "vms": hits} + + @mcp.tool() + def vm_bulk_action( + vmids: list[int], + action: str, + force: bool = False, + ) -> dict[str, Any]: + """Run ``start``/``stop``/``restart`` on many VMs/CTs in parallel. + + Locates each VMID across the cluster, fires the action, and collects + per-VM UPIDs (or errors) in one response. ``force`` applies to stop + and restart actions. Capped at 50 VMs per call to prevent runaway + fan-out; split larger lists client-side. + """ + valid_actions = {"start", "stop", "restart"} + if action not in valid_actions: + return {"error": f"Unsupported action {action!r}. Use one of {sorted(valid_actions)}."} + + # Hard cap. A typo like `vm_bulk_action(range(1, 10000), "stop")` should + # fail loud, not take down a cluster. 50 covers legit bulk ops on any + # homelab-scale setup. + _MAX_BULK = 50 + if len(vmids) > _MAX_BULK: + return { + "error": f"Too many VMIDs ({len(vmids)}); cap is {_MAX_BULK} per call. " + "Split into multiple calls.", + } + # Dedupe while preserving order -- repeated VMIDs are almost always a + # caller bug and doing the same stop/start twice is never what they want. + seen: set[int] = set() + unique_vmids = [v for v in vmids if not (v in seen or seen.add(v))] + + def _one(vmid: int) -> dict[str, Any]: + location = _find_vm_location(proxmox_client, vmid) + if not location: + return {"vmid": vmid, "error": "not found"} + n, vm_type = location + endpoint = f"nodes/{n}/{vm_type}/{vmid}/status/{action}" + params: dict[str, Any] = {} + if action in ("stop", "restart") and force: + params["forceStop"] = 1 + resp = proxmox_client.post(n, endpoint, **params) + if isinstance(resp, dict) and "error" in resp: + return {"vmid": vmid, "node": n, "error": resp["error"]} + upid = resp if isinstance(resp, str) else ( + resp.get("upid") if isinstance(resp, dict) else None + ) + return {"vmid": vmid, "node": n, "type": vm_type, "upid": upid} + + results = _parallel_map(_one, unique_vmids) + ok = sum(1 for r in results if "upid" in r) + return {"action": action, "total": len(results), "ok": ok, "results": results} diff --git a/src/beaconmcp/proxmox/apps/bridge.js b/src/beaconmcp/proxmox/apps/bridge.js new file mode 100644 index 0000000..008d6c1 --- /dev/null +++ b/src/beaconmcp/proxmox/apps/bridge.js @@ -0,0 +1,176 @@ +// Shared MCP Apps client bridge. +// +// JSON-RPC 2.0 over postMessage to the host, per the ext-apps spec. Injected +// into each ui:// document by panel.py in place of the +// marker, so the panels stay single self-contained resources. +// +// Deliberately small and dependency-free: an app resource is preloaded by the +// host before the tool even runs, so every kilobyte is paid up front. + +const MCPApp = (() => { + "use strict"; + + const PROTOCOL_VERSION = "2026-01-26"; + + let nextId = 0; + const pending = new Map(); + let hostCapabilities = {}; + let onToolResult = null; + let onHostContext = null; + let ready = false; + + function post(message) { + window.parent.postMessage(message, "*"); + } + + function request(method, params) { + const id = ++nextId; + post({ jsonrpc: "2.0", id, method, params }); + return new Promise((resolve, reject) => pending.set(id, { resolve, reject })); + } + + function notify(method, params) { + post({ jsonrpc: "2.0", method, params }); + } + + // The host owns the iframe height, so it has to be told when the content + // reflows -- otherwise the panel renders into a fixed sliver. + function reportSize() { + notify("ui/notifications/size-changed", { + width: document.documentElement.scrollWidth, + height: document.documentElement.scrollHeight, + }); + } + + // A CallToolResult carries the dict in structuredContent, but a host that + // strips it still sends the JSON as text -- fall back rather than blank out. + function unwrap(result) { + if (result && result.structuredContent) return result.structuredContent; + const first = result && result.content && result.content[0]; + if (first && first.type === "text") { + try { return JSON.parse(first.text); } catch { return null; } + } + return null; + } + + async function callTool(name, args) { + const result = await request("tools/call", { name, arguments: args }); + const data = unwrap(result); + if (result && result.isError) { + throw new Error((data && data.error) || "tool call failed"); + } + if (data && data.error) throw new Error(data.error); + return data; + } + + // Both of these are gated on host capabilities. Calling one the host did not + // advertise gets an error back, so check first and no-op quietly: a panel + // that works everywhere beats one that throws on a stricter host. + function updateModelContext(structuredContent, text) { + if (!hostCapabilities.updateModelContext) return Promise.resolve(false); + return request("ui/update-model-context", { + content: text ? [{ type: "text", text }] : undefined, + structuredContent, + }).then(() => true, () => false); + } + + function sendMessage(text) { + if (!hostCapabilities.message) return Promise.resolve(false); + return request("ui/message", { + role: "user", + content: [{ type: "text", text }], + }).then((r) => !(r && r.isError), () => false); + } + + function requestDisplayMode(mode) { + return request("ui/request-display-mode", { mode }).then( + (r) => (r && r.mode) || null, + () => null, + ); + } + + window.addEventListener("message", (event) => { + const message = event.data; + if (!message || message.jsonrpc !== "2.0") return; + + if (message.id != null && pending.has(message.id)) { + const { resolve, reject } = pending.get(message.id); + pending.delete(message.id); + if (message.error) reject(new Error(message.error.message || "request failed")); + else resolve(message.result); + return; + } + + switch (message.method) { + case "ui/notifications/tool-result": + if (onToolResult) onToolResult(unwrap(message.params)); + break; + case "ui/notifications/host-context-changed": + if (onHostContext) onHostContext(message.params); + break; + case "ui/resource-teardown": + post({ jsonrpc: "2.0", id: message.id, result: {} }); + break; + } + }); + + function applyHostContext(ctx) { + if (!ctx) return; + if (ctx.theme) document.documentElement.dataset.theme = ctx.theme; + const vars = ctx.styles && ctx.styles.variables; + if (vars) { + for (const [key, value] of Object.entries(vars)) { + document.documentElement.style.setProperty(key, value); + } + } + } + + /** + * Perform the ui/initialize handshake. + * + * Params are flat: appInfo / appCapabilities / protocolVersion. Nesting the + * capabilities or sending clientInfo instead of appInfo fails the host's + * schema check, and a rejected handshake is silent -- the host simply never + * answers, so the app never sends `initialized` and the host never delivers + * the tool result. See @modelcontextprotocol/ext-apps App.connect(). + * + * `onFail` is called if the host never completes the handshake, so the panel + * can say so instead of sitting on a spinner forever. + */ + function connect({ name, version = "1.0.0", onResult, onContext, onFail }) { + onToolResult = onResult; + onHostContext = (ctx) => { applyHostContext(ctx); if (onContext) onContext(ctx); }; + + setTimeout(() => { + if (!ready && onFail) { + onFail(new Error("The host did not complete the ui/initialize handshake.")); + } + }, 5000); + + return request("ui/initialize", { + appInfo: { name, version }, + appCapabilities: { availableDisplayModes: ["inline", "fullscreen"] }, + protocolVersion: PROTOCOL_VERSION, + }).then((result) => { + ready = true; + hostCapabilities = (result && result.hostCapabilities) || {}; + applyHostContext(result && result.hostContext); + notify("ui/notifications/initialized", {}); + reportSize(); + return result; + }).catch((err) => { + if (onFail) onFail(err); + throw err; + }); + } + + return { + connect, + callTool, + updateModelContext, + sendMessage, + requestDisplayMode, + reportSize, + hostSupports: (name) => Boolean(hostCapabilities[name]), + }; +})(); diff --git a/src/beaconmcp/proxmox/apps/cluster_panel.html b/src/beaconmcp/proxmox/apps/cluster_panel.html new file mode 100644 index 0000000..c89f853 --- /dev/null +++ b/src/beaconmcp/proxmox/apps/cluster_panel.html @@ -0,0 +1,292 @@ + + + +Cluster dashboard + + + +
Loading…
+ +
+
+
+

Cluster

+
+
+ +
+ +
+ +
Guests
+
+
+ +
+
+ +
+
+
+ + + + +
+ + diff --git a/src/beaconmcp/proxmox/apps/logs_panel.html b/src/beaconmcp/proxmox/apps/logs_panel.html new file mode 100644 index 0000000..907b565 --- /dev/null +++ b/src/beaconmcp/proxmox/apps/logs_panel.html @@ -0,0 +1,194 @@ + + + +Log viewer + + + +
Loading…
+ +
+
+
+

+
+
+ +
+ +
+
+ + +
+
+ + +
+
+ + +
+
+ + +
+ +
+ +
+ +
+ + diff --git a/src/beaconmcp/proxmox/apps/panel.css b/src/beaconmcp/proxmox/apps/panel.css new file mode 100644 index 0000000..5d31a81 --- /dev/null +++ b/src/beaconmcp/proxmox/apps/panel.css @@ -0,0 +1,96 @@ +/* Shared look for the ui:// panels. Injected alongside bridge.js by panel.py. + + Every colour is a host variable with a local fallback. A host that sends + hostContext.styles.variables (the spec's standardized names, which + bridge.js writes onto :root) drives the palette; one that sends nothing + but a theme falls through to the values below. Both halves of the theme + read the same variable so the host's choice wins either way -- the dark + block only changes what happens when it stays silent. */ + +:root { + --bg: var(--color-background-primary, #ffffff); + --fg: var(--color-text-primary, #1a1a1a); + --muted: var(--color-text-tertiary, #6b7280); + --line: var(--color-border-primary, #e5e7eb); + --accent: var(--color-ring-primary, #2563eb); + --danger: var(--color-text-danger, #dc2626); + --ok: var(--color-text-success, #16a34a); + --warn: var(--color-text-warning, #d97706); + --track: var(--color-background-tertiary, #f3f4f6); +} +:root[data-theme="dark"] { + --bg: var(--color-background-primary, #1c1c1c); + --fg: var(--color-text-primary, #ededed); + --muted: var(--color-text-tertiary, #9ca3af); + --line: var(--color-border-primary, #333333); + --accent: var(--color-ring-primary, #60a5fa); + --danger: var(--color-text-danger, #f87171); + --ok: var(--color-text-success, #4ade80); + --warn: var(--color-text-warning, #fbbf24); + --track: var(--color-background-tertiary, #2a2a2a); +} + +* { box-sizing: border-box; } + +body { + margin: 0; + padding: 14px; + background: var(--bg); + color: var(--fg); + font: 13px/1.45 var(--font-sans, system-ui, -apple-system, "Segoe UI", sans-serif); +} + +h1 { font-size: 15px; margin: 0; font-weight: 600; } +.sub { color: var(--muted); font-size: 12px; margin-top: 2px; } +header { display: flex; align-items: flex-start; justify-content: space-between; gap: 12px; } + +.badge { + padding: 2px 9px; border-radius: 999px; font-size: 11px; font-weight: 600; + text-transform: uppercase; letter-spacing: .04em; white-space: nowrap; + border: 1px solid var(--line); color: var(--muted); +} +.badge[data-status="running"], .badge[data-status="online"], .badge[data-status="OK"] { color: var(--ok); border-color: currentColor; } +.badge[data-status="paused"] { color: var(--accent); border-color: currentColor; } +.badge[data-status="unreachable"], .badge[data-status="unknown"] { color: var(--danger); border-color: currentColor; } + +.metrics { display: grid; grid-template-columns: repeat(auto-fit, minmax(140px, 1fr)); gap: 10px; margin: 14px 0; } +.metric { border: 1px solid var(--line); border-radius: 8px; padding: 9px 10px; } +.metric .label { color: var(--muted); font-size: 11px; text-transform: uppercase; letter-spacing: .04em; } +.metric .value { font-size: 15px; font-weight: 600; margin-top: 3px; font-variant-numeric: tabular-nums; } + +.bar { height: 4px; border-radius: 2px; background: var(--track); margin-top: 7px; overflow: hidden; } +.bar > i { display: block; height: 100%; background: var(--accent); } +.bar > i[data-level="warn"] { background: var(--warn); } +.bar > i[data-level="high"] { background: var(--danger); } + +fieldset { border: 0; margin: 0; padding: 0; } +fieldset[disabled] { opacity: .55; } +.row { display: flex; flex-wrap: wrap; gap: 7px; align-items: flex-end; } + +button { + font: inherit; padding: 6px 13px; border-radius: 6px; cursor: pointer; + border: 1px solid var(--line); background: transparent; color: var(--fg); +} +button:hover:not(:disabled) { border-color: var(--accent); } +button:disabled { cursor: not-allowed; opacity: .45; } +button.danger:hover:not(:disabled) { border-color: var(--danger); color: var(--danger); } + +label { display: block; font-size: 11px; color: var(--muted); margin-bottom: 3px; } +input, select { + font: inherit; padding: 5px 7px; border-radius: 6px; + border: 1px solid var(--line); background: var(--bg); color: var(--fg); +} + +hr { border: 0; border-top: 1px solid var(--line); margin: 14px 0; } +.section-title { font-size: 11px; text-transform: uppercase; letter-spacing: .04em; color: var(--muted); margin-bottom: 8px; } + +.note { margin-top: 12px; padding: 8px 10px; border-radius: 6px; font-size: 12px; border: 1px solid var(--line); } +.note[data-kind="error"], #loading[data-kind="error"] { color: var(--danger); border-color: currentColor; } +.note[data-kind="ok"] { color: var(--ok); border-color: currentColor; } + +table { border-collapse: collapse; width: 100%; font-size: 12px; } +th { text-align: left; font-weight: 600; color: var(--muted); font-size: 11px; text-transform: uppercase; letter-spacing: .04em; padding: 5px 8px; } +td { padding: 5px 8px; border-top: 1px solid var(--line); } +tbody tr:hover { background: var(--track); } + +[hidden] { display: none !important; } diff --git a/src/beaconmcp/proxmox/apps/vm_panel.html b/src/beaconmcp/proxmox/apps/vm_panel.html new file mode 100644 index 0000000..29fe698 --- /dev/null +++ b/src/beaconmcp/proxmox/apps/vm_panel.html @@ -0,0 +1,215 @@ + + + +VM control panel + + +
Loading…
+ +
+
+
+

+
+
+ +
+ +
+
+
CPU
+
+
+
+
+
Memory
+
+
+
+
+
Disk
+
+
+
+
+
Uptime
+
+
+
+ +
+
Power
+
+ + + + + +
+ +
+ +
Resources
+
+
+ + +
+
+ + +
+ +
+
+ + +
+ + diff --git a/src/beaconmcp/proxmox/client.py b/src/beaconmcp/proxmox/client.py new file mode 100644 index 0000000..da51c9a --- /dev/null +++ b/src/beaconmcp/proxmox/client.py @@ -0,0 +1,172 @@ +from __future__ import annotations + +import logging +import re +import threading +import time +from typing import Any + +from proxmoxer import ProxmoxAPI +from requests.exceptions import ConnectionError, Timeout + +from ..config import Config + +_logger = logging.getLogger("beaconmcp.proxmox") + +# Every Proxmox endpoint this server builds is an f-string with caller-supplied +# values spliced in as path segments (`.../snapshot/{snapname}/rollback`, +# `.../storage/{storage}/content`, ...). A value containing `/` or `..` would +# silently re-target the request at a different API endpoint than the tool +# advertises. Constrain each segment to the character set Proxmox itself +# allows for node names, storage ids, snapshot names and guest types. +# +# Deliberately excludes `/`, `..`, `.` and a leading `_` -- the latter because +# `api_call` walks the path with getattr() and proxmoxer's __getattr__ refuses +# (or, worse, resolves) dunder/private attribute names. +_SAFE_PATH_SEGMENT = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._@+-]*$") + + +class UnsafePathSegmentError(ValueError): + """Raised when a caller-supplied value would escape its path segment.""" + + +def _split_api_path(path: str) -> list[str]: + """Split an API path into segments, rejecting anything traversal-shaped. + + Raises :class:`UnsafePathSegmentError` with the offending segment so the + caller can surface an actionable message instead of issuing a request + against an endpoint nobody asked for. + """ + parts = [p for p in path.strip("/").split("/")] + for part in parts: + if not _SAFE_PATH_SEGMENT.match(part): + raise UnsafePathSegmentError( + f"illegal Proxmox API path segment {part!r} in {path!r}: " + "segments must match [A-Za-z0-9][A-Za-z0-9._@+-]* " + "(no slashes, no '..')" + ) + return parts + +# Transient-error retry: Proxmox API over the wire frequently hiccups on +# momentary network blips (TCP reset during cluster sync, TLS renegotiation +# behind a reverse proxy, etc). One quick retry with a short backoff covers +# the overwhelming majority without turning sustained outages into slow +# failures. Keep the numbers small and obvious -- callers already get a +# descriptive error dict back if retries don't help. +_RETRY_ATTEMPTS = 2 +_RETRY_BACKOFF_SECONDS = 0.5 + + +class ProxmoxClient: + """Manages connections to one or more Proxmox VE nodes via API tokens.""" + + def __init__(self, config: Config) -> None: + self._config = config + # A ProxmoxAPI wraps a requests.Session, which is not thread-safe. + # Sync tools now run on a worker-thread pool (see server._metric_tool), + # so a shared cache would hand the same Session to several threads at + # once. Each thread gets its own connection per node instead: no lock, + # no shared mutable state, and full parallelism across nodes. Worker + # threads are pooled and long-lived, so the cache still pays off. + self._local = threading.local() + + def _cache(self) -> dict[str, ProxmoxAPI]: + """Return the calling thread's node -> connection cache.""" + cache = getattr(self._local, "connections", None) + if cache is None: + cache = {} + self._local.connections = cache + return cache + + def _get_connection(self, node_name: str) -> ProxmoxAPI: + cache = self._cache() + conn = cache.get(node_name) + if conn is not None: + return conn + + pve_node = self._config.get_node(node_name) + if not pve_node: + raise NodeNotFoundError(node_name, [n.name for n in self._config.pve_nodes]) + + conn = ProxmoxAPI( + pve_node.host, + user=pve_node.token_id.split("!")[0], + token_name=pve_node.token_id.split("!")[1], + token_value=pve_node.token_secret, + verify_ssl=self._config.verify_ssl, + ) + cache[node_name] = conn + return conn + + def api_call(self, node_name: str, method: str, path: str, **kwargs: Any) -> Any: + """Execute an API call against a Proxmox node. + + Returns the result or a dict with 'error' key on failure. Transient + network errors get one quick retry; sustained unreachability returns + the descriptive error message. + """ + try: + parts = _split_api_path(path) + except UnsafePathSegmentError as e: + _logger.warning("rejected Proxmox API call on %s: %s", node_name, e) + return {"error": str(e)} + + last_exc: Exception | None = None + for attempt in range(_RETRY_ATTEMPTS): + try: + conn = self._get_connection(node_name) + obj = conn + for part in parts: + obj = getattr(obj, part) + return getattr(obj, method)(**kwargs) + except NodeNotFoundError: + raise + except (ConnectionError, Timeout) as e: + last_exc = e + # Drop this thread's cached connection so the retry rebuilds + # TLS state rather than re-using a half-broken socket. Other + # threads keep their own healthy sessions. + self._cache().pop(node_name, None) + if attempt + 1 < _RETRY_ATTEMPTS: + _logger.warning( + "transient error on %s %s (attempt %d/%d): %s", + node_name, path, attempt + 1, _RETRY_ATTEMPTS, e, + ) + time.sleep(_RETRY_BACKOFF_SECONDS) + continue + return { + "error": f"Node '{node_name}' is unreachable: {e}. " + "Try ssh_run to access the host directly, " + "or bmc_health_status if the server may be physically down." + } + except Exception as e: + return {"error": f"Proxmox API error on '{node_name}': {e}"} + # Defensive fallback -- loop should always return above. + return {"error": f"Node '{node_name}' is unreachable: {last_exc}"} + + def get(self, node_name: str, path: str, **kwargs: Any) -> Any: + return self.api_call(node_name, "get", path, **kwargs) + + def post(self, node_name: str, path: str, **kwargs: Any) -> Any: + return self.api_call(node_name, "post", path, **kwargs) + + def put(self, node_name: str, path: str, **kwargs: Any) -> Any: + return self.api_call(node_name, "put", path, **kwargs) + + def delete(self, node_name: str, path: str, **kwargs: Any) -> Any: + return self.api_call(node_name, "delete", path, **kwargs) + + @property + def configured_nodes(self) -> list[str]: + return [n.name for n in self._config.pve_nodes] + + +class NodeNotFoundError(Exception): + def __init__(self, node: str, available: list[str]) -> None: + self.node = node + self.available = available + super().__init__( + f"Node '{node}' is not configured. " + f"Available nodes: {', '.join(available)}. " + f"Check your .env file." + ) diff --git a/src/beaconmcp/proxmox/monitoring.py b/src/beaconmcp/proxmox/monitoring.py new file mode 100644 index 0000000..fe0bc82 --- /dev/null +++ b/src/beaconmcp/proxmox/monitoring.py @@ -0,0 +1,260 @@ +from __future__ import annotations + +from typing import Any + +from mcp.server.fastmcp import FastMCP + +from ..utils import filter_fields +from .client import ProxmoxClient + + +def register_monitoring_tools(mcp: FastMCP, client: ProxmoxClient) -> None: + """Register all Proxmox monitoring and diagnostic tools.""" + + @mcp.tool() + def proxmox_list_nodes(fields: list[str] | None = None) -> dict[str, Any]: + """List all Proxmox cluster nodes with their status (online/offline). + + Use this as the first step when diagnosing cluster health or checking which nodes are available. + Pass ``fields=[...]`` to trim each entry to only the keys you need + (e.g. ``["name", "status"]``). + Returns: {"nodes": [{name, status, cpu, mem_used_gb, mem_total_gb, uptime_h}]}. + If a node appears offline, use bmc_health_status to check if it's a hardware issue, + or ssh_run to try reaching it directly. + """ + # Query every configured node. In a joined cluster each member returns the + # same view (deduped by name); with standalone hosts each returns only + # itself, so we need to poll them all. Nodes we can't reach via their API + # are reported as "unreachable" so the caller knows which creds are stale. + results: dict[str, dict[str, Any]] = {} + unreachable: list[dict[str, Any]] = [] + + for node_name in client.configured_nodes: + data = client.get(node_name, "nodes") + if isinstance(data, dict) and "error" in data: + unreachable.append({"name": node_name, "status": "unreachable", "error": data["error"]}) + continue + if isinstance(data, list): + for node in data: + name = node.get("node") + if not name or name in results: + continue + results[name] = { + "name": name, + "status": node.get("status", "unknown"), + "cpu": round(node.get("cpu", 0) * 100, 1), + "mem_used_gb": round(node.get("mem", 0) / 1073741824, 1), + "mem_total_gb": round(node.get("maxmem", 0) / 1073741824, 1), + "uptime_h": round(node.get("uptime", 0) / 3600, 1), + } + continue + unreachable.append({"name": node_name, "status": "unknown", "raw": str(data)}) + + # Surface any configured-but-unreachable node that didn't appear in the + # cluster view (e.g., single-node setup where pve2 is down). + for entry in unreachable: + results.setdefault(entry["name"], entry) + + return {"nodes": filter_fields(list(results.values()), fields)} + + @mcp.tool() + def proxmox_node_status(node: str, fields: list[str] | None = None) -> dict[str, Any]: + """Get detailed status of a specific Proxmox node: CPU, RAM, disk, uptime, kernel, PVE version. + + Use after proxmox_list_nodes to drill into a specific node. + Provide the node name (e.g., 'pve1'). + Returns: {node, cpu_cores, cpu_model, cpu_pct, mem_used_gb, mem_total_gb, + swap_used_gb, swap_total_gb, rootfs_used_gb, rootfs_total_gb, uptime_h, + kernel_version, pve_version}. + """ + data = client.get(node, f"nodes/{node}/status") + if isinstance(data, dict) and "error" in data: + return data + result = { + "node": node, + "cpu_cores": data.get("cpuinfo", {}).get("cores"), + "cpu_model": data.get("cpuinfo", {}).get("model"), + "cpu_pct": round(data.get("cpu", 0) * 100, 1), + "mem_used_gb": round(data.get("memory", {}).get("used", 0) / 1073741824, 1), + "mem_total_gb": round(data.get("memory", {}).get("total", 0) / 1073741824, 1), + "swap_used_gb": round(data.get("swap", {}).get("used", 0) / 1073741824, 1), + "swap_total_gb": round(data.get("swap", {}).get("total", 0) / 1073741824, 1), + "rootfs_used_gb": round(data.get("rootfs", {}).get("used", 0) / 1073741824, 1), + "rootfs_total_gb": round(data.get("rootfs", {}).get("total", 0) / 1073741824, 1), + "uptime_h": round(data.get("uptime", 0) / 3600, 1), + "kernel_version": data.get("kversion"), + "pve_version": data.get("pveversion"), + } + return filter_fields(result, fields) + + @mcp.tool() + def proxmox_list_vms(node: str = "", fields: list[str] | None = None) -> dict[str, Any]: + """List all VMs and containers with their status and resource usage. + + Use to get an overview of what's running on the cluster. + Omit 'node' to list VMs across all configured nodes. + Provide a node name (e.g., 'pve1') to list only that node's VMs. + Returns: {"vms": {"": [{vmid, name, status, type, cpu_pct, + mem_used_mb, mem_max_mb, disk_used_gb, uptime_h}]}, "total": N}. + Per-node errors appear as {"error": "..."} entries in that node's list. + """ + target_nodes = [node] if node else client.configured_nodes + by_node: dict[str, list[dict[str, Any]]] = {} + total = 0 + + for n in target_nodes: + entries: list[dict[str, Any]] = [] + for vm_type in ("qemu", "lxc"): + data = client.get(n, f"nodes/{n}/{vm_type}") + if isinstance(data, dict) and "error" in data: + entries.append({"type": vm_type, "error": data["error"]}) + continue + if not isinstance(data, list): + continue + for vm in data: + entries.append({ + "vmid": vm.get("vmid"), + "name": vm.get("name", ""), + "status": vm.get("status"), + "type": vm_type, + "cpu_pct": round(vm.get("cpu", 0) * 100, 1), + "mem_used_mb": round(vm.get("mem", 0) / 1048576, 0), + "mem_max_mb": round(vm.get("maxmem", 0) / 1048576, 0), + "disk_used_gb": round(vm.get("disk", 0) / 1073741824, 1), + "uptime_h": round(vm.get("uptime", 0) / 3600, 1), + }) + entries.sort(key=lambda v: v.get("vmid", 0)) + by_node[n] = filter_fields(entries, fields) + total += sum(1 for e in entries if "vmid" in e) + + return {"vms": by_node, "total": total} + + @mcp.tool() + def proxmox_vm_status(node: str, vmid: int, fields: list[str] | None = None) -> dict[str, Any]: + """Get detailed status of a specific VM or container: CPU, RAM, disk I/O, network I/O, uptime. + + Use after proxmox_list_vms to drill into a specific VM. + Provide both the node name and VMID. + Auto-detects whether the target is a QEMU VM or LXC container. + Pass ``fields=[...]`` to trim the response to only the keys you need + (e.g. ``["name", "status", "cpu_pct"]``). + Returns: {node, vmid, type, name, status, cpu_pct, cpus, mem_used_mb, + mem_max_mb, disk_read_mb, disk_write_mb, net_in_mb, net_out_mb, uptime_h, + pid, config_summary: {cores, mem_mb, description}}. + """ + # Try qemu first, then lxc + for vm_type in ("qemu", "lxc"): + data = client.get(node, f"nodes/{node}/{vm_type}/{vmid}/status/current") + if isinstance(data, dict) and "error" in data: + if "does not exist" in str(data.get("error", "")).lower(): + continue + # Real error (network, auth) + return data + if isinstance(data, dict) and data.get("status"): + config_data = client.get(node, f"nodes/{node}/{vm_type}/{vmid}/config") + result: dict[str, Any] = { + "node": node, + "vmid": vmid, + "type": vm_type, + "name": data.get("name", ""), + "status": data.get("status"), + "cpu_pct": round(data.get("cpu", 0) * 100, 1), + "cpus": data.get("cpus"), + "mem_used_mb": round(data.get("mem", 0) / 1048576, 0), + "mem_max_mb": round(data.get("maxmem", 0) / 1048576, 0), + "disk_read_mb": round(data.get("diskread", 0) / 1048576, 1), + "disk_write_mb": round(data.get("diskwrite", 0) / 1048576, 1), + "net_in_mb": round(data.get("netin", 0) / 1048576, 1), + "net_out_mb": round(data.get("netout", 0) / 1048576, 1), + "uptime_h": round(data.get("uptime", 0) / 3600, 1), + "pid": data.get("pid"), + } + if isinstance(config_data, dict) and "error" not in config_data: + result["config_summary"] = { + "cores": config_data.get("cores"), + "mem_mb": config_data.get("memory"), + "description": config_data.get("description", ""), + } + return filter_fields(result, fields) + + return {"error": f"VM/CT {vmid} not found on node '{node}'. Check the VMID and node name."} + + @mcp.tool() + def proxmox_get_logs(node: str, source: str = "syslog", limit: int = 50) -> dict[str, Any]: + """Retrieve system logs from a Proxmox node. + + Use to diagnose system-level issues, crashes, or service failures. + Set 'source' to 'syslog' for system logs or 'tasks' for Proxmox task logs. + Adjust 'limit' to control how many log lines to return (default 50, max 500). + """ + limit = min(limit, 500) + + if source == "tasks": + data = client.get(node, f"nodes/{node}/tasks", limit=limit) + if isinstance(data, dict) and "error" in data: + return data + if isinstance(data, list): + return { + "node": node, + "source": "tasks", + "entries": [ + { + "upid": t.get("upid"), + "type": t.get("type"), + "status": t.get("status"), + "user": t.get("user"), + "starttime": t.get("starttime"), + "endtime": t.get("endtime"), + } + for t in data + ], + } + return {"node": node, "source": "tasks", "entries": [], "raw": str(data)} + + # syslog + data = client.get(node, f"nodes/{node}/syslog", limit=limit) + if isinstance(data, dict) and "error" in data: + return data + if isinstance(data, list): + return { + "node": node, + "source": "syslog", + "lines": [entry.get("t", "") for entry in data], + } + return {"node": node, "source": "syslog", "lines": [], "raw": str(data)} + + @mcp.tool() + def proxmox_get_tasks(node: str = "", limit: int = 20, fields: list[str] | None = None) -> dict[str, Any]: + """List recent Proxmox tasks across the cluster: migrations, backups, VM operations. + + Use to check what operations have been running or to investigate failed tasks. + Omit 'node' to list tasks from all configured nodes. + Pass ``fields=[...]`` to trim each entry to only the keys you need + (e.g. ``["upid", "status"]``). + Returns: {"tasks": {"": [{upid, type, status, user, starttime, + endtime}]}, "total": N}. Per-node errors appear as {"error": "..."} + entries in that node's list. + """ + target_nodes = [node] if node else client.configured_nodes + by_node: dict[str, list[dict[str, Any]]] = {} + total = 0 + + for n in target_nodes: + entries: list[dict[str, Any]] = [] + data = client.get(n, f"nodes/{n}/tasks", limit=limit) + if isinstance(data, dict) and "error" in data: + entries.append({"error": data["error"]}) + elif isinstance(data, list): + for t in data: + entries.append({ + "upid": t.get("upid"), + "type": t.get("type"), + "status": t.get("status"), + "user": t.get("user"), + "starttime": t.get("starttime"), + "endtime": t.get("endtime"), + }) + by_node[n] = filter_fields(entries, fields) + total += sum(1 for e in entries if "upid" in e) + + return {"tasks": by_node, "total": total} diff --git a/src/beaconmcp/proxmox/panel.py b/src/beaconmcp/proxmox/panel.py new file mode 100644 index 0000000..0dd1670 --- /dev/null +++ b/src/beaconmcp/proxmox/panel.py @@ -0,0 +1,222 @@ +"""MCP Apps panels. + +The MCP Apps extension (`io.modelcontextprotocol/ui`) lets a tool carry a +reference to an interactive UI: `_meta.ui.resourceUri` points at a `ui://` +resource served as `text/html;profile=mcp-app`, which the host renders in a +sandboxed iframe and talks to over JSON-RPC on `postMessage`. + +The wire format is all this module needs, so it runs on mcp 1.x. The `Apps` +extension class that wraps it lives in mcp 2.0 and requires `MCPServer`; the +two knobs it sets -- `meta=` on the tool and `mime_type=` on the resource -- +are already on `FastMCP`. + +Hosts that did not negotiate Apps ignore `_meta.ui` and just show the tool's +return value, which is why every panel tool returns its full snapshot as data +rather than a "see the panel" placeholder. + +Each panel is one HTML file under ``apps/``. They share ``bridge.js`` (the +JSON-RPC client) and ``panel.css`` (the look), spliced in at the +```` marker so what ships to the host stays a single +self-contained document. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +from mcp.server.fastmcp import FastMCP + +from .aggregators import ( + _collect_node_summaries, + _collect_storage_summaries, + _collect_vm_summaries, +) +from .client import ProxmoxClient + +APP_MIME_TYPE = "text/html;profile=mcp-app" + +VM_PANEL_URI = "ui://beaconmcp/vm-panel.html" +LOGS_PANEL_URI = "ui://beaconmcp/logs-panel.html" +CLUSTER_PANEL_URI = "ui://beaconmcp/cluster-panel.html" + +_APPS_DIR = Path(__file__).parent / "apps" +_RUNTIME_MARKER = "" + +_MB = 1048576 +_GB = 1073741824 + + +def _read_app(name: str) -> str: + """Load a panel document with the shared CSS and bridge spliced in.""" + html = (_APPS_DIR / name).read_text(encoding="utf-8") + runtime = ( + f"" + f"" + ) + return html.replace(_RUNTIME_MARKER, runtime) + + +def _vm_snapshot(client: ProxmoxClient, node: str, vmid: int) -> dict[str, Any]: + """Everything the VM panel renders, in one pass over both guest types.""" + for vm_type in ("qemu", "lxc"): + data = client.get(node, f"nodes/{node}/{vm_type}/{vmid}/status/current") + if isinstance(data, dict) and "error" in data: + if "does not exist" in str(data["error"]).lower(): + continue + return data + if not isinstance(data, dict) or not data.get("status"): + continue + + # QEMU reports disk=0 unless the guest agent is answering, so a zero + # here means "unknown", not "empty". LXC reports it for real. + disk_used = data.get("disk") or 0 + conf = client.get(node, f"nodes/{node}/{vm_type}/{vmid}/config") + if not isinstance(conf, dict) or "error" in conf: + conf = {} + + return { + "node": node, + "vmid": vmid, + "type": vm_type, + "name": data.get("name", ""), + "status": data.get("status"), + "cpu_pct": round(data.get("cpu", 0) * 100, 1), + "cpus": data.get("cpus"), + "mem_used_mb": round(data.get("mem", 0) / _MB), + "mem_max_mb": round(data.get("maxmem", 0) / _MB), + "disk_used_gb": round(disk_used / _GB, 1) if disk_used else None, + "disk_max_gb": round(data.get("maxdisk", 0) / _GB, 1), + "uptime_h": round(data.get("uptime", 0) / 3600, 1), + "cores": conf.get("cores"), + "memory_mb": conf.get("memory"), + } + + return {"error": f"VM/CT {vmid} not found on node '{node}'. Check the VMID and node name."} + + +def register_panel_tools(mcp: FastMCP, client: ProxmoxClient) -> None: + @mcp.tool( + meta={"ui": {"resourceUri": VM_PANEL_URI, "visibility": ["model", "app"]}}, + ) + def proxmox_vm_panel(node: str, vmid: int) -> dict[str, Any]: + """Open an interactive control panel for one VM or container. + + Renders live CPU / RAM / disk state with buttons for start, stop and + restart, and fields to change the CPU core count and memory. Use this + instead of proxmox_vm_status when the user wants to *act* on a guest + rather than just read its numbers, or when they ask to "manage", + "control" or "open" a VM. + + The panel drives the ordinary tools (proxmox_vm_start / _stop / + _restart / _config), so every action it takes goes through the same + approval the client applies to any other tool call. + + Returns: {node, vmid, type, name, status, cpu_pct, cpus, mem_used_mb, + mem_max_mb, disk_used_gb, disk_max_gb, uptime_h, cores, memory_mb}. + ``disk_used_gb`` is null when the guest does not report it. + """ + return _vm_snapshot(client, node, vmid) + + @mcp.tool( + meta={"ui": {"resourceUri": LOGS_PANEL_URI, "visibility": ["model", "app"]}}, + ) + def proxmox_logs_panel(node: str, source: str = "syslog", limit: int = 200) -> dict[str, Any]: + """Open a scrollable, filterable log and task viewer for one node. + + Prefer this over proxmox_get_logs whenever the user wants to *read* + logs rather than have them summarised: the panel keeps every line, + highlights errors and warnings, filters as you type, and can switch + between the syslog and the Proxmox task list without another turn. + + Args: + source: 'syslog' for system logs, 'tasks' for the task history. + limit: Lines to fetch, capped at 500 by the Proxmox API. + + Returns: {node, source, entries: [...]}. For syslog each entry is + {text, level} where level is error/warn/info, guessed from the line. + For tasks each entry is {upid, type, status, user, starttime, endtime, + level}. + """ + limit = min(limit, 500) + + if source == "tasks": + data = client.get(node, f"nodes/{node}/tasks", limit=limit) + if isinstance(data, dict) and "error" in data: + return data + entries = [ + { + "upid": t.get("upid"), + "type": t.get("type"), + "status": t.get("status"), + "user": t.get("user"), + "starttime": t.get("starttime"), + "endtime": t.get("endtime"), + # Proxmox writes "OK" for success and a message otherwise; + # a still-running task has no status yet. + "level": "info" if t.get("status") in ("OK", None) else "error", + } + for t in (data if isinstance(data, list) else []) + ] + return {"node": node, "source": "tasks", "entries": entries} + + data = client.get(node, f"nodes/{node}/syslog", limit=limit) + if isinstance(data, dict) and "error" in data: + return data + entries = [ + {"text": line, "level": _syslog_level(line)} + for line in (entry.get("t", "") for entry in (data if isinstance(data, list) else [])) + ] + return {"node": node, "source": "syslog", "entries": entries} + + @mcp.tool( + meta={"ui": {"resourceUri": CLUSTER_PANEL_URI, "visibility": ["model", "app"]}}, + ) + def cluster_overview_interactive(include_storage: bool = True) -> dict[str, Any]: + """Open an interactive cluster dashboard: nodes, guests and storage. + + The same data as cluster_overview, but rendered as a browsable panel: + nodes with CPU and memory pressure, a searchable guest table with + per-row start/stop, and storage pools with usage bars. Use it when the + user wants to look around the cluster rather than ask one question + about it. + + Returns: {nodes: [...], vms: [...], total_vms, storage: [...]}. + """ + nodes = _collect_node_summaries(client) + vms, total_vms = _collect_vm_summaries(client) + out: dict[str, Any] = {"nodes": nodes, "vms": vms, "total_vms": total_vms} + if include_storage: + out["storage"] = _collect_storage_summaries(client) + return out + + _register_app_resource(mcp, VM_PANEL_URI, "vm-panel", "VM control panel", "vm_panel.html") + _register_app_resource(mcp, LOGS_PANEL_URI, "logs-panel", "Log viewer", "logs_panel.html") + _register_app_resource( + mcp, CLUSTER_PANEL_URI, "cluster-panel", "Cluster dashboard", "cluster_panel.html", + ) + + +# Cheap keyword scan. Proxmox hands us raw journald text with no severity +# field, so the alternative to guessing is showing every line flat -- which is +# the thing this panel exists to fix. Over-flagging a line is harmless; the +# filter box is there for when it gets noisy. +_ERROR_WORDS = ("error", "fail", "fatal", "critical", "panic", "segfault", "refused", "timeout") +_WARN_WORDS = ("warn", "deprecat", "retry", "degraded", "unable") + + +def _syslog_level(line: str) -> str: + lowered = line.lower() + if any(word in lowered for word in _ERROR_WORDS): + return "error" + if any(word in lowered for word in _WARN_WORDS): + return "warn" + return "info" + + +def _register_app_resource( + mcp: FastMCP, uri: str, name: str, title: str, filename: str, +) -> None: + @mcp.resource(uri, name=name, title=title, mime_type=APP_MIME_TYPE) + def _app() -> str: + return _read_app(filename) diff --git a/src/beaconmcp/proxmox/system.py b/src/beaconmcp/proxmox/system.py new file mode 100644 index 0000000..b28d4ca --- /dev/null +++ b/src/beaconmcp/proxmox/system.py @@ -0,0 +1,863 @@ +from __future__ import annotations + +import base64 +import asyncio +import hashlib +import time +import uuid +import shlex +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +from mcp.server.fastmcp import FastMCP + +from .client import ProxmoxClient + + +@dataclass +class ExecSession: + exec_id: str + node: str + vmid: int + vm_type: str + command: str + status: str = "running" + stdout: str = "" + stderr: str = "" + exit_code: int | None = None + pid: int | None = None + started_at: float = field(default_factory=time.time) + + +_exec_sessions: dict[str, ExecSession] = {} +_EXEC_SESSION_TTL = 3600 + + +def _prune_exec_sessions() -> None: + now = time.time() + stale = [ + eid + for eid, s in _exec_sessions.items() + if s.status != "running" and now - s.started_at > _EXEC_SESSION_TTL + ] + for eid in stale: + del _exec_sessions[eid] + + +def _detect_vm_type(client: ProxmoxClient, node: str, vmid: int) -> str | None: + for vm_type in ("qemu", "lxc"): + data = client.get(node, f"nodes/{node}/{vm_type}/{vmid}/status/current") + if isinstance(data, dict) and "error" in data: + continue + if isinstance(data, dict) and data.get("status"): + return vm_type + return None + + +def _staging_dir(client: ProxmoxClient) -> Path: + """Return the resolved staging directory, creating it on first use.""" + base = Path(client._config.server.transfers_dir).expanduser().resolve() + base.mkdir(parents=True, exist_ok=True) + return base + + +def _staging_path(client: ProxmoxClient, name: str) -> Path: + """Resolve ``name`` against the staging dir, refusing path traversal. + + Only plain basenames are accepted: no slashes, no ``..``, no absolute + paths. The resolved target must remain inside the staging directory. + """ + if not isinstance(name, str) or not name or name in (".", ".."): + raise ValueError("staging filename must be a non-empty basename") + if "/" in name or "\\" in name or name.startswith(".."): + raise ValueError( + "staging filename must be a plain basename (no slashes, no '..')" + ) + base = _staging_dir(client) + target = (base / name).resolve() + try: + target.relative_to(base) + except ValueError as exc: + raise ValueError( + f"staging filename {name!r} resolves outside the transfers directory" + ) from exc + return target + + +def _sha256_file(path: Path) -> str: + h = hashlib.sha256() + with open(path, "rb") as f: + for chunk in iter(lambda: f.read(1024 * 1024), b""): + h.update(chunk) + return h.hexdigest() + + +# Hard cap for the guest-agent file-read path (matches proxmox_write_file). +_FILE_READ_MAX_BYTES = 1024 * 1024 + + +def _qemu_file_size(client: ProxmoxClient, node: str, vmid: int, path: str) -> int | None: + """Best-effort size (bytes) of ``path`` inside a QEMU VM via the guest agent. + + Returns ``None`` when the size can't be determined (no ``stat`` in the + guest, agent error, slow exec). Callers use this to reject oversized + reads *before* pulling the whole file into memory; ``None`` means + "couldn't check" so the caller falls back to its post-read guard rather + than blocking a legitimate read. + """ + start = client.post( + node, f"nodes/{node}/qemu/{vmid}/agent/exec", + command=["stat", "-c", "%s", path], + ) + if not isinstance(start, dict) or start.get("pid") is None: + return None + pid = start["pid"] + deadline = time.time() + 10 + while time.time() < deadline: + status = client.get( + node, f"nodes/{node}/qemu/{vmid}/agent/exec-status", pid=pid, + ) + if not isinstance(status, dict): + return None + if status.get("exited"): + if status.get("exitcode") != 0: + return None + out = status.get("out-data", "") + if status.get("out-data-encoding") == "base64" and out: + out = base64.b64decode(out).decode("utf-8", errors="replace") + try: + return int(str(out).strip()) + except (ValueError, AttributeError): + return None + time.sleep(0.5) + return None + + +def register_system_tools(mcp: FastMCP, client: ProxmoxClient, ssh_client: Any = None) -> None: + """Register Proxmox system administration and command execution tools.""" + + + @mcp.tool() + async def proxmox_read_file(node: str, vmid: int, path: str, binary: bool = False) -> dict[str, Any]: + """Read a file from a VM or container. + + For VMs, this uses the QEMU Guest Agent safely (file must be < 1MB). + For containers, this requires SSH to be configured. + """ + vm_type = _detect_vm_type(client, node, vmid) + if not vm_type: + return {"status": "error", "error": f"VM/CT {vmid} not found on node '{node}'."} + + if vm_type == "qemu": + # Pre-flight: refuse oversized files before the guest agent pulls + # the whole thing into memory. Best-effort -- if the guest has no + # `stat` we fall through to the post-read guard below. + size = await asyncio.get_running_loop().run_in_executor( + None, _qemu_file_size, client, node, vmid, path, + ) + if size is not None and size > _FILE_READ_MAX_BYTES: + return { + "status": "error", + "error": ( + f"File is {size} bytes, over the 1MB guest-agent read " + "limit. Use proxmox_download_file for large files." + ), + } + result = client.get(node, f"nodes/{node}/qemu/{vmid}/agent/file-read", file=path) + if isinstance(result, dict) and "error" in result: + return {"status": "error", "error": result["error"]} + # QEMU Guest Agent returns file content base64-encoded + try: + if isinstance(result, dict) and "content" in result: + import base64 + raw_b64 = result["content"] + if len(raw_b64) > 1398101: # ~1MB limit in base64 + return {"status": "error", "error": "File exceeds 1MB limit. Use SSH to download large files."} + + if binary: + return {"status": "success", "vmid": vmid, "node": node, "path": path, "content_base64": raw_b64} + + try: + content = base64.b64decode(raw_b64).decode("utf-8") + return {"status": "success", "vmid": vmid, "node": node, "path": path, "content": content} + except UnicodeDecodeError: + return {"status": "error", "error": "Binary data detected. Pass binary=True to retrieve as base64."} + return {"status": "success", "vmid": vmid, "node": node, "path": path, "content": str(result)} + except Exception as e: + return {"status": "error", "error": f"Failed to decode file content: {e}"} + + return {"status": "error", "error": "LXC file reading is currently unsupported via API. Please use ssh_run to cat the file."} + + @mcp.tool() + async def proxmox_write_file(node: str, vmid: int, path: str, content: str) -> dict[str, Any]: + """Write a file to a VM or container. + + For VMs, this uses the QEMU Guest Agent safely to avoid shell escaping issues. + For containers, this requires SSH to be configured. + """ + vm_type = _detect_vm_type(client, node, vmid) + if not vm_type: + return {"status": "error", "error": f"VM/CT {vmid} not found on node '{node}'."} + + if len(content.encode("utf-8")) > 1024 * 1024: + return {"status": "error", "error": "Content exceeds 1MB limit. Use SSH for large files."} + + if vm_type == "qemu": + import base64 + encoded = base64.b64encode(content.encode("utf-8")).decode("utf-8") + result = client.post(node, f"nodes/{node}/qemu/{vmid}/agent/file-write", file=path, content=encoded, encode=1) + if isinstance(result, dict) and "error" in result: + return {"status": "error", "error": result["error"]} + return {"status": "success", "vmid": vmid, "node": node, "path": path, "action": "file_write"} + + return {"status": "error", "error": "LXC file writing is currently unsupported via API. Please use ssh_exec_command to write the file."} + + @mcp.tool() + async def proxmox_upload_file( + node: str, + vmid: int, + source: str, + dest: str, + verify_checksum: bool = True, + ) -> dict[str, Any]: + """Upload a large file from the BeaconMCP staging dir into a VM/CT. + + ``source`` must be a plain basename of a file present in + ``server.transfers_dir`` (default ``~/.cache/beaconmcp/transfers``). + ``dest`` is an absolute path inside the guest. Capped by + ``server.transfers_max_mb`` (default 500 MB). + + - **LXC**: SFTP-streams to ``/tmp/`` on the Proxmox node, then + ``pct push `` into the container, and cleans up. + - **VM**: SFTP-streams directly into the VM. The VM must be + SSH-reachable — declared under ``ssh.hosts[]`` or matched by + ``ssh.vmid_to_ip``. For VMs without SSH, fall back to + ``proxmox_write_file`` (1 MB cap). + + When ``verify_checksum=True``, computes a SHA-256 locally and + re-checks it inside the guest after transfer. If the guest lacks + ``sha256sum`` (e.g. minimal Alpine CTs), the success response + carries ``checksum_verified=False`` and a ``warning`` field + rather than silently treating the transfer as verified. + """ + if not ssh_client: + return { + "status": "error", + "error": "Large-file transfer requires SSH. Add an ssh.hosts[] entry for the Proxmox node (or the VM).", + } + try: + local_path = _staging_path(client, source) + except ValueError as e: + return {"status": "error", "error": str(e)} + if not local_path.is_file(): + return { + "status": "error", + "error": ( + f"Staging file {source!r} not found. Place it in " + f"{_staging_dir(client)} first (SCP, dashboard, etc.)." + ), + } + size_bytes = local_path.stat().st_size + max_bytes = client._config.server.transfers_max_mb * 1024 * 1024 + if size_bytes > max_bytes: + return { + "status": "error", + "error": ( + f"File size {size_bytes} bytes exceeds the configured " + f"transfers_max_mb ({client._config.server.transfers_max_mb} MB)." + ), + } + if not isinstance(dest, str) or not dest.startswith("/"): + return {"status": "error", "error": "`dest` must be an absolute path inside the guest."} + + loop = asyncio.get_running_loop() + vm_type = await loop.run_in_executor(None, _detect_vm_type, client, node, vmid) + if not vm_type: + return {"status": "error", "error": f"VM/CT {vmid} not found on node '{node}'."} + + started = time.monotonic() + local_sha = await loop.run_in_executor(None, _sha256_file, local_path) if verify_checksum else None + + if vm_type == "lxc": + node_tmp = f"/tmp/beaconmcp-upload-{uuid.uuid4().hex[:12]}" + try: + await ssh_client.sftp_put(node, str(local_path), node_tmp) + except Exception as e: + return {"status": "error", "error": f"SFTP to node {node!r} failed: {e}"} + push_res = await ssh_client.exec_command( + node, + f"pct push {vmid} {shlex.quote(node_tmp)} {shlex.quote(dest)}", + timeout=600, + ) + await ssh_client.sftp_remove(node, node_tmp) + if push_res.get("exit_code") != 0: + return { + "status": "error", + "error": f"pct push failed: {push_res.get('stderr') or push_res.get('error') or 'unknown'}", + } + remote_sha = None + if verify_checksum: + sum_res = await ssh_client.exec_command( + node, + f"pct exec {vmid} -- sha256sum {shlex.quote(dest)}", + timeout=300, + ) + if sum_res.get("exit_code") == 0 and sum_res.get("stdout"): + remote_sha = sum_res["stdout"].split()[0] + if remote_sha and remote_sha != local_sha: + return { + "status": "error", + "error": f"Checksum mismatch after upload: local={local_sha} remote={remote_sha}", + "bytes": size_bytes, + } + result: dict[str, Any] = { + "status": "success", "vmid": vmid, "node": node, "dest": dest, + "bytes": size_bytes, "sha256": local_sha, + "duration_s": round(time.monotonic() - started, 2), + "transport": "sftp+pct_push", + "checksum_verified": remote_sha is not None, + } + if verify_checksum and remote_sha is None: + result["warning"] = ( + "sha256sum unavailable in guest; transfer was not checksum-verified." + ) + return result + + # QEMU: try direct SSH into the VM. + try: + ssh_client.resolve(str(vmid)) + except Exception as e: + return { + "status": "error", + "error": ( + f"VM {vmid} is not reachable via SSH ({e}). " + "Either declare it under ssh.hosts[] (or set ssh.vmid_to_ip), " + "or use proxmox_write_file for files ≤1MB." + ), + } + try: + await ssh_client.sftp_put(str(vmid), str(local_path), dest) + except Exception as e: + return {"status": "error", "error": f"SFTP to VM {vmid} failed: {e}"} + remote_sha = None + if verify_checksum: + sum_res = await ssh_client.exec_command( + str(vmid), f"sha256sum {shlex.quote(dest)}", timeout=300, + ) + if sum_res.get("exit_code") == 0 and sum_res.get("stdout"): + remote_sha = sum_res["stdout"].split()[0] + if remote_sha and remote_sha != local_sha: + return { + "status": "error", + "error": f"Checksum mismatch after upload: local={local_sha} remote={remote_sha}", + "bytes": size_bytes, + } + result: dict[str, Any] = { + "status": "success", "vmid": vmid, "node": node, "dest": dest, + "bytes": size_bytes, "sha256": local_sha, + "duration_s": round(time.monotonic() - started, 2), + "transport": "sftp", + "checksum_verified": remote_sha is not None, + } + if verify_checksum and remote_sha is None: + result["warning"] = ( + "sha256sum unavailable in guest; transfer was not checksum-verified." + ) + return result + + @mcp.tool() + async def proxmox_download_file( + node: str, + vmid: int, + source: str, + dest: str, + verify_checksum: bool = True, + overwrite: bool = False, + ) -> dict[str, Any]: + """Download a large file from a VM/CT into the BeaconMCP staging dir. + + ``source`` is an absolute path inside the guest. ``dest`` is a plain + basename written under ``server.transfers_dir`` (default + ``~/.cache/beaconmcp/transfers``). Capped by + ``server.transfers_max_mb`` (default 500 MB) — oversized sources + are rejected before any data is moved. + + Set ``overwrite=True`` to replace an existing file in the staging + directory. Without it, the tool refuses to clobber existing files. + Data is streamed to a sibling ``.part`` file and atomically + renamed on success, so an interrupted transfer never leaves a + half-written file at ``dest``. + + When ``verify_checksum=True`` but the guest lacks ``sha256sum``, + the success response carries ``checksum_verified=False`` and a + ``warning`` field rather than silently treating the transfer as + verified. + + - **LXC**: ``pct pull`` to ``/tmp/`` on the Proxmox node, then + SFTP-streams it back to the staging dir, and cleans up. + - **VM**: SFTP-streams directly from the VM. Requires the VM to be + SSH-reachable (see ``proxmox_upload_file`` for details). + """ + if not ssh_client: + return { + "status": "error", + "error": "Large-file transfer requires SSH. Add an ssh.hosts[] entry for the Proxmox node (or the VM).", + } + try: + local_path = _staging_path(client, dest) + except ValueError as e: + return {"status": "error", "error": str(e)} + if local_path.exists() and not overwrite: + return { + "status": "error", + "error": ( + f"Staging file {dest!r} already exists. " + "Pass overwrite=True to replace it." + ), + } + if not isinstance(source, str) or not source.startswith("/"): + return {"status": "error", "error": "`source` must be an absolute path inside the guest."} + + loop = asyncio.get_running_loop() + vm_type = await loop.run_in_executor(None, _detect_vm_type, client, node, vmid) + if not vm_type: + return {"status": "error", "error": f"VM/CT {vmid} not found on node '{node}'."} + max_bytes = client._config.server.transfers_max_mb * 1024 * 1024 + started = time.monotonic() + + if vm_type == "lxc": + # Pre-flight: refuse oversized files before any data movement. + size_res = await ssh_client.exec_command( + node, + f"pct exec {vmid} -- stat -c %s {shlex.quote(source)}", + timeout=60, + ) + if size_res.get("exit_code") != 0: + return { + "status": "error", + "error": f"Could not stat {source!r} in CT {vmid}: {size_res.get('stderr') or size_res.get('error')}", + } + try: + src_size = int((size_res.get("stdout") or "0").strip()) + except ValueError: + src_size = 0 + if src_size > max_bytes: + return { + "status": "error", + "error": ( + f"Source size {src_size} bytes exceeds the configured " + f"transfers_max_mb ({client._config.server.transfers_max_mb} MB)." + ), + } + remote_sha = None + if verify_checksum: + sum_res = await ssh_client.exec_command( + node, + f"pct exec {vmid} -- sha256sum {shlex.quote(source)}", + timeout=300, + ) + if sum_res.get("exit_code") == 0 and sum_res.get("stdout"): + remote_sha = sum_res["stdout"].split()[0] + node_tmp = f"/tmp/beaconmcp-download-{uuid.uuid4().hex[:12]}" + # Stream into a sibling .part file and rename on success so an + # interrupted transfer never leaves a half-written file at dest. + tmp_path = local_path.with_name(local_path.name + ".part") + pull_res = await ssh_client.exec_command( + node, + f"pct pull {vmid} {shlex.quote(source)} {shlex.quote(node_tmp)}", + timeout=600, + ) + if pull_res.get("exit_code") != 0: + await ssh_client.sftp_remove(node, node_tmp) + return { + "status": "error", + "error": f"pct pull failed: {pull_res.get('stderr') or pull_res.get('error') or 'unknown'}", + } + try: + await ssh_client.sftp_get(node, node_tmp, str(tmp_path)) + except Exception as e: + await ssh_client.sftp_remove(node, node_tmp) + tmp_path.unlink(missing_ok=True) + return {"status": "error", "error": f"SFTP from node {node!r} failed: {e}"} + await ssh_client.sftp_remove(node, node_tmp) + local_sha = await loop.run_in_executor(None, _sha256_file, tmp_path) if verify_checksum else None + if verify_checksum and remote_sha and local_sha != remote_sha: + size = tmp_path.stat().st_size + tmp_path.unlink(missing_ok=True) + return { + "status": "error", + "error": f"Checksum mismatch after download: remote={remote_sha} local={local_sha}", + "bytes": size, + } + tmp_path.replace(local_path) + result: dict[str, Any] = { + "status": "success", "vmid": vmid, "node": node, + "source": source, "staged_path": str(local_path), + "bytes": local_path.stat().st_size, "sha256": local_sha or remote_sha, + "duration_s": round(time.monotonic() - started, 2), + "transport": "pct_pull+sftp", + "checksum_verified": remote_sha is not None, + } + if verify_checksum and remote_sha is None: + result["warning"] = ( + "sha256sum unavailable in guest; transfer was not checksum-verified." + ) + return result + + # QEMU: direct SSH. + try: + ssh_client.resolve(str(vmid)) + except Exception as e: + return { + "status": "error", + "error": ( + f"VM {vmid} is not reachable via SSH ({e}). " + "Either declare it under ssh.hosts[] (or set ssh.vmid_to_ip), " + "or use proxmox_read_file for files ≤1MB." + ), + } + size_res = await ssh_client.exec_command( + str(vmid), f"stat -c %s {shlex.quote(source)}", timeout=60, + ) + if size_res.get("exit_code") != 0: + return { + "status": "error", + "error": f"Could not stat {source!r} in VM {vmid}: {size_res.get('stderr') or size_res.get('error')}", + } + try: + src_size = int((size_res.get("stdout") or "0").strip()) + except ValueError: + src_size = 0 + if src_size > max_bytes: + return { + "status": "error", + "error": ( + f"Source size {src_size} bytes exceeds the configured " + f"transfers_max_mb ({client._config.server.transfers_max_mb} MB)." + ), + } + remote_sha = None + if verify_checksum: + sum_res = await ssh_client.exec_command( + str(vmid), f"sha256sum {shlex.quote(source)}", timeout=300, + ) + if sum_res.get("exit_code") == 0 and sum_res.get("stdout"): + remote_sha = sum_res["stdout"].split()[0] + # Stream into a sibling .part file and rename on success so an + # interrupted transfer never leaves a half-written file at dest. + tmp_path = local_path.with_name(local_path.name + ".part") + try: + await ssh_client.sftp_get(str(vmid), source, str(tmp_path)) + except Exception as e: + tmp_path.unlink(missing_ok=True) + return {"status": "error", "error": f"SFTP from VM {vmid} failed: {e}"} + local_sha = await loop.run_in_executor(None, _sha256_file, tmp_path) if verify_checksum else None + if verify_checksum and remote_sha and local_sha != remote_sha: + size = tmp_path.stat().st_size + tmp_path.unlink(missing_ok=True) + return { + "status": "error", + "error": f"Checksum mismatch after download: remote={remote_sha} local={local_sha}", + "bytes": size, + } + tmp_path.replace(local_path) + result: dict[str, Any] = { + "status": "success", "vmid": vmid, "node": node, + "source": source, "staged_path": str(local_path), + "bytes": local_path.stat().st_size, "sha256": local_sha or remote_sha, + "duration_s": round(time.monotonic() - started, 2), + "transport": "sftp", + "checksum_verified": remote_sha is not None, + } + if verify_checksum and remote_sha is None: + result["warning"] = ( + "sha256sum unavailable in guest; transfer was not checksum-verified." + ) + return result + + @mcp.tool() + def proxmox_list_transfers() -> dict[str, Any]: + """List files currently in the BeaconMCP staging directory. + + Returns each file's basename, size in bytes, and last-modified epoch. + Use this to discover which ``source`` names are available for + ``proxmox_upload_file``, or to confirm a ``proxmox_download_file`` + landed. + """ + base = _staging_dir(client) + entries: list[dict[str, Any]] = [] + for entry in sorted(base.iterdir()): + if not entry.is_file(): + continue + stat = entry.stat() + entries.append({ + "name": entry.name, + "bytes": stat.st_size, + "modified": int(stat.st_mtime), + }) + return { + "transfers_dir": str(base), + "max_mb": client._config.server.transfers_max_mb, + "files": entries, + "total": len(entries), + } + + @mcp.tool() + def proxmox_delete_transfer(name: str) -> dict[str, Any]: + """Delete a file from the BeaconMCP staging directory. + + ``name`` must be a plain basename (no slashes). Use + ``proxmox_list_transfers`` to see available files. + """ + try: + target = _staging_path(client, name) + except ValueError as e: + return {"status": "error", "error": str(e)} + if not target.is_file(): + return {"status": "error", "error": f"File {name!r} not found in staging directory."} + size = target.stat().st_size + target.unlink() + return {"status": "success", "deleted": name, "freed_bytes": size} + + @mcp.tool() + def proxmox_storage_status(node: str = "") -> dict[str, Any]: + """Get storage status across the cluster: usage, type, content types. + + Use to check disk space, storage health, or find available storage. + Omit 'node' to list storage from all configured nodes. + """ + target_nodes = [node] if node else client.configured_nodes + by_node: dict[str, list[dict[str, Any]]] = {} + + for n in target_nodes: + entries: list[dict[str, Any]] = [] + data = client.get(n, f"nodes/{n}/storage") + if isinstance(data, dict) and "error" in data: + entries.append({"error": data["error"]}) + by_node[n] = entries + continue + if not isinstance(data, list): + by_node[n] = entries + continue + for s in data: + storage_name = s.get("storage") + if not storage_name: + continue + status = client.get(n, f"nodes/{n}/storage/{storage_name}/status") + used = status.get("used", 0) if isinstance(status, dict) and "error" not in status else 0 + total = status.get("total", 0) if isinstance(status, dict) and "error" not in status else 0 + + entries.append({ + "name": storage_name, + "type": s.get("type"), + "content": s.get("content"), + "enabled": s.get("enabled", 1) == 1, + "used_gb": round(used / 1073741824, 1), + "total_gb": round(total / 1073741824, 1), + "usage_pct": round(used / total * 100, 1) if total > 0 else 0, + }) + by_node[n] = entries + + return {"storage": by_node} + + @mcp.tool() + def proxmox_network_config(node: str) -> dict[str, Any]: + """Get network interface configuration of a Proxmox node.""" + data = client.get(node, f"nodes/{node}/network") + if isinstance(data, dict) and "error" in data: + return data + if not isinstance(data, list): + return {"node": node, "interfaces": [], "raw": str(data)} + + interfaces = [] + for iface in data: + interfaces.append({ + "name": iface.get("iface"), + "type": iface.get("type"), + "address": iface.get("address"), + "netmask": iface.get("netmask"), + "gateway": iface.get("gateway"), + "bridge_ports": iface.get("bridge_ports"), + "active": iface.get("active", False), + "method": iface.get("method"), + "cidr": iface.get("cidr"), + }) + return {"node": node, "interfaces": interfaces} + + def _start_async_qemu(node: str, vmid: int, command: str) -> dict[str, Any]: + _prune_exec_sessions() + exec_id = str(uuid.uuid4())[:8] + session = ExecSession( + exec_id=exec_id, + node=node, + vmid=vmid, + vm_type="qemu", + command=command, + ) + _exec_sessions[exec_id] = session + + parts = shlex.split(command) + result = client.post(node, f"nodes/{node}/qemu/{vmid}/agent/exec", command=parts) + if isinstance(result, dict) and "error" in result: + session.status = "failed" + session.stderr = str(result["error"]) + return {"exec_id": exec_id, "status": "failed", "error": result["error"]} + session.pid = result.get("pid") if isinstance(result, dict) else None + return {"exec_id": exec_id, "status": "running"} + + def _poll_session(exec_id: str) -> dict[str, Any]: + session = _exec_sessions.get(exec_id) + if not session: + return {"status": "error", "error": f"No command found with exec_id '{exec_id}'."} + + if session.status != "running": + return { + "exec_id": exec_id, + "status": "ok" if session.status == "completed" and session.exit_code == 0 else session.status, + "stdout": session.stdout, + "stderr": session.stderr, + "exit_code": session.exit_code, + "command": session.command, + "elapsed_s": round(time.time() - session.started_at, 1), + } + + if session.vm_type == "qemu" and session.pid is not None: + status_data = client.get( + session.node, + f"nodes/{session.node}/qemu/{session.vmid}/agent/exec-status", + pid=session.pid, + ) + if isinstance(status_data, dict) and status_data.get("exited"): + stdout = status_data.get("out-data", "") + stderr = status_data.get("err-data", "") + if status_data.get("out-data-encoding") == "base64" and stdout: + stdout = base64.b64decode(stdout).decode("utf-8", errors="replace") + if status_data.get("err-data-encoding") == "base64" and stderr: + stderr = base64.b64decode(stderr).decode("utf-8", errors="replace") + session.status = "completed" + session.stdout = stdout + session.stderr = stderr + session.exit_code = status_data.get("exitcode", -1) + + if time.time() - session.started_at > 600: + session.status = "timeout" + + elapsed = round(time.time() - session.started_at, 1) + if session.status == "running": + return {"status": "running", "exec_id": exec_id, "command": session.command, "elapsed_s": elapsed} + + return { + "status": "ok" if session.status == "completed" and session.exit_code == 0 else session.status, + "exec_id": exec_id, + "command": session.command, + "stdout": session.stdout, + "stderr": session.stderr, + "exit_code": session.exit_code, + "duration_s": elapsed, + } + + @mcp.tool() + async def proxmox_run( + node: str = "", + vmid: int = 0, + command: str = "", + timeout: int = 60, + wait: bool = True, + exec_id: str = "", + ) -> dict[str, Any]: + """Run a command inside a VM (QEMU Guest Agent) or container (LXC pct exec). + + Three call patterns: + - **Sync** (default): pass ``node``, ``vmid``, ``command``. Blocks up to ``timeout`` seconds (max 600). + - **Async start**: pass ``node``, ``vmid``, ``command``, ``wait=False``. Returns immediately. + - **Poll existing**: pass ``exec_id`` only. Returns the current status/output for that session. + """ + loop = asyncio.get_running_loop() + + def _format_ssh_result(xid: str, s) -> dict[str, Any]: + elapsed = round(time.time() - s.started_at, 1) + if s.status == "running": + return {"status": "running", "exec_id": xid, "command": s.command, "elapsed_s": elapsed} + return { + "status": "ok" if s.status == "completed" and s.exit_code == 0 else s.status, + "exec_id": xid, + "command": s.command, + "stdout": s.stdout, + "stderr": s.stderr, + "exit_code": s.exit_code, + "duration_s": elapsed, + } + + if exec_id: + if exec_id in _exec_sessions: + return await loop.run_in_executor(None, _poll_session, exec_id) + if ssh_client: + ssh_sess = ssh_client.get_session(exec_id) + if ssh_sess: + return _format_ssh_result(exec_id, ssh_sess) + return {"status": "error", "error": f"No command found with exec_id {exec_id!r}."} + + if not command: + return {"status": "error", "error": "`command` is required when `exec_id` is not provided."} + if not node or not vmid: + return {"status": "error", "error": "`node` and `vmid` are required to start a command."} + + vm_type = await loop.run_in_executor(None, _detect_vm_type, client, node, vmid) + if not vm_type: + return {"status": "error", "error": f"VM/CT {vmid} not found on node '{node}'."} + + max_timeout = min(max(timeout, 1), 600) + + if vm_type == "lxc": + if not ssh_client: + return {"status": "error", "error": "LXC execution requires SSH access to the Proxmox node, but SSH is not configured."} + try: + # Safe escaping to prevent shell injection (fixing the owner's feedback) + escaped_cmd = shlex.quote(command) + lxc_cmd = f"pct exec {vmid} -- sh -c {escaped_cmd}" + new_id = await ssh_client.exec_command_async(node, lxc_cmd) + except Exception as e: + return {"status": "error", "error": str(e)} + + if not wait: + return {"status": "running", "exec_id": new_id, "elapsed_s": 0} + + deadline = time.time() + max_timeout + while time.time() < deadline: + session = ssh_client.get_session(new_id) + if session and session.status != "running": + return _format_ssh_result(new_id, session) + await asyncio.sleep(1) + + session = ssh_client.get_session(new_id) + if session and session.status != "running": + return _format_ssh_result(new_id, session) + return { + "status": "running", + "exec_id": new_id, + "elapsed_s": int(time.time() - (session.started_at if session else time.time())), + "hint": "Command still running. Call proxmox_run(exec_id=...) to poll.", + } + else: + started = await loop.run_in_executor(None, _start_async_qemu, node, vmid, command) + if started.get("status") == "failed": + return started + new_id = started["exec_id"] + + if not wait: + return {"status": "running", "exec_id": new_id, "elapsed_s": 0} + + deadline = time.time() + max_timeout + while time.time() < deadline: + result = await loop.run_in_executor(None, _poll_session, new_id) + if result["status"] != "running": + return result + await asyncio.sleep(1) + + return { + "status": "running", + "exec_id": new_id, + "elapsed_s": int(time.time() - _exec_sessions[new_id].started_at), + "hint": "Command still running. Call proxmox_run(exec_id=...) to poll.", + } diff --git a/src/beaconmcp/proxmox/vms.py b/src/beaconmcp/proxmox/vms.py new file mode 100644 index 0000000..6d3c7bc --- /dev/null +++ b/src/beaconmcp/proxmox/vms.py @@ -0,0 +1,412 @@ +from __future__ import annotations + +import re +from typing import Any + +from mcp.server.fastmcp import FastMCP + +from ..utils import filter_fields +from .client import ProxmoxClient + + +def _detect_vm_type(client: ProxmoxClient, node: str, vmid: int) -> str | None: + """Detect whether a VMID is a QEMU VM or LXC container.""" + for vm_type in ("qemu", "lxc"): + data = client.get(node, f"nodes/{node}/{vm_type}/{vmid}/status/current") + if isinstance(data, dict) and "error" in data: + continue + if isinstance(data, dict) and data.get("status"): + return vm_type + return None + + +def _detect_backup_type(client: ProxmoxClient, node: str, archive: str) -> str | None: + """Return 'qemu' or 'lxc' for a backup volid, or None if undeterminable. + + Tries cheap string heuristics first (vzdump filenames + Proxmox Backup + Server's ``vm/`` / ``ct/`` namespaces), then falls back to a + storage-content lookup that reads the archive's ``subtype`` -- robust + against renamed or non-standard archive names. + """ + a = archive.lower() + if "vzdump-qemu-" in a or "backup/vm-" in a: + return "qemu" + if "vzdump-lxc-" in a or "backup/ct-" in a: + return "lxc" + # PBS namespaces: the guest-type segment is always followed by the + # numeric VMID (``backup/vm/100/...``, ``backup/ns//ct/200/...``). + # Anchoring on the id avoids misreading a *namespace* named "vm"/"ct". + m = re.search(r"(?:^|/)(vm|ct)/\d+(?:/|$)", a) + if m: + return "qemu" if m.group(1) == "vm" else "lxc" + + # Fallback: look the volid up in its source storage's content listing and + # read the type Proxmox itself reports. + src_storage = archive.split(":", 1)[0] if ":" in archive else "" + if not src_storage: + return None + content = client.get( + node, f"nodes/{node}/storage/{src_storage}/content", content="backup", + ) + if not isinstance(content, list): + return None + for item in content: + if not isinstance(item, dict) or item.get("volid") != archive: + continue + subtype = (item.get("subtype") or item.get("vmtype") or "").lower() + if subtype in ("qemu", "lxc"): + return subtype + fmt = (item.get("format") or "").lower() + if "vma" in fmt or "qemu" in fmt: + return "qemu" + if "tar" in fmt or "lxc" in fmt: + return "lxc" + return None + + +def register_vm_tools(mcp: FastMCP, client: ProxmoxClient) -> None: + """Register all Proxmox VM/CT lifecycle management tools.""" + + @mcp.tool() + def proxmox_vm_start(node: str, vmid: int) -> dict[str, Any]: + """Start a stopped VM or container. + + Use when a VM/CT needs to be powered on. + Provide the node name and VMID. Auto-detects VM vs container. + Returns the task UPID on success for tracking the operation. + """ + vm_type = _detect_vm_type(client, node, vmid) + if not vm_type: + return {"error": f"VM/CT {vmid} not found on node '{node}'. Check VMID and node name."} + + result = client.post(node, f"nodes/{node}/{vm_type}/{vmid}/status/start") + if isinstance(result, dict) and "error" in result: + return result + return {"vmid": vmid, "node": node, "action": "start", "upid": result} + + @mcp.tool() + def proxmox_vm_stop(node: str, vmid: int, force: bool = False) -> dict[str, Any]: + """Stop a running VM or container. + + Use to shut down a VM/CT. Set force=true for an immediate hard stop + (equivalent to pulling the power cord -- use only when a clean shutdown fails). + Default is a clean ACPI shutdown for VMs or clean stop for containers. + """ + vm_type = _detect_vm_type(client, node, vmid) + if not vm_type: + return {"error": f"VM/CT {vmid} not found on node '{node}'. Check VMID and node name."} + + endpoint = "stop" if force else "shutdown" + result = client.post(node, f"nodes/{node}/{vm_type}/{vmid}/status/{endpoint}") + if isinstance(result, dict) and "error" in result: + return result + return {"vmid": vmid, "node": node, "action": endpoint, "force": force, "upid": result} + + @mcp.tool() + def proxmox_vm_restart(node: str, vmid: int) -> dict[str, Any]: + """Restart a running VM or container (clean reboot). + + Use when a VM/CT needs to be rebooted. Sends an ACPI reboot signal for VMs + or a clean restart for containers. If the VM is unresponsive, stop it with force=true first, + then start it again. + """ + vm_type = _detect_vm_type(client, node, vmid) + if not vm_type: + return {"error": f"VM/CT {vmid} not found on node '{node}'. Check VMID and node name."} + + # Both QEMU and LXC use /status/reboot (PVE 7+). LXC does not expose /status/restart. + result = client.post(node, f"nodes/{node}/{vm_type}/{vmid}/status/reboot") + if isinstance(result, dict) and "error" in result: + return result + return {"vmid": vmid, "node": node, "action": "restart", "upid": result} + + @mcp.tool() + def proxmox_vm_create(node: str, vmid: int, vm_type: str = "qemu", config: dict[str, Any] | None = None) -> dict[str, Any]: + """Create a new VM or container on a Proxmox node. + + Use to provision new virtual machines or containers. + Set vm_type to 'qemu' for a VM or 'lxc' for a container. + Pass configuration as a dict (e.g., {"cores": 2, "memory": 4096, "net0": "virtio,bridge=vmbr0"}). + Refer to Proxmox API docs for available config options per VM type. + """ + if vm_type not in ("qemu", "lxc"): + return {"error": f"Invalid vm_type '{vm_type}'. Use 'qemu' for VMs or 'lxc' for containers."} + + create_params = config or {} + result = client.post(node, f"nodes/{node}/{vm_type}", vmid=vmid, **create_params) + if isinstance(result, dict) and "error" in result: + return result + return {"vmid": vmid, "node": node, "type": vm_type, "action": "create", "upid": result} + + @mcp.tool() + def proxmox_vm_clone(node: str, vmid: int, newid: int, name: str = "") -> dict[str, Any]: + """Clone an existing VM or container to create a copy. + + Use to duplicate a VM/CT. Provide the source VMID, the new VMID for the clone, + and optionally a name. The clone inherits the source configuration. + """ + vm_type = _detect_vm_type(client, node, vmid) + if not vm_type: + return {"error": f"VM/CT {vmid} not found on node '{node}'. Check VMID and node name."} + + kwargs: dict[str, Any] = {"newid": newid} + if name: + # PVE uses `name` for QEMU VMs and `hostname` for LXC containers. + kwargs["name" if vm_type == "qemu" else "hostname"] = name + + result = client.post(node, f"nodes/{node}/{vm_type}/{vmid}/clone", **kwargs) + if isinstance(result, dict) and "error" in result: + return result + return { + "src_vmid": vmid, + "new_vmid": newid, + "name": name, + "node": node, + "action": "clone", + "upid": result, + } + + @mcp.tool() + def proxmox_vm_migrate(node: str, vmid: int, target_node: str) -> dict[str, Any]: + """Migrate a VM or container to another Proxmox node. + + Use to move a VM/CT from one node to another (e.g., for maintenance or load balancing). + The VM can be running (live migration) or stopped. + Provide the current node, VMID, and the target node name. + """ + vm_type = _detect_vm_type(client, node, vmid) + if not vm_type: + return {"error": f"VM/CT {vmid} not found on node '{node}'. Check VMID and node name."} + + result = client.post(node, f"nodes/{node}/{vm_type}/{vmid}/migrate", target=target_node) + if isinstance(result, dict) and "error" in result: + return result + return { + "vmid": vmid, + "src_node": node, + "dst_node": target_node, + "action": "migrate", + "upid": result, + } + + @mcp.tool() + def proxmox_vm_config( + node: str, + vmid: int, + updates: dict[str, Any] | None = None, + fields: list[str] | None = None, + ) -> dict[str, Any]: + """Read or modify the configuration of a VM or container. + + Without 'updates': returns the full current configuration. + With 'updates': applies the provided config changes (e.g., {"memory": 4096, "cores": 4}). + Use to inspect or change VM settings like memory, CPU cores, network, disks, etc. + Pass ``fields=[...]`` (read-only mode) to trim the returned ``config`` + blob -- helpful because full VM configs can be large (dozens of + disk/net/hostpci keys). Ignored when ``updates`` is given. + """ + vm_type = _detect_vm_type(client, node, vmid) + if not vm_type: + return {"error": f"VM/CT {vmid} not found on node '{node}'. Check VMID and node name."} + + if updates is None: + data = client.get(node, f"nodes/{node}/{vm_type}/{vmid}/config") + if isinstance(data, dict) and "error" in data: + return data + return { + "vmid": vmid, + "node": node, + "type": vm_type, + "config": filter_fields(data, fields), + } + + result = client.put(node, f"nodes/{node}/{vm_type}/{vmid}/config", **updates) + if isinstance(result, dict) and "error" in result: + return result + return {"vmid": vmid, "node": node, "action": "config_update", "applied": updates} + + @mcp.tool() + def proxmox_snapshot_list(node: str, vmid: int) -> dict[str, Any]: + """List all snapshots for a VM or container. + + Returns the snapshot hierarchy including names, descriptions, and creation times. + """ + vm_type = _detect_vm_type(client, node, vmid) + if not vm_type: + return {"error": f"VM/CT {vmid} not found on node '{node}'. Check VMID and node name."} + + result = client.get(node, f"nodes/{node}/{vm_type}/{vmid}/snapshot") + if isinstance(result, dict) and "error" in result: + return result + + return {"vmid": vmid, "node": node, "snapshots": result} + + @mcp.tool() + def proxmox_snapshot_create( + node: str, + vmid: int, + snapname: str, + description: str = "", + vmstate: bool = False, + dry_run: bool = False, + ) -> dict[str, Any]: + """Create a new snapshot of a VM or container. + + Use this before risky operations (OS upgrades, risky commands) to establish a safe checkpoint. + Set vmstate=true to capture RAM (slower, only for running QEMU VMs). + """ + if dry_run: + return {"status": "dry_run", "message": f"Would create snapshot {snapname!r} for VM/CT {vmid} on {node}."} + + vm_type = _detect_vm_type(client, node, vmid) + if not vm_type: + return {"error": f"VM/CT {vmid} not found on node '{node}'. Check VMID and node name."} + + params = {"snapname": snapname} + if description: + params["description"] = description + if vmstate: + if vm_type == "lxc": + return {"error": "vmstate=True is only supported for QEMU VMs, not LXC containers."} + + status_data = client.get(node, f"nodes/{node}/{vm_type}/{vmid}/status/current") + is_running = isinstance(status_data, dict) and status_data.get("status") == "running" + if not is_running: + return {"error": "vmstate=True requires the VM to be running to capture RAM."} + + params["vmstate"] = 1 + + result = client.post(node, f"nodes/{node}/{vm_type}/{vmid}/snapshot", **params) + if isinstance(result, dict) and "error" in result: + return result + return {"vmid": vmid, "node": node, "action": "snapshot_create", "snapname": snapname, "upid": result} + + @mcp.tool() + def proxmox_snapshot_rollback(node: str, vmid: int, snapname: str, dry_run: bool = False) -> dict[str, Any]: + """Roll back a VM or container to a previous snapshot. + + Restores the guest to the exact state of the named snapshot. + Use proxmox_snapshot_list to find the correct snapname. + """ + if dry_run: + return {"status": "dry_run", "message": f"Would roll back VM/CT {vmid} on {node} to snapshot {snapname!r}."} + + vm_type = _detect_vm_type(client, node, vmid) + if not vm_type: + return {"error": f"VM/CT {vmid} not found on node '{node}'. Check VMID and node name."} + + result = client.post(node, f"nodes/{node}/{vm_type}/{vmid}/snapshot/{snapname}/rollback") + if isinstance(result, dict) and "error" in result: + return result + return {"vmid": vmid, "node": node, "action": "snapshot_rollback", "snapname": snapname, "upid": result} + + @mcp.tool() + def proxmox_snapshot_delete(node: str, vmid: int, snapname: str, dry_run: bool = False) -> dict[str, Any]: + """Delete a VM or container snapshot. + + Removes the snapshot from the storage backing the VM. + """ + if dry_run: + return {"status": "dry_run", "message": f"Would delete snapshot {snapname!r} from VM/CT {vmid} on {node}."} + + vm_type = _detect_vm_type(client, node, vmid) + if not vm_type: + return {"error": f"VM/CT {vmid} not found on node '{node}'. Check VMID and node name."} + + result = client.delete(node, f"nodes/{node}/{vm_type}/{vmid}/snapshot/{snapname}") + if isinstance(result, dict) and "error" in result: + return result + return {"vmid": vmid, "node": node, "action": "snapshot_delete", "snapname": snapname, "upid": result} + + @mcp.tool() + def proxmox_backup_create(node: str, vmid: int, storage: str, mode: str = "snapshot", compress: str = "zstd", notes: str = "") -> dict[str, Any]: + """Create a new backup of a VM or container. + + Args: + storage: The storage pool where the backup will be saved (e.g. 'local', 'pbs', 'nfs'). + mode: 'stop', 'suspend', or 'snapshot' (default is 'snapshot' for zero downtime). + compress: Compression algorithm ('zstd', 'lzo', 'gzip', or '0' for none). + notes: Optional description for the backup. + """ + vm_type = _detect_vm_type(client, node, vmid) + if not vm_type: + return {"error": f"VM/CT {vmid} not found on node '{node}'. Check VMID and node name."} + + params = { + "vmid": vmid, + "storage": storage, + "mode": mode, + "compress": compress, + } + if notes: + params["notes"] = notes + + result = client.post(node, f"nodes/{node}/vzdump", **params) + if isinstance(result, dict) and "error" in result: + return {"status": "error", "error": result["error"]} + return {"status": "success", "vmid": vmid, "node": node, "action": "backup_create", "storage": storage, "upid": result} + + @mcp.tool() + def proxmox_backup_list(node: str, storage: str, vmid: int | None = None) -> dict[str, Any]: + """List available backups (vzdump archives) on a specific storage pool. + + Args: + storage: The storage pool to query (e.g. 'local', 'pbs', 'nfs'). + vmid: Filter backups by a specific VM/CT. + """ + params = {"content": "backup"} + if vmid is not None: + params["vmid"] = vmid + + data = client.get(node, f"nodes/{node}/storage/{storage}/content", **params) + if isinstance(data, dict) and "error" in data: + return data + + if not isinstance(data, list): + return {"backups": []} + + backups = [] + for backup in data: + backups.append({ + "volid": backup.get("volid"), + "vmid": backup.get("vmid"), + "format": backup.get("format"), + "size_gb": round(backup.get("size", 0) / 1073741824, 2), + "ctime": backup.get("ctime"), + "notes": backup.get("notes", ""), + }) + + return {"node": node, "storage": storage, "backups": backups} + + @mcp.tool() + def proxmox_backup_restore(node: str, vmid: int, archive: str, force: bool = False, storage: str = "local-lvm") -> dict[str, Any]: + """Restore a VM or container from a backup archive. + + Args: + archive: The full volume ID of the backup archive (e.g. 'local:backup/vzdump-qemu-100-2023_10_25-00_00_00.vma.zst'). + Use proxmox_backup_list to find the correct 'volid'. + force: If True, overwrites an existing VM/CT if it already exists. + storage: Target storage for the restored disks (default 'local-lvm'). + """ + guest_type = _detect_backup_type(client, node, archive) + if guest_type is None: + return { + "status": "error", + "error": ( + "Could not determine if the backup is for 'qemu' or 'lxc'. " + "Pass a volid from proxmox_backup_list (its 'volid' field)." + ), + } + endpoint = f"nodes/{node}/{guest_type}" + + params = { + "vmid": vmid, + "archive": archive, + "force": 1 if force else 0, + "storage": storage, + } + + result = client.post(node, endpoint, **params) + if isinstance(result, dict) and "error" in result: + return {"status": "error", "error": result["error"]} + return {"status": "success", "vmid": vmid, "node": node, "action": "backup_restore", "archive": archive, "upid": result} diff --git a/src/beaconmcp/ratelimit.py b/src/beaconmcp/ratelimit.py new file mode 100644 index 0000000..9e48640 --- /dev/null +++ b/src/beaconmcp/ratelimit.py @@ -0,0 +1,223 @@ +"""Tiny in-memory sliding-window rate limiter. + +Covers the auth-adjacent endpoints (``/oauth/token``, ``/app/login``) so a +compromised or malicious client can't brute-force ``client_secret`` / TOTP +at line speed. The existing per-client TOTP lockout only triggers after a +valid-client-bad-TOTP pattern; this limiter fires earlier, on the *IP*, +regardless of which client_id is being tried. + +The bucket lives in-process: if you run multiple BeaconMCP instances +behind a load balancer each instance gets its own count. That's fine for +the single-host homelab target; deploy a real limiter (nginx, Traefik) in +front if you need global state. +""" + +from __future__ import annotations + +import ipaddress +import threading +import time +from collections import deque +from dataclasses import dataclass, field + + +@dataclass +class _Bucket: + events: deque[float] = field(default_factory=deque) + + +class RateLimiter: + """Sliding-window limiter: N events per ``window_seconds`` per key. + + ``check(key)`` returns True if the event is allowed (and records it), + False if it should be rejected. Keys are opaque strings -- we use the + client IP for auth endpoints. + """ + + def __init__(self, *, limit: int, window_seconds: float) -> None: + self._limit = limit + self._window = window_seconds + self._buckets: dict[str, _Bucket] = {} + self._lock = threading.Lock() + self._last_gc = 0.0 + + @staticmethod + def _prune_bucket(bucket: _Bucket, cutoff: float) -> None: + while bucket.events and bucket.events[0] <= cutoff: + bucket.events.popleft() + + def _collect_stale_buckets_locked(self, cutoff: float) -> None: + for key, bucket in list(self._buckets.items()): + self._prune_bucket(bucket, cutoff) + if not bucket.events: + del self._buckets[key] + + def check(self, key: str) -> bool: + now = time.monotonic() + cutoff = now - self._window + with self._lock: + bucket = self._buckets.get(key) + if bucket is None: + bucket = _Bucket() + self._buckets[key] = bucket + # Drop expired events. + self._prune_bucket(bucket, cutoff) + if len(bucket.events) >= self._limit: + return False + bucket.events.append(now) + # Opportunistic GC: once the map is large, reclaim stale keys whose + # events are all outside the window. + if len(self._buckets) > 1024 or (now - self._last_gc) >= self._window: + self._collect_stale_buckets_locked(cutoff) + self._last_gc = now + return True + + def retry_after(self, key: str) -> int: + """Seconds until ``key`` can make another request (0 if allowed now). + + Used to populate the ``Retry-After`` response header. + """ + with self._lock: + bucket = self._buckets.get(key) + if bucket is None or not bucket.events: + return 0 + cutoff = time.monotonic() - self._window + self._prune_bucket(bucket, cutoff) + if not bucket.events: + del self._buckets[key] + return 0 + oldest = bucket.events[0] + return max(0, int(self._window - (time.monotonic() - oldest)) + 1) + + +def _coerce_ip(value: object) -> str | None: + raw = str(value).strip() + if not raw: + return None + try: + return str(ipaddress.ip_address(raw)) + except ValueError: + pass + if raw.startswith("[") and "]" in raw: + try: + return str(ipaddress.ip_address(raw[1 : raw.index("]")])) + except ValueError: + pass + if raw.count(":") == 1: + host, _, port = raw.rpartition(":") + if host and port.isdigit(): + try: + return str(ipaddress.ip_address(host)) + except ValueError: + pass + return None + + +def _is_trusted_proxy(ip_value: str, trusted_proxies: tuple[str, ...]) -> bool: + try: + ip_obj = ipaddress.ip_address(ip_value) + except ValueError: + return False + + for raw_rule in trusted_proxies: + rule = raw_rule.strip() + if not rule: + continue + if "/" in rule: + try: + if ip_obj in ipaddress.ip_network(rule, strict=False): + return True + except ValueError: + continue + continue + rule_ip = _coerce_ip(rule) + if rule_ip is None: + continue + if ip_obj == ipaddress.ip_address(rule_ip): + return True + return False + + +def client_ip(request: object, trusted_proxies: tuple[str, ...] = ()) -> str: + """Best-effort client IP for a Starlette ``Request``. + + Honors ``X-Forwarded-For`` only when the direct peer is trusted. In that + case we walk the chain from right to left and return the first untrusted + hop, which avoids left-most spoofing when proxies append to the header. + """ + client = getattr(request, "client", None) + direct_peer = getattr(client, "host", None) if client is not None else None + direct_peer_raw = str(direct_peer) if direct_peer is not None else "" + direct_ip = _coerce_ip(direct_peer_raw) + + headers = getattr(request, "headers", None) + if headers is not None and direct_ip and _is_trusted_proxy(direct_ip, trusted_proxies): + fwd = headers.get("x-forwarded-for") if hasattr(headers, "get") else None + if fwd: + chain: list[str] = [ + ip for ip in (_coerce_ip(part) for part in fwd.split(",")) if ip is not None + ] + chain.append(direct_ip) + for hop in reversed(chain): + if not _is_trusted_proxy(hop, trusted_proxies): + return hop + + if direct_ip: + return direct_ip + if direct_peer_raw: + return direct_peer_raw + return "unknown" + + +def forwarded_host( + request: object, + trusted_proxies: tuple[str, ...] = (), + *, + default: str = "localhost", +) -> str: + """Client-facing Host for a Starlette ``Request``. + + Mirrors :func:`client_ip`'s trust model for the *host* dimension: the + ``X-Forwarded-Host`` header is attacker-controlled on a direct request, so + it is honored only when the direct peer is a declared trusted proxy. + Otherwise the request's own ``Host`` header is used (then ``default``). + + Like :func:`client_ip`, this keys on ``request.client.host`` being the real + TCP peer, so the server must run with uvicorn ``proxy_headers=False`` (see + ``__main__``): the default ``ProxyHeadersMiddleware`` would rewrite the peer + to the ``X-Forwarded-For`` client and the trusted-proxy branch below could + never open. + + Kept deliberately narrow -- it does NOT touch the scheme. ``X-Forwarded- + Proto`` is still read directly by the callers, because a TLS-terminating + edge (Cloudflare tunnel, nginx) legitimately needs it to report https even + when ``trusted_proxies`` is unset; gating it would silently downgrade the + Secure-cookie flag and the OAuth issuer to http. + """ + headers = getattr(request, "headers", None) + has_get = headers is not None and hasattr(headers, "get") + + def _hdr(name: str) -> str | None: + return headers.get(name) if has_get else None + + host_header = _hdr("host") or default + if not trusted_proxies: + return host_header + + client = getattr(request, "client", None) + direct_peer = getattr(client, "host", None) if client is not None else None + direct_ip = _coerce_ip(str(direct_peer)) if direct_peer is not None else None + if direct_ip and _is_trusted_proxy(direct_ip, trusted_proxies): + fwd = _hdr("x-forwarded-host") + if fwd: + # Take the last entry, not the first. A proxy that appends rather + # than overwrites puts its own value last, so the last entry is the + # one the nearest trusted proxy wrote; returning the first would + # hand back a client-supplied prefix. Symmetric with client_ip's + # right-to-left walk. Proxies that overwrite (the common case: + # nginx ``proxy_set_header X-Forwarded-Host $host``) leave a single + # entry, so first and last coincide. + parts = [p.strip() for p in fwd.split(",") if p.strip()] + if parts: + return parts[-1] + return host_header diff --git a/src/beaconmcp/security/__init__.py b/src/beaconmcp/security/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/beaconmcp/security/tools.py b/src/beaconmcp/security/tools.py new file mode 100644 index 0000000..1b610d9 --- /dev/null +++ b/src/beaconmcp/security/tools.py @@ -0,0 +1,41 @@ +"""Security-related MCP tools (session termination, etc.).""" + +from __future__ import annotations + +from mcp.server.fastmcp import FastMCP + +from ..auth import revoke_current_token + + +def register_security_tools(mcp: FastMCP) -> None: + @mcp.tool() + def security_end_session() -> dict: + """Invalidate the current bearer token. + + Call this as the very last step of a task when the caller has + finished using the server. After this returns, every subsequent + request made with the same bearer token will fail with 401, forcing + a fresh OAuth + 2FA round-trip. Use this to shrink the window during + which a stolen token could be replayed. + + Do not call this in the middle of a multi-step task; the next tool + call would be rejected. + + Returns ``{"revoked": true}`` on success, or + ``{"revoked": false, "reason": "..."}`` if no active token was found + (e.g. called outside an HTTP request context). + """ + if revoke_current_token(): + return { + "revoked": True, + "grace_s": 8, + "message": ( + "Bearer token scheduled for revocation. It remains valid " + "for ~8s so this response can reach the client; after " + "that the next request will need a fresh OAuth + 2FA." + ), + } + return { + "revoked": False, + "reason": "no active bearer token in request context", + } diff --git a/src/beaconmcp/server.py b/src/beaconmcp/server.py new file mode 100644 index 0000000..967a0ad --- /dev/null +++ b/src/beaconmcp/server.py @@ -0,0 +1,300 @@ +import base64 +import inspect +import os +from pathlib import Path + +import anyio + +from mcp.server.fastmcp import FastMCP +from mcp.server.transport_security import TransportSecuritySettings +from mcp.types import Icon + +from .bmc import build_registry as build_bmc_registry +from .bmc import register_bmc_tools +from .config import Config +from .maintenance import register_maintenance_tools +from .proxmox.aggregators import register_aggregator_tools +from .proxmox.client import ProxmoxClient +from .proxmox.monitoring import register_monitoring_tools +from .proxmox.panel import register_panel_tools +from .proxmox.system import register_system_tools +from .proxmox.vms import register_vm_tools +from .security.tools import register_security_tools +from .ssh.client import SSHClient +from .ssh.tools import register_ssh_tools + +from functools import partial, wraps +import time +from . import audit +from .auth import current_client_id +from .metrics import tool_calls, tool_latency_ms + +config = Config.load() +proxmox_client = ProxmoxClient(config) +ssh_client = SSHClient(config) +bmc_registry = build_bmc_registry(config) + + +def _csv_env(name: str, default: list[str]) -> list[str]: + raw = os.environ.get(name, "").strip() + if not raw: + return default + return [v.strip() for v in raw.split(",") if v.strip()] + + +# DNS-rebinding protection: the MCP SDK rejects any Host header that is not +# explicitly allowlisted. The public hostname behind the reverse proxy MUST +# appear either in ``server.allowed_hosts`` in beaconmcp.yaml or in the +# legacy BEACONMCP_ALLOWED_HOSTS env var. +_allowed_hosts = config.server.allowed_hosts or _csv_env( + "BEACONMCP_ALLOWED_HOSTS", + ["127.0.0.1:*", "localhost:*", "[::1]:*"], +) +_allowed_origins = config.server.allowed_origins or _csv_env( + "BEACONMCP_ALLOWED_ORIGINS", + ["https://assistant.ai", "https://chat.openai.com", "https://gemini.google.com"], +) + + +def _load_icons() -> list[Icon]: + """Load the bundled logo as a data-URL icon for MCP clients. + + Shipped inline so clients get the icon without needing a public static + route, and so nothing breaks when the server is hidden behind a tunnel + that only forwards /mcp. + """ + logo_path = Path(__file__).parent / "assets" / "logo.webp" + if not logo_path.is_file(): + return [] + data = base64.b64encode(logo_path.read_bytes()).decode("ascii") + return [ + Icon( + src=f"data:image/webp;base64,{data}", + mimeType="image/webp", + sizes=["512x512"], + ) + ] + + +def _build_instructions() -> str: + """Assemble the MCP server greeting from whatever capabilities are enabled. + + Each capability contributes its own blurb so a server with only SSH (e.g. + a couple of VPS) doesn't advertise Proxmox tools it can't expose, and a + Proxmox-only server doesn't pretend to have an SSH fallback. The greeting + is the first thing a client model reads — keeping it truthful is what + makes the conditional registration useful. + """ + blurbs: list[str] = [] + entry: list[str] = [] + if config.pve_nodes: + node_count = len(config.pve_nodes) + blurbs.append( + f"a Proxmox VE cluster ({node_count} node" + f"{'s' if node_count > 1 else ''}) via proxmox_* tools" + ) + entry.append("proxmox_list_nodes to see the cluster") + if config.ssh and config.ssh.hosts: + host_count = len(config.ssh.hosts) + blurbs.append( + f"{host_count} SSH host{'s' if host_count > 1 else ''} via " + "ssh_* tools for direct shell access" + ) + entry.append("ssh_list_sessions to track running commands") + if bmc_registry: + dev_count = len(bmc_registry) + blurbs.append( + f"{dev_count} BMC device{'s' if dev_count > 1 else ''} " + "(HP iLO / IPMI / iDRAC / Supermicro) via bmc_* tools for " + "hardware power and health" + ) + entry.append("bmc_list_devices to see hardware endpoints") + + if not blurbs: + # Defensive: the config loader refuses to start with no capability. + return "BeaconMCP running with no capabilities configured." + + body = "BeaconMCP exposes " + "; ".join(blurbs) + "." + entry_line = " Start with " + ", or ".join(entry) + "." if entry else "" + return body + entry_line + + +mcp = FastMCP( + "beaconmcp", + instructions=_build_instructions(), + website_url="https://github.com/Showdown76py/BeaconMCP", + icons=_load_icons(), + transport_security=TransportSecuritySettings( + enable_dns_rebinding_protection=True, + allowed_hosts=_allowed_hosts, + allowed_origins=_allowed_origins, + ), +) + + +# Wrap mcp.tool to inject metrics tracking. +# +# FastMCP (mcp 1.27) runs sync tool functions INLINE in the asyncio event +# loop -- it never offloads them to a thread. Most Proxmox tools are sync +# ``def`` wrappers around blocking proxmoxer calls (5s timeout x 2 retries), +# so a single unreachable node would otherwise freeze the entire server. +# We therefore make EVERY registered tool async from FastMCP's view and push +# sync bodies onto a worker thread, keeping blocking work off the event loop. +# ``@wraps(func)`` preserves the original signature so FastMCP still derives +# the tool input schema correctly. +_orig_tool = mcp.tool +def _metric_tool(*args, **kwargs): + def decorator(func): + tool_name = func.__name__ + is_coro = inspect.iscoroutinefunction(func) + + @wraps(func) + async def wrapper(*f_args, **f_kwargs): + start = time.monotonic() + status = "ok" + try: + if is_coro: + return await func(*f_args, **f_kwargs) + return await anyio.to_thread.run_sync( + partial(func, *f_args, **f_kwargs) + ) + except Exception: + status = "error" + raise + finally: + latency = (time.monotonic() - start) * 1000 + tool_calls.inc(tool=tool_name, status=status) + tool_latency_ms.observe(latency, tool=tool_name) + audit.emit( + "tool.call", tool=tool_name, status=status, + duration_ms=round(latency, 1), + client_id=current_client_id(), args=audit.compact_args(f_kwargs), + ) + + return _orig_tool(*args, **kwargs)(wrapper) + return decorator +mcp.tool = _metric_tool + + +@mcp.resource("beaconmcp://infrastructure") +def get_infrastructure() -> str: + """Infrastructure context: node topology, naming conventions, and access constraints.""" + if not config.infrastructure: + return "No infrastructure context configured." + + import yaml + + return yaml.dump(config.infrastructure, default_flow_style=False, allow_unicode=True) + + +@mcp.prompt() +def beaconmcp_context() -> str: + """Inject infrastructure context: topology and capability-aware diagnostic workflow.""" + topology_lines: list[str] = [] + if config.pve_nodes: + topology_lines.append( + "Proxmox nodes: " + ", ".join(n.name for n in config.pve_nodes) + ) + if config.ssh and config.ssh.hosts: + topology_lines.append( + "SSH hosts: " + ", ".join(h.name for h in config.ssh.hosts) + ) + if bmc_registry: + topology_lines.append( + "BMC devices: " + + ", ".join(f"{d.id} ({d.type})" for d in config.bmc_devices) + ) + topology = "\n".join(topology_lines) if topology_lines else "(no capabilities)" + + # Build a diagnostic workflow that only references tools that are + # actually registered. A VPS-only deployment gets a one-step workflow + # and no Proxmox/BMC references, which stops the model from suggesting + # tool calls that would 404. + steps: list[str] = [] + if config.pve_nodes: + steps.append( + "Start with cluster_overview for the whole cluster in one call, " + "or cluster_health(node=...) for node metrics + BMC + recent errors." + ) + steps.append( + "Drill in with proxmox_node_status / proxmox_list_vms as needed. " + "Pass fields=[...] on detail tools to trim the response." + ) + steps.append( + "For VM configuration issues, use proxmox_read_file to safely read configs via QEMU." + ) + steps.append( + "Find a VM by name with vm_find('web-*'); act on many at once with " + "vm_bulk_action(vmids=[...], action='stop'). " + "Consider taking a snapshot or proxmox_backup_create before a risky operation if relevant." + ) + if config.pve_nodes and config.ssh and config.ssh.hosts: + steps.append( + "If a Proxmox node is unreachable via API, try ssh_run against " + "the matching ssh.hosts entry." + ) + if bmc_registry: + steps.append( + "If a host is completely unresponsive, cluster_health already " + "includes BMC facts; otherwise use bmc_list_devices + " + "bmc_health_status / bmc_power_status." + ) + if config.pve_nodes and config.ssh and config.ssh.hosts: + steps.append( + "For in-VM issues, prefer proxmox_run (QEMU Guest Agent) or ssh_run. " + "Both auto-switch to async on timeout and accept exec_id for polling." + ) + elif config.pve_nodes: + steps.append("For in-VM issues, use proxmox_run (QEMU Guest Agent).") + elif config.ssh and config.ssh.hosts: + steps.append( + "Use ssh_run on declared hosts. Pass wait=False for long commands; " + "poll with ssh_run(exec_id=...)." + ) + workflow = "\n".join(f"{i}. {s}" for i, s in enumerate(steps, 1)) or "(no workflow: no capabilities configured)" + + infra = config.infrastructure + conventions = "" + if infra.get("conventions"): + conventions = "\n".join(f"- {k}: {v}" for k, v in infra["conventions"].items()) + + notes = "" + if infra.get("notes"): + notes = "\n".join(f"- {n}" for n in infra["notes"]) + + return f"""You are operating a BeaconMCP-managed infrastructure with the following topology: + +{topology} + +Conventions: +{conventions} + +Notes: +{notes} + +Diagnostic workflow: +{workflow}""" + + +# Register tool modules only for the capabilities that are actually +# configured. Each ``if`` gate here is what makes a VPS-only, Proxmox-only, +# or BMC-only deployment possible: tools the server cannot honor stay out +# of the exposed tool list. +if config.pve_nodes: + register_monitoring_tools(mcp, proxmox_client) + register_vm_tools(mcp, proxmox_client) + register_panel_tools(mcp, proxmox_client) + register_system_tools(mcp, proxmox_client, ssh_client if config.ssh and config.ssh.hosts else None) + # Aggregators ride on top of the Proxmox client and opportunistically + # pull BMC facts when the registry is non-empty. + register_aggregator_tools(mcp, proxmox_client, config, bmc_registry) +if config.ssh and config.ssh.hosts: + register_ssh_tools(mcp, ssh_client) +if bmc_registry: + register_bmc_tools(mcp, bmc_registry) +register_security_tools(mcp) +# Not tied to any infrastructure capability: keeping the server itself +# current is useful on every deployment shape. +register_maintenance_tools( + mcp, config.features.updates, config_path=config.source_path, +) diff --git a/src/beaconmcp/ssh/__init__.py b/src/beaconmcp/ssh/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/beaconmcp/ssh/client.py b/src/beaconmcp/ssh/client.py new file mode 100644 index 0000000..6af5abc --- /dev/null +++ b/src/beaconmcp/ssh/client.py @@ -0,0 +1,336 @@ +from __future__ import annotations + +import asyncio +import os +import time +import uuid +from dataclasses import dataclass, field +from typing import Any + +import asyncssh + +from ..config import Config, SSHHost + + +class SSHNotConfiguredError(Exception): + def __init__(self) -> None: + super().__init__( + "SSH is not configured. Add an 'ssh:' section with at least " + "one entry under 'ssh.hosts[]' in beaconmcp.yaml." + ) + + +class SSHHostResolutionError(Exception): + """Raised when a host identifier cannot be resolved to a declared SSH host.""" + + +@dataclass +class SSHExecSession: + exec_id: str + host: str + command: str + status: str = "running" + stdout: str = "" + stderr: str = "" + exit_code: int | None = None + started_at: float = field(default_factory=time.time) + + +_ssh_sessions: dict[str, SSHExecSession] = {} +_ssh_tasks: set[asyncio.Task[None]] = set() +_connection_cache: dict[str, tuple[asyncssh.SSHClientConnection, float]] = {} +_CONNECTION_TTL = 300 # 5 minutes +_SSH_SESSION_TTL = 3600 # drop completed sessions older than this + + +def _prune_ssh_sessions() -> None: + now = time.time() + stale = [ + eid + for eid, s in _ssh_sessions.items() + if s.status != "running" and now - s.started_at > _SSH_SESSION_TTL + ] + for eid in stale: + del _ssh_sessions[eid] + + +async def _connect_to_host( + spec: SSHHost, + *, + known_hosts: str | None = None, + strict_host_key_checking: bool = False, +) -> asyncssh.SSHClientConnection: + """Open an asyncssh connection to a declared host using its auth method. + + Exposed at module level so BMC jump-host tunneling in ``bmc/hp_ilo.py`` + can reuse the same auth plumbing (password vs. key_file, port override, + trusted host keys) instead of duplicating it. + + Host-key verification: + * ``known_hosts`` (path): asyncssh loads the file and refuses unknown keys. + * ``strict_host_key_checking=True`` with no ``known_hosts``: use the + caller's ``~/.ssh/known_hosts`` (asyncssh's default when the kwarg + is omitted entirely). + * Neither: pass ``known_hosts=None`` -- accept any key. Default to keep + existing trusted-LAN deployments working unchanged. + """ + connect_kwargs: dict[str, Any] = { + "host": spec.host, + "port": spec.port, + "username": spec.user, + } + if known_hosts: + connect_kwargs["known_hosts"] = os.path.expanduser(known_hosts) + elif not strict_host_key_checking: + # Trusted-LAN default. + connect_kwargs["known_hosts"] = None + # else: omit the kwarg -> asyncssh uses ~/.ssh/known_hosts automatically. + if spec.password: + connect_kwargs["password"] = spec.password + elif spec.key_file: + connect_kwargs["client_keys"] = [os.path.expanduser(spec.key_file)] + return await asyncssh.connect(**connect_kwargs) + + +class SSHClient: + """Async SSH client with declarative host resolution and connection caching. + + Every SSH target must be declared under ``ssh.hosts[]`` in beaconmcp.yaml. + The client looks up an identifier (a host name, a numeric VMID, or a raw + IP/hostname) against that declaration to recover full connect params + (host, port, user, password or key_file). There is no implicit + passthrough to an arbitrary host with shared credentials — this keeps + credentials a declarative concern, not a runtime guess. + """ + + def __init__(self, config: Config) -> None: + self._config = config + + def resolve(self, identifier: str) -> SSHHost: + """Resolve an identifier to a declared :class:`SSHHost`. + + Resolution order: + 1. Match ``identifier`` against ``ssh.hosts[].name``. + 2. If numeric and ``ssh.vmid_to_ip`` is set, apply the template and + match the resulting IP against ``ssh.hosts[].host``. + 3. Otherwise, match ``identifier`` directly against + ``ssh.hosts[].host``. + + Raises :class:`SSHNotConfiguredError` if SSH has no hosts declared, + or :class:`SSHHostResolutionError` with an actionable message when + no host matches. + """ + if not self._config.ssh or not self._config.ssh.hosts: + raise SSHNotConfiguredError() + + # 1. By declared name + by_name = self._config.get_ssh_host(identifier) + if by_name is not None: + return by_name + + # 2. Numeric VMID → template → address match + if identifier.isdigit(): + template = self._config.ssh.vmid_to_ip + if not template: + raise SSHHostResolutionError( + f"Identifier {identifier!r} looks like a numeric VMID " + "but no 'ssh.vmid_to_ip' template is configured. Either " + "set the template (e.g. '192.168.1.{id}') or reference " + "one of the declared host names under ssh.hosts[]." + ) + try: + resolved_ip = template.format(id=identifier) + except (KeyError, IndexError) as exc: + raise SSHHostResolutionError( + f"ssh.vmid_to_ip template {template!r} is invalid: {exc}. " + "Use '{id}' as the only placeholder." + ) from exc + by_addr = self._config.get_ssh_host_by_address(resolved_ip) + if by_addr is not None: + return by_addr + raise SSHHostResolutionError( + f"VMID {identifier!r} resolves to {resolved_ip!r} via " + "ssh.vmid_to_ip, but no ssh.hosts[] entry has that address. " + f"Declare one (e.g. name: vm-{identifier}, host: " + f"{resolved_ip}, user: ..., password or key_file: ...)." + ) + + # 3. Direct IP/hostname + by_addr = self._config.get_ssh_host_by_address(identifier) + if by_addr is not None: + return by_addr + + declared = ", ".join(h.name for h in self._config.ssh.hosts) or "" + hint = "" + # When the identifier matches a Proxmox node that wasn't declared as + # an SSH host, point the caller at the two common fixes instead of + # just reporting "not declared". This is the single most common + # foot-gun — pre-2.0 code let you SSH into a Proxmox node by name + # implicitly. + if any(n.name == identifier for n in self._config.pve_nodes): + hint = ( + f" Note: {identifier!r} is a Proxmox node. To reach it via " + "SSH, either add it under ssh.hosts[] explicitly, or set " + "'ssh.inherit_proxmox_nodes: true' with 'ssh.defaults:' so " + "every node is auto-declared. If you meant to run something " + "*inside* a VM/LXC on that node, use proxmox_run(node=..., " + "vmid=..., command=...) instead — it goes through QEMU Guest " + "Agent / pct exec and doesn't need SSH." + ) + elif identifier.isdigit(): + hint = ( + f" Note: {identifier!r} looks like a VMID. To run a command " + "inside that guest, prefer proxmox_run(node=..., " + f"vmid={identifier}, command=...)." + ) + raise SSHHostResolutionError( + f"Host {identifier!r} is not declared in ssh.hosts[]. Add an " + "entry (name, host, user, password or key_file) to enable SSH " + f"to this target. Declared hosts: {declared}.{hint}" + ) + + def resolve_host(self, identifier: str) -> str: + """Return the connect-target address for an identifier. + + Helper used by ``ssh_run`` (and its ``wait=False`` / ``exec_id=…`` + polling paths) to surface the resolved IP/hostname in the tool + response. Prefer :meth:`resolve` when the full host spec (port, + user, auth) is needed. + """ + return self.resolve(identifier).host + + async def _get_connection(self, identifier: str) -> asyncssh.SSHClientConnection: + """Get or create an SSH connection with caching. + + Cache key is the declared host *name*, so distinct declarations + sharing the same address still get distinct cached connections + (useful when one address has multiple user accounts). + """ + host_spec = self.resolve(identifier) + cache_key = host_spec.name + + if cache_key in _connection_cache: + conn, created_at = _connection_cache[cache_key] + if time.time() - created_at < _CONNECTION_TTL: + try: + if not conn.is_closed(): + return conn + except Exception: + pass + try: + conn.close() + except Exception: + pass + del _connection_cache[cache_key] + + global_kh = self._config.ssh.known_hosts if self._config.ssh else None + global_strict = ( + self._config.ssh.strict_host_key_checking if self._config.ssh else False + ) + # Per-host overrides win; ``None`` means "inherit the global setting". + kh = host_spec.known_hosts if host_spec.known_hosts is not None else global_kh + strict = ( + host_spec.strict_host_key_checking + if host_spec.strict_host_key_checking is not None + else global_strict + ) + conn = await _connect_to_host( + host_spec, known_hosts=kh, strict_host_key_checking=strict, + ) + _connection_cache[cache_key] = (conn, time.time()) + return conn + + async def exec_command(self, host: str, command: str, timeout: int = 60) -> dict[str, Any]: + """Execute a command via SSH and wait for the result.""" + try: + conn = await self._get_connection(host) + result = await asyncio.wait_for( + conn.run(command, check=False), + timeout=timeout, + ) + return { + "stdout": result.stdout or "", + "stderr": result.stderr or "", + "exit_code": result.exit_status, + } + except asyncio.TimeoutError: + return { + "stdout": "", + "stderr": "", + "exit_code": None, + "status": "timeout", + "error": f"Command timed out after {timeout}s. Use ssh_run(..., wait=False) to start async and poll with ssh_run(exec_id=...).", + } + except SSHNotConfiguredError: + raise + except Exception as e: + return {"error": f"SSH connection to '{host}' failed: {e}. Check SSH credentials and host accessibility."} + + async def exec_command_async(self, host: str, command: str) -> str: + """Start a long-running command and return an exec_id.""" + _prune_ssh_sessions() + exec_id = str(uuid.uuid4())[:8] + session = SSHExecSession(exec_id=exec_id, host=host, command=command) + _ssh_sessions[exec_id] = session + + async def _run() -> None: + try: + conn = await self._get_connection(host) + result = await asyncio.wait_for(conn.run(command, check=False), timeout=600) + session.stdout = result.stdout or "" + session.stderr = result.stderr or "" + session.exit_code = result.exit_status + session.status = "completed" + except asyncio.TimeoutError: + session.status = "timeout" + except Exception as e: + session.status = "failed" + session.stderr = str(e) + + # Keep a reference to the task so it isn't garbage collected mid-run. + task = asyncio.create_task(_run()) + _ssh_tasks.add(task) + task.add_done_callback(_ssh_tasks.discard) + return exec_id + + async def sftp_put(self, host: str, local_path: str, remote_path: str) -> None: + """Stream a local file to ``remote_path`` on ``host`` via SFTP. + + Used by ``proxmox_upload_file`` for large transfers (asyncssh streams + in chunks, so RAM usage stays bounded regardless of file size). + """ + conn = await self._get_connection(host) + async with conn.start_sftp_client() as sftp: + await sftp.put(local_path, remote_path) + + async def sftp_get(self, host: str, remote_path: str, local_path: str) -> None: + """Stream ``remote_path`` from ``host`` into ``local_path`` via SFTP.""" + conn = await self._get_connection(host) + async with conn.start_sftp_client() as sftp: + await sftp.get(remote_path, local_path) + + async def sftp_remove(self, host: str, remote_path: str) -> None: + """Best-effort SFTP unlink, used to clean up staging files on a node.""" + try: + conn = await self._get_connection(host) + async with conn.start_sftp_client() as sftp: + await sftp.remove(remote_path) + except Exception: + pass + + @staticmethod + def get_session(exec_id: str) -> SSHExecSession | None: + return _ssh_sessions.get(exec_id) + + @staticmethod + def list_sessions() -> list[dict[str, Any]]: + return [ + { + "exec_id": s.exec_id, + "host": s.host, + "command": s.command, + "status": s.status, + "elapsed_s": round(time.time() - s.started_at), + } + for s in _ssh_sessions.values() + ] diff --git a/src/beaconmcp/ssh/tools.py b/src/beaconmcp/ssh/tools.py new file mode 100644 index 0000000..7411c21 --- /dev/null +++ b/src/beaconmcp/ssh/tools.py @@ -0,0 +1,114 @@ +from __future__ import annotations + +import asyncio +import time +from typing import Any + +from mcp.server.fastmcp import FastMCP + +from .client import SSHClient, SSHHostResolutionError, SSHNotConfiguredError + + +def _session_to_result(exec_id: str, session) -> dict[str, Any]: + """Turn an SSH session row into a ``proxmox_run``-shaped response.""" + elapsed = round(time.time() - session.started_at, 1) + if session.status == "running": + return { + "status": "running", + "exec_id": exec_id, + "host": session.host, + "command": session.command, + "elapsed_s": elapsed, + } + return { + "status": "ok" if session.status == "completed" and session.exit_code == 0 else session.status, + "exec_id": exec_id, + "host": session.host, + "command": session.command, + "stdout": session.stdout, + "stderr": session.stderr, + "exit_code": session.exit_code, + "duration_s": elapsed, + } + + +def register_ssh_tools(mcp: FastMCP, ssh_client: SSHClient) -> None: + """Register SSH command execution tools.""" + + @mcp.tool() + async def ssh_run( + host: str = "", + command: str = "", + timeout: int = 60, + wait: bool = True, + exec_id: str = "", + ) -> dict[str, Any]: + """Run a command on a host via SSH. Handles sync + async in one tool. + + Three call patterns: + + - **Sync** (default): pass ``host`` + ``command``. Blocks up to + ``timeout`` seconds (max 600). Completes -> returns + ``stdout``/``stderr``/``exit_code``. Times out -> auto-switches to + async and returns ``{status: "running", exec_id}``. + - **Async start**: ``host`` + ``command`` + ``wait=False``. + Returns ``{status: "running", exec_id}`` immediately. + - **Poll existing**: pass ``exec_id`` only. Returns the current + status/output for that session. + + ``host`` must resolve to a declared ``ssh.hosts[]`` entry. Accepts: + an entry ``name``; a numeric VMID when ``ssh.vmid_to_ip`` is set + (e.g. ``"110"`` -> ``"192.168.1.110"``); or a declared ``host`` + address. If ``ssh.inherit_proxmox_nodes: true``, every Proxmox node + is auto-declared as an SSH host under its own name, so reaching the + hypervisor reuses the same identifier as ``proxmox_run(node=…)``. + + To run **inside a VM or LXC** managed by Proxmox, prefer + ``proxmox_run`` (QEMU Guest Agent / ``pct exec``) — no SSH is needed + and it works even when the guest has no inbound network reachability. + """ + if exec_id: + session = SSHClient.get_session(exec_id) + if not session: + return {"status": "error", "error": f"No SSH command with exec_id {exec_id!r}."} + return _session_to_result(exec_id, session) + + if not host or not command: + return {"status": "error", "error": "`host` and `command` are required when `exec_id` is not provided."} + + max_timeout = min(max(timeout, 1), 600) + + try: + new_id = await ssh_client.exec_command_async(host, command) + except (SSHNotConfiguredError, SSHHostResolutionError) as e: + return {"status": "error", "error": str(e)} + + if not wait: + return {"status": "running", "exec_id": new_id, "host": host, "elapsed_s": 0} + + deadline = time.time() + max_timeout + while time.time() < deadline: + session = SSHClient.get_session(new_id) + if session and session.status != "running": + return _session_to_result(new_id, session) + await asyncio.sleep(0.5) + + session = SSHClient.get_session(new_id) + if session and session.status != "running": + return _session_to_result(new_id, session) + return { + "status": "running", + "exec_id": new_id, + "host": host, + "elapsed_s": int(time.time() - (session.started_at if session else time.time())), + "hint": "Command still running. Call ssh_run(exec_id=...) to poll.", + } + + @mcp.tool() + def ssh_list_sessions() -> dict[str, Any]: + """List all active and recent SSH command sessions. + + Returns exec_id, host, command, status, and elapsed time for each session. + """ + sessions = SSHClient.list_sessions() + return {"sessions": sessions, "total": len(sessions)} diff --git a/src/beaconmcp/updates.py b/src/beaconmcp/updates.py new file mode 100644 index 0000000..47f9c1e --- /dev/null +++ b/src/beaconmcp/updates.py @@ -0,0 +1,821 @@ +"""Update detection and self-update for BeaconMCP. + +BeaconMCP ships no PyPI package and cuts no releases: the canonical install +is a ``git clone`` at ``/opt/beaconmcp`` with a venv and a systemd unit (see +``deploy/install.sh``). So "is there an update?" means *is this checkout +behind the remote default branch?*, not "is there a newer version string". + +Three things live here: + +* :func:`detect_installation` -- how this server was installed, so the + advice we give matches reality instead of assuming everyone ran the + install script. +* :func:`check_for_update` -- a cached, fail-soft, read-only check. It also + diffs the *new* ``.env.example`` / ``beaconmcp.yaml.example`` against the + operator's actual files, which is how we can say "this update wants a + variable you haven't set" before they apply it. +* :func:`apply_update` -- pull, reinstall dependencies, **validate the + config**, and roll back if the new revision cannot load it. Restarting + into a config that refuses to parse would take the server down with no + one at the keyboard, so validation is a hard gate, not a warning. + +Nothing here ever raises into a caller: an air-gapped box, a missing git +binary or a detached HEAD all degrade to "couldn't check", never to a +broken dashboard or a failed tool call. +""" + +from __future__ import annotations + +import os +import re +import shutil +import subprocess +import sys +import threading +import time +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +_REPO_URL = "https://github.com/Showdown76py/BeaconMCP" + +#: How long a successful check stays fresh. Updates are not urgent and the +#: check shells out to git, so once every few hours is plenty. +CHECK_TTL_SECONDS = 6 * 3600 +#: Failures are retried sooner -- a transient DNS blip shouldn't mean six +#: hours of "unknown". +FAILED_CHECK_TTL_SECONDS = 15 * 60 + +#: Ceiling on any git/pip subprocess. `pip install -e .` on a cold cache is +#: the slow one; the rest finish in well under a second. +_GIT_TIMEOUT = 60 +_PIP_TIMEOUT = 900 + + +def current_version() -> str: + """Installed version string, preferring package metadata.""" + try: + from importlib.metadata import version + + return version("beaconmcp") + except Exception: # noqa: BLE001 - not installed as a distribution + from . import __version__ + + return __version__ + + +# --------------------------------------------------------------------------- +# Installation shape +# --------------------------------------------------------------------------- + +@dataclass +class Installation: + """How this particular server was installed.""" + + #: "git" (clone, the documented install), "pip" (installed as a + #: distribution from a URL/wheel), "docker", or "unknown". + kind: str + #: Root of the git checkout, when there is one. + root: Path | None + #: Interpreter running us -- also the venv's python when there is a venv. + python: str + #: Virtualenv prefix, or None when running against a system interpreter. + venv: Path | None + #: True when the package is imported straight from the checkout. + editable: bool + #: True when the process was started by systemd. + under_systemd: bool + #: systemd unit to restart, when we can name one. + service: str | None + #: True when running inside a container. + in_container: bool + + def to_json(self) -> dict[str, Any]: + return { + "kind": self.kind, + "root": str(self.root) if self.root else None, + "python": self.python, + "venv": str(self.venv) if self.venv else None, + "editable": self.editable, + "under_systemd": self.under_systemd, + "service": self.service, + "in_container": self.in_container, + } + + +def _package_root() -> Path: + """Directory holding ``src/beaconmcp`` -- i.e. the repo root when cloned.""" + # ...//src/beaconmcp/updates.py -> parents[2] == + return Path(__file__).resolve().parents[2] + + +def _detect_service() -> str | None: + """Name the systemd unit, if one is installed for us.""" + for candidate in ( + "/etc/systemd/system/beaconmcp.service", + "/lib/systemd/system/beaconmcp.service", + "/usr/lib/systemd/system/beaconmcp.service", + ): + if Path(candidate).is_file(): + return "beaconmcp" + return None + + +def detect_installation() -> Installation: + """Inspect the runtime to work out how BeaconMCP got here.""" + root = _package_root() + is_git = (root / ".git").exists() + venv = Path(sys.prefix) if sys.prefix != sys.base_prefix else None + # systemd exports INVOCATION_ID to every unit it starts; it is the one + # signal that does not require guessing at pid 1 or parsing /proc. + under_systemd = bool(os.environ.get("INVOCATION_ID")) + in_container = ( + Path("/.dockerenv").exists() + or os.environ.get("container") is not None + ) + + if is_git: + kind = "git" + elif in_container: + kind = "docker" + else: + try: + from importlib.metadata import distribution + + distribution("beaconmcp") + kind = "pip" + except Exception: # noqa: BLE001 + kind = "unknown" + + return Installation( + kind=kind, + root=root if is_git else None, + python=sys.executable, + venv=venv, + editable=is_git, + under_systemd=under_systemd, + service=_detect_service(), + in_container=in_container, + ) + + +# --------------------------------------------------------------------------- +# git plumbing +# --------------------------------------------------------------------------- + +def _git(root: Path, *args: str, timeout: int = _GIT_TIMEOUT) -> tuple[int, str, str]: + """Run a git command in ``root``. Never raises.""" + if not shutil.which("git"): + return 127, "", "git is not installed" + try: + proc = subprocess.run( + ["git", *args], + cwd=str(root), + capture_output=True, + text=True, + timeout=timeout, + # Never let git try to prompt for credentials: on a private + # remote it would hang until the timeout instead of failing. + env={**os.environ, "GIT_TERMINAL_PROMPT": "0", "GIT_ASKPASS": ""}, + ) + return proc.returncode, proc.stdout.strip(), proc.stderr.strip() + except subprocess.TimeoutExpired: + return 124, "", f"git {' '.join(args)} timed out" + except OSError as exc: + return 1, "", str(exc) + + +def _default_branch(root: Path) -> str: + """Remote default branch name, falling back to ``main``.""" + code, out, _ = _git(root, "symbolic-ref", "--short", "refs/remotes/origin/HEAD") + if code == 0 and out.startswith("origin/"): + return out.split("/", 1)[1] + # Not every clone has origin/HEAD set (shallow clones, older git). + code, out, _ = _git(root, "remote", "show", "origin") + if code == 0: + match = re.search(r"HEAD branch:\s*(\S+)", out) + if match: + return match.group(1) + return "main" + + +def working_tree_dirty(root: Path) -> bool: + """True when tracked files have uncommitted modifications.""" + code, out, _ = _git(root, "status", "--porcelain", "--untracked-files=no") + return code == 0 and bool(out) + + +# --------------------------------------------------------------------------- +# Config drift: what the new revision wants that the operator hasn't set +# --------------------------------------------------------------------------- + +_ENV_ASSIGNMENT = re.compile(r"^\s*(?:export\s+)?([A-Z][A-Z0-9_]*)\s*=") + + +def _env_names(text: str) -> list[str]: + """Variable names assigned in a dotenv-style file (comments included). + + Comments count on purpose: ``.env.example`` documents optional settings + as ``# GEMINI_API_KEY=`` and those are exactly the ones an operator + wants to hear about after an update. + """ + names: list[str] = [] + for raw in text.splitlines(): + line = raw.lstrip() + if line.startswith("#"): + line = line.lstrip("#").lstrip() + match = _ENV_ASSIGNMENT.match(line) + if match: + names.append(match.group(1)) + return names + + +def _yaml_paths(text: str) -> set[str]: + """Dotted key paths in a YAML document, list items collapsed away.""" + try: + import yaml + + data = yaml.safe_load(text) + except Exception: # noqa: BLE001 - malformed example, nothing to diff + return set() + + paths: set[str] = set() + + def walk(node: Any, prefix: str) -> None: + if isinstance(node, dict): + for key, value in node.items(): + path = f"{prefix}.{key}" if prefix else str(key) + paths.add(path) + walk(value, path) + elif isinstance(node, list): + # Sequence entries are instances (nodes, hosts, devices), not + # settings -- their *shape* is what matters, so recurse without + # adding an index to the path. + for item in node: + walk(item, prefix) + + walk(data, "") + return paths + + +@dataclass +class ConfigDrift: + """Settings the incoming revision knows about and this install does not.""" + + new_env_vars: list[str] = field(default_factory=list) + new_config_keys: list[str] = field(default_factory=list) + + @property + def empty(self) -> bool: + return not self.new_env_vars and not self.new_config_keys + + def to_json(self) -> dict[str, Any]: + return { + "new_env_vars": self.new_env_vars, + "new_config_keys": self.new_config_keys, + } + + +def _read_local(root: Path, *names: str) -> str | None: + for name in names: + path = root / name + if path.is_file(): + try: + return path.read_text(encoding="utf-8", errors="replace") + except OSError: + return None + return None + + +def config_drift(root: Path, ref: str, config_path: Path | None = None) -> ConfigDrift: + """Diff the example files at ``ref`` against what this install actually has. + + Deliberately compares against the operator's *real* files rather than + the local examples: someone who set ``GEMINI_API_KEY`` before it was + documented should not be told to set it again. + """ + drift = ConfigDrift() + + code, new_env, _ = _git(root, "show", f"{ref}:.env.example") + if code == 0: + local_env = _read_local(root, ".env") or "" + known = set(_env_names(local_env)) | set(os.environ) + for name in _env_names(new_env): + if name not in known and name not in drift.new_env_vars: + drift.new_env_vars.append(name) + + code, new_yaml, _ = _git(root, "show", f"{ref}:beaconmcp.yaml.example") + if code == 0: + local_yaml = None + if config_path and config_path.is_file(): + try: + local_yaml = config_path.read_text(encoding="utf-8", errors="replace") + except OSError: + local_yaml = None + if local_yaml is None: + local_yaml = _read_local(root, "beaconmcp.yaml") or "" + have = _yaml_paths(local_yaml) + # Anything the operator already configured, plus its ancestors, is + # "known"; only genuinely new leaves are worth reporting. + for path in sorted(_yaml_paths(new_yaml) - have): + if any(p.startswith(path + ".") for p in have): + continue # a parent of something already configured + drift.new_config_keys.append(path) + + return drift + + +# --------------------------------------------------------------------------- +# Update check +# --------------------------------------------------------------------------- + +@dataclass +class UpdateInfo: + """Result of one update check. Always renderable, even on failure.""" + + checked_at: float + available: bool = False + error: str | None = None + version: str = "" + install_kind: str = "unknown" + branch: str | None = None + current_ref: str | None = None + latest_ref: str | None = None + behind: int = 0 + commits: list[dict[str, str]] = field(default_factory=list) + drift: ConfigDrift = field(default_factory=ConfigDrift) + instructions: list[str] = field(default_factory=list) + can_self_update: bool = False + blockers: list[str] = field(default_factory=list) + repo_url: str = _REPO_URL + + def to_json(self) -> dict[str, Any]: + return { + "checked_at": self.checked_at, + "available": self.available, + "error": self.error, + "version": self.version, + "install_kind": self.install_kind, + "branch": self.branch, + "current_ref": self.current_ref, + "latest_ref": self.latest_ref, + "behind": self.behind, + "commits": self.commits, + "config": self.drift.to_json(), + "instructions": self.instructions, + "can_self_update": self.can_self_update, + "blockers": self.blockers, + "repo_url": self.repo_url, + "compare_url": ( + f"{self.repo_url}/compare/{self.current_ref}...{self.latest_ref}" + if self.current_ref and self.latest_ref and self.available + else None + ), + } + + +def manual_instructions(install: Installation) -> list[str]: + """Shell commands that update *this* install, in order.""" + if install.kind == "git" and install.root: + root = install.root + pip = ( + str(install.venv / "bin" / "pip") + if install.venv and (install.venv / "bin" / "pip").exists() + else f"{install.python} -m pip" + ) + steps = [f"cd {root}", "git pull --ff-only", f"{pip} install -e ."] + if install.service: + steps.append(f"systemctl restart {install.service}") + return steps + if install.kind == "docker": + return [ + "docker compose pull", + "docker compose up -d", + "# (or: docker pull && docker compose up -d)", + ] + if install.kind == "pip": + pip = f"{install.python} -m pip" + steps = [f"{pip} install --upgrade 'beaconmcp @ git+{_REPO_URL}.git'"] + if install.service: + steps.append(f"systemctl restart {install.service}") + return steps + return [ + "# Could not determine how BeaconMCP was installed here.", + "# Re-run the installer from a checkout: bash deploy/install.sh", + ] + + +def _self_update_blockers(install: Installation, root: Path | None) -> list[str]: + """Reasons ``apply_update`` would refuse, as operator-facing sentences.""" + blockers: list[str] = [] + if install.kind != "git" or root is None: + blockers.append( + f"this is a {install.kind} install, and automatic updates only " + "support a git checkout" + ) + return blockers + if not shutil.which("git"): + blockers.append("the git binary is not on PATH") + if working_tree_dirty(root): + blockers.append( + "the checkout has uncommitted changes -- commit or stash them " + "first so the update cannot discard your work" + ) + return blockers + + +_cache_lock = threading.Lock() +_cached: UpdateInfo | None = None + +#: Serializes the git work. Two entry points can reach this concurrently -- +#: the dashboard button and the MCP tool -- and two ``git pull`` / +#: ``pip install`` runs in one checkout would fight over index.lock and +#: could leave a half-applied tree. ``_apply_lock`` is never waited on: a +#: second updater is told one is already running rather than queueing +#: behind a pip that may take minutes. +_apply_lock = threading.Lock() +#: Held across an uncached check so N dashboard tabs opening at once cause +#: one ``git fetch``, not N. Waiters get the result the winner cached. +_check_lock = threading.RLock() + + +def check_for_update( + *, + force: bool = False, + config_path: Path | None = None, + install: Installation | None = None, +) -> UpdateInfo: + """Return update status, using a cached result when it is still fresh. + + Read-only: it fetches git objects (which never touches the working tree) + and shells out to ``git show``. Failures are captured in + :attr:`UpdateInfo.error`, never raised. + + ``install`` overrides autodetection; passing one also bypasses the + cache, since the cache is keyed on "this server" and nothing else. + """ + global _cached + + if install is not None: + return _check_uncached(config_path=config_path, install=install) + + def _fresh() -> UpdateInfo | None: + with _cache_lock: + cached = _cached + if cached is None: + return None + ttl = FAILED_CHECK_TTL_SECONDS if cached.error else CHECK_TTL_SECONDS + return cached if time.time() - cached.checked_at < ttl else None + + if not force: + hit = _fresh() + if hit is not None: + return hit + + with _check_lock: + # Someone may have refreshed it while we waited for the lock; a + # forced check still runs, since that is the point of forcing. + if not force: + hit = _fresh() + if hit is not None: + return hit + info = _check_uncached(config_path=config_path) + with _cache_lock: + _cached = info + return info + + +def cached_update() -> UpdateInfo | None: + """Last check result, without triggering a new one.""" + with _cache_lock: + return _cached + + +def invalidate_cache() -> None: + global _cached + with _cache_lock: + _cached = None + + +def _check_uncached( + *, config_path: Path | None = None, install: Installation | None = None, +) -> UpdateInfo: + install = install or detect_installation() + info = UpdateInfo( + checked_at=time.time(), + version=current_version(), + install_kind=install.kind, + instructions=manual_instructions(install), + ) + + root = install.root + if install.kind != "git" or root is None: + info.error = ( + f"cannot check automatically: this is a {install.kind} install, " + "not a git checkout" + ) + info.blockers = _self_update_blockers(install, root) + return info + + code, head, err = _git(root, "rev-parse", "--short", "HEAD") + if code != 0: + info.error = f"could not read the local revision ({err or 'git failed'})" + return info + info.current_ref = head + + branch = _default_branch(root) + info.branch = branch + + code, _, err = _git(root, "fetch", "--quiet", "origin", branch) + if code != 0: + info.error = f"could not reach the remote ({err or 'git fetch failed'})" + info.blockers = _self_update_blockers(install, root) + return info + + remote_ref = f"origin/{branch}" + code, latest, err = _git(root, "rev-parse", "--short", remote_ref) + if code != 0: + info.error = f"could not read {remote_ref} ({err or 'git failed'})" + return info + info.latest_ref = latest + + code, count, _ = _git(root, "rev-list", "--count", f"HEAD..{remote_ref}") + info.behind = int(count) if code == 0 and count.isdigit() else 0 + info.available = info.behind > 0 + + if not info.available: + return info + + code, log, _ = _git( + root, "log", "--no-merges", "--max-count=20", + "--pretty=format:%h\x1f%s\x1f%aI", f"HEAD..{remote_ref}", + ) + if code == 0 and log: + for line in log.splitlines(): + parts = line.split("\x1f") + if len(parts) == 3: + info.commits.append( + {"sha": parts[0], "subject": parts[1], "date": parts[2]} + ) + + info.drift = config_drift(root, remote_ref, config_path) + info.blockers = _self_update_blockers(install, root) + info.can_self_update = not info.blockers + return info + + +# --------------------------------------------------------------------------- +# Applying an update +# --------------------------------------------------------------------------- + +@dataclass +class UpdateStep: + name: str + ok: bool + detail: str = "" + + def to_json(self) -> dict[str, Any]: + return {"step": self.name, "ok": self.ok, "detail": self.detail} + + +@dataclass +class UpdateResult: + ok: bool + steps: list[UpdateStep] = field(default_factory=list) + from_ref: str | None = None + to_ref: str | None = None + rolled_back: bool = False + restart_scheduled: bool = False + restart_in_seconds: int = 0 + message: str = "" + drift: ConfigDrift = field(default_factory=ConfigDrift) + + def to_json(self) -> dict[str, Any]: + return { + "ok": self.ok, + "steps": [s.to_json() for s in self.steps], + "from_ref": self.from_ref, + "to_ref": self.to_ref, + "rolled_back": self.rolled_back, + "restart_scheduled": self.restart_scheduled, + "restart_in_seconds": self.restart_in_seconds, + "message": self.message, + "config": self.drift.to_json(), + } + + +def _pip_command(install: Installation) -> list[str]: + if install.venv: + for candidate in ( + install.venv / "bin" / "pip", + install.venv / "Scripts" / "pip.exe", + ): + if candidate.exists(): + return [str(candidate)] + return [install.python, "-m", "pip"] + + +def _run(cmd: list[str], cwd: Path, timeout: int) -> tuple[int, str]: + """Run a command, returning ``(returncode, combined output)``.""" + try: + proc = subprocess.run( + cmd, cwd=str(cwd), capture_output=True, text=True, timeout=timeout, + ) + except subprocess.TimeoutExpired: + return 124, f"{' '.join(cmd)} timed out after {timeout}s" + except OSError as exc: + return 1, str(exc) + output = (proc.stdout or "") + (proc.stderr or "") + return proc.returncode, output.strip() + + +def _tail(text: str, limit: int = 1500) -> str: + """Keep the end of a command's output -- that's where errors are.""" + text = text.strip() + return text if len(text) <= limit else "…" + text[-limit:] + + +def _schedule_restart(service: str, delay: int) -> bool: + """Restart the unit after ``delay`` seconds, detached from this process. + + A direct ``systemctl restart`` would kill us mid-response, so the caller + would never learn whether the update worked. Detaching and sleeping lets + the tool result (or the HTTP response) reach the client first. + """ + if not shutil.which("systemctl"): + return False + try: + subprocess.Popen( + # Values go through argv, never interpolated into the script: + # `service` is a literal today, but a future change that made it + # configurable must not turn this into a shell injection. + [ + "sh", "-c", 'sleep "$1"; systemctl restart "$2"', + "sh", str(int(delay)), service, + ], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + start_new_session=True, + ) + return True + except OSError: + return False + + +def apply_update( + *, + restart: bool = True, + restart_delay: int = 5, + config_path: Path | None = None, + install: Installation | None = None, +) -> UpdateResult: + """Pull, reinstall dependencies, validate the config, then restart. + + The config validation is a **gate**: if the new revision cannot load the + operator's configuration (a newly required setting, a renamed key), the + checkout is rolled back to where it started and nothing is restarted. + An unattended update that bricks the server is worse than no update. + """ + if not _apply_lock.acquire(blocking=False): + busy = UpdateResult(ok=False) + busy.message = ( + "An update is already running on this server. Wait for it to " + "finish before starting another one." + ) + busy.steps.append(UpdateStep("preflight", False, busy.message)) + return busy + try: + return _apply_update_locked( + restart=restart, + restart_delay=restart_delay, + config_path=config_path, + install=install, + ) + finally: + _apply_lock.release() + + +def _apply_update_locked( + *, + restart: bool, + restart_delay: int, + config_path: Path | None, + install: Installation | None, +) -> UpdateResult: + install = install or detect_installation() + result = UpdateResult(ok=False) + + root = install.root + blockers = _self_update_blockers(install, root) + if blockers or root is None: + result.message = "Refusing to update: " + "; ".join(blockers) + result.steps.append(UpdateStep("preflight", False, result.message)) + return result + result.steps.append(UpdateStep("preflight", True, "git checkout is clean")) + + code, from_ref, _ = _git(root, "rev-parse", "HEAD") + if code != 0: + result.message = "Could not read the current revision." + result.steps.append(UpdateStep("read-head", False, result.message)) + return result + result.from_ref = from_ref[:12] + + branch = _default_branch(root) + code, out, err = _git(root, "pull", "--ff-only", "origin", branch) + if code != 0: + detail = _tail(err or out) + result.message = ( + f"git pull failed: {detail}. Nothing was changed." + ) + result.steps.append(UpdateStep("git-pull", False, detail)) + return result + code, to_ref, _ = _git(root, "rev-parse", "HEAD") + result.to_ref = to_ref[:12] if code == 0 else None + result.steps.append( + UpdateStep("git-pull", True, f"{result.from_ref} -> {result.to_ref}") + ) + + if result.from_ref == result.to_ref: + result.ok = True + result.message = "Already up to date; nothing to do." + return result + + def _rollback(reason: str) -> UpdateResult: + code, out, err = _git(root, "reset", "--hard", from_ref) + rolled = code == 0 + if rolled: + # Put the dependency set back too, so a half-applied update + # doesn't leave newer libraries against older code. + _run([*_pip_command(install), "install", "-e", "."], root, _PIP_TIMEOUT) + result.rolled_back = rolled + result.steps.append( + UpdateStep( + "rollback", rolled, + f"restored {result.from_ref}" if rolled else _tail(err or out), + ) + ) + result.ok = False + result.message = reason + ( + " The checkout was rolled back and the server was NOT restarted." + if rolled + else " ROLLBACK FAILED -- fix the checkout by hand before restarting." + ) + return result + + code, out = _run( + [*_pip_command(install), "install", "-e", "."], root, _PIP_TIMEOUT, + ) + if code != 0: + result.steps.append(UpdateStep("pip-install", False, _tail(out))) + return _rollback(f"Dependency install failed: {_tail(out, 400)}.") + result.steps.append(UpdateStep("pip-install", True, "dependencies up to date")) + + # Config gate. Run in a subprocess so the *new* code parses the config, + # not the copy this process imported at boot. + validate = [install.python, "-m", "beaconmcp", "validate-config"] + if config_path: + validate += ["--config", str(config_path)] + code, out = _run(validate, root, _GIT_TIMEOUT) + if code != 0: + result.steps.append(UpdateStep("validate-config", False, _tail(out))) + return _rollback( + f"The new revision cannot load your configuration: {_tail(out, 600)}" + ) + result.steps.append(UpdateStep("validate-config", True, "config still loads")) + + result.drift = config_drift(root, "HEAD", config_path) + result.ok = True + + if restart and install.service: + scheduled = _schedule_restart(install.service, restart_delay) + result.restart_scheduled = scheduled + result.restart_in_seconds = restart_delay if scheduled else 0 + result.steps.append( + UpdateStep( + "restart", scheduled, + f"systemctl restart {install.service} in {restart_delay}s" + if scheduled else "could not schedule a restart (no systemctl)", + ) + ) + + bits = [f"Updated {result.from_ref} -> {result.to_ref}."] + if result.restart_scheduled: + bits.append( + f"The service restarts in {restart_delay}s to run the new code." + ) + elif install.service: + bits.append(f"Restart it with: systemctl restart {install.service}") + else: + bits.append("Restart the server process to run the new code.") + if result.drift.new_env_vars: + bits.append( + "New environment variables you may need to set in .env: " + + ", ".join(result.drift.new_env_vars) + ) + if result.drift.new_config_keys: + bits.append( + "New beaconmcp.yaml settings are available: " + + ", ".join(result.drift.new_config_keys[:10]) + ) + result.message = " ".join(bits) + return result diff --git a/src/beaconmcp/utils.py b/src/beaconmcp/utils.py new file mode 100644 index 0000000..457d1b0 --- /dev/null +++ b/src/beaconmcp/utils.py @@ -0,0 +1,103 @@ +"""Shared response-shaping helpers for BeaconMCP tools. + +These helpers exist so individual tool modules don't each reimplement the two +cross-cutting patterns BeaconMCP relies on to stay token-efficient: + +* ``filter_fields`` lets callers trim tool output to the keys they need, cutting + the payload on the wire without forcing the server to ship a separate tool for + every projection. +* ``parse_since`` lets any time-windowed tool (``proxmox_get_tasks`` etc.) + accept either a relative duration (``"15m"``, ``"2h"``) or an absolute + epoch/ISO timestamp. +""" + +from __future__ import annotations + +import re +import time +from datetime import datetime, timezone +from typing import Any + + +def filter_fields(data: Any, fields: list[str] | None) -> Any: + """Return ``data`` trimmed to only the keys listed in ``fields``. + + - ``fields`` None / empty -> data is returned unchanged. + - dict -> returns a new dict with only the requested keys + (missing keys are skipped silently). + - list of dicts -> applies the same filter to every element. + - everything else -> returned unchanged (ints, strings, None, ...). + + Design notes: + * Missing keys are silently dropped rather than raising so callers can share + one ``fields`` list across tools that return slightly different shapes. + * Nested dicts/lists are kept as-is; this is a single-level projection on + purpose so callers keep predictable output shape. + """ + if not fields: + return data + keep = set(fields) + if isinstance(data, dict): + return {k: v for k, v in data.items() if k in keep} + if isinstance(data, list): + return [ + {k: v for k, v in item.items() if k in keep} + if isinstance(item, dict) + else item + for item in data + ] + return data + + +_SINCE_RE = re.compile(r"^\s*(\d+)\s*([smhd])\s*$", re.IGNORECASE) + + +def parse_since(value: Any, now: float | None = None) -> int | None: + """Parse a ``since`` argument into an epoch-seconds lower bound. + + Accepted forms: + * None / "" / 0 -> returns None (no lower bound). + * "" -> duration relative to ``now``. Units: s/m/h/d. + e.g. ``"15m"`` -> now - 900. + * int or numeric str -> treated as a unix epoch in seconds. + * ISO-8601 string -> parsed via ``datetime.fromisoformat``; naive values + are interpreted as UTC. + + Raises ``ValueError`` on anything else, so tools can surface a clean error + to the caller instead of silently misinterpreting input. + """ + if value in (None, "", 0): + return None + + current = now if now is not None else time.time() + + if isinstance(value, (int, float)): + return int(value) + + if isinstance(value, str): + match = _SINCE_RE.match(value) + if match: + n = int(match.group(1)) + unit = match.group(2).lower() + mult = {"s": 1, "m": 60, "h": 3600, "d": 86400}[unit] + return int(current - n * mult) + + # Numeric epoch as a string. + if value.strip().isdigit(): + return int(value.strip()) + + try: + dt = datetime.fromisoformat(value.replace("Z", "+00:00")) + except ValueError as exc: + raise ValueError( + f"Unrecognized 'since' value {value!r}. " + "Expected a duration like '15m'/'2h'/'1d', a unix epoch, or an ISO-8601 timestamp." + ) from exc + if dt.tzinfo is None: + dt = dt.replace(tzinfo=timezone.utc) + return int(dt.timestamp()) + + raise ValueError( + f"Unsupported 'since' type {type(value).__name__}. " + "Expected str, int, or float." + ) diff --git a/src/beaconmcp/wizard.py b/src/beaconmcp/wizard.py new file mode 100644 index 0000000..dbbc985 --- /dev/null +++ b/src/beaconmcp/wizard.py @@ -0,0 +1,1624 @@ +"""Interactive TUI config wizard for BeaconMCP (`beaconmcp init`). + +Three-pane layout: section menu on the left, section-specific form in the +middle, live ``beaconmcp.yaml`` preview on the right. The draft stays in +memory until the user saves — at which point the YAML gets written to +disk and any referenced ``${VAR}`` placeholders are appended to ``.env`` +with empty values for the user to fill in. + +The wizard is intentionally a **bootstrap** tool, not a full config +editor. It covers the capabilities (Proxmox, SSH, BMC), the critical +server fields (allowed_hosts / allowed_origins / trusted_proxies), and nothing else — +tweaks to dashboard settings or obscure fields happen by editing the +resulting YAML directly. Keeping the scope small means the preview pane +stays honest: what you see is the whole file. +""" + +from __future__ import annotations + +import os +import re +import sys +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Callable + +try: + from textual.app import App, ComposeResult + from textual.binding import Binding + from textual.containers import Horizontal, Vertical, VerticalScroll + from textual.screen import ModalScreen + from textual.widgets import ( + Button, + DataTable, + Footer, + Header, + Input, + Label, + ListItem, + ListView, + Static, + Switch, + TextArea, + ) +except ImportError as _exc: # pragma: no cover - import guard + _WIZARD_IMPORT_ERROR = _exc + + # Stubs so the class definitions below can be loaded without textual + # installed. Actual use is gated by `_WIZARD_IMPORT_ERROR` in + # `run_wizard`, which prints an install hint and exits. + class _Stub: + def __init__(self, *args: Any, **kwargs: Any) -> None: ... + def __class_getitem__(cls, item: Any) -> type: # noqa: D401 + return cls + + App = ComposeResult = Binding = _Stub # type: ignore[assignment,misc] + Horizontal = Vertical = VerticalScroll = _Stub # type: ignore[assignment,misc] + ModalScreen = _Stub # type: ignore[assignment,misc] + Button = DataTable = Footer = Header = Input = Label = _Stub # type: ignore[assignment,misc] + ListItem = ListView = Static = Switch = TextArea = _Stub # type: ignore[assignment,misc] +else: + _WIZARD_IMPORT_ERROR = None + + +# --------------------------------------------------------------------------- +# Draft data model (lenient mirror of beaconmcp.config dataclasses) +# --------------------------------------------------------------------------- + + +@dataclass +class PVENodeDraft: + name: str = "" + host: str = "" + token_id: str = "" + token_secret_env: str = "" # env var name, rendered as ${NAME} + # Raw (non-${VAR}) secret read from an existing YAML. Round-trip + # preservation for users who inlined their token — we re-emit it + # verbatim if token_secret_env is empty. + token_secret_literal: str = "" + + +@dataclass +class SSHDefaultsDraft: + user: str = "root" + port: int = 22 + # Exactly one of these two should be set in a valid draft. Wizard + # enforces it via the form, but the model accepts both empty so the + # user can start typing. + password_env: str = "" + key_file: str = "" + password_literal: str = "" + + +@dataclass +class SSHHostDraft: + name: str = "" + host: str = "" + user: str = "root" + port: int = 22 + password_env: str = "" + key_file: str = "" + password_literal: str = "" + # Per-host host-key overrides. Empty string means "inherit the global + # ssh-level setting"; strict is tri-state ("", "true", "false") for the + # same reason. + known_hosts: str = "" + strict_host_key_checking: str = "" + + +@dataclass +class SSHDraft: + enabled: bool = True + vmid_to_ip: str = "" + inherit_proxmox_nodes: bool = True + defaults: SSHDefaultsDraft = field(default_factory=SSHDefaultsDraft) + hosts: list[SSHHostDraft] = field(default_factory=list) + # Global host-key verification (see beaconmcp.yaml.example). Empty + # known_hosts + strict off keeps the permissive accept-any-key default. + known_hosts: str = "" + strict_host_key_checking: bool = False + + +@dataclass +class BMCDeviceDraft: + id: str = "" + type: str = "hp_ilo" + host: str = "" + user: str = "" + password_env: str = "" + jump_host: str = "" # references ssh.hosts[].name + password_literal: str = "" + + +@dataclass +class ServerDraft: + host: str = "0.0.0.0" + port: int = 8420 + allowed_hosts: list[str] = field(default_factory=lambda: ["127.0.0.1:*", "localhost:*", "[::1]:*"]) + allowed_origins: list[str] = field( + default_factory=lambda: [ + "https://assistant.ai", + "https://chatgpt.com", + "https://chat.mistral.ai", + "https://gemini.google.com", + ] + ) + trusted_proxies: list[str] = field(default_factory=lambda: ["127.0.0.1", "::1"]) + session_key_env: str = "" # env var name + allow_dynamic_registration: bool = False + # Optional paths; empty string means "use the built-in default" + # (tokens.db beside clients_file / /opt/beaconmcp/audit.log). + tokens_db: str = "" + audit_log: str = "" + # Named-token lifetime in seconds; empty means the 30-day default. + named_token_ttl: str = "" + + +@dataclass +class DashboardDraft: + enabled: bool = True + gemini_api_key_env: str = "" # env var name + limit_5h_usd: float = 2.0 + limit_week_usd: float = 10.0 + public_url: str = "" + mcp_mode: str = "local" # "local" | "remote" + + +@dataclass +class ConfigDraft: + server: ServerDraft = field(default_factory=ServerDraft) + pve_nodes: list[PVENodeDraft] = field(default_factory=list) + verify_ssl: bool = False # proxmox.verify_ssl + ssh: SSHDraft = field(default_factory=SSHDraft) + bmc_devices: list[BMCDeviceDraft] = field(default_factory=list) + dashboard: DashboardDraft = field(default_factory=DashboardDraft) + + def referenced_env_vars(self) -> list[str]: + """Collect every ``${VAR}`` name the draft references. + + Used when saving to append placeholders to ``.env`` so the user has + one file to fill in after the wizard exits. + """ + names: list[str] = [] + if self.server.session_key_env: + names.append(self.server.session_key_env) + for n in self.pve_nodes: + if n.token_secret_env: + names.append(n.token_secret_env) + if self.ssh.enabled: + if self.ssh.defaults.password_env: + names.append(self.ssh.defaults.password_env) + for h in self.ssh.hosts: + if h.password_env: + names.append(h.password_env) + for d in self.bmc_devices: + if d.password_env: + names.append(d.password_env) + if self.dashboard.enabled and self.dashboard.gemini_api_key_env: + names.append(self.dashboard.gemini_api_key_env) + # Dedupe while preserving order + seen: set[str] = set() + out: list[str] = [] + for name in names: + if name in seen: + continue + seen.add(name) + out.append(name) + return out + + +# --------------------------------------------------------------------------- +# YAML rendering — hand-rolled so we control comments and quoting precisely. +# --------------------------------------------------------------------------- + + +def _q(value: str) -> str: + """Quote a YAML scalar when it contains reserved characters.""" + if not value: + return '""' + if any(ch in value for ch in "!@:#&*`{}[]|>?,%"): + return f'"{value}"' + # A bare "-" (or "- foo") would parse as a sequence entry, not a scalar. + if value == "-" or value.startswith("- "): + return f'"{value}"' + if value.lower() in {"true", "false", "yes", "no", "on", "off", "null", "~"}: + return f'"{value}"' + return value + + +def render_yaml(draft: ConfigDraft) -> str: + """Render the draft as a ``beaconmcp.yaml`` string.""" + lines: list[str] = [] + lines.append("# Generated by `beaconmcp init`. Edit freely once saved.") + lines.append("version: 1") + lines.append("") + + # Server + lines.append("server:") + lines.append(f" host: {_q(draft.server.host)}") + lines.append(f" port: {draft.server.port}") + if draft.server.allowed_hosts: + lines.append(" allowed_hosts:") + for h in draft.server.allowed_hosts: + lines.append(f" - {_q(h)}") + if draft.server.allowed_origins: + lines.append(" allowed_origins:") + for o in draft.server.allowed_origins: + lines.append(f" - {o}") + if draft.server.trusted_proxies: + lines.append(" trusted_proxies:") + for p in draft.server.trusted_proxies: + lines.append(f" - {_q(p)}") + if draft.server.session_key_env: + lines.append(f" session_key: ${{{draft.server.session_key_env}}}") + if draft.server.allow_dynamic_registration: + lines.append(" allow_dynamic_registration: true") + if draft.server.tokens_db: + lines.append(f" tokens_db: {_q(draft.server.tokens_db)}") + if draft.server.audit_log: + lines.append(f" audit_log: {_q(draft.server.audit_log)}") + if draft.server.named_token_ttl: + lines.append(f" named_token_ttl: {draft.server.named_token_ttl}") + lines.append("") + + # Proxmox + if draft.pve_nodes: + lines.append("proxmox:") + lines.append(f" verify_ssl: {'true' if draft.verify_ssl else 'false'}") + lines.append(" nodes:") + for n in draft.pve_nodes: + lines.append(f" - name: {_q(n.name)}") + lines.append(f" host: {_q(n.host)}") + lines.append(f" token_id: {_q(n.token_id)}") + if n.token_secret_env: + secret = f"${{{n.token_secret_env}}}" + elif n.token_secret_literal: + secret = _q(n.token_secret_literal) + else: + secret = '""' + lines.append(f" token_secret: {secret}") + lines.append("") + + # SSH + if draft.ssh.enabled and ( + draft.ssh.hosts + or draft.ssh.inherit_proxmox_nodes + or draft.ssh.vmid_to_ip + ): + lines.append("ssh:") + if draft.ssh.vmid_to_ip: + lines.append(f" vmid_to_ip: {_q(draft.ssh.vmid_to_ip)}") + d = draft.ssh.defaults + if draft.ssh.inherit_proxmox_nodes or d.password_env or d.key_file: + lines.append(" defaults:") + lines.append(f" user: {_q(d.user)}") + if d.port and d.port != 22: + lines.append(f" port: {d.port}") + if d.key_file: + lines.append(f" key_file: {_q(d.key_file)}") + elif d.password_env: + lines.append(f" password: ${{{d.password_env}}}") + elif d.password_literal: + lines.append(f" password: {_q(d.password_literal)}") + if draft.ssh.inherit_proxmox_nodes: + lines.append(" inherit_proxmox_nodes: true") + if draft.ssh.known_hosts: + lines.append(f" known_hosts: {_q(draft.ssh.known_hosts)}") + if draft.ssh.strict_host_key_checking: + lines.append(" strict_host_key_checking: true") + if draft.ssh.hosts: + lines.append(" hosts:") + for h in draft.ssh.hosts: + lines.append(f" - name: {_q(h.name)}") + lines.append(f" host: {_q(h.host)}") + lines.append(f" user: {_q(h.user)}") + if h.port and h.port != 22: + lines.append(f" port: {h.port}") + if h.key_file: + lines.append(f" key_file: {_q(h.key_file)}") + elif h.password_env: + lines.append(f" password: ${{{h.password_env}}}") + elif h.password_literal: + lines.append(f" password: {_q(h.password_literal)}") + if h.known_hosts: + lines.append(f" known_hosts: {_q(h.known_hosts)}") + if h.strict_host_key_checking in ("true", "false"): + lines.append( + f" strict_host_key_checking: {h.strict_host_key_checking}" + ) + lines.append("") + + # BMC + if draft.bmc_devices: + lines.append("bmc:") + lines.append(" devices:") + for b in draft.bmc_devices: + lines.append(f" - id: {_q(b.id)}") + lines.append(f" type: {b.type}") + lines.append(f" host: {_q(b.host)}") + lines.append(f" user: {_q(b.user)}") + if b.password_env: + secret = f"${{{b.password_env}}}" + elif b.password_literal: + secret = _q(b.password_literal) + else: + secret = '""' + lines.append(f" password: {secret}") + if b.jump_host: + lines.append(f" jump_host: {_q(b.jump_host)}") + lines.append("") + + # Features (dashboard). Only emit when the user has departed from the + # defaults — keeps the generated file readable. + dash = draft.dashboard + non_default = ( + not dash.enabled + or dash.gemini_api_key_env + or dash.public_url + or dash.mcp_mode != "local" + or dash.limit_5h_usd != 2.0 + or dash.limit_week_usd != 10.0 + ) + if non_default: + lines.append("features:") + lines.append(" dashboard:") + lines.append(f" enabled: {'true' if dash.enabled else 'false'}") + if dash.gemini_api_key_env: + lines.append(f" gemini_api_key: ${{{dash.gemini_api_key_env}}}") + if dash.public_url: + lines.append(f" public_url: {_q(dash.public_url)}") + if dash.mcp_mode and dash.mcp_mode != "local": + lines.append(f" mcp_mode: {_q(dash.mcp_mode)}") + if dash.limit_5h_usd != 2.0 or dash.limit_week_usd != 10.0: + lines.append(" limits:") + lines.append(f" per_5h_usd: {dash.limit_5h_usd}") + lines.append(f" per_week_usd: {dash.limit_week_usd}") + lines.append("") + + return "\n".join(lines).rstrip() + "\n" + + +# --------------------------------------------------------------------------- +# YAML loading — inverse of render_yaml. Preserves ${VAR} placeholders +# rather than resolving them so the draft round-trips cleanly. +# --------------------------------------------------------------------------- + + +# Accept any shell-valid identifier, case-insensitive. Earlier versions +# of the loader required uppercase and silently dropped secrets that +# didn't match — a lossy round-trip we're not repeating. +_ENV_REF = re.compile(r"^\$\{([A-Za-z_][A-Za-z0-9_]*)\}$") + + +def _split_secret(value: Any) -> tuple[str, str]: + """Classify a secret scalar. + + Returns ``(env_name, literal)`` where exactly one side is populated: + - ``(NAME, "")`` when ``value`` is ``${NAME}`` + - ``("", raw)`` when ``value`` is a non-empty string that isn't a ``${VAR}`` + - ``("", "")`` when ``value`` is empty / missing / non-string + """ + if not isinstance(value, str): + return "", "" + s = value.strip() + if not s: + return "", "" + m = _ENV_REF.match(s) + if m: + return m.group(1), "" + return "", s + + +def load_yaml_into_draft(path: Path) -> ConfigDraft: + """Parse ``beaconmcp.yaml`` into a ``ConfigDraft`` for the wizard. + + Unknown or malformed sections are skipped rather than raised — this is + an editing convenience, not a validating loader. ``beaconmcp + validate-config`` remains the source of truth. + + Secret fields (``token_secret``, ``password``) are read as raw strings; + if they match ``${VAR}`` the env var name is stored in the ``*_env`` + draft field so saving renders the same placeholder back. + """ + import yaml # lazy: keep import cost off the module load path + + draft = ConfigDraft() + # Deliberately not swallowing yaml.YAMLError / OSError here: a + # parse failure used to silently return an empty draft, which — + # combined with the wizard overwriting on save — wiped users' + # configs. run_wizard() now catches and refuses to start. + raw = yaml.safe_load(path.read_text(encoding="utf-8")) or {} + if not isinstance(raw, dict): + return draft + + server = raw.get("server") or {} + if isinstance(server, dict): + if server.get("host"): + draft.server.host = str(server["host"]) + if isinstance(server.get("port"), int): + draft.server.port = int(server["port"]) + hosts = server.get("allowed_hosts") + if isinstance(hosts, list): + draft.server.allowed_hosts = [str(h) for h in hosts if h] + origins = server.get("allowed_origins") + if isinstance(origins, list): + draft.server.allowed_origins = [str(o) for o in origins if o] + proxies = server.get("trusted_proxies") + if isinstance(proxies, list): + draft.server.trusted_proxies = [str(p) for p in proxies if p] + sk_env, _ = _split_secret(server.get("session_key")) + draft.server.session_key_env = sk_env + draft.server.allow_dynamic_registration = bool( + server.get("allow_dynamic_registration", False) + ) + draft.server.tokens_db = str(server.get("tokens_db") or "") + draft.server.audit_log = str(server.get("audit_log") or "") + # ``0`` is meaningful (never expires) -- only absence means default. + draft.server.named_token_ttl = ( + str(server.get("named_token_ttl")) + if server.get("named_token_ttl") is not None + else "" + ) + + proxmox = raw.get("proxmox") or {} + if isinstance(proxmox, dict): + draft.verify_ssl = bool(proxmox.get("verify_ssl", False)) + for node in proxmox.get("nodes") or []: + if not isinstance(node, dict): + continue + env_name, literal = _split_secret(node.get("token_secret")) + draft.pve_nodes.append(PVENodeDraft( + name=str(node.get("name") or ""), + host=str(node.get("host") or ""), + token_id=str(node.get("token_id") or ""), + token_secret_env=env_name, + token_secret_literal=literal, + )) + + ssh = raw.get("ssh") + if isinstance(ssh, dict): + draft.ssh.enabled = True + draft.ssh.vmid_to_ip = str(ssh.get("vmid_to_ip") or "") + draft.ssh.inherit_proxmox_nodes = bool(ssh.get("inherit_proxmox_nodes", False)) + draft.ssh.known_hosts = str(ssh.get("known_hosts") or "") + draft.ssh.strict_host_key_checking = bool( + ssh.get("strict_host_key_checking", False) + ) + defaults = ssh.get("defaults") or {} + if isinstance(defaults, dict): + draft.ssh.defaults.user = str(defaults.get("user") or "root") + port = defaults.get("port") + draft.ssh.defaults.port = int(port) if isinstance(port, int) else 22 + draft.ssh.defaults.key_file = str(defaults.get("key_file") or "") + env_name, literal = _split_secret(defaults.get("password")) + draft.ssh.defaults.password_env = env_name + draft.ssh.defaults.password_literal = literal + for host in ssh.get("hosts") or []: + if not isinstance(host, dict): + continue + port = host.get("port") + env_name, literal = _split_secret(host.get("password")) + strict_raw = host.get("strict_host_key_checking") + draft.ssh.hosts.append(SSHHostDraft( + name=str(host.get("name") or ""), + host=str(host.get("host") or ""), + user=str(host.get("user") or "root"), + port=int(port) if isinstance(port, int) else 22, + password_env=env_name, + password_literal=literal, + key_file=str(host.get("key_file") or ""), + known_hosts=str(host.get("known_hosts") or ""), + strict_host_key_checking=( + "" if strict_raw is None else ("true" if strict_raw else "false") + ), + )) + else: + # No ssh block means SSH is disabled in the saved config. + draft.ssh.enabled = False + + bmc = raw.get("bmc") or {} + if isinstance(bmc, dict): + for dev in bmc.get("devices") or []: + if not isinstance(dev, dict): + continue + env_name, literal = _split_secret(dev.get("password")) + draft.bmc_devices.append(BMCDeviceDraft( + id=str(dev.get("id") or ""), + type=str(dev.get("type") or "hp_ilo"), + host=str(dev.get("host") or ""), + user=str(dev.get("user") or ""), + password_env=env_name, + password_literal=literal, + jump_host=str(dev.get("jump_host") or ""), + )) + + features = raw.get("features") or {} + if isinstance(features, dict): + dash_raw = features.get("dashboard") or {} + if isinstance(dash_raw, dict): + draft.dashboard.enabled = bool(dash_raw.get("enabled", True)) + gk_env, _ = _split_secret(dash_raw.get("gemini_api_key")) + draft.dashboard.gemini_api_key_env = gk_env + if dash_raw.get("public_url"): + draft.dashboard.public_url = str(dash_raw["public_url"]) + if dash_raw.get("mcp_mode"): + draft.dashboard.mcp_mode = str(dash_raw["mcp_mode"]).strip().lower() + limits = dash_raw.get("limits") or {} + if isinstance(limits, dict): + try: + if "per_5h_usd" in limits: + draft.dashboard.limit_5h_usd = float(limits["per_5h_usd"]) + if "per_week_usd" in limits: + draft.dashboard.limit_week_usd = float(limits["per_week_usd"]) + except (TypeError, ValueError): + pass + + return draft + + +# --------------------------------------------------------------------------- +# Textual app +# --------------------------------------------------------------------------- + +SECTIONS = [ + ("proxmox", "Proxmox nodes"), + ("ssh", "SSH"), + ("bmc", "BMC devices"), + ("server", "Server"), + ("dashboard", "Dashboard"), + ("save", "Save & exit"), +] + + +CSS = """ +Screen { + layout: vertical; +} + +#body { + layout: horizontal; + height: 1fr; +} + +#sidebar { + width: 24; + border-right: solid $primary-background; + padding: 1; +} + +#sidebar ListView { + background: $surface; + height: auto; +} + +#main { + width: 1fr; + padding: 1 2; +} + +#preview { + width: 55; + padding: 1; + border-left: solid $primary-background; +} + +#preview-title { + color: $text-muted; + text-style: bold; + margin-bottom: 1; +} + +#preview-area { + background: $surface; + border: solid $primary-background; + height: 1fr; +} + +.section-heading { + text-style: bold; + color: $accent; + margin-bottom: 1; +} + +.hint { + color: $text-muted; + margin-bottom: 1; +} + +DataTable { + height: auto; + max-height: 12; + margin: 1 0; +} + +/* In-panel list editors (allowed_hosts, allowed_origins, ...). Without + a bound, TextArea defaults to 1fr and several of them in the same + VerticalScroll fight each other into unusable thin strips. */ +.list-area { + height: 7; + margin-bottom: 1; +} + +.field-label { + color: $text-muted; + text-style: bold; + margin-top: 1; +} + +.form-row { + layout: horizontal; + height: auto; + margin: 0 0 1 0; +} + +.form-row Label { + width: 16; + padding: 1 1 0 0; +} + +.form-row Input { + width: 1fr; +} + +.form-actions { + layout: horizontal; + height: auto; + margin-top: 1; +} + +.form-actions Button { + margin-right: 1; +} + +Switch { + margin-right: 1; +} +""" + + +# --------------------------------------------------------------------------- +# Modals — forms for add/edit flows +# --------------------------------------------------------------------------- + + +class _FormModal(ModalScreen[dict[str, str] | None]): + """Generic modal with a list of (label, key, initial) fields. + + Returns a dict of entered values on Save, or None on Cancel. + """ + + BINDINGS = [ + Binding("escape", "cancel", "Cancel"), + Binding("ctrl+s", "save", "Save"), + ] + + def __init__( + self, + title: str, + fields: list[tuple[str, str, str]], + hint: str = "", + ) -> None: + super().__init__() + self._title = title + self._fields = fields + self._hint = hint + + def compose(self) -> ComposeResult: + with Vertical(id="modal-box"): + yield Static(self._title, classes="section-heading") + if self._hint: + yield Static(self._hint, classes="hint") + for label, key, initial in self._fields: + with Horizontal(classes="form-row"): + yield Label(label + ":") + yield Input(value=initial, id=f"f-{key}") + with Horizontal(classes="form-actions"): + yield Button("Save", id="ok", variant="primary") + yield Button("Cancel", id="cancel") + + def on_mount(self) -> None: + first = self.query(Input).first() + if first is not None: + first.focus() + + def on_button_pressed(self, event: Button.Pressed) -> None: + if event.button.id == "cancel": + self.dismiss(None) + elif event.button.id == "ok": + self.action_save() + + def action_save(self) -> None: + out: dict[str, str] = {} + for _label, key, _initial in self._fields: + out[key] = self.query_one(f"#f-{key}", Input).value.strip() + self.dismiss(out) + + def action_cancel(self) -> None: + self.dismiss(None) + + +# --------------------------------------------------------------------------- +# Section panels — each renders into the centre pane +# --------------------------------------------------------------------------- + + +class _ProxmoxPanel(Static): + def __init__(self, draft: ConfigDraft, on_change: Callable[[], None]) -> None: + super().__init__() + self.draft = draft + self.on_change = on_change + + def compose(self) -> ComposeResult: + yield Static("Proxmox nodes", classes="section-heading") + yield Static( + "One entry per Proxmox node. Use LAN IPs in `host:` — same " + "address will be reused for SSH inheritance.", + classes="hint", + ) + with Horizontal(classes="form-row"): + yield Label("verify_ssl:") + yield Switch(value=self.draft.verify_ssl, id="pve-verify-ssl") + yield DataTable(id="pve-table", cursor_type="row", zebra_stripes=True) + with Horizontal(classes="form-actions"): + yield Button("Add", id="pve-add", variant="primary") + yield Button("Edit", id="pve-edit") + yield Button("Delete", id="pve-delete", variant="error") + + def on_mount(self) -> None: + table = self.query_one("#pve-table", DataTable) + table.add_columns("name", "host", "token_id", "secret env") + self._refresh_table() + + def on_switch_changed(self, event: Switch.Changed) -> None: + if event.switch.id == "pve-verify-ssl": + self.draft.verify_ssl = event.value + self.on_change() + + def _refresh_table(self) -> None: + table = self.query_one("#pve-table", DataTable) + table.clear() + for n in self.draft.pve_nodes: + table.add_row( + n.name or "—", + n.host or "—", + n.token_id or "—", + f"${{{n.token_secret_env}}}" if n.token_secret_env else "—", + ) + + def on_button_pressed(self, event: Button.Pressed) -> None: + if event.button.id == "pve-add": + self._open_form(None) + elif event.button.id == "pve-edit": + idx = self._selected_row() + if idx is not None: + self._open_form(idx) + elif event.button.id == "pve-delete": + idx = self._selected_row() + if idx is not None: + del self.draft.pve_nodes[idx] + self._refresh_table() + self.on_change() + + def _selected_row(self) -> int | None: + table = self.query_one("#pve-table", DataTable) + if table.cursor_row is None or not self.draft.pve_nodes: + return None + idx = table.cursor_row + if 0 <= idx < len(self.draft.pve_nodes): + return idx + return None + + def _open_form(self, idx: int | None) -> None: + existing = self.draft.pve_nodes[idx] if idx is not None else PVENodeDraft() + default_env = existing.token_secret_env or ( + f"PVE{len(self.draft.pve_nodes) + 1}_TOKEN_SECRET" if idx is None else "" + ) + modal = _FormModal( + title="Proxmox node" if idx is None else f"Edit {existing.name or 'node'}", + hint=( + "host: LAN IP of the node (e.g. 10.0.0.1). token_id: the " + "Proxmox API token ID in user@realm!tokenname shape. The " + "secret itself lives in .env — type the env-var name here." + ), + fields=[ + ("name", "name", existing.name), + ("host", "host", existing.host), + ("token id", "token_id", existing.token_id or "root@pam!beaconmcp"), + ("secret env", "token_secret_env", default_env), + ], + ) + + def after(result: dict[str, str] | None) -> None: + if result is None: + return + entry = PVENodeDraft( + name=result["name"], + host=result["host"], + token_id=result["token_id"], + token_secret_env=result["token_secret_env"], + ) + if idx is None: + self.draft.pve_nodes.append(entry) + else: + self.draft.pve_nodes[idx] = entry + self._refresh_table() + self.on_change() + + self.app.push_screen(modal, after) + + +class _SSHPanel(Static): + def __init__(self, draft: ConfigDraft, on_change: Callable[[], None]) -> None: + super().__init__() + self.draft = draft + self.on_change = on_change + + def compose(self) -> ComposeResult: + yield Static("SSH capability", classes="section-heading") + yield Static( + "Flip inheritance on to reach every Proxmox node via SSH using " + "the `defaults` creds — no per-node duplication.", + classes="hint", + ) + + ssh = self.draft.ssh + with Horizontal(classes="form-row"): + yield Label("Enable SSH:") + yield Switch(value=ssh.enabled, id="ssh-enabled") + with Horizontal(classes="form-row"): + yield Label("vmid_to_ip:") + yield Input( + value=ssh.vmid_to_ip, + placeholder="e.g. 192.168.1.{id} (leave empty to disable)", + id="ssh-vmid", + ) + with Horizontal(classes="form-row"): + yield Label("Inherit PVE nodes:") + yield Switch(value=ssh.inherit_proxmox_nodes, id="ssh-inherit") + + yield Static("Host-key verification", classes="section-heading") + yield Static( + "Defaults to accept-any-key (trusted LAN). Point known_hosts at " + "an OpenSSH file to pin keys, or flip strict to use " + "~/.ssh/known_hosts. Hosts can override per-entry in their form.", + classes="hint", + ) + with Horizontal(classes="form-row"): + yield Label("known_hosts:") + yield Input( + value=ssh.known_hosts, + placeholder="/etc/beaconmcp/known_hosts (empty = accept any)", + id="ssh-known-hosts", + ) + with Horizontal(classes="form-row"): + yield Label("Strict host keys:") + yield Switch(value=ssh.strict_host_key_checking, id="ssh-strict") + + yield Static("Default credentials", classes="section-heading") + yield Static( + "Used for inherited Proxmox entries. Provide exactly one of " + "key_file OR password (env var name).", + classes="hint", + ) + with Horizontal(classes="form-row"): + yield Label("Default user:") + yield Input(value=ssh.defaults.user, id="ssh-def-user") + with Horizontal(classes="form-row"): + yield Label("Key file:") + yield Input( + value=ssh.defaults.key_file, + placeholder="~/.ssh/beaconmcp", + id="ssh-def-key", + ) + with Horizontal(classes="form-row"): + yield Label("Password env:") + yield Input( + value=ssh.defaults.password_env, + placeholder="(only if no key_file)", + id="ssh-def-pw", + ) + + yield Static("Explicit hosts", classes="section-heading") + yield Static( + "Targets outside your Proxmox cluster (VPS, bastion, remote " + "node with its own creds). Names may match a Proxmox node — " + "the explicit entry shadows inheritance.", + classes="hint", + ) + yield DataTable(id="ssh-table", cursor_type="row", zebra_stripes=True) + with Horizontal(classes="form-actions"): + yield Button("Add host", id="ssh-add", variant="primary") + yield Button("Edit", id="ssh-edit") + yield Button("Delete", id="ssh-delete", variant="error") + + def on_mount(self) -> None: + table = self.query_one("#ssh-table", DataTable) + table.add_columns("name", "host", "user", "auth") + self._refresh_table() + + def _refresh_table(self) -> None: + table = self.query_one("#ssh-table", DataTable) + table.clear() + for h in self.draft.ssh.hosts: + auth = h.key_file or (f"${{{h.password_env}}}" if h.password_env else "—") + table.add_row(h.name or "—", h.host or "—", h.user or "—", auth) + + def on_switch_changed(self, event: Switch.Changed) -> None: + if event.switch.id == "ssh-enabled": + self.draft.ssh.enabled = event.value + elif event.switch.id == "ssh-inherit": + self.draft.ssh.inherit_proxmox_nodes = event.value + elif event.switch.id == "ssh-strict": + self.draft.ssh.strict_host_key_checking = event.value + self.on_change() + + def on_input_changed(self, event: Input.Changed) -> None: + ssh = self.draft.ssh + if event.input.id == "ssh-vmid": + ssh.vmid_to_ip = event.value.strip() + elif event.input.id == "ssh-known-hosts": + ssh.known_hosts = event.value.strip() + elif event.input.id == "ssh-def-user": + ssh.defaults.user = event.value.strip() or "root" + elif event.input.id == "ssh-def-key": + ssh.defaults.key_file = event.value.strip() + if ssh.defaults.key_file: + ssh.defaults.password_env = "" + elif event.input.id == "ssh-def-pw": + ssh.defaults.password_env = event.value.strip() + if ssh.defaults.password_env: + ssh.defaults.key_file = "" + self.on_change() + + def on_button_pressed(self, event: Button.Pressed) -> None: + if event.button.id == "ssh-add": + self._open_form(None) + elif event.button.id == "ssh-edit": + idx = self._selected_row() + if idx is not None: + self._open_form(idx) + elif event.button.id == "ssh-delete": + idx = self._selected_row() + if idx is not None: + del self.draft.ssh.hosts[idx] + self._refresh_table() + self.on_change() + + def _selected_row(self) -> int | None: + table = self.query_one("#ssh-table", DataTable) + if table.cursor_row is None or not self.draft.ssh.hosts: + return None + idx = table.cursor_row + if 0 <= idx < len(self.draft.ssh.hosts): + return idx + return None + + def _open_form(self, idx: int | None) -> None: + existing = self.draft.ssh.hosts[idx] if idx is not None else SSHHostDraft() + modal = _FormModal( + title="SSH host" if idx is None else f"Edit {existing.name or 'host'}", + hint=( + "key_file or password env — one of the two, not both. " + "Leave port empty for 22. known_hosts / strict host keys " + "(true/false) are optional overrides of the global SSH " + "settings; leave empty to inherit." + ), + fields=[ + ("name", "name", existing.name), + ("host", "host", existing.host), + ("user", "user", existing.user), + ("port", "port", str(existing.port) if existing.port and existing.port != 22 else ""), + ("key_file", "key_file", existing.key_file), + ("password env", "password_env", existing.password_env), + ("known_hosts", "known_hosts", existing.known_hosts), + ( + "strict host keys", + "strict_host_key_checking", + existing.strict_host_key_checking, + ), + ], + ) + + def after(result: dict[str, str] | None) -> None: + if result is None: + return + port = int(result["port"]) if result["port"].isdigit() else 22 + key = result["key_file"] + pw = result["password_env"] + # Enforce mutual exclusion + if key and pw: + pw = "" + strict = result["strict_host_key_checking"].strip().lower() + if strict not in ("true", "false"): + strict = "" # inherit the global setting + entry = SSHHostDraft( + name=result["name"], + host=result["host"], + user=result["user"] or "root", + port=port, + key_file=key, + password_env=pw, + known_hosts=result["known_hosts"].strip(), + strict_host_key_checking=strict, + ) + if idx is None: + self.draft.ssh.hosts.append(entry) + else: + self.draft.ssh.hosts[idx] = entry + self._refresh_table() + self.on_change() + + self.app.push_screen(modal, after) + + +class _BMCPanel(Static): + def __init__(self, draft: ConfigDraft, on_change: Callable[[], None]) -> None: + super().__init__() + self.draft = draft + self.on_change = on_change + + def compose(self) -> ComposeResult: + yield Static("BMC devices", classes="section-heading") + yield Static( + "HP iLO, IPMI, iDRAC or Supermicro. `jump_host` (optional) " + "references an ssh.hosts[] entry by name.", + classes="hint", + ) + yield DataTable(id="bmc-table", cursor_type="row", zebra_stripes=True) + with Horizontal(classes="form-actions"): + yield Button("Add", id="bmc-add", variant="primary") + yield Button("Edit", id="bmc-edit") + yield Button("Delete", id="bmc-delete", variant="error") + + def on_mount(self) -> None: + table = self.query_one("#bmc-table", DataTable) + table.add_columns("id", "type", "host", "jump_host") + self._refresh_table() + + def _refresh_table(self) -> None: + table = self.query_one("#bmc-table", DataTable) + table.clear() + for d in self.draft.bmc_devices: + table.add_row(d.id or "—", d.type, d.host or "—", d.jump_host or "—") + + def on_button_pressed(self, event: Button.Pressed) -> None: + if event.button.id == "bmc-add": + self._open_form(None) + elif event.button.id == "bmc-edit": + idx = self._selected_row() + if idx is not None: + self._open_form(idx) + elif event.button.id == "bmc-delete": + idx = self._selected_row() + if idx is not None: + del self.draft.bmc_devices[idx] + self._refresh_table() + self.on_change() + + def _selected_row(self) -> int | None: + table = self.query_one("#bmc-table", DataTable) + if table.cursor_row is None or not self.draft.bmc_devices: + return None + idx = table.cursor_row + if 0 <= idx < len(self.draft.bmc_devices): + return idx + return None + + def _open_form(self, idx: int | None) -> None: + existing = self.draft.bmc_devices[idx] if idx is not None else BMCDeviceDraft() + default_env = existing.password_env or ( + f"BMC{len(self.draft.bmc_devices) + 1}_PASSWORD" if idx is None else "" + ) + modal = _FormModal( + title="BMC device" if idx is None else f"Edit {existing.id or 'device'}", + hint=( + "type: hp_ilo | ipmi | idrac | supermicro. jump_host is the " + "name of an ssh.hosts[] entry used to tunnel into a private " + "management VLAN (leave empty for direct access)." + ), + fields=[ + ("id", "id", existing.id), + ("type", "type", existing.type), + ("host", "host", existing.host), + ("user", "user", existing.user or "Administrator"), + ("password env", "password_env", default_env), + ("jump_host", "jump_host", existing.jump_host), + ], + ) + + def after(result: dict[str, str] | None) -> None: + if result is None: + return + entry = BMCDeviceDraft( + id=result["id"], + type=result["type"] or "hp_ilo", + host=result["host"], + user=result["user"], + password_env=result["password_env"], + jump_host=result["jump_host"], + ) + if idx is None: + self.draft.bmc_devices.append(entry) + else: + self.draft.bmc_devices[idx] = entry + self._refresh_table() + self.on_change() + + self.app.push_screen(modal, after) + + +class _ServerPanel(Static): + def __init__(self, draft: ConfigDraft, on_change: Callable[[], None]) -> None: + super().__init__() + self.draft = draft + self.on_change = on_change + + def compose(self) -> ComposeResult: + yield Static("Server", classes="section-heading") + yield Static( + "Bind address/port, DNS-rebinding allowlist + CORS origins. " + "One entry per line for the list fields. trusted_proxies accepts " + "IPs/CIDRs and the token 'cloudflare'.", + classes="hint", + ) + srv = self.draft.server + with Horizontal(classes="form-row"): + yield Label("host:") + yield Input(value=srv.host, id="srv-host", placeholder="0.0.0.0") + with Horizontal(classes="form-row"): + yield Label("port:") + yield Input(value=str(srv.port), id="srv-port", placeholder="8420") + yield Static("allowed_hosts (one per line)", classes="field-label") + yield TextArea( + "\n".join(srv.allowed_hosts), + id="srv-hosts", + show_line_numbers=False, + classes="list-area", + ) + yield Static("allowed_origins (one per line)", classes="field-label") + yield TextArea( + "\n".join(srv.allowed_origins), + id="srv-origins", + show_line_numbers=False, + classes="list-area", + ) + yield Static("trusted_proxies (one per line)", classes="field-label") + yield TextArea( + "\n".join(srv.trusted_proxies), + id="srv-trusted-proxies", + show_line_numbers=False, + classes="list-area", + ) + yield Static( + "Session key env (${VAR} name) — leave empty to auto-generate", + classes="field-label", + ) + yield Input( + value=srv.session_key_env, + id="srv-sessionkey", + placeholder="BEACONMCP_SESSION_KEY", + ) + with Horizontal(classes="form-row"): + yield Label("Dynamic reg:") + yield Switch( + value=srv.allow_dynamic_registration, id="srv-dynreg" + ) + yield Static( + "Dynamic registration lets clients without a pre-provisioned " + "client_id (notably ChatGPT) self-register via a dashboard-minted " + "slug. Off by default.", + classes="hint", + ) + yield Static( + "tokens_db — SQLite file persisting named API tokens across " + "restarts. Empty = tokens.db beside clients_file.", + classes="field-label", + ) + yield Input( + value=srv.tokens_db, + id="srv-tokens-db", + placeholder="/opt/beaconmcp/tokens.db", + ) + yield Static( + "audit_log — JSON-lines audit file. Empty = " + "/opt/beaconmcp/audit.log; '-' keeps stderr only.", + classes="field-label", + ) + yield Input( + value=srv.audit_log, + id="srv-audit-log", + placeholder="/opt/beaconmcp/audit.log", + ) + yield Static( + "named_token_ttl — lifetime (seconds) of named API tokens. " + "Empty = 30 days; 0 = never expires (revoke-only). Internal " + "session bearers stay 24 h.", + classes="field-label", + ) + yield Input( + value=srv.named_token_ttl, + id="srv-named-ttl", + placeholder="2592000", + ) + + def on_input_changed(self, event: Input.Changed) -> None: + srv = self.draft.server + if event.input.id == "srv-host": + srv.host = event.value.strip() or "0.0.0.0" + elif event.input.id == "srv-port": + raw = event.value.strip() + if raw.isdigit(): + srv.port = int(raw) + elif event.input.id == "srv-sessionkey": + srv.session_key_env = event.value.strip() + elif event.input.id == "srv-tokens-db": + srv.tokens_db = event.value.strip() + elif event.input.id == "srv-audit-log": + srv.audit_log = event.value.strip() + elif event.input.id == "srv-named-ttl": + raw = event.value.strip() + srv.named_token_ttl = raw if raw.isdigit() else "" + self.on_change() + + def on_switch_changed(self, event: Switch.Changed) -> None: + if event.switch.id == "srv-dynreg": + self.draft.server.allow_dynamic_registration = event.value + self.on_change() + + def on_text_area_changed(self, event: TextArea.Changed) -> None: + lines = [line.strip() for line in event.text_area.text.splitlines() if line.strip()] + if event.text_area.id == "srv-hosts": + self.draft.server.allowed_hosts = lines + elif event.text_area.id == "srv-origins": + self.draft.server.allowed_origins = lines + elif event.text_area.id == "srv-trusted-proxies": + self.draft.server.trusted_proxies = lines + self.on_change() + + +class _DashboardPanel(Static): + def __init__(self, draft: ConfigDraft, on_change: Callable[[], None]) -> None: + super().__init__() + self.draft = draft + self.on_change = on_change + + def compose(self) -> ComposeResult: + yield Static("Dashboard", classes="section-heading") + yield Static( + "Optional web panel (/app/login, /app/chat, /app/tokens). " + "Runs an AI chat backed by Gemini; leave the key empty to " + "disable the chat while keeping the panel for token management.", + classes="hint", + ) + dash = self.draft.dashboard + with Horizontal(classes="form-row"): + yield Label("Enabled:") + yield Switch(value=dash.enabled, id="dash-enabled") + yield Static("Gemini API key env (${VAR} name)") + yield Input( + value=dash.gemini_api_key_env, + id="dash-gemini", + placeholder="GEMINI_API_KEY (leave empty to disable chat)", + ) + with Horizontal(classes="form-row"): + yield Label("Public URL:") + yield Input( + value=dash.public_url, + id="dash-url", + placeholder="https://beacon.example.com (for OAuth redirects)", + ) + with Horizontal(classes="form-row"): + yield Label("MCP mode:") + yield Input( + value=dash.mcp_mode, + id="dash-mode", + placeholder="local | remote", + ) + yield Static( + "Spending caps for Gemini chat (USD). Dashboard stops answering " + "when either threshold is hit.", + classes="hint", + ) + with Horizontal(classes="form-row"): + yield Label("5h limit $:") + yield Input( + value=str(dash.limit_5h_usd), + id="dash-5h", + placeholder="2.0", + ) + with Horizontal(classes="form-row"): + yield Label("Weekly limit $:") + yield Input( + value=str(dash.limit_week_usd), + id="dash-week", + placeholder="10.0", + ) + + def on_switch_changed(self, event: Switch.Changed) -> None: + if event.switch.id == "dash-enabled": + self.draft.dashboard.enabled = event.value + self.on_change() + + def on_input_changed(self, event: Input.Changed) -> None: + dash = self.draft.dashboard + if event.input.id == "dash-gemini": + dash.gemini_api_key_env = event.value.strip() + elif event.input.id == "dash-url": + dash.public_url = event.value.strip() + elif event.input.id == "dash-mode": + mode = event.value.strip().lower() + if mode in ("local", "remote"): + dash.mcp_mode = mode + elif not mode: + dash.mcp_mode = "local" + elif event.input.id in ("dash-5h", "dash-week"): + try: + val = float(event.value.strip()) + except ValueError: + return + if event.input.id == "dash-5h": + dash.limit_5h_usd = val + else: + dash.limit_week_usd = val + self.on_change() + + +class _SavePanel(Static): + def __init__( + self, + draft: ConfigDraft, + yaml_path: Path, + env_path: Path, + on_save: Callable[[Path, Path], None], + ) -> None: + super().__init__() + self.draft = draft + self.yaml_path = yaml_path + self.env_path = env_path + self.on_save = on_save + + def compose(self) -> ComposeResult: + yield Static("Save & exit", classes="section-heading") + yield Static( + f"YAML will be written to: {self.yaml_path}\n" + f".env will be extended at: {self.env_path}", + classes="hint", + ) + yield Static("Referenced env vars (need values in .env):", classes="section-heading") + refs = self.draft.referenced_env_vars() + yield Static("\n".join(f" - {n}" for n in refs) if refs else "(none)") + with Horizontal(classes="form-actions"): + yield Button("Save config", id="save", variant="primary") + yield Button("Cancel", id="cancel") + yield Static("", id="save-status", classes="hint") + + def on_button_pressed(self, event: Button.Pressed) -> None: + if event.button.id == "cancel": + self.app.exit() + elif event.button.id == "save": + try: + self.on_save(self.yaml_path, self.env_path) + self.query_one("#save-status", Static).update( + f"Saved. Edit {self.env_path} to fill in the secrets, then " + "run `beaconmcp validate-config`." + ) + except Exception as exc: # noqa: BLE001 + self.query_one("#save-status", Static).update( + f"[red]Save failed: {exc}[/red]" + ) + + +# --------------------------------------------------------------------------- +# Main app +# --------------------------------------------------------------------------- + + +class ConfigWizardApp(App[None]): + CSS = CSS + TITLE = "BeaconMCP — config wizard" + SUB_TITLE = "beaconmcp init" + + BINDINGS = [ + Binding("q", "quit", "Quit"), + Binding("ctrl+s", "quick_save", "Save"), + ] + + def __init__( + self, + yaml_path: Path, + env_path: Path, + draft: ConfigDraft | None = None, + ) -> None: + super().__init__() + self.draft = draft if draft is not None else ConfigDraft() + self.yaml_path = yaml_path + self.env_path = env_path + self._current_section = "proxmox" + + def compose(self) -> ComposeResult: + yield Header(show_clock=False) + with Horizontal(id="body"): + with Vertical(id="sidebar"): + yield Static("Sections", classes="section-heading") + yield ListView( + *[ListItem(Label(label), id=f"sect-{key}") for key, label in SECTIONS], + id="sections", + ) + yield Static("", classes="hint") + yield Static( + "Tip: arrow keys to move, enter to open a section, " + "tab to jump between panes.", + classes="hint", + ) + with VerticalScroll(id="main"): + yield Static("Select a section on the left.", id="main-content") + with Vertical(id="preview"): + yield Static("beaconmcp.yaml (live preview)", id="preview-title") + yield TextArea("", id="preview-area", read_only=True, show_line_numbers=False) + yield Footer() + + def on_mount(self) -> None: + lv = self.query_one("#sections", ListView) + lv.focus() + self._show_section("proxmox") + self._refresh_preview() + + def on_list_view_selected(self, event: ListView.Selected) -> None: + item_id = event.item.id or "" + if not item_id.startswith("sect-"): + return + self._show_section(item_id[len("sect-"):]) + + def on_list_view_highlighted(self, event: ListView.Highlighted) -> None: + # Highlight on arrow keys also swaps the panel, so the user doesn't + # have to press enter to preview each section. + if event.item is None: + return + item_id = event.item.id or "" + if item_id.startswith("sect-"): + self._show_section(item_id[len("sect-"):]) + + def _show_section(self, key: str) -> None: + self._current_section = key + container = self.query_one("#main", VerticalScroll) + container.remove_children() + panel: Static + if key == "proxmox": + panel = _ProxmoxPanel(self.draft, self._refresh_preview) + elif key == "ssh": + panel = _SSHPanel(self.draft, self._refresh_preview) + elif key == "bmc": + panel = _BMCPanel(self.draft, self._refresh_preview) + elif key == "server": + panel = _ServerPanel(self.draft, self._refresh_preview) + elif key == "dashboard": + panel = _DashboardPanel(self.draft, self._refresh_preview) + elif key == "save": + panel = _SavePanel( + self.draft, self.yaml_path, self.env_path, self._write_files + ) + else: + panel = Static("Unknown section.") + container.mount(panel) + + def _refresh_preview(self) -> None: + self.query_one("#preview-area", TextArea).text = render_yaml(self.draft) + + def action_quick_save(self) -> None: + # Triggered by Ctrl+S anywhere in the app. Doesn't exit — user can + # keep editing. Status gets reflected on the save panel if open. + try: + self._write_files(self.yaml_path, self.env_path) + except Exception: # noqa: BLE001 + pass + + def _write_files(self, yaml_path: Path, env_path: Path) -> None: + yaml_path.parent.mkdir(parents=True, exist_ok=True) + _backup_existing(yaml_path) + yaml_path.write_text(render_yaml(self.draft), encoding="utf-8") + _merge_env_placeholders(env_path, self.draft.referenced_env_vars()) + + +def _backup_existing(path: Path) -> None: + """Copy ``path`` to a timestamped sibling before overwriting. + + The wizard rewrites the YAML from the draft, which means any field + the loader didn't understand is lost on save. A timestamped backup + makes that recoverable instead of catastrophic. + """ + if not path.exists(): + return + from datetime import datetime + stamp = datetime.now().strftime("%Y%m%d-%H%M%S") + backup = path.with_name(f"{path.name}.bak.{stamp}") + try: + backup.write_bytes(path.read_bytes()) + except OSError: + # A failed backup shouldn't block saving, but it shouldn't + # silently succeed either — surface it via preview later. + pass + + +def _merge_env_placeholders(env_path: Path, names: list[str]) -> None: + """Ensure every referenced env var has a line in ``.env``. + + Existing values are preserved. Missing names get an empty placeholder + with a comment noting the wizard added them. Passing an empty list is + a no-op. + """ + if not names: + return + env_path.parent.mkdir(parents=True, exist_ok=True) + existing = {} + if env_path.exists(): + for line in env_path.read_text(encoding="utf-8").splitlines(): + if "=" in line and not line.lstrip().startswith("#"): + key = line.split("=", 1)[0].strip() + if key: + existing[key] = True + to_add = [n for n in names if n not in existing] + if not to_add: + return + with env_path.open("a", encoding="utf-8") as f: + if env_path.stat().st_size and not env_path.read_text(encoding="utf-8").endswith("\n"): + f.write("\n") + f.write("\n# Added by `beaconmcp init` — fill these in.\n") + for name in to_add: + f.write(f"{name}=\n") + # This file is where every BMC password, SSH password and Proxmox API + # token ends up. `open("a")` creates it with the process umask (0644 on + # a stock Debian box), so lock it down explicitly -- same treatment as + # clients.json / tokens.db. Best-effort: never block a save on chmod. + try: + os.chmod(env_path, 0o600) + except OSError: + pass + + +# --------------------------------------------------------------------------- +# Entry point used by the CLI +# --------------------------------------------------------------------------- + + +def run_wizard( + yaml_path: Path | None = None, + env_path: Path | None = None, + *, + start_blank: bool = False, +) -> int: + """Launch the wizard. Returns a process exit code. + + If ``yaml_path`` already exists and ``start_blank`` is False, the file + is parsed into the draft so the user edits the existing config + in-place. ``start_blank=True`` discards whatever's on disk — the + caller is responsible for confirming this is safe. + """ + if _WIZARD_IMPORT_ERROR is not None: + print( + "The interactive wizard needs the optional 'textual' dependency.\n" + "Install it with:\n" + " pip install 'beaconmcp[wizard]'\n" + f"Import failed with: {_WIZARD_IMPORT_ERROR}", + ) + return 1 + + yaml_path = yaml_path or Path(os.environ.get("BEACONMCP_CONFIG", "beaconmcp.yaml")) + env_path = env_path or Path(".env") + + draft: ConfigDraft | None = None + if yaml_path.exists() and not start_blank: + try: + draft = load_yaml_into_draft(yaml_path) + except Exception as exc: # noqa: BLE001 - YAML, OS, encoding, ... + print( + f"ERROR: could not parse existing {yaml_path}: {exc}\n" + f"Refusing to start the wizard because saving would " + f"overwrite the file with an empty config.\n" + f"Fix the YAML by hand, or rerun with --blank to start " + f"fresh (back up the file first).", + file=sys.stderr, + ) + return 1 + + ConfigWizardApp(yaml_path=yaml_path, env_path=env_path, draft=draft).run() + return 0 diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..646b098 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,29 @@ +"""Shared pytest configuration. + +``test_integration.py`` is a standalone script with its own ``TestRunner`` +and is designed to be executed as ``python tests/test_integration.py`` +against live infrastructure. Its ``test_*`` functions take a ``runner`` +and ``tools`` positional argument, which pytest tries (and fails) to +resolve as fixtures -- producing a pile of ERRORs on every ``pytest`` +run even when nothing destructive would have happened. + +We tell pytest to skip that file at collection time unless the opt-in +environment variable ``BEACONMCP_RUN_INTEGRATION=1`` is set. The script +path stays runnable as a plain Python program. +""" + +from __future__ import annotations + +import os + + +def _run_integration_enabled() -> bool: + return os.environ.get("BEACONMCP_RUN_INTEGRATION", "").strip().lower() in ( + "1", "true", "yes", "on", + ) + + +# Collected by pytest; a relative path listed here is skipped entirely. +collect_ignore: list[str] = [] +if not _run_integration_enabled(): + collect_ignore.append("test_integration.py") diff --git a/tests/test_audit.py b/tests/test_audit.py new file mode 100644 index 0000000..cb96d65 --- /dev/null +++ b/tests/test_audit.py @@ -0,0 +1,69 @@ +"""Tests for the JSON-lines audit logger.""" + +from __future__ import annotations + +import json +import logging + +from beaconmcp import audit + + +def test_emit_writes_json_line(caplog) -> None: + with caplog.at_level(logging.INFO, logger="beaconmcp.audit"): + audit.emit("auth.login", client_id="c1", outcome="ok") + rec = caplog.records[-1] + data = json.loads(rec.getMessage()) + assert data["event"] == "auth.login" + assert data["client_id"] == "c1" + assert data["outcome"] == "ok" + assert "ts" in data + + +def test_redacts_sensitive_fields(caplog) -> None: + with caplog.at_level(logging.INFO, logger="beaconmcp.audit"): + audit.emit( + "tool.call", + tool="ssh_run", + args={"host": "pve1", "password": "hunter2", "nested": {"token": "abc"}}, + ) + data = json.loads(caplog.records[-1].getMessage()) + assert data["args"]["host"] == "pve1" + assert data["args"]["password"] == "***" + assert data["args"]["nested"]["token"] == "***" + + +def test_emit_never_raises(monkeypatch) -> None: + def boom(_msg: str) -> None: + raise RuntimeError("sink died") + + monkeypatch.setattr(audit._logger, "info", boom) + # Should swallow the exception -- audit must never break a request. + audit.emit("anything", x=1) + + +def test_redacts_top_level_fields(caplog) -> None: + with caplog.at_level(logging.INFO, logger="beaconmcp.audit"): + audit.emit("auth.authorize.ok", client_id="c1", totp="123456") + data = json.loads(caplog.records[-1].getMessage()) + assert data["client_id"] == "c1" + assert data["totp"] == "***" + + +def test_compact_args_collapses_long_strings() -> None: + out = audit.compact_args({"node": "pve1", "blob": "x" * 500}) + assert out["node"] == "pve1" + assert out["blob"] == "" + + +def test_compact_args_always_collapses_content_keys() -> None: + # Short command lines / file payloads can embed secrets that key-based + # redaction cannot see -- they must never reach the sink verbatim. + out = audit.compact_args({ + "command": "mysql -u root -phunter2 db", + "content": "DB_PASSWORD=hunter2", + "path": "/root/.env", + }) + assert "hunter2" not in str(out) + assert out["command"] == "" + assert out["content"] == "" + assert out["path"] == "/root/.env" diff --git a/tests/test_audit_log_wiring.py b/tests/test_audit_log_wiring.py new file mode 100644 index 0000000..d992016 --- /dev/null +++ b/tests/test_audit_log_wiring.py @@ -0,0 +1,64 @@ +"""Tests for the audit-log file wiring in the CLI entrypoint.""" + +from __future__ import annotations + +import logging +import stat +from pathlib import Path + +import pytest + +from beaconmcp.__main__ import _configure_audit_log + + +@pytest.fixture(autouse=True) +def _clean_audit_handlers(): + """Detach any file handler the test wired so runs stay independent.""" + logger = logging.getLogger("beaconmcp.audit") + before = list(logger.handlers) + yield + for h in logger.handlers[:]: + if h not in before: + logger.removeHandler(h) + h.close() + + +def test_audit_file_created_owner_only(tmp_path: Path, monkeypatch) -> None: + target = tmp_path / "audit.log" + monkeypatch.setenv("BEACONMCP_AUDIT_LOG", str(target)) + _configure_audit_log() + assert target.exists() + assert stat.S_IMODE(target.stat().st_mode) == 0o600 + + +def test_env_var_wins_over_config_value(tmp_path: Path, monkeypatch) -> None: + env_target = tmp_path / "from-env.log" + monkeypatch.setenv("BEACONMCP_AUDIT_LOG", str(env_target)) + _configure_audit_log(str(tmp_path / "from-config.log")) + assert env_target.exists() + assert not (tmp_path / "from-config.log").exists() + + +def test_config_value_used_without_env(tmp_path: Path, monkeypatch) -> None: + monkeypatch.delenv("BEACONMCP_AUDIT_LOG", raising=False) + target = tmp_path / "from-config.log" + _configure_audit_log(str(target)) + assert target.exists() + + +def test_dash_disables_file(tmp_path: Path, monkeypatch) -> None: + monkeypatch.setenv("BEACONMCP_AUDIT_LOG", "-") + logger = logging.getLogger("beaconmcp.audit") + before = [h for h in logger.handlers if isinstance(h, logging.FileHandler)] + _configure_audit_log(str(tmp_path / "ignored.log")) + after = [h for h in logger.handlers if isinstance(h, logging.FileHandler)] + assert before == after + assert not (tmp_path / "ignored.log").exists() + + +def test_unwritable_path_degrades_gracefully(tmp_path: Path, monkeypatch) -> None: + blocker = tmp_path / "blocked" + blocker.write_text("not a dir") + monkeypatch.setenv("BEACONMCP_AUDIT_LOG", str(blocker / "audit.log")) + # Must warn and keep going, never raise. + _configure_audit_log() diff --git a/tests/test_bmc_ipmi.py b/tests/test_bmc_ipmi.py new file mode 100644 index 0000000..8050f25 --- /dev/null +++ b/tests/test_bmc_ipmi.py @@ -0,0 +1,142 @@ +"""Unit tests for the generic IPMI backend (mocked subprocess).""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, patch + +import pytest + +from beaconmcp.bmc.ipmi import GenericIPMIBackend +from beaconmcp.config import ( + BMCDevice, + Config, + FeaturesConfig, + PVENode, + ServerConfig, +) + + +def _cfg() -> Config: + return Config( + server=ServerConfig(), + pve_nodes=[ + PVENode( + name="pve1", + host="pve1.example.com", + token_id="root@pam!beaconmcp", + token_secret="x", + ) + ], + bmc_devices=[], + ssh=None, + features=FeaturesConfig(), + verify_ssl=False, + infrastructure={}, + ) + + +def _backend() -> GenericIPMIBackend: + device = BMCDevice( + id="rack1-ipmi", + type="ipmi", + host="10.0.0.11", + user="admin", + password="pw", + ) + return GenericIPMIBackend(device, _cfg()) + + +class _FakeProc: + def __init__(self, stdout: bytes, stderr: bytes = b"", rc: int = 0) -> None: + self._stdout = stdout + self._stderr = stderr + self.returncode = rc + + async def communicate(self) -> tuple[bytes, bytes]: + return self._stdout, self._stderr + + +@pytest.mark.asyncio +async def test_power_on_calls_ipmitool_with_correct_argv() -> None: + backend = _backend() + fake_proc = _FakeProc(b"Chassis Power Control: Up/On\n") + mock = AsyncMock(return_value=fake_proc) + + with patch("asyncio.create_subprocess_exec", mock): + result = await backend.power_on() + + mock.assert_awaited_once() + assert mock.await_args is not None + argv = mock.await_args.args + assert argv[0] == "ipmitool" + assert "-H" in argv and "10.0.0.11" in argv + assert "-U" in argv and "admin" in argv + assert argv[-3:] == ("chassis", "power", "on") + assert result["action"] == "power_on" + assert result["result"] == "success" + + +@pytest.mark.asyncio +async def test_password_never_appears_in_argv() -> None: + """The BMC password must travel via IPMI_PASSWORD, not `-P` on argv. + + argv is readable by every local user through /proc//cmdline. + """ + backend = _backend() + mock = AsyncMock(return_value=_FakeProc(b"Chassis Power is on\n")) + + with patch("asyncio.create_subprocess_exec", mock): + await backend.power_status() + + argv = mock.await_args.args + assert "-P" not in argv + assert "pw" not in argv + assert "-E" in argv + assert mock.await_args.kwargs["env"]["IPMI_PASSWORD"] == "pw" + + +@pytest.mark.asyncio +async def test_power_status_parses_on_off() -> None: + backend = _backend() + fake_proc = _FakeProc(b"Chassis Power is on\n") + with patch( + "asyncio.create_subprocess_exec", + AsyncMock(return_value=fake_proc), + ): + result = await backend.power_status() + assert result["power_status"] == "on" + + fake_proc = _FakeProc(b"Chassis Power is off\n") + with patch( + "asyncio.create_subprocess_exec", + AsyncMock(return_value=fake_proc), + ): + result = await backend.power_status() + assert result["power_status"] == "off" + + +@pytest.mark.asyncio +async def test_missing_ipmitool_binary_returns_error() -> None: + backend = _backend() + with patch( + "asyncio.create_subprocess_exec", + AsyncMock(side_effect=FileNotFoundError()), + ): + result = await backend.power_on() + assert "error" in result + assert "ipmitool" in result["error"] + + +@pytest.mark.asyncio +async def test_event_log_limits_output() -> None: + backend = _backend() + lines = "\n".join(f"event {i}" for i in range(60)).encode() + fake_proc = _FakeProc(lines) + with patch( + "asyncio.create_subprocess_exec", + AsyncMock(return_value=fake_proc), + ): + result = await backend.event_log(limit=10) + assert result["total"] == 10 + assert result["events"][0] == "event 50" + assert result["events"][-1] == "event 59" diff --git a/tests/test_bmc_registry.py b/tests/test_bmc_registry.py new file mode 100644 index 0000000..0d843e3 --- /dev/null +++ b/tests/test_bmc_registry.py @@ -0,0 +1,99 @@ +"""Unit tests for the BMC registry.""" + +from __future__ import annotations + +import pytest + +from beaconmcp.bmc import build_registry +from beaconmcp.bmc.hp_ilo import HPILOBackend +from beaconmcp.bmc.ipmi import GenericIPMIBackend +from beaconmcp.bmc.redfish import RedfishBackend +from beaconmcp.config import ( + BMCDevice, + Config, + FeaturesConfig, + PVENode, + ServerConfig, +) + + +def _make_config(devices: list[BMCDevice]) -> Config: + return Config( + server=ServerConfig(), + pve_nodes=[ + PVENode( + name="pve1", + host="pve1.example.com", + token_id="root@pam!beaconmcp", + token_secret="x", + ) + ], + bmc_devices=devices, + ssh=None, + features=FeaturesConfig(), + verify_ssl=False, + infrastructure={}, + ) + + +def test_empty_registry() -> None: + registry = build_registry(_make_config([])) + assert registry == {} + + +def test_single_hp_ilo_device() -> None: + dev = BMCDevice(id="rack1-ilo", type="hp_ilo", host="10.0.0.10", user="admin", password="pw") + registry = build_registry(_make_config([dev])) + assert list(registry.keys()) == ["rack1-ilo"] + assert isinstance(registry["rack1-ilo"], HPILOBackend) + assert registry["rack1-ilo"].type == "hp_ilo" + + +def test_multiple_mixed_devices() -> None: + cfg = _make_config( + [ + BMCDevice(id="ilo", type="hp_ilo", host="10.0.0.10", user="a", password="x"), + BMCDevice(id="ipmi", type="ipmi", host="10.0.0.11", user="a", password="y"), + BMCDevice(id="dell", type="idrac", host="10.0.0.12", user="a", password="z"), + BMCDevice(id="smci", type="supermicro", host="10.0.0.13", user="a", password="w"), + ] + ) + registry = build_registry(cfg) + + assert set(registry.keys()) == {"ilo", "ipmi", "dell", "smci"} + assert isinstance(registry["ilo"], HPILOBackend) + assert isinstance(registry["ipmi"], GenericIPMIBackend) + # idrac and supermicro device types both build to the universal Redfish backend. + assert isinstance(registry["dell"], RedfishBackend) + assert isinstance(registry["smci"], RedfishBackend) + + +def test_unknown_type_raises_at_startup() -> None: + cfg = _make_config( + [BMCDevice(id="x", type="nope", host="10.0.0.10", user="a", password="b")] + ) + with pytest.raises(ValueError, match="Unknown BMC type 'nope'"): + build_registry(cfg) + + +@pytest.mark.asyncio +async def test_redfish_backend_returns_error_when_unreachable() -> None: + # An idrac device builds to RedfishBackend; with no real network the + # underlying httpx request fails and power_status() surfaces an error dict + # rather than raising. + cfg = _make_config( + [ + BMCDevice( + id="dell", + type="idrac", + host="127.0.0.1:1", # unreachable: nothing listening here + user="a", + password="z", + ) + ] + ) + registry = build_registry(cfg) + assert isinstance(registry["dell"], RedfishBackend) + result = await registry["dell"].power_status() + assert isinstance(result, dict) + assert "error" in result diff --git a/tests/test_cloudflare_diagnostic.py b/tests/test_cloudflare_diagnostic.py new file mode 100644 index 0000000..2846428 --- /dev/null +++ b/tests/test_cloudflare_diagnostic.py @@ -0,0 +1,73 @@ +"""Cloudflare 401 diagnostic tests. + +When an MCP request reaches BeaconMCP without a usable ``Authorization`` +header but *with* Cloudflare's ``cf-ray`` edge header, the 401 body gains an +actionable ``hint`` (and the server logs a throttled warning) because a +Cloudflare WAF / Access / Bot-Fight-Mode rule is the overwhelmingly likely +cause. See ``docs/cloudflare.md``. + +The 401 status and ``WWW-Authenticate`` header are the caller's job and are +covered where ``auth_middleware`` is exercised; these tests own the body and +the logging contract. +""" + +from __future__ import annotations + +import logging +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).parent.parent / "src")) + +from beaconmcp.__main__ import _CF_EDGE_HEADER, _build_unauthorized_body, _cf_log_state + +_HINT_MARKER = "docs/cloudflare.md" + + +@pytest.fixture(autouse=True) +def reset_log_throttle(): + """The Cloudflare warning is rate-limited via module state; start clean so + tests don't suppress each other's log lines.""" + _cf_log_state.update(last=0.0, suppressed=0) + yield + _cf_log_state.update(last=0.0, suppressed=0) + + +@pytest.mark.parametrize("error", ["unauthorized", "invalid_token"]) +def test_cf_ray_adds_hint_and_logs(error: str, caplog) -> None: + with caplog.at_level(logging.WARNING, logger="beaconmcp"): + body = _build_unauthorized_body( + {_CF_EDGE_HEADER: "7d9f0c2a1b3e4f56-AMS"}, error=error + ) + + assert body["error"] == error + assert _HINT_MARKER in body["hint"] + assert any("Cloudflare" in r.message for r in caplog.records) + + +def test_no_cf_ray_keeps_body_minimal_and_quiet(caplog) -> None: + with caplog.at_level(logging.WARNING, logger="beaconmcp"): + body = _build_unauthorized_body({}, error="unauthorized") + + assert body == {"error": "unauthorized"} + assert not any("Cloudflare" in r.message for r in caplog.records) + + +def test_repeated_unauthorized_hits_do_not_flood_the_log(caplog) -> None: + """A public /mcp is scanned continuously and every hit carries a cf-ray. + Each one must still get its hint, but the journal must not grow by a line + per anonymous request. + """ + headers = {_CF_EDGE_HEADER: "abc-AMS"} + with caplog.at_level(logging.WARNING, logger="beaconmcp"): + bodies = [ + _build_unauthorized_body(headers, error="unauthorized") + for _ in range(200) + ] + + assert all(_HINT_MARKER in b["hint"] for b in bodies) + cf_records = [r for r in caplog.records if "Cloudflare" in r.message] + assert len(cf_records) == 1 + assert _cf_log_state["suppressed"] == 199 diff --git a/tests/test_config_yaml.py b/tests/test_config_yaml.py new file mode 100644 index 0000000..2af9c18 --- /dev/null +++ b/tests/test_config_yaml.py @@ -0,0 +1,741 @@ +"""Unit tests for the YAML-first config loader.""" + +from __future__ import annotations + +import textwrap +import warnings +from pathlib import Path + +import pytest + +from beaconmcp.config import Config, ConfigError + + +def _write(path: Path, yaml_text: str) -> Path: + path.write_text(textwrap.dedent(yaml_text).lstrip()) + return path + + +def test_yaml_happy_path(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("PVE1_TOKEN_SECRET", "secret1") + monkeypatch.setenv("RACK1_ILO_PASSWORD", "ilopw") + monkeypatch.setenv("PVE1_SSH_PW", "ssh-secret") + path = _write( + tmp_path / "beaconmcp.yaml", + """ + version: 1 + proxmox: + verify_ssl: false + nodes: + - name: pve1 + host: pve1.example.com + token_id: root@pam!beaconmcp + token_secret: ${PVE1_TOKEN_SECRET} + ssh: + hosts: + - name: pve1-ssh + host: pve1.example.com + user: root + password: ${PVE1_SSH_PW} + bmc: + devices: + - id: rack1-ilo + type: hp_ilo + host: 10.0.0.10 + user: Administrator + password: ${RACK1_ILO_PASSWORD} + jump_host: pve1-ssh + """, + ) + + cfg = Config.load(config_path=path) + + assert [n.name for n in cfg.pve_nodes] == ["pve1"] + assert cfg.pve_nodes[0].token_secret == "secret1" + assert len(cfg.bmc_devices) == 1 + assert cfg.bmc_devices[0].password == "ilopw" + assert cfg.bmc_devices[0].jump_host == "pve1-ssh" + assert cfg.ssh is not None and len(cfg.ssh.hosts) == 1 + assert cfg.ssh.hosts[0].password == "ssh-secret" + assert cfg.verify_ssl is False + + +def test_missing_env_ref_raises_with_path(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("PVE1_TOKEN_SECRET", raising=False) + path = _write( + tmp_path / "beaconmcp.yaml", + """ + version: 1 + proxmox: + nodes: + - name: pve1 + host: pve1.example.com + token_id: root@pam!beaconmcp + token_secret: ${PVE1_TOKEN_SECRET} + """, + ) + + with pytest.raises(ConfigError) as exc: + Config.load(config_path=path) + message = str(exc.value) + assert "PVE1_TOKEN_SECRET" in message + assert "proxmox.nodes.[0].token_secret" in message + + +def test_duplicate_bmc_device_id_raises(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("PVE1_TOKEN_SECRET", "x") + path = _write( + tmp_path / "beaconmcp.yaml", + """ + version: 1 + proxmox: + nodes: + - name: pve1 + host: pve1.example.com + token_id: root@pam!beaconmcp + token_secret: ${PVE1_TOKEN_SECRET} + bmc: + devices: + - id: dup + type: hp_ilo + host: 10.0.0.10 + user: admin + password: x + - id: dup + type: ipmi + host: 10.0.0.11 + user: admin + password: y + """, + ) + + with pytest.raises(ConfigError, match="Duplicate BMC device id"): + Config.load(config_path=path) + + +def test_at_least_one_capability_required(tmp_path: Path) -> None: + """A config with no Proxmox, no SSH, no BMC should be refused.""" + path = _write( + tmp_path / "beaconmcp.yaml", + """ + version: 1 + proxmox: + nodes: [] + """, + ) + with pytest.raises(ConfigError, match="at least one capability"): + Config.load(config_path=path) + + +def test_ssh_only_minimal_config(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """A VPS-only user: no Proxmox, no BMC, just SSH hosts. Must load fine.""" + monkeypatch.setenv("VPS1_PW", "pw1") + monkeypatch.setenv("VPS2_PW", "pw2") + path = _write( + tmp_path / "beaconmcp.yaml", + """ + version: 1 + ssh: + hosts: + - name: vps1 + host: 198.51.100.10 + user: root + password: ${VPS1_PW} + - name: vps2 + host: 198.51.100.11 + port: 2222 + user: admin + password: ${VPS2_PW} + """, + ) + cfg = Config.load(config_path=path) + assert cfg.pve_nodes == [] + assert cfg.bmc_devices == [] + assert cfg.ssh is not None + assert [h.name for h in cfg.ssh.hosts] == ["vps1", "vps2"] + assert cfg.ssh.hosts[1].port == 2222 + + +def test_proxmox_only_config(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """Proxmox-only: no ssh: and no bmc: section. Must load.""" + monkeypatch.setenv("PVE1_TOKEN_SECRET", "x") + path = _write( + tmp_path / "beaconmcp.yaml", + """ + version: 1 + proxmox: + nodes: + - name: pve1 + host: pve1.example.com + token_id: root@pam!beaconmcp + token_secret: ${PVE1_TOKEN_SECRET} + """, + ) + cfg = Config.load(config_path=path) + assert len(cfg.pve_nodes) == 1 + assert cfg.ssh is None + assert cfg.bmc_devices == [] + + +def test_ssh_legacy_shape_rejected(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """Legacy ssh.user/ssh.password flat shape must produce an actionable error.""" + monkeypatch.setenv("SSH_PW", "x") + path = _write( + tmp_path / "beaconmcp.yaml", + """ + version: 1 + ssh: + user: root + password: ${SSH_PW} + vmid_to_ip: "192.168.1.{id}" + """, + ) + with pytest.raises(ConfigError, match="legacy shape"): + Config.load(config_path=path) + + +def test_ssh_host_name_can_match_proxmox_node( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """An ssh.hosts[].name may match a proxmox.nodes[].name. + + Pre-2.0 the loader rejected this collision. The two names live in + separate tool namespaces (``ssh_*`` vs ``proxmox_*``) so there is no + routing ambiguity, and matching names removes the need for synthetic + ``*-ssh`` suffixes. + """ + monkeypatch.setenv("PVE1_TOKEN_SECRET", "x") + monkeypatch.setenv("SSH_PW", "y") + path = _write( + tmp_path / "beaconmcp.yaml", + """ + version: 1 + proxmox: + nodes: + - name: pve1 + host: pve1.example.com + token_id: root@pam!beaconmcp + token_secret: ${PVE1_TOKEN_SECRET} + ssh: + hosts: + - name: pve1 + host: pve1.example.com + user: root + password: ${SSH_PW} + """, + ) + cfg = Config.load(config_path=path) + assert cfg.ssh is not None + assert [h.name for h in cfg.ssh.hosts] == ["pve1"] + assert cfg.pve_nodes[0].name == "pve1" + + +def test_ssh_inherit_proxmox_nodes_synthesizes_hosts( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """With ``inherit_proxmox_nodes: true`` + defaults, every Proxmox node + that isn't covered by an explicit ssh.hosts[] entry gets one synthesized. + Restores the pre-2.0 single-credential-block ergonomic.""" + monkeypatch.setenv("PVE1_TOKEN_SECRET", "x") + monkeypatch.setenv("PVE2_TOKEN_SECRET", "x") + path = _write( + tmp_path / "beaconmcp.yaml", + """ + version: 1 + proxmox: + nodes: + - name: pve1 + host: 10.0.0.1 + token_id: root@pam!beaconmcp + token_secret: ${PVE1_TOKEN_SECRET} + - name: pve2 + host: 10.0.0.2 + token_id: root@pam!beaconmcp + token_secret: ${PVE2_TOKEN_SECRET} + ssh: + defaults: + user: root + key_file: ~/.ssh/homelab + inherit_proxmox_nodes: true + """, + ) + cfg = Config.load(config_path=path) + assert cfg.ssh is not None + names = sorted(h.name for h in cfg.ssh.hosts) + assert names == ["pve1", "pve2"] + pve1 = next(h for h in cfg.ssh.hosts if h.name == "pve1") + assert pve1.host == "10.0.0.1" + assert pve1.user == "root" + assert pve1.key_file == "~/.ssh/homelab" + assert pve1.password is None + + +def test_ssh_inherit_proxmox_nodes_explicit_override_wins( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """An explicit ssh.hosts[] entry shadows inheritance by name or address. + + Here pve2 is declared explicitly; inheritance must only synthesize pve1. + """ + monkeypatch.setenv("PVE1_TOKEN_SECRET", "x") + monkeypatch.setenv("PVE2_TOKEN_SECRET", "x") + monkeypatch.setenv("PVE2_PW", "specific") + path = _write( + tmp_path / "beaconmcp.yaml", + """ + version: 1 + proxmox: + nodes: + - name: pve1 + host: 10.0.0.1 + token_id: root@pam!beaconmcp + token_secret: ${PVE1_TOKEN_SECRET} + - name: pve2 + host: 10.0.0.2 + token_id: root@pam!beaconmcp + token_secret: ${PVE2_TOKEN_SECRET} + ssh: + defaults: + user: root + key_file: ~/.ssh/homelab + inherit_proxmox_nodes: true + hosts: + - name: pve2 + host: 10.0.0.2 + user: admin + password: ${PVE2_PW} + """, + ) + cfg = Config.load(config_path=path) + assert cfg.ssh is not None + pve2 = next(h for h in cfg.ssh.hosts if h.name == "pve2") + assert pve2.user == "admin" + assert pve2.password == "specific" + # pve1 still inherits. + pve1 = next(h for h in cfg.ssh.hosts if h.name == "pve1") + assert pve1.key_file == "~/.ssh/homelab" + + +def test_ssh_inherit_proxmox_nodes_requires_defaults( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """``inherit_proxmox_nodes: true`` without ``defaults:`` is rejected.""" + monkeypatch.setenv("PVE1_TOKEN_SECRET", "x") + path = _write( + tmp_path / "beaconmcp.yaml", + """ + version: 1 + proxmox: + nodes: + - name: pve1 + host: 10.0.0.1 + token_id: root@pam!beaconmcp + token_secret: ${PVE1_TOKEN_SECRET} + ssh: + inherit_proxmox_nodes: true + """, + ) + with pytest.raises(ConfigError, match=r"requires 'ssh\.defaults"): + Config.load(config_path=path) + + +def test_ssh_defaults_requires_one_auth_method( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Defaults block follows the same rule as an explicit host entry.""" + monkeypatch.setenv("PVE1_TOKEN_SECRET", "x") + path = _write( + tmp_path / "beaconmcp.yaml", + """ + version: 1 + proxmox: + nodes: + - name: pve1 + host: 10.0.0.1 + token_id: root@pam!beaconmcp + token_secret: ${PVE1_TOKEN_SECRET} + ssh: + defaults: + user: root + inherit_proxmox_nodes: true + """, + ) + with pytest.raises(ConfigError, match="ssh.defaults.*neither"): + Config.load(config_path=path) + + +def test_ssh_host_requires_one_auth_method( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Each SSH host must provide exactly one of password or key_file.""" + # Case 1: neither provided + path1 = _write( + tmp_path / "neither.yaml", + """ + version: 1 + ssh: + hosts: + - name: vps1 + host: 198.51.100.10 + user: root + """, + ) + with pytest.raises(ConfigError, match="neither"): + Config.load(config_path=path1) + + # Case 2: both provided + monkeypatch.setenv("VPS_PW", "x") + path2 = _write( + tmp_path / "both.yaml", + """ + version: 1 + ssh: + hosts: + - name: vps1 + host: 198.51.100.10 + user: root + password: ${VPS_PW} + key_file: ~/.ssh/id_ed25519 + """, + ) + with pytest.raises(ConfigError, match="both"): + Config.load(config_path=path2) + + +def test_ssh_empty_hosts_list_rejected(tmp_path: Path) -> None: + """An ssh: section that resolves to zero hosts is a config mistake. + + Empty ``hosts: []`` with no ``inherit_proxmox_nodes`` → the SSH section + contributes nothing, which is almost certainly an oversight. Require + the user to either declare a host or flip inheritance on. + """ + path = _write( + tmp_path / "beaconmcp.yaml", + """ + version: 1 + ssh: + hosts: [] + """, + ) + with pytest.raises(ConfigError, match="inherit_proxmox_nodes"): + Config.load(config_path=path) + + +def test_ssh_duplicate_host_names_rejected( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("PW", "x") + path = _write( + tmp_path / "beaconmcp.yaml", + """ + version: 1 + ssh: + hosts: + - name: dup + host: 198.51.100.10 + user: root + password: ${PW} + - name: dup + host: 198.51.100.11 + user: root + password: ${PW} + """, + ) + with pytest.raises(ConfigError, match="duplicate name"): + Config.load(config_path=path) + + +def test_bmc_jump_host_must_reference_ssh_host( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """bmc.devices[].jump_host must match an ssh.hosts[].name — caught at load.""" + monkeypatch.setenv("ILO_PW", "x") + path = _write( + tmp_path / "beaconmcp.yaml", + """ + version: 1 + bmc: + devices: + - id: ilo1 + type: hp_ilo + host: 10.0.0.10 + user: Administrator + password: ${ILO_PW} + jump_host: nonexistent + """, + ) + with pytest.raises(ConfigError, match="jump_host"): + Config.load(config_path=path) + + +def test_legacy_env_fallback_drops_ssh( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Legacy env path still synthesizes Proxmox and BMC, but drops SSH with a warning.""" + monkeypatch.chdir(tmp_path) + monkeypatch.delenv("BEACONMCP_CONFIG", raising=False) + monkeypatch.setenv("PVE1_HOST", "pve1.example.com") + monkeypatch.setenv("PVE1_TOKEN_ID", "root@pam!beaconmcp") + monkeypatch.setenv("PVE1_TOKEN_SECRET", "legacy-secret") + monkeypatch.setenv("SSH_USER", "root") + monkeypatch.setenv("SSH_PASSWORD", "legacy-ssh") + monkeypatch.setenv("ILO_HOST", "10.0.0.10") + monkeypatch.setenv("ILO_USER", "Administrator") + monkeypatch.setenv("ILO_PASSWORD", "legacy-ilo") + # Deliberately don't set ILO_JUMP_HOST so BMC loads without needing + # an SSH host (legacy path can no longer synthesize one). + + with warnings.catch_warnings(record=True) as captured: + warnings.simplefilter("always") + cfg = Config.load() + + messages = [str(w.message) for w in captured] + # Two deprecations: top-level legacy-env + SSH-drop. + assert any("deprecated" in m.lower() for m in messages) + assert any("SSH" in m and "no longer supported" in m for m in messages) + assert len(cfg.pve_nodes) == 1 + assert cfg.pve_nodes[0].token_secret == "legacy-secret" + assert len(cfg.bmc_devices) == 1 + assert cfg.bmc_devices[0].type == "hp_ilo" + assert cfg.bmc_devices[0].jump_host is None # jump not synthesized in legacy + assert cfg.ssh is None + + +def test_get_node_host_strips_port(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """proxmox.nodes[].host often carries the API port (e.g. :443 behind a reverse + proxy). SSH and BMC-over-SSH-tunnel need the bare hostname.""" + monkeypatch.setenv("PVE1_TOKEN_SECRET", "x") + monkeypatch.setenv("PVE2_TOKEN_SECRET", "y") + path = _write( + tmp_path / "beaconmcp.yaml", + """ + version: 1 + proxmox: + nodes: + - name: pve1 + host: pve1.example.com:443 + token_id: "root@pam!beaconmcp" + token_secret: ${PVE1_TOKEN_SECRET} + - name: pve6 + host: "[::1]:8006" + token_id: "root@pam!beaconmcp" + token_secret: ${PVE2_TOKEN_SECRET} + """, + ) + cfg = Config.load(config_path=path) + assert cfg.get_node_host("pve1") == "pve1.example.com" + assert cfg.get_node_host("pve6") == "[::1]" + assert cfg.get_node_host("missing") is None + # The raw .host value is preserved for proxmoxer which accepts host:port. + assert cfg.pve_nodes[0].host == "pve1.example.com:443" + + +def test_redacted_masks_secrets(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("PVE1_TOKEN_SECRET", "abcdefghij") + monkeypatch.setenv("VPS_PW", "1234567890") + path = _write( + tmp_path / "beaconmcp.yaml", + """ + version: 1 + proxmox: + nodes: + - name: pve1 + host: pve1.example.com + token_id: root@pam!beaconmcp + token_secret: ${PVE1_TOKEN_SECRET} + ssh: + hosts: + - name: vps1 + host: 198.51.100.10 + user: root + password: ${VPS_PW} + """, + ) + cfg = Config.load(config_path=path) + redacted = cfg.redacted() + s = str(redacted) + assert "abcdefghij" not in s + assert "1234567890" not in s + assert "***" in redacted["proxmox"]["nodes"][0]["token_secret"] + assert "***" in redacted["ssh"]["hosts"][0]["password"] + + +def test_server_trusted_proxies_cloudflare_macro_expands( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("VPS_PW", "pw") + path = _write( + tmp_path / "beaconmcp.yaml", + """ + version: 1 + server: + trusted_proxies: + - cloudflare + - 127.0.0.1 + - 127.0.0.1 + ssh: + hosts: + - name: vps1 + host: 198.51.100.10 + user: root + password: ${VPS_PW} + """, + ) + + cfg = Config.load(config_path=path) + assert "173.245.48.0/20" in cfg.server.trusted_proxies + assert "2a06:98c0::/29" in cfg.server.trusted_proxies + assert cfg.server.trusted_proxies.count("127.0.0.1") == 1 + assert "trusted_proxies" in cfg.redacted()["server"] + + +def test_server_trusted_proxies_must_be_list( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("VPS_PW", "pw") + path = _write( + tmp_path / "beaconmcp.yaml", + """ + version: 1 + server: + trusted_proxies: cloudflare + ssh: + hosts: + - name: vps1 + host: 198.51.100.10 + user: root + password: ${VPS_PW} + """, + ) + + with pytest.raises(ConfigError, match="server.trusted_proxies"): + Config.load(config_path=path) + + +def test_get_ssh_host_accessors( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("VPS_PW", "pw") + path = _write( + tmp_path / "beaconmcp.yaml", + """ + version: 1 + ssh: + hosts: + - name: vps1 + host: 198.51.100.10 + user: root + password: ${VPS_PW} + """, + ) + cfg = Config.load(config_path=path) + h = cfg.get_ssh_host("vps1") + assert h is not None + assert h.host == "198.51.100.10" + assert cfg.get_ssh_host("missing") is None + assert cfg.get_ssh_host_by_address("198.51.100.10") is not None + assert cfg.get_ssh_host_by_address("10.0.0.1") is None + + +def test_ssh_per_host_host_key_overrides( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("VPS_PW", "pw") + path = _write( + tmp_path / "beaconmcp.yaml", + """ + version: 1 + ssh: + known_hosts: /etc/beaconmcp/known_hosts + strict_host_key_checking: true + hosts: + - name: lan-node + host: 10.0.0.1 + user: root + password: ${VPS_PW} + - name: vps + host: 198.51.100.10 + user: root + password: ${VPS_PW} + known_hosts: /etc/beaconmcp/vps_known_hosts + strict_host_key_checking: false + """, + ) + cfg = Config.load(config_path=path) + lan = cfg.get_ssh_host("lan-node") + assert lan is not None + # Unset per-host fields stay None so the SSH client inherits the globals. + assert lan.known_hosts is None + assert lan.strict_host_key_checking is None + vps = cfg.get_ssh_host("vps") + assert vps is not None + assert vps.known_hosts == "/etc/beaconmcp/vps_known_hosts" + assert vps.strict_host_key_checking is False + + +def test_server_tokens_db_and_audit_log_keys( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("VPS_PW", "pw") + path = _write( + tmp_path / "beaconmcp.yaml", + """ + version: 1 + server: + tokens_db: /var/lib/beaconmcp/tokens.db + audit_log: "-" + ssh: + hosts: + - name: vps1 + host: 198.51.100.10 + user: root + password: ${VPS_PW} + """, + ) + cfg = Config.load(config_path=path) + assert cfg.server.tokens_db == Path("/var/lib/beaconmcp/tokens.db") + assert cfg.server.audit_log == "-" + + +def test_server_tokens_db_and_audit_log_default_none( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("VPS_PW", "pw") + path = _write( + tmp_path / "beaconmcp.yaml", + """ + version: 1 + ssh: + hosts: + - name: vps1 + host: 198.51.100.10 + user: root + password: ${VPS_PW} + """, + ) + cfg = Config.load(config_path=path) + assert cfg.server.tokens_db is None + assert cfg.server.audit_log is None + + +def test_server_named_token_ttl_zero_and_negative( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("VPS_PW", "pw") + base = """ + version: 1 + server: + named_token_ttl: {value} + ssh: + hosts: + - name: vps1 + host: 198.51.100.10 + user: root + password: ${{VPS_PW}} + """ + # 0 is a deliberate setting (never expires) and must survive parsing. + cfg = Config.load(config_path=_write(tmp_path / "zero.yaml", base.format(value=0))) + assert cfg.server.named_token_ttl == 0 + # Negatives are nonsense -- fail fast with the offending key. + with pytest.raises(ConfigError, match="named_token_ttl"): + Config.load(config_path=_write(tmp_path / "neg.yaml", base.format(value=-5))) diff --git a/tests/test_dashboard_chat.py b/tests/test_dashboard_chat.py new file mode 100644 index 0000000..351374d --- /dev/null +++ b/tests/test_dashboard_chat.py @@ -0,0 +1,998 @@ +"""Integration tests for Stage 2: conversations + chat stream with FakeChatEngine.""" + +from __future__ import annotations + +import json +import os +import sys +import time +from pathlib import Path + +import pytest +from starlette.applications import Starlette +from starlette.testclient import TestClient + +sys.path.insert(0, str(Path(__file__).parent.parent / "src")) + +from beaconmcp.auth import TotpResult +from beaconmcp.dashboard.app import BEARER_TTL_SECONDS, DashboardDeps, build_dashboard_routes +from beaconmcp.dashboard.chat import ( + ErrorEvent, + FakeChatEngine, + FakeScript, + TextDelta, + ToolCallEnd, + ToolCallStart, + ToolConfirmRequired, +) +from beaconmcp.dashboard.confirmations import ConfirmationStore +from beaconmcp.dashboard.conversations import ConversationStore +from beaconmcp.dashboard.csrf import CSRF_COOKIE +from beaconmcp.dashboard.db import Database +from beaconmcp.dashboard.session import SessionStore + + +class FakeClientStore: + def verify(self, cid, sec): return cid == "c" and sec == "s" + def check_totp(self, cid, code): + return TotpResult.OK if code == "123456" else TotpResult.INVALID + def get_name(self, cid): return "Test" + + +class _FakeTokenCapExceeded(Exception): + pass + + +class FakeTokenStore: + NAMED_TOKEN_CAP = 3 + + def __init__(self): + self.revoked = [] + self._n = 0 + self._live: dict[str, str] = {} + # token -> (name, created_at, expires_at) for named tokens only + self._named: dict[str, tuple[str, float, float]] = {} + + def issue(self, cid, *, name=None): + if name is not None: + active = sum( + 1 for t, meta in self._named.items() + if self._live.get(t) == cid + ) + if active >= self.NAMED_TOKEN_CAP: + err = _FakeTokenCapExceeded(f"max {self.NAMED_TOKEN_CAP}") + err.__class__.__name__ = "TokenCapExceeded" + raise err + self._n += 1 + token = f"b_{self._n}" + ("x" * 60) # pad so prefix[:12] is stable + self._live[token] = cid + if name is not None: + self._named[token] = (name, time.time(), time.time() + BEARER_TTL_SECONDS) + return token, BEARER_TTL_SECONDS + + def validate(self, token): + return self._live.get(token) + + def revoke(self, token): + self.revoked.append(token) + self._live.pop(token, None) + self._named.pop(token, None) + return True + + # Named-token APIs ------------------------------------------------ + def list_named(self, cid): + from dataclasses import dataclass + + @dataclass + class _Row: + token: str + name: str + created_at: float + expires_at: float + + return [ + _Row(token=t, name=meta[0], created_at=meta[1], expires_at=meta[2]) + for t, meta in self._named.items() + if self._live.get(t) == cid + ] + + def count_named(self, cid): + return len(self.list_named(cid)) + + def revoke_named(self, token_prefix, cid): + if len(token_prefix) < 6: + return False + matches = [ + t for t in self._named + if t.startswith(token_prefix) and self._live.get(t) == cid + ] + if len(matches) != 1: + return False + return self.revoke(matches[0]) + + +@pytest.fixture() +def engine(): + return FakeChatEngine(FakeScript(events=[TextDelta(text="pong")], title_text="Un titre")) + + +@pytest.fixture() +def deps(tmp_path, engine): + db = Database(tmp_path / "dashboard.db") + return DashboardDeps( + database=db, + session_store=SessionStore(db, key=os.urandom(32)), + client_store=FakeClientStore(), + token_store=FakeTokenStore(), + totp_locked=lambda cid: False, + totp_record_failure=lambda cid: None, + totp_record_success=lambda cid: None, + conversations=ConversationStore(db), + engine=engine, + confirmations=ConfirmationStore(), + mcp_public_url="https://mcp.example/", + ) + + +@pytest.fixture() +def app_and_client(deps): + app = Starlette(routes=build_dashboard_routes(deps)) + return app, TestClient(app, follow_redirects=False) + + +def _login(client): + r = client.get("/app/login") + csrf = r.cookies.get(CSRF_COOKIE) + r = client.post("/app/login", data={ + "csrf_token": csrf, "client_id": "c", + "client_secret": "s", "totp": "123456", "remember": "on", + }) + assert r.status_code == 303 + return client.cookies.get(CSRF_COOKIE) + + +# --------------------------------------------------------------------------- +# Conversations API +# --------------------------------------------------------------------------- + +def test_conv_list_empty(app_and_client): + _, client = app_and_client + _login(client) + r = client.get("/app/api/conversations") + assert r.status_code == 200 + assert r.json() == {"conversations": []} + + +def test_conv_list_requires_auth(app_and_client): + _, client = app_and_client + r = client.get("/app/api/conversations") + assert r.status_code == 401 + + +def test_conv_create_and_list(app_and_client): + _, client = app_and_client + csrf = _login(client) + r = client.post( + "/app/api/conversations", + headers={"X-CSRF-Token": csrf, "Content-Type": "application/json"}, + content=json.dumps({"model": "gemini-3.5-flash-lite", "effort": "medium"}), + ) + assert r.status_code == 201 + conv = r.json()["conversation"] + assert conv["model"] == "gemini-3.5-flash-lite" + assert conv["thinking_effort"] == "medium" + + r = client.get("/app/api/conversations") + assert r.status_code == 200 + assert len(r.json()["conversations"]) == 1 + + +def test_conv_create_csrf(app_and_client): + _, client = app_and_client + _login(client) + r = client.post( + "/app/api/conversations", + headers={"Content-Type": "application/json"}, + content="{}", + ) + assert r.status_code == 403 + + +def test_conv_patch(app_and_client): + _, client = app_and_client + csrf = _login(client) + r = client.post( + "/app/api/conversations", + headers={"X-CSRF-Token": csrf, "Content-Type": "application/json"}, + content="{}", + ) + cid = r.json()["conversation"]["id"] + r = client.patch( + f"/app/api/conversations/{cid}", + headers={"X-CSRF-Token": csrf, "Content-Type": "application/json"}, + content=json.dumps({"title": "Nouveau titre", "effort": "high"}), + ) + assert r.status_code == 200 + conv = r.json()["conversation"] + assert conv["title"] == "Nouveau titre" + assert conv["thinking_effort"] == "high" + + +def test_conv_patch_rejects_invalid_effort(app_and_client): + _, client = app_and_client + csrf = _login(client) + r = client.post("/app/api/conversations", + headers={"X-CSRF-Token": csrf, "Content-Type": "application/json"}, + content="{}") + cid = r.json()["conversation"]["id"] + r = client.patch( + f"/app/api/conversations/{cid}", + headers={"X-CSRF-Token": csrf, "Content-Type": "application/json"}, + content=json.dumps({"effort": "ULTRA"}), + ) + assert r.status_code == 200 + assert r.json()["conversation"]["thinking_effort"] == "low" # unchanged + + +def test_conv_delete(app_and_client): + _, client = app_and_client + csrf = _login(client) + r = client.post("/app/api/conversations", + headers={"X-CSRF-Token": csrf, "Content-Type": "application/json"}, + content="{}") + cid = r.json()["conversation"]["id"] + r = client.delete(f"/app/api/conversations/{cid}", + headers={"X-CSRF-Token": csrf}) + assert r.status_code == 204 + r = client.get(f"/app/api/conversations/{cid}") + assert r.status_code == 404 + + +def test_conv_scoped_to_client(deps, app_and_client): + _, client = app_and_client + csrf = _login(client) + # Create a conversation for client c + r = client.post("/app/api/conversations", + headers={"X-CSRF-Token": csrf, "Content-Type": "application/json"}, + content="{}") + cid = r.json()["conversation"]["id"] + + # Another client's conversation -- should not be visible + other = deps.conversations.create(client_id="otherclient", model="gemini-3-flash", effort="low") + + r = client.get("/app/api/conversations") + ids = [c["id"] for c in r.json()["conversations"]] + assert cid in ids + assert other.id not in ids + + +# --------------------------------------------------------------------------- +# Chat stream (SSE) with FakeChatEngine +# --------------------------------------------------------------------------- + +def _parse_sse(text): + events = [] + for frame in text.strip().split("\n\n"): + ev = None + data = [] + for line in frame.split("\n"): + if line.startswith("event:"): + ev = line[6:].strip() + elif line.startswith("data:"): + data.append(line[5:].strip()) + payload = json.loads("\n".join(data)) if data else {} + events.append((ev, payload)) + return events + + +def test_chat_stream_simple_text(app_and_client, engine, deps): + _, client = app_and_client + csrf = _login(client) + r = client.post("/app/api/conversations", + headers={"X-CSRF-Token": csrf, "Content-Type": "application/json"}, + content="{}") + cid = r.json()["conversation"]["id"] + + r = client.post( + "/app/api/chat/stream", + headers={"X-CSRF-Token": csrf, "Content-Type": "application/json"}, + content=json.dumps({"conversation_id": cid, "content": "ping"}), + ) + assert r.status_code == 200 + assert r.headers["content-type"].startswith("text/event-stream") + + events = _parse_sse(r.text) + names = [e[0] for e in events] + assert "text_delta" in names + assert "done" in names + assert "title_updated" in names + + text_events = [e[1] for e in events if e[0] == "text_delta"] + assert text_events[0]["text"] == "pong" + + # Engine saw the call + assert len(engine.calls) == 1 + turn = engine.calls[0] + assert turn.user_text == "ping" + assert turn.bearer.startswith("b_") + # Local mode (default) ignores mcp_public_url and uses loopback so + # the dashboard never round-trips through its own reverse proxy. + assert turn.mcp_url == "http://127.0.0.1:8420/mcp" + assert turn.history == [] + + # Message persisted in DB + msgs = deps.conversations.list_messages(cid) + assert [m.role for m in msgs] == ["user", "assistant"] + assert msgs[0].content == "ping" + assert msgs[1].content == "pong" + + +def test_chat_stream_tool_call(app_and_client, engine, deps): + _, client = app_and_client + engine.script = FakeScript(events=[ + ToolCallStart(id="tc1", name="proxmox_list_vms", args={"node": "pve1"}), + ToolCallEnd(id="tc1", status="ok", preview="2 VMs", duration_ms=42), + TextDelta(text="Voici la liste."), + ]) + csrf = _login(client) + r = client.post("/app/api/conversations", + headers={"X-CSRF-Token": csrf, "Content-Type": "application/json"}, + content="{}") + cid = r.json()["conversation"]["id"] + + r = client.post( + "/app/api/chat/stream", + headers={"X-CSRF-Token": csrf, "Content-Type": "application/json"}, + content=json.dumps({"conversation_id": cid, "content": "liste"}), + ) + events = _parse_sse(r.text) + names = [e[0] for e in events] + assert "tool_call" in names + assert "tool_result" in names + + msgs = deps.conversations.list_messages(cid) + assistant = msgs[1] + assert len(assistant.tool_calls) == 1 + tc = assistant.tool_calls[0] + assert tc.name == "proxmox_list_vms" + assert tc.status == "ok" + assert tc.preview == "2 VMs" + assert tc.duration_ms == 42 + + +def test_refresh_page_shows_totp_when_token_wiped(tmp_path, engine): + """After a restart, /app/refresh must render the TOTP form instead + of redirecting back to /app/chat and creating a redirect loop. + """ + class WipedTokenStore(FakeTokenStore): + def validate(self, token): + return None # TokenStore wiped by restart + + db = Database(tmp_path / "dashboard.db") + deps = DashboardDeps( + database=db, + session_store=SessionStore(db, key=os.urandom(32)), + client_store=FakeClientStore(), + token_store=WipedTokenStore(), + totp_locked=lambda cid: False, + totp_record_failure=lambda cid: None, + totp_record_success=lambda cid: None, + conversations=ConversationStore(db), + engine=engine, + ) + app = Starlette(routes=build_dashboard_routes(deps)) + client = TestClient(app, follow_redirects=False) + + _login(client) + r = client.get("/app/refresh") + # Must render the TOTP form, NOT 302 back to /app/chat. + assert r.status_code == 200 + assert "totp" in r.text.lower() + + +def test_chat_page_redirects_to_refresh_when_token_wiped(tmp_path, engine): + """/app/chat must send the user to /app/refresh when the bearer + is wiped, to break the refresh<->chat redirect loop. + """ + class WipedTokenStore(FakeTokenStore): + def validate(self, token): + return None + + db = Database(tmp_path / "dashboard.db") + deps = DashboardDeps( + database=db, + session_store=SessionStore(db, key=os.urandom(32)), + client_store=FakeClientStore(), + token_store=WipedTokenStore(), + totp_locked=lambda cid: False, + totp_record_failure=lambda cid: None, + totp_record_success=lambda cid: None, + conversations=ConversationStore(db), + engine=engine, + ) + app = Starlette(routes=build_dashboard_routes(deps)) + client = TestClient(app, follow_redirects=False) + + _login(client) + r = client.get("/app/chat") + assert r.status_code == 302 + assert r.headers["location"] == "/app/refresh" + + +def test_chat_stream_detects_token_wiped_after_restart(tmp_path, engine): + """After a service restart TokenStore is empty but the session's + bearer is still there; we must emit session_expired before reaching + the MCP server with a stale token. + + We simulate the "restart mid-session" flow: login + create + conversation succeed while the token is live, then we flip the + store into its post-restart state and try to stream. + """ + class SwitchableTokenStore(FakeTokenStore): + wiped = False + def validate(self, token): + if self.wiped: + return None + return super().validate(token) + + db = Database(tmp_path / "dashboard.db") + token_store = SwitchableTokenStore() + deps = DashboardDeps( + database=db, + session_store=SessionStore(db, key=os.urandom(32)), + client_store=FakeClientStore(), + token_store=token_store, + totp_locked=lambda cid: False, + totp_record_failure=lambda cid: None, + totp_record_success=lambda cid: None, + conversations=ConversationStore(db), + engine=engine, + ) + app = Starlette(routes=build_dashboard_routes(deps)) + client = TestClient(app, follow_redirects=False) + + csrf = _login(client) + r = client.post("/app/api/conversations", + headers={"X-CSRF-Token": csrf, "Content-Type": "application/json"}, + content="{}") + cid = r.json()["conversation"]["id"] + + # Simulate `systemctl restart beaconmcp` wiping every issued token. + token_store.wiped = True + + r = client.post( + "/app/api/chat/stream", + headers={"X-CSRF-Token": csrf, "Content-Type": "application/json"}, + content=json.dumps({"conversation_id": cid, "content": "ping"}), + ) + events = _parse_sse(r.text) + assert events[0][0] == "session_expired" + # Engine must NOT have been called with a doomed bearer. + assert len(engine.calls) == 0 + + +def test_chat_stream_session_expired(app_and_client, deps): + _, client = app_and_client + csrf = _login(client) + r = client.post("/app/api/conversations", + headers={"X-CSRF-Token": csrf, "Content-Type": "application/json"}, + content="{}") + cid = r.json()["conversation"]["id"] + + # Expire the bearer + sessions = deps.session_store.list_for_client("c") + deps.session_store._db.conn().execute( + "UPDATE sessions SET mcp_bearer_expires_at = 0 WHERE session_id = ?", + (sessions[0].session_id,), + ) + + r = client.post( + "/app/api/chat/stream", + headers={"X-CSRF-Token": csrf, "Content-Type": "application/json"}, + content=json.dumps({"conversation_id": cid, "content": "ping"}), + ) + events = _parse_sse(r.text) + assert events[0][0] == "session_expired" + + +def test_chat_stream_invalid_body(app_and_client): + _, client = app_and_client + csrf = _login(client) + r = client.post( + "/app/api/chat/stream", + headers={"X-CSRF-Token": csrf, "Content-Type": "application/json"}, + content=json.dumps({}), + ) + assert r.status_code == 400 + + +def test_chat_stream_not_found(app_and_client): + _, client = app_and_client + csrf = _login(client) + r = client.post( + "/app/api/chat/stream", + headers={"X-CSRF-Token": csrf, "Content-Type": "application/json"}, + content=json.dumps({"conversation_id": "nope", "content": "x"}), + ) + assert r.status_code == 404 + + +def test_chat_stream_error_event(app_and_client, engine): + _, client = app_and_client + engine.script = FakeScript(events=[ + TextDelta(text="partial "), + ErrorEvent(code="boom", message="tool failure"), + ]) + csrf = _login(client) + r = client.post("/app/api/conversations", + headers={"X-CSRF-Token": csrf, "Content-Type": "application/json"}, + content="{}") + cid = r.json()["conversation"]["id"] + r = client.post( + "/app/api/chat/stream", + headers={"X-CSRF-Token": csrf, "Content-Type": "application/json"}, + content=json.dumps({"conversation_id": cid, "content": "ping"}), + ) + events = _parse_sse(r.text) + codes = [e[0] for e in events] + assert "error" in codes + # Stream should still emit a done event so client state is coherent. + assert "done" in codes + + +def test_chat_persists_history_for_second_turn(app_and_client, engine, deps): + _, client = app_and_client + engine.script = FakeScript(events=[TextDelta(text="ack")], title_text="t") + csrf = _login(client) + r = client.post("/app/api/conversations", + headers={"X-CSRF-Token": csrf, "Content-Type": "application/json"}, + content="{}") + cid = r.json()["conversation"]["id"] + + client.post( + "/app/api/chat/stream", + headers={"X-CSRF-Token": csrf, "Content-Type": "application/json"}, + content=json.dumps({"conversation_id": cid, "content": "first"}), + ) + client.post( + "/app/api/chat/stream", + headers={"X-CSRF-Token": csrf, "Content-Type": "application/json"}, + content=json.dumps({"conversation_id": cid, "content": "second"}), + ) + # Second turn: history should contain the first user + first assistant. + history = engine.calls[1].history + assert [m.role for m in history] == ["user", "assistant"] + assert history[0].content == "first" + assert history[1].content == "ack" + + +def test_chat_stream_ssh_tool_emits_confirm_event(app_and_client, engine, deps): + """ssh_run triggers a tool_confirm_required SSE frame and the + engine must wait for a decision via /app/api/chat/confirm. + """ + import threading + import time as _t + + engine.script = FakeScript(events=[ + ToolCallStart(id="tc1", name="ssh_run", args={"host": "pve1", "command": "ls"}), + ToolConfirmRequired(id="tc1", name="ssh_run", args={"host": "pve1", "command": "ls"}), + ToolCallEnd(id="tc1", status="ok", preview="ok", duration_ms=50), + TextDelta(text="done"), + ]) + + _, client = app_and_client + csrf = _login(client) + r = client.post("/app/api/conversations", + headers={"X-CSRF-Token": csrf, "Content-Type": "application/json"}, + content="{}") + cid = r.json()["conversation"]["id"] + + # Approve the confirmation from a parallel thread so the SSE stream + # is free to unblock and finish. + def approve_when_ready(): + deadline = _t.time() + 5 + while _t.time() < deadline: + pending = deps.confirmations.pending_for( + deps.session_store.list_for_client("c")[0].session_id + ) + if pending: + client.post( + "/app/api/chat/confirm", + headers={"X-CSRF-Token": csrf, "Content-Type": "application/json"}, + content=json.dumps({"call_id": pending[0], "approve": True}), + ) + return + _t.sleep(0.05) + + t = threading.Thread(target=approve_when_ready) + t.start() + try: + r = client.post( + "/app/api/chat/stream", + headers={"X-CSRF-Token": csrf, "Content-Type": "application/json"}, + content=json.dumps({"conversation_id": cid, "content": "run ls"}), + ) + finally: + t.join(timeout=5) + + events = _parse_sse(r.text) + names = [e[0] for e in events] + assert "tool_confirm_required" in names + assert "tool_result" in names + # done must still fire so the client state settles. + assert "done" in names + + +def test_confirm_endpoint_requires_csrf(app_and_client): + _, client = app_and_client + _login(client) + r = client.post( + "/app/api/chat/confirm", + headers={"Content-Type": "application/json"}, + content=json.dumps({"call_id": "x", "approve": True}), + ) + assert r.status_code == 403 + + +def test_confirm_endpoint_rejects_unknown_call_id(app_and_client): + _, client = app_and_client + csrf = _login(client) + r = client.post( + "/app/api/chat/confirm", + headers={"X-CSRF-Token": csrf, "Content-Type": "application/json"}, + content=json.dumps({"call_id": "nope", "approve": True}), + ) + assert r.status_code == 404 + + +# --------------------------------------------------------------------------- +# Tokens page +# --------------------------------------------------------------------------- + +def test_tokens_page_requires_auth(app_and_client): + _, client = app_and_client + r = client.get("/app/tokens") + assert r.status_code == 302 + assert r.headers["location"] == "/app/login" + + +def test_tokens_page_lists_empty(app_and_client): + _, client = app_and_client + _login(client) + r = client.get("/app/tokens") + assert r.status_code == 200 + assert "No active tokens" in r.text + assert "0 / 3" in r.text # count indicator + + +def test_tokens_create_requires_name(app_and_client): + _, client = app_and_client + csrf = _login(client) + r = client.post( + "/app/tokens", + headers={"X-CSRF-Token": csrf, "Content-Type": "application/x-www-form-urlencoded"}, + data={"csrf_token": csrf, "name": "", "totp": "123456"}, + ) + assert r.status_code == 200 + assert "Name is required" in r.text + + +def test_tokens_create_requires_totp(app_and_client): + _, client = app_and_client + csrf = _login(client) + r = client.post( + "/app/tokens", + data={"csrf_token": csrf, "name": "Gemini Web", "totp": "999999"}, + ) + assert r.status_code == 200 + assert "Invalid 2FA code" in r.text + + +def test_tokens_create_success(app_and_client, deps): + _, client = app_and_client + csrf = _login(client) + r = client.post( + "/app/tokens", + data={"csrf_token": csrf, "name": "Gemini Web", "totp": "123456"}, + ) + assert r.status_code == 200 + assert "New token: Gemini Web" in r.text + assert deps.token_store.count_named("c") == 1 + + +def test_tokens_create_enforces_cap(app_and_client, deps): + _, client = app_and_client + csrf = _login(client) + for i in range(3): + r = client.post( + "/app/tokens", + data={"csrf_token": csrf, "name": f"Client {i}", "totp": "123456"}, + ) + assert r.status_code == 200, f"issue {i} failed" + # Fourth attempt should fail + r = client.post( + "/app/tokens", + data={"csrf_token": csrf, "name": "Client 4", "totp": "123456"}, + ) + assert r.status_code == 200 + assert "Limit reached" in r.text or "3 active tokens" in r.text + assert deps.token_store.count_named("c") == 3 + + +def test_tokens_revoke_removes_token(app_and_client, deps): + _, client = app_and_client + csrf = _login(client) + # Create two tokens + client.post("/app/tokens", + data={"csrf_token": csrf, "name": "A", "totp": "123456"}) + client.post("/app/tokens", + data={"csrf_token": csrf, "name": "B", "totp": "123456"}) + assert deps.token_store.count_named("c") == 2 + + rows = deps.token_store.list_named("c") + prefix = rows[0].token[:12] + + r = client.post( + "/app/tokens/revoke", + data={"csrf_token": csrf, "token_prefix": prefix}, + ) + assert r.status_code == 303 + assert deps.token_store.count_named("c") == 1 + + +def test_tokens_revoke_csrf_required(app_and_client): + _, client = app_and_client + _login(client) + r = client.post( + "/app/tokens/revoke", + data={"token_prefix": "abcdef"}, + ) + assert r.status_code == 403 + + +def test_tokens_revoke_scoped_to_client(app_and_client, deps): + """A client cannot revoke another client's token even with the right prefix.""" + _, client = app_and_client + csrf = _login(client) + client.post("/app/tokens", + data={"csrf_token": csrf, "name": "A", "totp": "123456"}) + + # Craft a fake token belonging to another client + deps.token_store._live["foreign_token_xxxxxxxxxxxx"] = "other" + deps.token_store._named["foreign_token_xxxxxxxxxxxx"] = ("ForeignApp", 0, 1e12) + + r = client.post( + "/app/tokens/revoke", + data={"csrf_token": csrf, "token_prefix": "foreign_toke"}, + ) + # Route still redirects (no enumeration), but the token stays alive. + assert r.status_code == 303 + assert deps.token_store.validate("foreign_token_xxxxxxxxxxxx") == "other" + + +def test_chat_page_auth_required(app_and_client): + _, client = app_and_client + r = client.get("/app/chat") + assert r.status_code == 302 + + +def test_chat_stream_remote_mode_still_routes_public_url(tmp_path, engine): + """In remote mode the public URL is still the resolved target. + + The engine itself refuses to drive a remote turn (the SDK's + backend-driven MCP mode is broken under our auth), but the URL + resolution logic is independent and continues to honour the + configured public hostname. That lets us keep remote-mode + configuration around for a future re-enablement without changing + the routing layer. + """ + db = Database(tmp_path / "dashboard.db") + deps = DashboardDeps( + database=db, + session_store=SessionStore(db, key=os.urandom(32)), + client_store=FakeClientStore(), + token_store=FakeTokenStore(), + totp_locked=lambda cid: False, + totp_record_failure=lambda cid: None, + totp_record_success=lambda cid: None, + conversations=ConversationStore(db), + engine=engine, + mcp_public_url="https://mcp.example/", + mcp_mode="remote", + ) + app = Starlette(routes=build_dashboard_routes(deps)) + client = TestClient(app, follow_redirects=False) + + csrf = _login(client) + r = client.post("/app/api/conversations", + headers={"X-CSRF-Token": csrf, "Content-Type": "application/json"}, + content="{}") + cid = r.json()["conversation"]["id"] + client.post( + "/app/api/chat/stream", + headers={"X-CSRF-Token": csrf, "Content-Type": "application/json"}, + content=json.dumps({"conversation_id": cid, "content": "ping"}), + ) + assert engine.calls[-1].mcp_url == "https://mcp.example/mcp" + assert engine.calls[-1].mcp_mode == "remote" + + +def test_gemini_engine_rejects_remote_mode(): + """GeminiChatEngine yields an ErrorEvent instead of calling the SDK.""" + import asyncio + from beaconmcp.dashboard.chat import ( + ErrorEvent, + GeminiChatEngine, + TurnInput, + ) + + engine_real = GeminiChatEngine(api_key="test") + turn = TurnInput( + history=[], user_text="x", model="gemini-3-flash-preview", + effort="low", bearer="b", + mcp_url="https://mcp.example/mcp", mcp_mode="remote", + ) + + async def _drain(): + return [e async for e in engine_real.run(turn)] + + events = asyncio.run(_drain()) + assert len(events) == 1 + assert isinstance(events[0], ErrorEvent) + assert events[0].code == "remote_mode_disabled" + + +def test_chat_page_renders_after_login(app_and_client): + _, client = app_and_client + _login(client) + r = client.get("/app/chat") + assert r.status_code == 200 + assert "chat-root" in r.text + assert "gemini-3.6-flash" in r.text + assert "gemini-3.5-flash-lite" in r.text + assert "gemini-3.1-pro-preview" in r.text + + +def test_needs_confirmation_includes_run_tools(): + """Every tool that can fire arbitrary shell on a host/VM must require + human approval -- both the SSH and the QEMU Guest Agent exec paths, + now unified as ``ssh_run`` / ``proxmox_run``. Legacy ``*_exec_command*`` + names stay in the allow-list defensively in case an older MCP server + is still wired up. + """ + from beaconmcp.dashboard.chat import ( + _NEEDS_CONFIRMATION, + _tool_call_requires_confirmation, + ) + + # Unified names: required. + assert "ssh_run" in _NEEDS_CONFIRMATION + assert "proxmox_run" in _NEEDS_CONFIRMATION + # Legacy names: still guarded. + assert "ssh_exec_command" in _NEEDS_CONFIRMATION + assert "ssh_exec_command_async" in _NEEDS_CONFIRMATION + assert "proxmox_exec_command" in _NEEDS_CONFIRMATION + assert "proxmox_exec_command_async" in _NEEDS_CONFIRMATION + + # Sync + async-start (command present) must confirm on unified tools. + assert _tool_call_requires_confirmation( + "ssh_run", {"host": "pve1", "command": "ls"} + ) + assert _tool_call_requires_confirmation( + "proxmox_run", {"node": "pve1", "vmid": 101, "command": "ls"} + ) + assert _tool_call_requires_confirmation( + "ssh_run", {"host": "pve1", "command": "ls", "wait": False} + ) + + # Poll-only call (exec_id, no command) is read-only: no modal. + assert not _tool_call_requires_confirmation("ssh_run", {"exec_id": "abc"}) + assert not _tool_call_requires_confirmation("proxmox_run", {"exec_id": "abc"}) + + # Legacy sync tools still prompt (no poll-exempt shortcut -- they + # always carry a ``command``). + assert _tool_call_requires_confirmation( + "ssh_exec_command", {"host": "pve1", "command": "ls"} + ) + + # Read-only result-fetchers and unrelated tools never confirm. + assert "proxmox_exec_get_result" not in _NEEDS_CONFIRMATION + assert "ssh_exec_get_result" not in _NEEDS_CONFIRMATION + assert not _tool_call_requires_confirmation("proxmox_list_nodes", {}) + assert not _tool_call_requires_confirmation("cluster_overview", {}) + + +def test_self_update_needs_confirmation_when_applied(): + """Applying an update is the most consequential call this server has. + + ``beaconmcp_self_update(confirm=True)`` pulls new code, reinstalls + dependencies and restarts the service. Nothing stops an injected + instruction from asking for it, so it must reach the modal -- while the + ``confirm=False`` preview shape stays a plain read. + """ + from beaconmcp.dashboard.chat import _tool_call_requires_confirmation + + assert _tool_call_requires_confirmation("beaconmcp_self_update", {"confirm": True}) + assert not _tool_call_requires_confirmation("beaconmcp_self_update", {"confirm": False}) + assert not _tool_call_requires_confirmation("beaconmcp_self_update", {}) + # Checking is read-only and must not raise a modal. + assert not _tool_call_requires_confirmation("beaconmcp_check_update", {}) + + +def test_vm_create_with_config_needs_confirmation(): + """A ``config`` dict on vm-create can carry code-execution keys. + + ``proxmox_vm_create`` forwards its ``config`` straight to the PVE API, so + an injected instruction could smuggle a ``hookscript`` or raw QEMU + ``args`` that runs code on the host. The create-with-config shape must + reach the modal, while a bare create (an empty shell VM) stays a plain + unattended add like clone/start. + """ + from beaconmcp.dashboard.chat import _tool_call_requires_confirmation + + assert _tool_call_requires_confirmation( + "proxmox_vm_create", + {"node": "pve1", "vmid": 100, "config": {"hookscript": "local:snippets/x.sh"}}, + ) + # A benign config still gets the modal -- we gate on presence, not on + # inspecting keys, mirroring the proxmox_vm_config ``updates`` treatment. + assert _tool_call_requires_confirmation( + "proxmox_vm_create", {"node": "pve1", "vmid": 100, "config": {"cores": 2}}, + ) + # No config -> empty shell, nothing to smuggle, no modal. + assert not _tool_call_requires_confirmation( + "proxmox_vm_create", {"node": "pve1", "vmid": 100}, + ) + assert not _tool_call_requires_confirmation( + "proxmox_vm_create", {"node": "pve1", "vmid": 100, "config": None}, + ) + + +def test_every_registered_tool_is_gated_or_deliberately_not(): + """Guard against a new tool quietly landing outside the gate. + + The list below is the reviewed set of tools that may run unattended: + reads, and calls that only add (create/clone/start/backup). A new tool + name showing up here means someone must decide which side it is on -- + which is exactly how ``beaconmcp_self_update`` slipped through when the + self-update tools were added. + """ + import pathlib + import re + + from beaconmcp.dashboard.chat import ( + _CONFIRM_WHEN_ARG_PRESENT, + _NEEDS_CONFIRMATION, + ) + + src_root = pathlib.Path(__file__).parent.parent / "src" / "beaconmcp" + registered: set[str] = set() + for path in src_root.rglob("*.py"): + registered.update( + re.findall( + r"@mcp\.tool\((?:[^)]*|\s*\n(?:.*\n)*?\s*)\)\s*\n\s*(?:async )?def (\w+)", + path.read_text(), + ) + ) + assert "ssh_run" in registered, "tool discovery regex stopped matching" + + ungated_by_design = { + "beaconmcp_check_update", + "bmc_get_event_log", "bmc_health_status", "bmc_list_devices", + "bmc_power_on", "bmc_power_status", "bmc_server_info", + "cluster_health", "cluster_overview", "cluster_overview_interactive", + "proxmox_backup_create", "proxmox_backup_list", + "proxmox_get_logs", "proxmox_get_tasks", + "proxmox_list_nodes", "proxmox_list_transfers", "proxmox_list_vms", + "proxmox_logs_panel", "proxmox_network_config", "proxmox_node_status", + "proxmox_read_file", "proxmox_snapshot_create", "proxmox_snapshot_list", + "proxmox_storage_status", "proxmox_vm_clone", + "proxmox_vm_panel", "proxmox_vm_start", "proxmox_vm_status", + "security_end_session", "ssh_list_sessions", "vm_find", + } + gated = set(_NEEDS_CONFIRMATION) | set(_CONFIRM_WHEN_ARG_PRESENT) + unclassified = registered - gated - ungated_by_design + assert not unclassified, ( + "new tool(s) outside the confirmation gate; decide and list them: " + f"{sorted(unclassified)}" + ) diff --git a/tests/test_dashboard_integration.py b/tests/test_dashboard_integration.py new file mode 100644 index 0000000..ea0ab99 --- /dev/null +++ b/tests/test_dashboard_integration.py @@ -0,0 +1,485 @@ +"""Dashboard integration tests using Starlette's TestClient. + +Mocks ClientStore + TokenStore so we don't need the real auth backend. +Run with:: + + pytest tests/test_dashboard_integration.py -v +""" + +from __future__ import annotations + +import os +import sys +import time +from pathlib import Path + +import pytest +from starlette.applications import Starlette +from starlette.testclient import TestClient + +sys.path.insert(0, str(Path(__file__).parent.parent / "src")) + +from beaconmcp.auth import TOTP_REPLAY_MESSAGE, TotpResult +from beaconmcp.dashboard.app import ( + BEARER_TTL_SECONDS, + DashboardDeps, + SESSION_COOKIE, + build_dashboard_routes, +) +from beaconmcp.dashboard.csrf import CSRF_COOKIE +from beaconmcp.dashboard.db import Database +from beaconmcp.dashboard.session import SessionStore + + +# --------------------------------------------------------------------------- +# Mocks +# --------------------------------------------------------------------------- + +class FakeClientStore: + def __init__(self): + self.clients = { + "beaconmcp_test": { + "secret": "sk_test", + "name": "Test Client", + "totp": "123456", + } + } + + def verify(self, client_id, secret): + c = self.clients.get(client_id) + return bool(c and c["secret"] == secret) + + def check_totp(self, client_id, code): + c = self.clients.get(client_id) + return TotpResult.OK if (c and c["totp"] == code) else TotpResult.INVALID + + def get_name(self, client_id): + c = self.clients.get(client_id) + return c["name"] if c else None + + +class FakeTokenStore: + def __init__(self): + self._tokens: dict[str, str] = {} + self.next_ttl = BEARER_TTL_SECONDS + self.revoked: list[str] = [] + self._counter = 0 + + def issue(self, client_id): + self._counter += 1 + token = f"bearer_{client_id}_{self._counter}" + self._tokens[token] = client_id + return token, self.next_ttl + + def validate(self, token): + return self._tokens.get(token) + + def revoke(self, token): + self.revoked.append(token) + self._tokens.pop(token, None) + return True + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + +@pytest.fixture() +def deps(tmp_path, monkeypatch): + monkeypatch.setenv("BEACONMCP_DASHBOARD_DB", str(tmp_path / "dashboard.db")) + db = Database(tmp_path / "dashboard.db") + session_store = SessionStore(db, key=os.urandom(32)) + failures: dict[str, tuple[int, float]] = {} + + def totp_locked(cid): + e = failures.get(cid) + if not e: + return False + count, until = e + return count >= 5 and time.time() < until + + def totp_record_failure(cid): + c, _ = failures.get(cid, (0, 0.0)) + failures[cid] = (c + 1, time.time() + 300) + + def totp_record_success(cid): + failures.pop(cid, None) + + # A sentinel non-None engine so the post-login landing stays /app/chat. + # These integration tests exercise chat-mode routing; the tokens-only + # mode (engine=None) is covered separately. + return DashboardDeps( + database=db, + session_store=session_store, + client_store=FakeClientStore(), + token_store=FakeTokenStore(), + totp_locked=totp_locked, + totp_record_failure=totp_record_failure, + totp_record_success=totp_record_success, + engine=object(), # type: ignore[arg-type] + ) + + +@pytest.fixture() +def client(deps): + app = Starlette(routes=build_dashboard_routes(deps)) + return TestClient(app, follow_redirects=False) + + +@pytest.fixture() +def tokens_only_client(tmp_path): + """A fixture mirroring ``deps``/``client`` but with engine=None. + + Exercises the tokens-only mode of the dashboard: no Gemini key set, + so ``/app/chat`` redirects to ``/app/tokens`` and post-login lands + there directly. + """ + db = Database(tmp_path / "dashboard-tokens-only.db") + session_store = SessionStore(db, key=os.urandom(32)) + failures: dict[str, tuple[int, float]] = {} + + deps_local = DashboardDeps( + database=db, + session_store=session_store, + client_store=FakeClientStore(), + token_store=FakeTokenStore(), + totp_locked=lambda cid: ( + failures.get(cid, (0, 0.0))[0] >= 5 + and time.time() < failures.get(cid, (0, 0.0))[1] + ), + totp_record_failure=lambda cid: failures.__setitem__( + cid, + (failures.get(cid, (0, 0.0))[0] + 1, time.time() + 300), + ), + totp_record_success=lambda cid: (failures.pop(cid, None), None)[1], + engine=None, + ) + app = Starlette(routes=build_dashboard_routes(deps_local)) + return TestClient(app, follow_redirects=False) + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + +def _login_form(csrf_token: str, **overrides) -> dict: + data = { + "csrf_token": csrf_token, + "client_id": "beaconmcp_test", + "client_secret": "sk_test", + "totp": "123456", + "remember": "on", + } + data.update(overrides) + return data + + +def _csrf(client: TestClient) -> str: + """Hit the login page to obtain a CSRF cookie.""" + r = client.get("/app/login") + assert r.status_code == 200 + token = r.cookies.get(CSRF_COOKIE) + assert token, "CSRF cookie missing" + return token + + +def test_index_redirects_to_login(client): + r = client.get("/") + assert r.status_code == 302 + assert r.headers["location"] == "/app/login" + + +def test_login_page_renders(client): + r = client.get("/app/login") + assert r.status_code == 200 + assert "Sign in" in r.text + assert "Client ID" in r.text + assert r.cookies.get(CSRF_COOKIE) + + +def test_login_post_csrf_required(client): + r = client.post( + "/app/login", + data={ + "client_id": "beaconmcp_test", + "client_secret": "sk_test", + "totp": "123456", + }, + ) + assert r.status_code == 403 + assert r.json() == {"error": "csrf"} + + +def test_login_post_wrong_credentials(client): + token = _csrf(client) + r = client.post( + "/app/login", + data=_login_form(token, client_secret="wrong"), + ) + assert r.status_code == 401 + assert "Invalid credentials" in r.text + + +def test_login_post_wrong_totp(client): + token = _csrf(client) + r = client.post( + "/app/login", + data=_login_form(token, totp="000000"), + ) + assert r.status_code == 401 + assert "Invalid 2FA code" in r.text + + +def test_login_post_success(client, deps): + token = _csrf(client) + r = client.post("/app/login", data=_login_form(token)) + assert r.status_code == 303 + assert r.headers["location"] == "/app/chat" + assert r.cookies.get(SESSION_COOKIE) + # Session persisted + sessions = deps.session_store.list_for_client("beaconmcp_test") + assert len(sessions) == 1 + assert sessions[0].mcp_bearer.startswith("bearer_beaconmcp_test_") + + +def test_chat_requires_session(client): + r = client.get("/app/chat") + assert r.status_code == 302 + assert r.headers["location"] == "/app/login" + + +def test_chat_accessible_after_login(client): + token = _csrf(client) + client.post("/app/login", data=_login_form(token)) + r = client.get("/app/chat") + assert r.status_code == 200 + assert "beaconmcp_test" in r.text + + +def test_logout_revokes_bearer(client, deps): + token = _csrf(client) + client.post("/app/login", data=_login_form(token)) + sessions = deps.session_store.list_for_client("beaconmcp_test") + bearer = sessions[0].mcp_bearer + + # CSRF cookie is rotated on login, fetch the fresh one. + new_token = client.cookies.get(CSRF_COOKIE) + r = client.post("/app/logout", data={"csrf_token": new_token}) + assert r.status_code == 303 + assert r.headers["location"] == "/app/login" + assert bearer in deps.token_store.revoked + assert deps.session_store.list_for_client("beaconmcp_test") == [] + + +def test_refresh_requires_session(client): + r = client.get("/app/refresh") + assert r.status_code == 302 + assert r.headers["location"] == "/app/login" + + +def test_refresh_when_bearer_still_valid_redirects_to_chat(client): + token = _csrf(client) + client.post("/app/login", data=_login_form(token)) + r = client.get("/app/refresh") + assert r.status_code == 302 + assert r.headers["location"] == "/app/chat" + + +def test_refresh_when_bearer_expired_renders_form(client, deps): + token = _csrf(client) + client.post("/app/login", data=_login_form(token)) + sessions = deps.session_store.list_for_client("beaconmcp_test") + deps.session_store._db.conn().execute( + "UPDATE sessions SET mcp_bearer_expires_at = ? WHERE session_id = ?", + (0, sessions[0].session_id), + ) + + r = client.get("/app/refresh") + assert r.status_code == 200 + assert "Test Client" in r.text + # New UI replaces the "2FA code" label with the 6-digit boxes + copy. + assert "6-digit code" in r.text + + +def test_refresh_post_re_issues_bearer(client, deps): + token = _csrf(client) + client.post("/app/login", data=_login_form(token)) + sessions = deps.session_store.list_for_client("beaconmcp_test") + sid = sessions[0].session_id + old_bearer = sessions[0].mcp_bearer + deps.session_store._db.conn().execute( + "UPDATE sessions SET mcp_bearer_expires_at = ? WHERE session_id = ?", + (0, sid), + ) + + new_token = client.cookies.get(CSRF_COOKIE) + r = client.post( + "/app/refresh", + data={"csrf_token": new_token, "totp": "123456"}, + ) + assert r.status_code == 303 + assert r.headers["location"] == "/app/chat" + + refreshed = deps.session_store.load(sid) + assert refreshed is not None + assert refreshed.mcp_bearer != old_bearer + assert refreshed.bearer_valid() + + +def test_refresh_wrong_totp(client, deps): + token = _csrf(client) + client.post("/app/login", data=_login_form(token)) + sessions = deps.session_store.list_for_client("beaconmcp_test") + deps.session_store._db.conn().execute( + "UPDATE sessions SET mcp_bearer_expires_at = ? WHERE session_id = ?", + (0, sessions[0].session_id), + ) + + new_token = client.cookies.get(CSRF_COOKIE) + r = client.post( + "/app/refresh", + data={"csrf_token": new_token, "totp": "999999"}, + ) + assert r.status_code == 401 + assert "Invalid 2FA code" in r.text + + +def test_login_after_5_failed_totp_locks_out(client, deps): + token = _csrf(client) + for _ in range(5): + r = client.post("/app/login", data=_login_form(token, totp="000000")) + assert r.status_code == 401 + + r = client.post("/app/login", data=_login_form(token)) + assert r.status_code == 429 + assert "Too many attempts" in r.text + + +def test_replayed_totp_never_locks_the_operator_out(deps, monkeypatch): + """Re-submitting an already-spent code is an ordinary mistake, not a + failed auth attempt: the dashboard asks for 2FA on several actions, so an + operator hitting two of them inside one 30 s step would otherwise burn + through the 5-strike lockout without ever typing a wrong code. + """ + # The operator already spent this code on a previous action. + spent = {"123456"} + + def check_totp(client_id, code): + if code != "123456": + return TotpResult.INVALID + return TotpResult.REPLAY if code in spent else TotpResult.OK + + monkeypatch.setattr(deps.client_store, "check_totp", check_totp) + client = TestClient( + Starlette(routes=build_dashboard_routes(deps)), follow_redirects=False + ) + token = _csrf(client) + + # Well past the 5-strike lockout, if replays counted. + for _ in range(10): + r = client.post("/app/login", data=_login_form(token)) + assert r.status_code == 401 + assert TOTP_REPLAY_MESSAGE in r.text + # Not "wrong code" -- the operator must not be sent clock-hunting. + assert "clock is in sync" not in r.text + + assert not deps.totp_locked("beaconmcp_test") + # And the next code rolls over into a normal login. + spent.clear() + assert client.post("/app/login", data=_login_form(token)).status_code == 303 + + +def test_logout_csrf_required(client): + r = client.post("/app/logout", data={}) + assert r.status_code == 403 + + +def test_existing_session_skips_login_page(client): + token = _csrf(client) + client.post("/app/login", data=_login_form(token)) + r = client.get("/app/login") + assert r.status_code == 302 + assert r.headers["location"] == "/app/chat" + + +def test_existing_session_with_stale_bearer_redirects_to_refresh(client, deps): + token = _csrf(client) + client.post("/app/login", data=_login_form(token)) + sessions = deps.session_store.list_for_client("beaconmcp_test") + deps.session_store._db.conn().execute( + "UPDATE sessions SET mcp_bearer_expires_at = ? WHERE session_id = ?", + (0, sessions[0].session_id), + ) + + r = client.get("/app/login") + assert r.status_code == 302 + assert r.headers["location"] == "/app/refresh" + + +def test_security_headers_present(client): + r = client.get("/app/login") + assert r.headers.get("X-Frame-Options") == "DENY" + assert r.headers.get("X-Content-Type-Options") == "nosniff" + assert "Referrer-Policy" in r.headers + assert "Content-Security-Policy" in r.headers + + +# --------------------------------------------------------------------------- +# Tokens-only mode (engine=None): chat redirects to tokens +# --------------------------------------------------------------------------- + + +def test_tokens_only_login_lands_on_tokens(tokens_only_client): + r = tokens_only_client.get("/app/login") + token = r.cookies.get(CSRF_COOKIE) + assert token + r = tokens_only_client.post( + "/app/login", + data={ + "csrf_token": token, + "client_id": "beaconmcp_test", + "client_secret": "sk_test", + "totp": "123456", + "remember": "on", + }, + ) + assert r.status_code == 303 + assert r.headers["location"] == "/app/tokens" + + +def test_tokens_only_chat_redirects_to_tokens(tokens_only_client): + r = tokens_only_client.get("/app/login") + token = r.cookies.get(CSRF_COOKIE) + tokens_only_client.post( + "/app/login", + data={ + "csrf_token": token, + "client_id": "beaconmcp_test", + "client_secret": "sk_test", + "totp": "123456", + "remember": "on", + }, + ) + r = tokens_only_client.get("/app/chat") + assert r.status_code == 302 + assert r.headers["location"] == "/app/tokens" + + +def test_tokens_only_login_page_redirects_when_authenticated(tokens_only_client): + r = tokens_only_client.get("/app/login") + token = r.cookies.get(CSRF_COOKIE) + tokens_only_client.post( + "/app/login", + data={ + "csrf_token": token, + "client_id": "beaconmcp_test", + "client_secret": "sk_test", + "totp": "123456", + "remember": "on", + }, + ) + # Already authenticated: /app/login should bounce to the tokens page + # since there is no chat engine configured. + r = tokens_only_client.get("/app/login") + assert r.status_code == 302 + assert r.headers["location"] == "/app/tokens" diff --git a/tests/test_dashboard_mcp_apps.py b/tests/test_dashboard_mcp_apps.py new file mode 100644 index 0000000..86722d9 --- /dev/null +++ b/tests/test_dashboard_mcp_apps.py @@ -0,0 +1,560 @@ +"""MCP Apps support in the dashboard chat (#35). + +The dashboard is both the MCP client that negotiates the extension and the +host that renders the ``ui://`` documents. These tests pin the parts a +refactor could quietly break without any visible symptom: the capability on +the wire, the headers that decide whether a panel can be framed at all and +what it may reach from inside the frame, and the allow-list that says which +tool calls a panel may make on its own. +""" + +from __future__ import annotations + +import json +import os +import sys +from contextlib import asynccontextmanager +from pathlib import Path + +import pytest +from starlette.applications import Starlette +from starlette.testclient import TestClient + +sys.path.insert(0, str(Path(__file__).parent.parent / "src")) + +from beaconmcp.dashboard import mcp_bridge +from beaconmcp.dashboard.app import DashboardDeps, build_dashboard_routes +from beaconmcp.dashboard.chat import ( + FakeChatEngine, + FakeScript, + ToolCallEnd, + ToolCallStart, + assemble_assistant_message, + format_app_context, + panel_call_allowed, +) +from beaconmcp.dashboard.confirmations import ConfirmationStore +from beaconmcp.dashboard.conversations import ConversationStore +from beaconmcp.dashboard.csrf import CSRF_COOKIE +from beaconmcp.dashboard.db import Database +from beaconmcp.dashboard.session import SessionStore + +from test_dashboard_chat import FakeClientStore, FakeTokenStore # noqa: E402 + + +PANEL_URI = "ui://beaconmcp/vm-panel.html" +PANEL_HTML = "vm" + + +# --------------------------------------------------------------------------- +# Client half: capability negotiation and _meta.ui plumbing +# --------------------------------------------------------------------------- + +@pytest.mark.anyio +async def test_initialize_declares_the_apps_extension(monkeypatch): + """The extension rides on ``ClientCapabilities`` even on mcp 1.x. + + ``extensions`` only becomes a typed field in mcp 2.0, which is what made + this look blocked behind the ``<2`` pin. It is not: the model allows + extras, so the field serialises under the name the spec gives it. If a + future SDK bump makes ``ClientCapabilities`` strict, this fails here + rather than silently dropping the capability at runtime. + """ + from mcp import types + from mcp.client.session import ClientSession + + captured: dict = {} + + async def fake_send_request(self, request, result_type, **kwargs): + captured["request"] = request + return None + + monkeypatch.setattr(ClientSession, "send_request", fake_send_request) + + session = object.__new__(mcp_bridge.AppsClientSession) + request = types.ClientRequest( + types.InitializeRequest( + params=types.InitializeRequestParams( + protocolVersion=types.LATEST_PROTOCOL_VERSION, + capabilities=types.ClientCapabilities(), + clientInfo=types.Implementation(name="t", version="1"), + ), + ) + ) + await session.send_request(request, types.InitializeResult) + + wire = captured["request"].model_dump(by_alias=True, exclude_none=True) + extensions = wire["params"]["capabilities"]["extensions"] + assert extensions == { + "io.modelcontextprotocol/ui": { + "mimeTypes": ["text/html;profile=mcp-app"], + } + } + + +@pytest.mark.anyio +async def test_non_initialize_requests_pass_through_untouched(monkeypatch): + """Only the handshake is rewritten; every other request goes as-is.""" + from mcp import types + from mcp.client.session import ClientSession + + captured: dict = {} + + async def fake_send_request(self, request, result_type, **kwargs): + captured["request"] = request + return None + + monkeypatch.setattr(ClientSession, "send_request", fake_send_request) + + session = object.__new__(mcp_bridge.AppsClientSession) + original = types.ClientRequest(types.ListToolsRequest(method="tools/list")) + await session.send_request(original, types.ListToolsResult) + + assert captured["request"] is original + assert captured["request"].model_dump(by_alias=True, exclude_none=True) == { + "method": "tools/list", + } + + +def test_ui_resource_uri_extraction(): + assert mcp_bridge.ui_resource_uri({"ui": {"resourceUri": PANEL_URI}}) == PANEL_URI + assert mcp_bridge.ui_resource_uri(None) is None + assert mcp_bridge.ui_resource_uri({"ui": {}}) is None + # A non-ui:// target is not a panel. Honouring one would let a tool + # point the host's iframe at an arbitrary URL. + assert mcp_bridge.ui_resource_uri({"ui": {"resourceUri": "https://evil/x"}}) is None + assert mcp_bridge.ui_resource_uri({"ui": {"resourceUri": "file:///etc/passwd"}}) is None + + +def test_ui_resource_uris_by_tool(): + class _Tool: + def __init__(self, name, meta): + self.name = name + self.meta = meta + + mapping = mcp_bridge.ui_resource_uris_by_tool([ + _Tool("proxmox_vm_panel", {"ui": {"resourceUri": PANEL_URI}}), + _Tool("proxmox_list_vms", None), + ]) + assert mapping == {"proxmox_vm_panel": PANEL_URI} + + +class _FakeResourceContent: + def __init__(self, text, mime): + self.text = text + self.mimeType = mime + + +class _FakeReadResult: + def __init__(self, contents): + self.contents = contents + + +class _FakeSession: + def __init__(self, contents): + self._contents = contents + + async def read_resource(self, uri): + return _FakeReadResult(self._contents) + + +@pytest.mark.anyio +async def test_read_ui_resource_accepts_an_app_document(): + session = _FakeSession([_FakeResourceContent(PANEL_HTML, "text/html;profile=mcp-app")]) + assert await mcp_bridge.read_ui_resource(session, PANEL_URI) == PANEL_HTML + + +@pytest.mark.anyio +async def test_read_ui_resource_refuses_a_plain_resource(): + """The MIME check is what stops this being a generic resource proxy.""" + session = _FakeSession([_FakeResourceContent("nodes: []", "text/plain")]) + with pytest.raises(mcp_bridge.UiResourceError): + await mcp_bridge.read_ui_resource(session, PANEL_URI) + + +@pytest.mark.anyio +async def test_read_ui_resource_refuses_a_non_ui_scheme(): + session = _FakeSession([_FakeResourceContent("x", "text/html;profile=mcp-app")]) + with pytest.raises(mcp_bridge.UiResourceError): + await mcp_bridge.read_ui_resource(session, "beaconmcp://infrastructure") + + +# --------------------------------------------------------------------------- +# The panel allow-list -- the decision #35 asked to be made explicitly +# --------------------------------------------------------------------------- + +@pytest.mark.parametrize("name,args", [ + ("proxmox_vm_panel", {"node": "pve1", "vmid": 104}), + ("cluster_overview_interactive", {}), + ("proxmox_vm_start", {"node": "pve1", "vmid": 104}), + ("proxmox_vm_stop", {"node": "pve1", "vmid": 104}), + ("proxmox_vm_restart", {"node": "pve1", "vmid": 104}), + ("proxmox_vm_config", {"node": "pve1", "vmid": 104, "updates": {"cores": 4}}), +]) +def test_panel_may_drive_guest_lifecycle(name, args): + """A labelled button on one guest is the click; a modal would restate it.""" + assert panel_call_allowed(name, args) is True + + +@pytest.mark.parametrize("name,args", [ + ("ssh_run", {"host": "pve1", "command": "rm -rf /"}), + ("proxmox_run", {"node": "pve1", "vmid": 104, "command": "id"}), + ("proxmox_write_file", {"node": "pve1", "vmid": 104, "path": "/root/.ssh/authorized_keys"}), + ("vm_bulk_action", {"vmids": [1, 2, 3], "action": "stop"}), + ("proxmox_snapshot_rollback", {"node": "pve1", "vmid": 104, "snapname": "s"}), + ("proxmox_backup_restore", {"node": "pve1", "vmid": 104}), + ("bmc_power_off", {"device_id": "rack1"}), + # updates is an open-ended guest config. Exempting the tool wholesale + # would exempt hookscript, raw QEMU args and device passthrough with it. + ("proxmox_vm_config", {"node": "pve1", "vmid": 104, + "updates": {"hookscript": "local:snippets/x.sh"}}), + ("proxmox_vm_config", {"node": "pve1", "vmid": 104, + "updates": {"cores": 4, "args": "-device x"}}), + # Pulls new code, reinstalls dependencies and restarts the service. + ("beaconmcp_self_update", {"confirm": True}), +]) +def test_panel_may_not_reach_the_gated_tools(name, args): + """The exemption is a closed list, not "iframe calls skip the gate". + + A ui:// document is HTML the server wrote. Letting one through the gate + by virtue of being in a frame would hand every connected MCP server a + way around the approval it is documented to be subject to. + """ + assert panel_call_allowed(name, args) is False + + +# --------------------------------------------------------------------------- +# Model context from panels +# --------------------------------------------------------------------------- + +def test_format_app_context_labels_the_source(): + block = format_app_context([ + {"tool": "proxmox_vm_panel", "text": "VM 104 is now stopped.", "structured": {"status": "stopped"}}, + ]) + assert "proxmox_vm_panel" in block + assert "VM 104 is now stopped." in block + assert '"status": "stopped"' in block + # The model has to know this did not come from the operator's keyboard. + assert "did not go through you" in block + + +def test_format_app_context_empty(): + assert format_app_context([]) == "" + assert format_app_context(None) == "" + assert format_app_context([{"tool": "x"}]) == "" + + +def test_format_app_context_truncates(): + block = format_app_context([ + {"tool": "cluster", "text": "x" * 10_000, "structured": None}, + ]) + assert len(block) < 5_000 + + +def test_assemble_persists_only_the_panel_uri(): + """The snapshot is not stored -- reopening refetches instead.""" + content, tool_calls, _ = assemble_assistant_message([ + ToolCallStart(id="1", name="proxmox_vm_panel", args={"node": "pve1", "vmid": 104}), + ToolCallEnd( + id="1", status="ok", preview="{}", duration_ms=12, + ui={"resourceUri": PANEL_URI, "result": {"structuredContent": {"vmid": 104}}}, + ), + ]) + assert tool_calls[0].ui_resource_uri == PANEL_URI + assert "result" not in tool_calls[0].to_json() + + +# --------------------------------------------------------------------------- +# Host routes +# --------------------------------------------------------------------------- + +@pytest.fixture() +def engine(): + return FakeChatEngine(FakeScript(events=[ + ToolCallStart(id="fc_0", name="proxmox_vm_panel", args={"node": "pve1", "vmid": 104}), + ToolCallEnd( + id="fc_0", status="ok", preview='{"vmid": 104}', duration_ms=7, + ui={ + "resourceUri": PANEL_URI, + "result": {"content": [], "isError": False, + "structuredContent": {"vmid": 104, "status": "running"}}, + }, + ), + ])) + + +@pytest.fixture() +def deps(tmp_path, engine): + db = Database(tmp_path / "dashboard.db") + return DashboardDeps( + database=db, + session_store=SessionStore(db, key=os.urandom(32)), + client_store=FakeClientStore(), + token_store=FakeTokenStore(), + totp_locked=lambda cid: False, + totp_record_failure=lambda cid: None, + totp_record_success=lambda cid: None, + conversations=ConversationStore(db), + engine=engine, + confirmations=ConfirmationStore(), + ) + + +@pytest.fixture() +def client(deps): + mcp_bridge.cache_clear() + app = Starlette(routes=build_dashboard_routes(deps)) + return TestClient(app, follow_redirects=False) + + +def _login(client): + r = client.get("/app/login") + csrf = r.cookies.get(CSRF_COOKIE) + r = client.post("/app/login", data={ + "csrf_token": csrf, "client_id": "c", + "client_secret": "s", "totp": "123456", "remember": "on", + }) + assert r.status_code == 303 + return client.cookies.get(CSRF_COOKIE) + + +class _StubSession: + """Stands in for a live MCP session in the two host routes.""" + + def __init__(self): + self.calls: list[tuple[str, dict]] = [] + + async def call_tool(self, name, args): + self.calls.append((name, args)) + + class _Result: + content = [] + isError = False + structuredContent = {"ok": True} + + return _Result() + + +@pytest.fixture() +def stub_mcp(monkeypatch): + stub = _StubSession() + + @asynccontextmanager + async def fake_open_session(url, bearer, **kwargs): + yield stub + + async def fake_read(session, uri): + if uri != PANEL_URI: + raise mcp_bridge.UiResourceError("unknown resource") + return PANEL_HTML + + monkeypatch.setattr(mcp_bridge, "open_session", fake_open_session) + monkeypatch.setattr(mcp_bridge, "read_ui_resource", fake_read) + return stub + + +# --- panel document --------------------------------------------------------- + +def test_panel_route_requires_a_session(client, stub_mcp): + r = client.get(f"/app/api/mcp/panel?uri={PANEL_URI}") + assert r.status_code == 401 + + +def test_panel_route_rejects_a_non_ui_uri(client, stub_mcp): + _login(client) + r = client.get("/app/api/mcp/panel?uri=beaconmcp://infrastructure") + assert r.status_code == 400 + assert r.json()["error"] == "invalid_uri" + + +def test_panel_route_serves_the_document_with_its_own_csp(client, stub_mcp): + """Framable by us, inline script allowed, and no way out of the frame. + + The default dashboard headers are the opposite on both counts -- + ``X-Frame-Options: DENY`` would stop the panel rendering at all and + ``script-src 'self'`` would kill the inline bridge -- so this route + setting its own is load-bearing, not tidiness. + """ + _login(client) + r = client.get(f"/app/api/mcp/panel?uri={PANEL_URI}") + assert r.status_code == 200 + assert r.text == PANEL_HTML + + csp = r.headers["content-security-policy"] + assert "frame-ancestors 'self'" in csp + assert "script-src 'unsafe-inline'" in csp + # The panel reaches the cluster through its parent, never directly. + assert "connect-src 'none'" in csp + assert "default-src 'none'" in csp + assert r.headers["x-frame-options"] == "SAMEORIGIN" + assert "DENY" not in r.headers["x-frame-options"] + + +def test_panel_route_reports_an_unreadable_resource(client, stub_mcp): + _login(client) + r = client.get("/app/api/mcp/panel?uri=ui://beaconmcp/nope.html") + assert r.status_code == 502 + assert r.json()["error"] == "resource_unavailable" + + +def test_panel_documents_are_cached(client, stub_mcp, monkeypatch): + _login(client) + reads: list[str] = [] + + async def counting_read(session, uri): + reads.append(uri) + return PANEL_HTML + + monkeypatch.setattr(mcp_bridge, "read_ui_resource", counting_read) + for _ in range(3): + assert client.get(f"/app/api/mcp/panel?uri={PANEL_URI}").status_code == 200 + assert len(reads) == 1 + + +# --- tools/call relay ------------------------------------------------------- + +def test_relay_requires_a_session(client, stub_mcp): + r = client.post( + "/app/api/mcp/call", + headers={"Content-Type": "application/json"}, + content=json.dumps({"name": "proxmox_vm_stop", "arguments": {}}), + ) + assert r.status_code == 401 + + +def test_relay_requires_csrf(client, stub_mcp): + _login(client) + r = client.post( + "/app/api/mcp/call", + headers={"Content-Type": "application/json"}, + content=json.dumps({"name": "proxmox_vm_stop", "arguments": {}}), + ) + assert r.status_code == 403 + assert r.json()["error"] == "csrf" + + +def test_relay_runs_an_allowed_tool(client, stub_mcp): + csrf = _login(client) + r = client.post( + "/app/api/mcp/call", + headers={"X-CSRF-Token": csrf, "Content-Type": "application/json"}, + content=json.dumps({ + "name": "proxmox_vm_stop", "arguments": {"node": "pve1", "vmid": 104}, + }), + ) + assert r.status_code == 200 + assert r.json()["result"]["structuredContent"] == {"ok": True} + assert stub_mcp.calls == [("proxmox_vm_stop", {"node": "pve1", "vmid": 104})] + + +def test_relay_refuses_a_gated_tool_without_calling_it(client, stub_mcp): + csrf = _login(client) + r = client.post( + "/app/api/mcp/call", + headers={"X-CSRF-Token": csrf, "Content-Type": "application/json"}, + content=json.dumps({ + "name": "ssh_run", "arguments": {"host": "pve1", "command": "id"}, + }), + ) + assert r.status_code == 403 + assert r.json()["error"] == "confirmation_required" + assert stub_mcp.calls == [] + + +def test_relay_rejects_a_malformed_body(client, stub_mcp): + csrf = _login(client) + r = client.post( + "/app/api/mcp/call", + headers={"X-CSRF-Token": csrf, "Content-Type": "application/json"}, + content=json.dumps({"name": "", "arguments": {}}), + ) + assert r.status_code == 400 + + +# --- turn plumbing ---------------------------------------------------------- + +def _parse_sse(text): + events = [] + for frame in text.strip().split("\n\n"): + ev, data = None, [] + for line in frame.split("\n"): + if line.startswith("event:"): + ev = line[6:].strip() + elif line.startswith("data:"): + data.append(line[5:].strip()) + events.append((ev, json.loads("\n".join(data)) if data else {})) + return events + + +def test_tool_result_frame_carries_the_panel(client, deps): + """Without this the browser never learns there is a frame to mount.""" + csrf = _login(client) + r = client.post( + "/app/api/conversations", + headers={"X-CSRF-Token": csrf, "Content-Type": "application/json"}, + content="{}", + ) + conv_id = r.json()["conversation"]["id"] + + r = client.post( + "/app/api/chat/stream", + headers={"X-CSRF-Token": csrf, "Content-Type": "application/json"}, + content=json.dumps({"conversation_id": conv_id, "content": "open vm 104"}), + ) + results = [d for e, d in _parse_sse(r.text) if e == "tool_result"] + assert results and results[0]["ui"]["resourceUri"] == PANEL_URI + assert results[0]["ui"]["result"]["structuredContent"]["status"] == "running" + + stored = deps.conversations.list_messages(conv_id)[1].tool_calls[0] + assert stored.ui_resource_uri == PANEL_URI + + +def test_app_context_reaches_the_turn(client, engine): + csrf = _login(client) + r = client.post( + "/app/api/conversations", + headers={"X-CSRF-Token": csrf, "Content-Type": "application/json"}, + content="{}", + ) + conv_id = r.json()["conversation"]["id"] + + client.post( + "/app/api/chat/stream", + headers={"X-CSRF-Token": csrf, "Content-Type": "application/json"}, + content=json.dumps({ + "conversation_id": conv_id, + "content": "et maintenant ?", + "app_context": [ + {"tool": "proxmox_vm_panel", "text": "VM 104 is now stopped.", + "structured": {"status": "stopped"}}, + "junk", + {"tool": "x"}, + ], + }), + ) + assert engine.calls[0].app_context == [ + {"tool": "proxmox_vm_panel", "text": "VM 104 is now stopped.", + "structured": {"status": "stopped"}}, + ] + + +def test_app_context_entry_count_is_capped(client, engine): + csrf = _login(client) + r = client.post( + "/app/api/conversations", + headers={"X-CSRF-Token": csrf, "Content-Type": "application/json"}, + content="{}", + ) + conv_id = r.json()["conversation"]["id"] + client.post( + "/app/api/chat/stream", + headers={"X-CSRF-Token": csrf, "Content-Type": "application/json"}, + content=json.dumps({ + "conversation_id": conv_id, + "content": "ping", + "app_context": [ + {"tool": f"p{i}", "text": "x"} for i in range(50) + ], + }), + ) + assert len(engine.calls[0].app_context) == 8 diff --git a/tests/test_dashboard_unit.py b/tests/test_dashboard_unit.py new file mode 100644 index 0000000..13f9908 --- /dev/null +++ b/tests/test_dashboard_unit.py @@ -0,0 +1,785 @@ +"""Unit tests for the BeaconMCP dashboard module. + +Run with:: + + pytest tests/test_dashboard_unit.py -v +""" + +from __future__ import annotations + +import os +import sys +import time +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).parent.parent / "src")) + +from beaconmcp.dashboard.db import Database +from beaconmcp.dashboard.session import ( + SESSION_TTL_SECONDS, + SessionStore, + load_session_key, +) + + +@pytest.fixture() +def db(tmp_path): + return Database(tmp_path / "dashboard.db") + + +@pytest.fixture() +def store(db): + key = os.urandom(32) + return SessionStore(db, key=key) + + +# --------------------------------------------------------------------------- +# load_session_key +# --------------------------------------------------------------------------- + +def test_load_session_key_missing(monkeypatch): + monkeypatch.delenv("BEACONMCP_SESSION_KEY", raising=False) + with pytest.raises(RuntimeError, match="BEACONMCP_SESSION_KEY"): + load_session_key() + + +def test_load_session_key_invalid_base64(monkeypatch): + monkeypatch.setenv("BEACONMCP_SESSION_KEY", "not!!!base64@@@") + with pytest.raises(RuntimeError, match="not valid base64"): + load_session_key() + + +def test_load_session_key_wrong_length(monkeypatch): + import base64 + + monkeypatch.setenv("BEACONMCP_SESSION_KEY", base64.b64encode(b"too short").decode()) + with pytest.raises(RuntimeError, match="32 bytes"): + load_session_key() + + +def test_load_session_key_ok(monkeypatch): + import base64 + + raw = os.urandom(32) + monkeypatch.setenv("BEACONMCP_SESSION_KEY", base64.b64encode(raw).decode()) + assert load_session_key() == raw + + +# --------------------------------------------------------------------------- +# SessionStore lifecycle +# --------------------------------------------------------------------------- + +def test_create_and_load(store): + s = store.create( + client_id="beaconmcp_abc", + client_secret="sk_supersecret", + mcp_bearer="bearer_xyz", + bearer_ttl_seconds=3600, + user_agent="pytest", + ) + assert s.session_id + assert s.client_id == "beaconmcp_abc" + assert s.mcp_bearer == "bearer_xyz" + assert s.bearer_valid() + assert not s.is_expired() + + loaded = store.load(s.session_id) + assert loaded is not None + assert loaded.session_id == s.session_id + assert loaded.client_id == "beaconmcp_abc" + + +def test_load_unknown_returns_none(store): + assert store.load("nope") is None + assert store.load("") is None + + +def test_client_secret_round_trip(store): + s = store.create( + client_id="c", client_secret="sk_top_secret", + mcp_bearer="b", bearer_ttl_seconds=60, user_agent=None, + ) + assert store.get_client_secret(s.session_id) == "sk_top_secret" + + +def test_client_secret_with_wrong_key_fails(db): + key1 = os.urandom(32) + key2 = os.urandom(32) + s1 = SessionStore(db, key=key1) + sess = s1.create( + client_id="c", client_secret="sk_secret", + mcp_bearer="b", bearer_ttl_seconds=60, user_agent=None, + ) + s2 = SessionStore(db, key=key2) + assert s2.get_client_secret(sess.session_id) is None + + +def test_update_bearer(store): + s = store.create( + client_id="c", client_secret="sk", mcp_bearer="old", + bearer_ttl_seconds=10, user_agent=None, + ) + store.update_bearer(s.session_id, mcp_bearer="new", bearer_ttl_seconds=86400) + loaded = store.load(s.session_id) + assert loaded.mcp_bearer == "new" + assert loaded.bearer_valid() + + +def test_session_expiry(store): + s = store.create( + client_id="c", client_secret="sk", mcp_bearer="b", + bearer_ttl_seconds=60, user_agent=None, + ) + # Force expiry + store._db.conn().execute( + "UPDATE sessions SET expires_at = ? WHERE session_id = ?", + (time.time() - 1, s.session_id), + ) + assert store.load(s.session_id) is None + # Auto-deleted + row = store._db.conn().execute( + "SELECT 1 FROM sessions WHERE session_id = ?", (s.session_id,), + ).fetchone() + assert row is None + + +def test_bearer_invalid_when_expired(store): + s = store.create( + client_id="c", client_secret="sk", mcp_bearer="b", + bearer_ttl_seconds=10, user_agent=None, + ) + assert s.bearer_valid() + s.mcp_bearer_expires_at = time.time() - 1 + assert not s.bearer_valid() + + +def test_delete_returns_bearer(store): + s = store.create( + client_id="c", client_secret="sk", mcp_bearer="bearer_to_revoke", + bearer_ttl_seconds=60, user_agent=None, + ) + bearer = store.delete(s.session_id) + assert bearer == "bearer_to_revoke" + assert store.load(s.session_id) is None + assert store.delete(s.session_id) is None + + +def test_delete_all_for_client(store): + sessions = [ + store.create( + client_id="cA", client_secret="sk", mcp_bearer=f"b{i}", + bearer_ttl_seconds=60, user_agent=None, + ) + for i in range(3) + ] + other = store.create( + client_id="cB", client_secret="sk", mcp_bearer="b_other", + bearer_ttl_seconds=60, user_agent=None, + ) + revoked = store.delete_all_for_client("cA") + assert sorted(revoked) == sorted([s.mcp_bearer for s in sessions]) + assert store.load(other.session_id) is not None + + +def test_list_for_client_orders_by_last_seen(store): + s_old = store.create( + client_id="c", client_secret="sk", mcp_bearer="b1", + bearer_ttl_seconds=60, user_agent=None, + ) + time.sleep(0.01) + s_new = store.create( + client_id="c", client_secret="sk", mcp_bearer="b2", + bearer_ttl_seconds=60, user_agent=None, + ) + listed = store.list_for_client("c") + assert [s.session_id for s in listed] == [s_new.session_id, s_old.session_id] + + +def test_session_ttl_is_90_days(): + assert SESSION_TTL_SECONDS == 90 * 24 * 3600 + + +def test_cleanup_expired(store): + s_old = store.create( + client_id="c", client_secret="sk", mcp_bearer="b1", + bearer_ttl_seconds=60, user_agent=None, + ) + s_new = store.create( + client_id="c", client_secret="sk", mcp_bearer="b2", + bearer_ttl_seconds=60, user_agent=None, + ) + store._db.conn().execute( + "UPDATE sessions SET expires_at = ? WHERE session_id = ?", + (time.time() - 1, s_old.session_id), + ) + deleted = store.cleanup_expired() + assert deleted == 1 + assert store.load(s_old.session_id) is None + assert store.load(s_new.session_id) is not None + + +def test_unwrap_exception_single(): + from beaconmcp.dashboard.chat import _unwrap_exception + + err = ValueError("boom") + assert _unwrap_exception(err) is err + + +def test_unwrap_exception_simple_group(): + from beaconmcp.dashboard.chat import _unwrap_exception + + inner = RuntimeError("real cause") + group = ExceptionGroup("task group", [inner]) + assert _unwrap_exception(group) is inner + + +def test_unwrap_exception_nested_groups(): + from beaconmcp.dashboard.chat import _unwrap_exception + + inner = ConnectionError("network") + nested = ExceptionGroup("inner", [inner]) + outer = ExceptionGroup("outer", [nested]) + assert _unwrap_exception(outer) is inner + + +def test_unwrap_exception_prefers_leaf_over_group(): + from beaconmcp.dashboard.chat import _unwrap_exception + + leaf = TypeError("t") + sibling_group = ExceptionGroup("sibling", [RuntimeError("deep")]) + outer = ExceptionGroup("outer", [sibling_group, leaf]) + assert _unwrap_exception(outer) is leaf + + +def test_classify_error_preview_model_permission_denied(): + from beaconmcp.dashboard.chat import _classify_error + + err = Exception( + "403 PERMISSION_DENIED. The caller does not have permission" + ) + code, msg = _classify_error(err, "gemini-3-flash-preview") + assert code == "model_access_denied" + # The way out must name models the picker still offers. + assert "gemini-3.6-flash" in msg + assert "gemini-3-flash-preview" in msg + + +def test_classify_error_stable_model_permission_denied(): + from beaconmcp.dashboard.chat import _classify_error + + err = Exception("403 PERMISSION_DENIED. caller issue") + code, msg = _classify_error(err, "gemini-2.5-flash") + assert code == "permission_denied" + assert "gemini-2.5-flash" in msg + + +def test_classify_error_model_not_found(): + from beaconmcp.dashboard.chat import _classify_error + + err = Exception("404 NOT_FOUND. models/foo is not found") + code, msg = _classify_error(err, "foo") + assert code == "model_not_found" + + +def test_classify_error_rate_limit(): + from beaconmcp.dashboard.chat import _classify_error + + err = Exception("429 RESOURCE_EXHAUSTED. Quota exceeded") + code, _msg = _classify_error(err, "gemini-2.5-flash") + assert code == "rate_limited" + + +def test_classify_error_generic(): + from beaconmcp.dashboard.chat import _classify_error + + err = RuntimeError("boom") + code, msg = _classify_error(err, "gemini-2.5-flash") + assert code == "gemini_error" + assert "RuntimeError" in msg + + +def test_classify_error_upstream_internal(): + from beaconmcp.dashboard.chat import _classify_error + + err = Exception( + "500 INTERNAL. {'error': {'code': 500, 'message': 'Internal error encountered.', 'status': 'INTERNAL'}}" + ) + code, msg = _classify_error(err, "gemini-3-flash-preview") + assert code == "upstream_internal" + assert "500" in msg + assert "gemini-3-flash-preview" in msg + + +def test_classify_error_upstream_unavailable(): + from beaconmcp.dashboard.chat import _classify_error + + err = Exception("503 UNAVAILABLE. The service is temporarily unavailable") + code, _msg = _classify_error(err, "gemini-2.5-flash") + assert code == "upstream_unavailable" + + +def test_classify_error_upstream_timeout(): + from beaconmcp.dashboard.chat import _classify_error + + err = Exception("504 DEADLINE_EXCEEDED") + code, _msg = _classify_error(err, "gemini-2.5-flash") + assert code == "upstream_timeout" + + +def test_token_store_named_issue_and_list(): + from beaconmcp.auth import TokenStore + + ts = TokenStore() + t1, _ = ts.issue("cid", name="Gemini Web") + t2, _ = ts.issue("cid", name="ChatGPT") + + rows = ts.list_named("cid") + assert len(rows) == 2 + # Newest first + assert rows[0].name == "ChatGPT" + assert rows[1].name == "Gemini Web" + assert ts.count_named("cid") == 2 + + # Unnamed dashboard session token must not appear + ts.issue("cid") + assert ts.count_named("cid") == 2 + + +def test_token_store_cap_is_three(): + from beaconmcp.auth import TokenCapExceeded, TokenStore + + ts = TokenStore() + for i in range(3): + ts.issue("cid", name=f"t{i}") + with pytest.raises(TokenCapExceeded): + ts.issue("cid", name="overflow") + # Cap is per-client: another client is unaffected + ts.issue("other", name="ok") + + +def test_token_store_cap_frees_after_revoke(): + from beaconmcp.auth import TokenStore + + ts = TokenStore() + t1, _ = ts.issue("cid", name="a") + ts.issue("cid", name="b") + ts.issue("cid", name="c") + assert ts.count_named("cid") == 3 + + # Expire (not just schedule-revoke) the first token so the cap count drops. + ts._tokens[t1].expires_at = 0 + # Re-issue should now succeed. + ts.issue("cid", name="d") + assert ts.count_named("cid") == 3 + + +def test_token_store_revoke_named_by_prefix_scoped_to_client(): + from beaconmcp.auth import TokenStore + + ts = TokenStore() + t_mine, _ = ts.issue("me", name="Mine") + t_other, _ = ts.issue("other", name="Theirs") + prefix = t_other[:12] + + # I can't revoke someone else's token even with the right prefix. + assert ts.revoke_named(prefix, "me") is False + # Owner can revoke their own. + assert ts.revoke_named(prefix, "other") is True + + +def test_token_store_revoke_named_requires_min_prefix(): + from beaconmcp.auth import TokenStore + + ts = TokenStore() + ts.issue("cid", name="a") + assert ts.revoke_named("ab", "cid") is False # too short + + +def test_mcp_tool_to_declaration_passes_input_schema(): + from google.genai import types + + from beaconmcp.dashboard.chat import _mcp_tool_to_declaration + + class FakeMCPTool: + name = "proxmox_list_vms" + description = "List VMs on a Proxmox node" + inputSchema = { + "type": "object", + "properties": {"node": {"type": "string"}}, + "required": ["node"], + } + + decl = _mcp_tool_to_declaration(FakeMCPTool(), types) + assert decl.name == "proxmox_list_vms" + assert decl.description == "List VMs on a Proxmox node" + assert decl.parameters_json_schema["required"] == ["node"] + + +def test_mcp_tool_to_declaration_defaults_schema_when_missing(): + from google.genai import types + + from beaconmcp.dashboard.chat import _mcp_tool_to_declaration + + class FakeMCPTool: + name = "ping" + description = "" + inputSchema = None + + decl = _mcp_tool_to_declaration(FakeMCPTool(), types) + # Gemini rejects empty/missing schemas; helper must substitute an + # empty object schema so the tool still registers. + assert decl.parameters_json_schema == {"type": "object", "properties": {}} + +def test_build_google_search_tool_supports_google_search_class(): + from beaconmcp.dashboard.chat import _build_google_search_tool + + class _Tool: + def __init__(self, **kwargs): + self.kwargs = kwargs + + class _GoogleSearch: + pass + + class _Types: + Tool = _Tool + GoogleSearch = _GoogleSearch + + tool = _build_google_search_tool(_Types) + assert tool is not None + assert "google_search" in tool.kwargs + assert isinstance(tool.kwargs["google_search"], _GoogleSearch) + + +def test_build_google_search_tool_supports_legacy_toolgooglesearch(): + from beaconmcp.dashboard.chat import _build_google_search_tool + + class _Tool: + def __init__(self, **kwargs): + self.kwargs = kwargs + + class _ToolGoogleSearch: + pass + + class _Types: + Tool = _Tool + ToolGoogleSearch = _ToolGoogleSearch + + tool = _build_google_search_tool(_Types) + assert tool is not None + assert "google_search" in tool.kwargs + assert isinstance(tool.kwargs["google_search"], _ToolGoogleSearch) + + +def test_tool_name_from_server_tool_type_mapping(): + from beaconmcp.dashboard.chat import _tool_name_from_server_tool_type + + assert _tool_name_from_server_tool_type("GOOGLE_SEARCH_WEB") == "google_search_web" + assert _tool_name_from_server_tool_type("ToolType.URL_CONTEXT") == "url_context" + assert _tool_name_from_server_tool_type(None) == "unspecified" + + +def test_tool_response_is_error_helper(): + from beaconmcp.dashboard.chat import _tool_response_is_error + + assert _tool_response_is_error({"status": "error"}) is True + assert _tool_response_is_error({"error": {"message": "boom"}}) is True + assert _tool_response_is_error({"url_retrieval_status": "URL_RETRIEVAL_STATUS_ERROR"}) is True + assert _tool_response_is_error({"status": "ok"}) is False + + +def test_compose_system_instruction_includes_interleaving_and_server_context(): + from beaconmcp.dashboard.chat import _compose_system_instruction + + text = _compose_system_instruction("Tools available: ssh_run, proxmox_run") + assert "never batch multiple MCP function calls" in text + assert "Before each tool call" in text + assert "Tools available: ssh_run, proxmox_run" in text + + +def test_compose_system_instruction_without_server_context(): + from beaconmcp.dashboard.chat import _compose_system_instruction + + text = _compose_system_instruction(None) + assert "never batch multiple MCP function calls" in text + assert "function_call" in text + assert "Tools available:" not in text + + +def test_mcp_call_result_to_response_flattens_text_content(): + from beaconmcp.dashboard.chat import _mcp_call_result_to_response + + class FakeText: + text = "hello" + + class FakeResult: + content = [FakeText()] + isError = False + structuredContent = None + + payload = _mcp_call_result_to_response(FakeResult()) + assert payload == {"content": [{"type": "text", "text": "hello"}]} + + +def test_mcp_call_result_to_response_marks_error(): + from beaconmcp.dashboard.chat import _mcp_call_result_to_response + + class FakeText: + text = "boom" + + class FakeResult: + content = [FakeText()] + isError = True + structuredContent = None + + payload = _mcp_call_result_to_response(FakeResult()) + assert payload["error"] is True + + +def test_is_transient_error_matches_5xx(): + from beaconmcp.dashboard.chat import _is_transient_error + + assert _is_transient_error(Exception("500 INTERNAL. Internal error")) + assert _is_transient_error(Exception("503 UNAVAILABLE")) + assert _is_transient_error(Exception("504 DEADLINE_EXCEEDED")) + assert not _is_transient_error(Exception("403 PERMISSION_DENIED")) + assert not _is_transient_error(RuntimeError("boom")) + + +def _run_retry_scenario(monkeypatch, fake_run): + """Helper: swap in ``fake_run`` on a GeminiChatEngine and drain events.""" + import asyncio as _asyncio + from beaconmcp.dashboard import chat as chat_mod + from beaconmcp.dashboard.chat import GeminiChatEngine, TurnInput + + async def _noop_sleep(_): + return None + + monkeypatch.setattr(chat_mod.asyncio, "sleep", _noop_sleep) + engine = GeminiChatEngine(api_key="test") + engine._run = fake_run # type: ignore[assignment] + turn = TurnInput( + history=[], user_text="x", model="gemini-2.5-flash", + effort="low", bearer="b", mcp_url="http://localhost/mcp", + ) + + async def _drain(): + return [e async for e in engine.run(turn)] + + return _asyncio.run(_drain()) + + +def test_gemini_retry_recovers_from_transient_500(monkeypatch): + """run() retries transient 5xx errors before surfacing them.""" + from beaconmcp.dashboard.chat import TextDelta + + attempts = {"n": 0} + + async def fake_run(_turn): + attempts["n"] += 1 + if attempts["n"] < 3: + raise Exception( + "500 INTERNAL. {'error': {'code': 500, 'status': 'INTERNAL'}}" + ) + yield TextDelta(text="ok after retries") + + events = _run_retry_scenario(monkeypatch, fake_run) + assert attempts["n"] == 3 + assert len(events) == 1 + assert isinstance(events[0], TextDelta) + assert events[0].text == "ok after retries" + + +def test_gemini_retry_surfaces_after_max_attempts(monkeypatch): + from beaconmcp.dashboard.chat import ErrorEvent + + attempts = {"n": 0} + + async def fake_run(_turn): + attempts["n"] += 1 + raise Exception("500 INTERNAL. persistent") + yield # pragma: no cover -- makes fake_run an async generator + + events = _run_retry_scenario(monkeypatch, fake_run) + # initial attempt + 2 retries = 3 total + assert attempts["n"] == 3 + assert len(events) == 1 + assert isinstance(events[0], ErrorEvent) + assert events[0].code == "upstream_internal" + + +def test_gemini_retry_skips_on_non_transient(monkeypatch): + from beaconmcp.dashboard.chat import ErrorEvent + + attempts = {"n": 0} + + async def fake_run(_turn): + attempts["n"] += 1 + raise Exception("403 PERMISSION_DENIED. caller") + yield # pragma: no cover + + events = _run_retry_scenario(monkeypatch, fake_run) + assert attempts["n"] == 1 # no retries for 403 + assert isinstance(events[0], ErrorEvent) + assert events[0].code == "permission_denied" + + +def test_thinking_config_for_gemini_3(): + from beaconmcp.dashboard.chat import GeminiChatEngine + + cfg = GeminiChatEngine._build_thinking_config("gemini-3-flash-preview", "high") + assert cfg.thinking_level is not None + assert cfg.thinking_budget is None + assert cfg.include_thoughts is False + + +def test_thinking_config_for_gemini_2_5(): + from beaconmcp.dashboard.chat import GeminiChatEngine + + cfg = GeminiChatEngine._build_thinking_config("gemini-2.5-flash", "medium") + assert cfg.thinking_level is None + assert cfg.thinking_budget == 4096 + + +def test_thinking_config_clamps_gemini_2_5_pro_minimum(): + """2.5 Pro cannot disable thinking; budget must clamp to 128+.""" + from beaconmcp.dashboard.chat import GeminiChatEngine + + cfg = GeminiChatEngine._build_thinking_config("gemini-2.5-pro", "minimal") + assert cfg.thinking_budget == 128 # clamped up from 0 + + # 2.5 Flash has no minimum -- "minimal" means disable thinking. + cfg_flash = GeminiChatEngine._build_thinking_config("gemini-2.5-flash", "minimal") + assert cfg_flash.thinking_budget == 0 + + +def test_thinking_config_unknown_effort_defaults_to_low(): + from beaconmcp.dashboard.chat import GeminiChatEngine + + cfg3 = GeminiChatEngine._build_thinking_config("gemini-3.1-pro-preview", "nonsense") + # Enum repr contains LOW + assert "LOW" in str(cfg3.thinking_level) + cfg25 = GeminiChatEngine._build_thinking_config("gemini-2.5-pro", "nonsense") + assert cfg25.thinking_budget == 1024 # low + + +def test_migration_v1_to_v2_renames_gemini_models(tmp_path): + """A DB created under schema v1 with bare model names must be upgraded.""" + import sqlite3 + + path = tmp_path / "legacy.db" + # Hand-build a v1 database that predates the migration. + conn = sqlite3.connect(path) + conn.executescript( + """ + CREATE TABLE conversations ( + id TEXT PRIMARY KEY, client_id TEXT NOT NULL, title TEXT, + model TEXT NOT NULL DEFAULT 'gemini-3-flash', + thinking_effort TEXT NOT NULL DEFAULT 'low', + created_at REAL NOT NULL, updated_at REAL NOT NULL + ); + CREATE TABLE messages ( + id TEXT PRIMARY KEY, conversation_id TEXT NOT NULL, + role TEXT NOT NULL, content TEXT, tool_calls TEXT, + thinking_summary TEXT, model TEXT, effort TEXT, + created_at REAL NOT NULL + ); + """ + ) + conn.execute("INSERT INTO conversations VALUES ('c1', 'cli', NULL, 'gemini-3-flash', 'low', 0, 0)") + conn.execute("INSERT INTO conversations VALUES ('c2', 'cli', NULL, 'gemini-3.1-pro', 'medium', 0, 0)") + conn.execute("INSERT INTO messages VALUES ('m1', 'c1', 'assistant', 'hi', NULL, NULL, 'gemini-3-flash', 'low', 0)") + conn.execute("PRAGMA user_version = 1") + conn.commit() + conn.close() + + # Opening via Database() runs every pending migration, so c1 is first + # renamed to gemini-3-flash-preview by v2 and then moved forward to the + # current Flash by v6. + db = Database(path) + rows = db.conn().execute( + "SELECT id, model FROM conversations ORDER BY id" + ).fetchall() + assert dict(rows[0]) == {"id": "c1", "model": "gemini-3.6-flash"} + assert dict(rows[1]) == {"id": "c2", "model": "gemini-3.1-pro-preview"} + # v2 renamed the same model; v6 substitutes a different one, so the + # message keeps naming whichever model actually wrote the reply. + msg = db.conn().execute("SELECT model FROM messages WHERE id='m1'").fetchone() + assert msg["model"] == "gemini-3-flash-preview" + + # user_version reflects the migration (latest schema version). + from beaconmcp.dashboard.db import _LATEST_VERSION + + ver = db.conn().execute("PRAGMA user_version").fetchone()[0] + assert ver == _LATEST_VERSION + + +def test_migration_v6_moves_retired_models_forward(tmp_path): + """A conversation left on a retired model must still be usable. + + ``conversations.model`` is what the next turn runs on. Leaving a + retired id there would fail ``VALID_MODELS`` and silently fall back, + which reads as the picker forgetting the operator's choice. + """ + from beaconmcp.dashboard.conversations import VALID_MODELS + + db = Database(tmp_path / "d.db") + conn = db.conn() + for cid, model in ( + ("c1", "gemini-2.5-flash"), + ("c2", "gemini-2.5-pro"), + ("c3", "gemini-3-flash-preview"), + ("c4", "gemini-3.1-pro-preview"), + ): + conn.execute( + "INSERT INTO conversations (id, client_id, title, model, " + "thinking_effort, created_at, updated_at) VALUES (?,?,?,?,?,?,?)", + (cid, "cli", None, model, "low", 0, 0), + ) + conn.execute("PRAGMA user_version = 5") + conn.commit() + + from beaconmcp.dashboard.db import _migrate + + _migrate(conn) + rows = conn.execute( + "SELECT id, model FROM conversations ORDER BY id" + ).fetchall() + assert [r["model"] for r in rows] == [ + "gemini-3.6-flash", + "gemini-3.1-pro-preview", + "gemini-3.6-flash", + "gemini-3.1-pro-preview", + ] + assert all(r["model"] in VALID_MODELS for r in rows) + + +def test_every_offered_model_has_a_price(): + """A model in the picker with no rate would bill at the fallback's.""" + from beaconmcp.dashboard.conversations import DEFAULT_MODEL, VALID_MODELS + from beaconmcp.dashboard.usage import _PRICING + + assert DEFAULT_MODEL in VALID_MODELS + for model in VALID_MODELS: + assert model in _PRICING, model + + +def test_retired_models_keep_their_rates(): + """Old turns must not be re-priced at the current model's rate.""" + from beaconmcp.dashboard.usage import _PRICING + + for model in ("gemini-2.5-flash", "gemini-2.5-pro", "gemini-3-flash-preview"): + assert model in _PRICING, model + + +def test_short_ciphertext_decryption_returns_none(store): + s = store.create( + client_id="c", client_secret="sk", mcp_bearer="b", + bearer_ttl_seconds=60, user_agent=None, + ) + store._db.conn().execute( + "UPDATE sessions SET client_secret_enc = ? WHERE session_id = ?", + (b"xx", s.session_id), + ) + assert store.get_client_secret(s.session_id) is None diff --git a/tests/test_dashboard_usage.py b/tests/test_dashboard_usage.py new file mode 100644 index 0000000..acbf1b0 --- /dev/null +++ b/tests/test_dashboard_usage.py @@ -0,0 +1,434 @@ +"""Unit + integration tests for per-client usage tracking and budget caps.""" + +from __future__ import annotations + +import os +import sys +import time +from pathlib import Path + +import pytest +from starlette.applications import Starlette +from starlette.testclient import TestClient + +sys.path.insert(0, str(Path(__file__).parent.parent / "src")) + +from beaconmcp.auth import TotpResult +from beaconmcp.dashboard.app import DashboardDeps, build_dashboard_routes +from beaconmcp.dashboard.chat import ( + FakeChatEngine, + FakeScript, + TextDelta, + UsageAccumulated, +) +from beaconmcp.dashboard.confirmations import ConfirmationStore +from beaconmcp.dashboard.conversations import ConversationStore +from beaconmcp.dashboard.csrf import CSRF_COOKIE +from beaconmcp.dashboard.db import Database +from beaconmcp.dashboard.session import SessionStore +from beaconmcp.dashboard.usage import Budget, UsageMeter, UsageStore + + +# --------------------------------------------------------------------------- +# Helpers copied from test_dashboard_chat (kept local so this file is +# self-contained and can run even if that one is skipped). +# --------------------------------------------------------------------------- + + +class FakeClientStore: + def verify(self, cid, sec): return cid == "c" and sec == "s" + def check_totp(self, cid, code): + return TotpResult.OK if code == "123456" else TotpResult.INVALID + def get_name(self, cid): return "Test" + + +class FakeTokenStore: + NAMED_TOKEN_CAP = 3 + + def __init__(self): + self._live: dict[str, str] = {} + + def issue(self, cid, *, name=None): + token = f"b_{len(self._live) + 1}" + ("x" * 60) + self._live[token] = cid + return token, 24 * 3600 + + def validate(self, token): + return self._live.get(token) + + def revoke(self, token): + self._live.pop(token, None) + return True + + +def _login(client) -> str: + r = client.get("/app/login") + csrf = r.cookies.get(CSRF_COOKIE) + assert csrf is not None + r = client.post("/app/login", data={ + "csrf_token": csrf, "client_id": "c", + "client_secret": "s", "totp": "123456", "remember": "on", + }) + assert r.status_code == 303 + out = client.cookies.get(CSRF_COOKIE) + assert out is not None + return out + + +# --------------------------------------------------------------------------- +# Unit: UsageMeter +# --------------------------------------------------------------------------- + + +def test_cost_flash_no_cache(): + # 1M prompt + 1M output on 2.5-flash = $0.30 + $2.50 = $2.80. + cost = UsageMeter.cost_usd( + "gemini-2.5-flash", + prompt_tokens=1_000_000, cached_tokens=0, output_tokens=1_000_000, + ) + assert cost == pytest.approx(2.80) + + +def test_cost_flash_with_cache_hit(): + # Half the input comes from cache: 500k * $0.30 + 500k * $0.03 + 0 out. + cost = UsageMeter.cost_usd( + "gemini-2.5-flash", + prompt_tokens=1_000_000, cached_tokens=500_000, output_tokens=0, + ) + assert cost == pytest.approx(0.30 * 0.5 + 0.03 * 0.5) + + +def test_cost_pro_high_tier(): + # Pro model crosses the 200k threshold -> high-tier pricing applies. + cost = UsageMeter.cost_usd( + "gemini-2.5-pro", + prompt_tokens=300_000, cached_tokens=0, output_tokens=0, + ) + # 300k * $2.50 / 1M = $0.75 + assert cost == pytest.approx(0.75) + + +def test_cost_unknown_model_falls_back_to_flash(): + cost = UsageMeter.cost_usd( + "gemini-unknown-9", + prompt_tokens=1_000_000, cached_tokens=0, output_tokens=0, + ) + # Same rate as gemini-3.6-flash input, the current default. + assert cost == pytest.approx(1.50) + + +def test_cost_retired_model_keeps_its_own_rate(): + """A turn billed on 2.5 Flash must not be re-priced at 3.6 Flash.""" + cost = UsageMeter.cost_usd( + "gemini-2.5-flash", + prompt_tokens=1_000_000, cached_tokens=0, output_tokens=0, + ) + assert cost == pytest.approx(0.30) + + +def test_cost_new_flash_models(): + assert UsageMeter.cost_usd( + "gemini-3.6-flash", + prompt_tokens=1_000_000, cached_tokens=0, output_tokens=1_000_000, + ) == pytest.approx(1.50 + 7.50) + assert UsageMeter.cost_usd( + "gemini-3.5-flash-lite", + prompt_tokens=1_000_000, cached_tokens=0, output_tokens=1_000_000, + ) == pytest.approx(0.30 + 2.50) + + +def test_cost_cached_over_prompt_is_clamped(): + # Defensive: if Gemini ever reports more cached tokens than prompt + # tokens, the billable_input floor is 0 (not negative). + cost = UsageMeter.cost_usd( + "gemini-2.5-flash", + prompt_tokens=100, cached_tokens=200, output_tokens=0, + ) + # Billable input is 0, cached is 200 at cached rate. + assert cost == pytest.approx(200 * 0.03 / 1_000_000) + + +# --------------------------------------------------------------------------- +# Unit: UsageStore (5h session + rolling week) +# --------------------------------------------------------------------------- + + +@pytest.fixture() +def store(tmp_path): + db = Database(tmp_path / "usage.db") + return UsageStore(db, Budget(limit_5h_usd=1.0, limit_week_usd=5.0)) + + +def test_record_turn_opens_session(store): + now = 1000.0 + store.record_turn( + client_id="c", conversation_id="x", message_id="m1", + model="gemini-2.5-flash", + prompt_tokens=1000, cached_tokens=0, output_tokens=500, + cost_usd=0.10, now=now, + ) + snap = store.snapshot("c", now=now) + assert snap.spent_5h_usd == pytest.approx(0.10) + assert snap.session_5h_started_at == now + assert snap.session_5h_reset_at == now + 5 * 3600 + assert snap.spent_week_usd == pytest.approx(0.10) + + +def test_five_hour_window_resets_after_expiry(store): + t0 = 1000.0 + store.record_turn( + client_id="c", conversation_id="x", message_id="m1", + model="gemini-2.5-flash", + prompt_tokens=1000, cached_tokens=0, output_tokens=500, + cost_usd=0.60, now=t0, + ) + # Just inside 5h: still part of the same session. + store.record_turn( + client_id="c", conversation_id="x", message_id="m2", + model="gemini-2.5-flash", + prompt_tokens=500, cached_tokens=0, output_tokens=500, + cost_usd=0.10, now=t0 + 1000, + ) + snap_mid = store.snapshot("c", now=t0 + 1000) + assert snap_mid.spent_5h_usd == pytest.approx(0.70) + + # Past 5h: next turn opens a fresh session at its own ts. + t1 = t0 + 6 * 3600 + store.record_turn( + client_id="c", conversation_id="x", message_id="m3", + model="gemini-2.5-flash", + prompt_tokens=1000, cached_tokens=0, output_tokens=500, + cost_usd=0.05, now=t1, + ) + snap_new = store.snapshot("c", now=t1) + assert snap_new.spent_5h_usd == pytest.approx(0.05) + assert snap_new.session_5h_started_at == t1 + + +def test_snapshot_with_expired_session_shows_zero(store): + t0 = 1000.0 + store.record_turn( + client_id="c", conversation_id="x", message_id="m1", + model="gemini-2.5-flash", + prompt_tokens=1000, cached_tokens=0, output_tokens=500, + cost_usd=0.40, now=t0, + ) + # Look up the snapshot after the 5h window has already lapsed but + # WITHOUT a new turn in between -- should report zero and no active + # session, so the UI footer stops showing stale data. + snap = store.snapshot("c", now=t0 + 6 * 3600) + assert snap.spent_5h_usd == 0.0 + assert snap.session_5h_started_at is None + + +def test_check_budget_blocks_when_five_hour_exceeded(store): + t0 = 1000.0 + # Spend exactly the cap in one go. + store.record_turn( + client_id="c", conversation_id="x", message_id="m1", + model="gemini-2.5-flash", + prompt_tokens=1000, cached_tokens=0, output_tokens=500, + cost_usd=1.0, now=t0, + ) + block = store.check_budget("c", now=t0 + 60) + assert block is not None + assert block.window == "5h" + assert block.limit_usd == 1.0 + assert block.reset_at == t0 + 5 * 3600 + + +def test_check_budget_blocks_when_weekly_exceeded(tmp_path): + db = Database(tmp_path / "usage.db") + # 5h cap disabled so we test the weekly path in isolation. + s = UsageStore(db, Budget(limit_5h_usd=0.0, limit_week_usd=1.0)) + t0 = 1_000_000.0 + # Accumulate enough cost across two turns spread out in time. + s.record_turn( + client_id="c", conversation_id="x", message_id="m1", + model="gemini-2.5-flash", + prompt_tokens=1, cached_tokens=0, output_tokens=1, + cost_usd=0.60, now=t0, + ) + s.record_turn( + client_id="c", conversation_id="x", message_id="m2", + model="gemini-2.5-flash", + prompt_tokens=1, cached_tokens=0, output_tokens=1, + cost_usd=0.50, now=t0 + 3 * 24 * 3600, + ) + block = s.check_budget("c", now=t0 + 4 * 24 * 3600) + assert block is not None + assert block.window == "week" + + +def test_check_budget_allows_when_disabled(tmp_path): + db = Database(tmp_path / "usage.db") + s = UsageStore(db, Budget(limit_5h_usd=0.0, limit_week_usd=0.0)) + # Huge spend but no limit configured -> always allow. + s.record_turn( + client_id="c", conversation_id="x", message_id="m1", + model="gemini-2.5-flash", + prompt_tokens=1, cached_tokens=0, output_tokens=1, + cost_usd=999.0, now=time.time(), + ) + assert s.check_budget("c") is None + + +def test_weekly_window_is_rolling(tmp_path): + """Events older than 7 days must fall out of the weekly sum.""" + db = Database(tmp_path / "usage.db") + s = UsageStore(db, Budget(limit_5h_usd=0.0, limit_week_usd=10.0)) + t_old = 1000.0 + s.record_turn( + client_id="c", conversation_id="x", message_id="m1", + model="gemini-2.5-flash", + prompt_tokens=1, cached_tokens=0, output_tokens=1, + cost_usd=5.0, now=t_old, + ) + # >7 days later: old event drops out, new one is the only contributor. + t_now = t_old + 8 * 24 * 3600 + s.record_turn( + client_id="c", conversation_id="x", message_id="m2", + model="gemini-2.5-flash", + prompt_tokens=1, cached_tokens=0, output_tokens=1, + cost_usd=1.0, now=t_now, + ) + snap = s.snapshot("c", now=t_now) + assert snap.spent_week_usd == pytest.approx(1.0) + + +# --------------------------------------------------------------------------- +# Integration: chat stream enforces caps + persists cost +# --------------------------------------------------------------------------- + + +def _build_deps(tmp_path, *, engine, budget): + db = Database(tmp_path / "dashboard.db") + return DashboardDeps( + database=db, + session_store=SessionStore(db, key=os.urandom(32)), + client_store=FakeClientStore(), + token_store=FakeTokenStore(), + totp_locked=lambda cid: False, + totp_record_failure=lambda cid: None, + totp_record_success=lambda cid: None, + conversations=ConversationStore(db), + engine=engine, + confirmations=ConfirmationStore(), + usage=UsageStore(db, budget), + ) + + +def test_stream_records_cost_after_turn(tmp_path): + # Engine script: yields a text delta + a usage event reporting + # 100k prompt (0 cached) + 50k output on flash => 100000*0.30/1M + + # 50000*2.50/1M = 0.03 + 0.125 = 0.155. + engine = FakeChatEngine(FakeScript(events=[ + TextDelta(text="hello"), + UsageAccumulated( + model="gemini-2.5-flash", + prompt_tokens=100_000, + cached_tokens=0, + output_tokens=50_000, + ), + ], title_text="t")) + + deps = _build_deps( + tmp_path, engine=engine, + budget=Budget(limit_5h_usd=1.0, limit_week_usd=5.0), + ) + app = Starlette(routes=build_dashboard_routes(deps)) + client = TestClient(app, follow_redirects=False) + csrf = _login(client) + + # Create conversation. + r = client.post( + "/app/api/conversations", + json={"model": "gemini-2.5-flash", "effort": "low"}, + headers={"X-CSRF-Token": csrf}, + ) + conv_id = r.json()["conversation"]["id"] + + # Stream a turn. + r = client.post( + "/app/api/chat/stream", + json={"conversation_id": conv_id, "content": "hi"}, + headers={"X-CSRF-Token": csrf}, + ) + assert r.status_code == 200 + body = r.text + assert "event: usage_update" in body + + # Snapshot API should reflect the recorded cost. + r = client.get("/app/api/usage") + assert r.status_code == 200 + usage = r.json()["usage"] + expected_cost = 100_000 * 0.30 / 1_000_000 + 50_000 * 2.50 / 1_000_000 + assert usage["spent_5h_usd"] == pytest.approx(expected_cost, rel=1e-6) + assert usage["spent_week_usd"] == pytest.approx(expected_cost, rel=1e-6) + assert usage["limit_5h_usd"] == 1.0 + assert usage["limit_week_usd"] == 5.0 + + +def test_stream_rejects_when_over_5h_cap(tmp_path): + engine = FakeChatEngine(FakeScript(events=[TextDelta(text="unused")])) + deps = _build_deps( + tmp_path, engine=engine, + budget=Budget(limit_5h_usd=0.01, limit_week_usd=5.0), + ) + # Pre-load a spend that already exceeds the 5h cap. + assert deps.usage is not None + deps.usage.record_turn( + client_id="c", conversation_id=None, message_id=None, + model="gemini-2.5-flash", + prompt_tokens=0, cached_tokens=0, output_tokens=0, + cost_usd=0.02, + ) + + app = Starlette(routes=build_dashboard_routes(deps)) + client = TestClient(app, follow_redirects=False) + csrf = _login(client) + + r = client.post( + "/app/api/conversations", + json={"model": "gemini-2.5-flash", "effort": "low"}, + headers={"X-CSRF-Token": csrf}, + ) + conv_id = r.json()["conversation"]["id"] + + r = client.post( + "/app/api/chat/stream", + json={"conversation_id": conv_id, "content": "hi"}, + headers={"X-CSRF-Token": csrf}, + ) + assert r.status_code == 200 + body = r.text + assert '"code": "quota_exceeded"' in body + # Engine must not have been called -- the cap is enforced before + # we reach google-genai. + assert engine.calls == [] + + +def test_usage_endpoint_returns_disabled_snapshot_when_store_missing(tmp_path): + db = Database(tmp_path / "dashboard.db") + deps = DashboardDeps( + database=db, + session_store=SessionStore(db, key=os.urandom(32)), + client_store=FakeClientStore(), + token_store=FakeTokenStore(), + totp_locked=lambda cid: False, + totp_record_failure=lambda cid: None, + totp_record_success=lambda cid: None, + conversations=ConversationStore(db), + engine=FakeChatEngine(FakeScript()), + confirmations=ConfirmationStore(), + usage=None, # explicitly disabled + ) + app = Starlette(routes=build_dashboard_routes(deps)) + client = TestClient(app, follow_redirects=False) + _login(client) + + r = client.get("/app/api/usage") + assert r.status_code == 200 + u = r.json()["usage"] + assert u["limit_5h_usd"] == 0.0 + assert u["limit_week_usd"] == 0.0 diff --git a/tests/test_dynamic_registration.py b/tests/test_dynamic_registration.py new file mode 100644 index 0000000..1eb5571 --- /dev/null +++ b/tests/test_dynamic_registration.py @@ -0,0 +1,288 @@ +"""Tests for OAuth Dynamic Client Registration (ChatGPT connector flow). + +Covers the two correctness-critical paths: + +1. :class:`DynamicSlugStore.consume` is atomic & single-use — concurrent + claims deterministically produce exactly one winner. +2. Derived :class:`Client` rows delegate TOTP verification to their + owner, and owner revocation cascades to all derived rows. +""" + +from __future__ import annotations + +import sys +import threading +import time +from pathlib import Path + +import pyotp +import pytest + +sys.path.insert(0, str(Path(__file__).parent.parent / "src")) + +from beaconmcp.auth import ClientStore +from beaconmcp.dashboard.db import Database +from beaconmcp.dashboard.dyn_reg import ( + SLUG_TTL_SECONDS, + DynamicSlugStore, + SlugAlreadyConsumed, +) + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture() +def db(tmp_path: Path) -> Database: + return Database(tmp_path / "dash.db") + + +@pytest.fixture() +def slug_store(db: Database) -> DynamicSlugStore: + return DynamicSlugStore(db) + + +@pytest.fixture() +def clients(tmp_path: Path) -> ClientStore: + return ClientStore(tmp_path / "clients.json") + + +# --------------------------------------------------------------------------- +# Slug lifecycle +# --------------------------------------------------------------------------- + + +def test_mint_then_consume_records_resulting_client(slug_store: DynamicSlugStore) -> None: + row = slug_store.mint(owner_client_id="owner_1", label="ChatGPT iPhone") + assert row.used_at is None + + claimed = slug_store.consume(row.slug, resulting_client_id="new_client_1") + assert claimed.used_at is not None + assert claimed.resulting_client_id == "new_client_1" + + reloaded = slug_store.load(row.slug) + assert reloaded is not None and reloaded.resulting_client_id == "new_client_1" + + +def test_consume_twice_raises(slug_store: DynamicSlugStore) -> None: + row = slug_store.mint(owner_client_id="owner_1", label="x") + slug_store.consume(row.slug, resulting_client_id="c1") + with pytest.raises(SlugAlreadyConsumed): + slug_store.consume(row.slug, resulting_client_id="c2") + + +def test_consume_expired_slug_rejected( + slug_store: DynamicSlugStore, db: Database, +) -> None: + row = slug_store.mint(owner_client_id="owner_1", label="x") + # Fast-forward expiry by rewriting the row directly. + db.conn().execute( + "UPDATE oauth_dynamic_slugs SET expires_at = ? WHERE slug = ?", + (time.time() - 1, row.slug), + ) + with pytest.raises(SlugAlreadyConsumed): + slug_store.consume(row.slug, resulting_client_id="c1") + + +def test_consume_unknown_slug_rejected(slug_store: DynamicSlugStore) -> None: + with pytest.raises(SlugAlreadyConsumed): + slug_store.consume("not-a-real-slug", resulting_client_id="c1") + + +def test_concurrent_consume_has_exactly_one_winner( + tmp_path: Path, +) -> None: + """Two threads race to claim the same slug; exactly one wins.""" + # Each thread uses its own Database instance (fresh connection) to + # mimic the real HTTP server where handlers run on a thread pool. + db_path = tmp_path / "dash.db" + bootstrap = Database(db_path) + store = DynamicSlugStore(bootstrap) + row = store.mint(owner_client_id="owner_1", label="race") + + barrier = threading.Barrier(10) + wins: list[str] = [] + losses: list[str] = [] + lock = threading.Lock() + + def worker(i: int) -> None: + local_store = DynamicSlugStore(Database(db_path)) + barrier.wait() + try: + local_store.consume(row.slug, resulting_client_id=f"c{i}") + with lock: + wins.append(f"c{i}") + except SlugAlreadyConsumed: + with lock: + losses.append(f"c{i}") + + threads = [threading.Thread(target=worker, args=(i,)) for i in range(10)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert len(wins) == 1, f"expected exactly one winner, got {wins}" + assert len(losses) == 9 + + +def test_delete_unused_only_removes_unconsumed( + slug_store: DynamicSlugStore, +) -> None: + row = slug_store.mint(owner_client_id="owner_1", label="pending") + assert slug_store.delete_unused(row.slug, "owner_1") is True + # Idempotent: second call is a no-op, returns False. + assert slug_store.delete_unused(row.slug, "owner_1") is False + + # Consumed slugs survive as audit trail. + row2 = slug_store.mint(owner_client_id="owner_1", label="used") + slug_store.consume(row2.slug, resulting_client_id="c1") + assert slug_store.delete_unused(row2.slug, "owner_1") is False + + +def test_delete_unused_scoped_to_owner(slug_store: DynamicSlugStore) -> None: + row = slug_store.mint(owner_client_id="owner_A", label="x") + assert slug_store.delete_unused(row.slug, "owner_B") is False + assert slug_store.load(row.slug) is not None + + +def test_prune_expired_drops_only_unused_expired_rows( + slug_store: DynamicSlugStore, db: Database, +) -> None: + # 1 fresh, 1 expired-unused, 1 expired-used. + fresh = slug_store.mint(owner_client_id="o", label="fresh") + expired_unused = slug_store.mint(owner_client_id="o", label="exp-unused") + expired_used = slug_store.mint(owner_client_id="o", label="exp-used") + + slug_store.consume(expired_used.slug, resulting_client_id="c1") + + now = time.time() + db.conn().execute( + "UPDATE oauth_dynamic_slugs SET expires_at = ? WHERE slug IN (?, ?)", + (now - 1, expired_unused.slug, expired_used.slug), + ) + + removed = slug_store.prune_expired() + assert removed == 1 + assert slug_store.load(fresh.slug) is not None + assert slug_store.load(expired_unused.slug) is None + assert slug_store.load(expired_used.slug) is not None + + +def test_ttl_matches_design(slug_store: DynamicSlugStore) -> None: + """Guard against accidental TTL changes — the dashboard UI copy depends + on this value, and a sudden bump would surprise users.""" + row = slug_store.mint(owner_client_id="o", label="x") + delta = row.expires_at - row.created_at + assert abs(delta - SLUG_TTL_SECONDS) < 1.0 + + +# --------------------------------------------------------------------------- +# Dynamic-client TOTP delegation +# --------------------------------------------------------------------------- + + +def test_dynamic_client_totp_delegates_to_owner(clients: ClientStore) -> None: + owner_id, _, owner_seed = clients.create("human") + derived_id, derived_secret = clients.create_dynamic( + owner_client_id=owner_id, + name="ChatGPT (derived)", + registration_source="chatgpt:slug1", + ) + + # Derived client authenticates with its own secret but NOT its own TOTP. + assert clients.verify(derived_id, derived_secret) is True + now_code = pyotp.TOTP(owner_seed).now() + assert clients.verify_totp(derived_id, now_code) is True + # Wrong codes still fail. + assert clients.verify_totp(derived_id, "000000") is False + + +def test_dynamic_client_without_owner_secret_cannot_verify_totp( + clients: ClientStore, +) -> None: + """If the owner is revoked after the derived client is created, TOTP + verification fails closed — a derived client can't outlive its owner.""" + owner_id, _, owner_seed = clients.create("human") + derived_id, _ = clients.create_dynamic( + owner_client_id=owner_id, + name="derived", + registration_source="chatgpt:slug1", + ) + + # Revoke owner; derived row cascades away (verify_totp returns False + # because the client is gone). + clients.revoke(owner_id) + assert clients.verify_totp(derived_id, pyotp.TOTP(owner_seed).now()) is False + + +def test_revoke_owner_cascades_to_derived(clients: ClientStore) -> None: + owner_id, _, _ = clients.create("human") + d1, _ = clients.create_dynamic( + owner_client_id=owner_id, name="chatgpt-1", registration_source="chatgpt:s1", + ) + d2, _ = clients.create_dynamic( + owner_client_id=owner_id, name="chatgpt-2", registration_source="chatgpt:s2", + ) + + assert clients.exists(d1) and clients.exists(d2) + clients.revoke(owner_id) + assert not clients.exists(d1) + assert not clients.exists(d2) + + +def test_revoke_derived_leaves_owner_intact(clients: ClientStore) -> None: + owner_id, _, _ = clients.create("human") + d1, _ = clients.create_dynamic( + owner_client_id=owner_id, name="chatgpt", registration_source="chatgpt:s1", + ) + clients.revoke(d1) + assert clients.exists(owner_id) + assert not clients.exists(d1) + + +def test_list_derived_scoped_to_owner(clients: ClientStore) -> None: + a_id, _, _ = clients.create("owner_A") + b_id, _, _ = clients.create("owner_B") + clients.create_dynamic( + owner_client_id=a_id, name="A-chatgpt", registration_source="chatgpt:s1", + ) + clients.create_dynamic( + owner_client_id=b_id, name="B-chatgpt", registration_source="chatgpt:s2", + ) + a_derived = clients.list_derived(a_id) + assert len(a_derived) == 1 and a_derived[0].name == "A-chatgpt" + + +def test_create_dynamic_requires_existing_owner(clients: ClientStore) -> None: + with pytest.raises(ValueError): + clients.create_dynamic( + owner_client_id="does-not-exist", + name="x", + registration_source="chatgpt:s1", + ) + + +def test_derived_client_round_trips_through_disk( + tmp_path: Path, +) -> None: + """Owner + derived survive a ClientStore reopen (JSON serialization).""" + path = tmp_path / "clients.json" + first = ClientStore(path) + owner_id, _, owner_seed = first.create("human") + derived_id, derived_secret = first.create_dynamic( + owner_client_id=owner_id, + name="chatgpt", + registration_source="chatgpt:s1", + ) + + second = ClientStore(path) + derived = second.get(derived_id) + assert derived is not None + assert derived.owner_client_id == owner_id + assert derived.registration_source == "chatgpt:s1" + # TOTP delegation still works after reload. + assert second.verify_totp(derived_id, pyotp.TOTP(owner_seed).now()) is True + assert second.verify(derived_id, derived_secret) is True diff --git a/tests/test_integration.py b/tests/test_integration.py new file mode 100644 index 0000000..89d747f --- /dev/null +++ b/tests/test_integration.py @@ -0,0 +1,917 @@ +""" +BeaconMCP -- Integration test suite +=================================== + +Run against real infrastructure once services are back online. +Requires a filled .env file with real credentials. + +Usage: + # Run all tests (requires all services: PVE1 + SSH + iLO) + python tests/test_integration.py + + # Run a specific section + python tests/test_integration.py --section proxmox + python tests/test_integration.py --section ssh + python tests/test_integration.py --section bmc + + # Run with a test VM (for destructive tests: start/stop/clone) + python tests/test_integration.py --test-vmid 9999 + +Prerequisites: + - API token created on pve1 (see README.md) + - QEMU Guest Agent installed in at least one VM + - SSH access to pve1 with password auth + - iLO accessible via pve1 tunnel (for iLO tests) +""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import os +import sys +import time +from dataclasses import dataclass +from pathlib import Path + +# Ensure the project is importable +sys.path.insert(0, str(Path(__file__).parent.parent / "src")) + +from dotenv import load_dotenv + +load_dotenv(Path(__file__).parent.parent / ".env") + + +# --------------------------------------------------------------------------- +# Test infrastructure +# --------------------------------------------------------------------------- + +@dataclass +class TestResult: + name: str + passed: bool + message: str + data: dict | None = None + + +class TestRunner: + def __init__(self) -> None: + self.results: list[TestResult] = [] + self._section = "" + + def section(self, name: str) -> None: + self._section = name + print(f"\n{'=' * 60}") + print(f" {name}") + print(f"{'=' * 60}") + + def record(self, name: str, passed: bool, message: str, data: dict | None = None) -> None: + full_name = f"[{self._section}] {name}" if self._section else name + result = TestResult(full_name, passed, message, data) + self.results.append(result) + icon = "PASS" if passed else "FAIL" + print(f" [{icon}] {name}") + if not passed: + print(f" {message}") + if data and not passed: + preview = json.dumps(data, indent=2, default=str)[:300] + print(f" {preview}") + + def summary(self) -> None: + passed = sum(1 for r in self.results if r.passed) + failed = sum(1 for r in self.results if not r.passed) + total = len(self.results) + print(f"\n{'=' * 60}") + print(f" RESULTS: {passed}/{total} passed, {failed} failed") + print(f"{'=' * 60}") + if failed: + print("\n Failed tests:") + for r in self.results: + if not r.passed: + print(f" - {r.name}: {r.message}") + print() + + +# --------------------------------------------------------------------------- +# Import tools (done lazily so .env is loaded first) +# --------------------------------------------------------------------------- + +def get_tools() -> dict: + """Import and return all registered MCP tools.""" + from beaconmcp.server import mcp + return mcp._tool_manager._tools + + +def call_tool(tools: dict, name: str, **kwargs): + """Call an MCP tool function directly, handling both sync and async.""" + tool = tools[name] + result = tool.fn(**kwargs) + if asyncio.iscoroutine(result): + result = asyncio.get_event_loop().run_until_complete(result) + return result + + +# --------------------------------------------------------------------------- +# Test sections +# --------------------------------------------------------------------------- + +def test_proxmox_monitoring(runner: TestRunner, tools: dict) -> None: + runner.section("Proxmox Monitoring") + + # T1: List nodes + result = call_tool(tools, "proxmox_list_nodes") + has_nodes = isinstance(result, dict) and "nodes" in result and len(result["nodes"]) > 0 + runner.record( + "proxmox_list_nodes returns nodes", + has_nodes, + "Expected a list of nodes with at least 1 entry", + result, + ) + + if has_nodes: + pve1_node = next((n for n in result["nodes"] if n.get("name") == "pve1"), None) + runner.record( + "pve1 is present in node list", + pve1_node is not None, + "pve1 should appear in the cluster node list", + result, + ) + runner.record( + "pve1 is online", + pve1_node is not None and pve1_node.get("status") == "online", + f"pve1 status: {pve1_node.get('status') if pve1_node else 'missing'}", + pve1_node, + ) + + # T2: Node status + result = call_tool(tools, "proxmox_node_status", node="pve1") + has_cpu = isinstance(result, dict) and "cpu_cores" in result + runner.record( + "proxmox_node_status returns CPU/RAM/disk info", + has_cpu and "mem_total_gb" in result and "rootfs_total_gb" in result, + "Expected cpu_cores, mem_total_gb, rootfs_total_gb", + result, + ) + if has_cpu: + runner.record( + "CPU usage is a percentage (0-100)", + 0 <= result.get("cpu_pct", -1) <= 100, + f"cpu_pct = {result.get('cpu_pct')}", + ) + runner.record( + "Uptime is positive", + result.get("uptime_h", 0) > 0, + f"uptime_h = {result.get('uptime_h')}", + ) + runner.record( + "PVE version is present", + result.get("pve_version") is not None, + f"pve_version = {result.get('pve_version')}", + ) + + # T3: List VMs + result = call_tool(tools, "proxmox_list_vms", node="pve1") + has_vms = isinstance(result, dict) and "vms" in result + runner.record( + "proxmox_list_vms returns VM list", + has_vms, + "Expected a 'vms' key with a dict keyed by node", + result, + ) + # New shape: {"vms": {"pve1": [...]}, "total": N} + all_vm_entries = [] + if has_vms and isinstance(result["vms"], dict): + for _node, entries in result["vms"].items(): + for e in entries: + if "vmid" in e: + all_vm_entries.append(e) + if all_vm_entries: + first_vm = all_vm_entries[0] + runner.record( + "VMs have required fields (vmid, name, status, type)", + all(k in first_vm for k in ("vmid", "name", "status", "type")), + f"First VM keys: {list(first_vm.keys())}", + first_vm, + ) + + # T4: VM status (use first running VM found) + running_vm = next((v for v in all_vm_entries if v.get("status") == "running"), None) + if running_vm: + result = call_tool(tools, "proxmox_vm_status", node="pve1", vmid=running_vm["vmid"]) + runner.record( + f"proxmox_vm_status for VMID {running_vm['vmid']} returns details", + isinstance(result, dict) and "status" in result and "cpu_pct" in result, + "Expected status, cpu_pct, memory fields", + result, + ) + else: + runner.record( + "proxmox_vm_status (skipped: no running VM found)", + True, + "Need at least one running VM to test vm_status", + ) + + # T5: Get logs + result = call_tool(tools, "proxmox_get_logs", node="pve1", source="syslog", limit=10) + runner.record( + "proxmox_get_logs (syslog) returns log lines", + isinstance(result, dict) and "lines" in result and len(result.get("lines", [])) > 0, + "Expected non-empty 'lines' list", + result if isinstance(result, dict) and "error" in result else None, + ) + + result = call_tool(tools, "proxmox_get_logs", node="pve1", source="tasks", limit=5) + runner.record( + "proxmox_get_logs (tasks) returns task entries", + isinstance(result, dict) and "entries" in result, + "Expected 'entries' list", + result if isinstance(result, dict) and "error" in result else None, + ) + + # T6: Get tasks + result = call_tool(tools, "proxmox_get_tasks", node="pve1", limit=5) + runner.record( + "proxmox_get_tasks returns task list", + isinstance(result, dict) and "tasks" in result, + "Expected 'tasks' key", + result, + ) + + +def test_proxmox_system(runner: TestRunner, tools: dict) -> None: + runner.section("Proxmox System") + + # T7: Storage status -- new shape: {"storage": {"pve1": [{name, type, ...}]}} + result = call_tool(tools, "proxmox_storage_status", node="pve1") + storage_entries = [] + if isinstance(result, dict) and isinstance(result.get("storage"), dict): + for _node, entries in result["storage"].items(): + storage_entries.extend(e for e in entries if "name" in e) + has_storage = len(storage_entries) > 0 + runner.record( + "proxmox_storage_status returns storage list", + has_storage, + "Expected at least one storage entry", + result, + ) + if has_storage: + first = storage_entries[0] + runner.record( + "Storage entries have name/type/usage fields", + all(k in first for k in ("name", "type", "used_gb", "total_gb")), + f"First storage keys: {list(first.keys())}", + first, + ) + + # T8: Network config + result = call_tool(tools, "proxmox_network_config", node="pve1") + has_ifaces = isinstance(result, dict) and "interfaces" in result and len(result.get("interfaces", [])) > 0 + runner.record( + "proxmox_network_config returns interface list", + has_ifaces, + "Expected at least one network interface", + result, + ) + if has_ifaces: + bridge = next((i for i in result["interfaces"] if i.get("type") == "bridge"), None) + runner.record( + "At least one bridge interface exists", + bridge is not None, + "Proxmox nodes should have at least one bridge (vmbr0)", + ) + + +def test_proxmox_exec(runner: TestRunner, tools: dict) -> None: + runner.section("Proxmox Command Execution (QEMU Guest Agent)") + + # Find a running QEMU VM with guest agent (new vms shape: dict-by-node) + vms_result = call_tool(tools, "proxmox_list_vms", node="pve1") + all_entries = [] + if isinstance(vms_result, dict) and isinstance(vms_result.get("vms"), dict): + for _n, entries in vms_result["vms"].items(): + all_entries.extend(e for e in entries if "vmid" in e) + running_qemu = next( + (v for v in all_entries + if v.get("status") == "running" and v.get("type") == "qemu"), + None, + ) + + if not running_qemu: + runner.record( + "proxmox_run (skipped: no running QEMU VM)", + True, + "Need a running QEMU VM with guest agent to test exec", + ) + return + + vmid = running_qemu["vmid"] + print(f" Using VMID {vmid} ({running_qemu.get('name', '?')}) for exec tests") + + # T9: Sync exec -- simple command + result = call_tool(tools, "proxmox_run", node="pve1", vmid=vmid, + command="echo BeaconMCP-test", timeout=30) + runner.record( + f"proxmox_run 'echo' in VM {vmid}", + isinstance(result, dict) and result.get("status") == "ok" and "BeaconMCP-test" in result.get("stdout", ""), + f"Expected stdout containing 'BeaconMCP-test', got: {result}", + result, + ) + + # T10: Sync exec -- exit code + result = call_tool(tools, "proxmox_run", node="pve1", vmid=vmid, + command="cat /etc/hostname", timeout=30) + runner.record( + "proxmox_run 'cat /etc/hostname' returns exit_code 0", + isinstance(result, dict) and result.get("exit_code") == 0, + f"exit_code = {result.get('exit_code')}, stdout = {result.get('stdout', '')[:100]}", + result, + ) + + # T11: Async start + poll via unified proxmox_run + result = call_tool(tools, "proxmox_run", node="pve1", vmid=vmid, + command="sleep 3 && echo async-done", wait=False) + has_exec_id = isinstance(result, dict) and "exec_id" in result + runner.record( + "proxmox_run(wait=False) returns exec_id", + has_exec_id and result.get("status") == "running", + "Expected status=running with exec_id", + result, + ) + + if has_exec_id: + exec_id = result["exec_id"] + # Poll until done (max 30s) + deadline = time.time() + 30 + final_result = None + while time.time() < deadline: + final_result = call_tool(tools, "proxmox_run", exec_id=exec_id) + if isinstance(final_result, dict) and final_result.get("status") != "running": + break + time.sleep(2) + + runner.record( + "proxmox_run(exec_id=...) returns completed result", + isinstance(final_result, dict) and final_result.get("status") == "ok", + f"Final status: {final_result.get('status') if final_result else 'none'}", + final_result, + ) + if isinstance(final_result, dict) and final_result.get("status") == "ok": + runner.record( + "Async stdout contains 'async-done'", + "async-done" in final_result.get("stdout", ""), + f"stdout = {final_result.get('stdout', '')[:100]}", + ) + + # T12: Exec on non-existent VM + result = call_tool(tools, "proxmox_run", node="pve1", vmid=99999, + command="echo test", timeout=10) + runner.record( + "proxmox_run on invalid VMID returns error", + isinstance(result, dict) and result.get("status") == "error", + f"Expected error, got: {result}", + result, + ) + + +def test_proxmox_exec_lxc(runner: TestRunner, tools: dict) -> None: + runner.section("Proxmox Command Execution (LXC)") + + vms_result = call_tool(tools, "proxmox_list_vms", node="pve1") + all_entries = [] + if isinstance(vms_result, dict) and isinstance(vms_result.get("vms"), dict): + for _n, entries in vms_result["vms"].items(): + all_entries.extend(e for e in entries if "vmid" in e) + running_lxc = next( + (v for v in all_entries + if v.get("status") == "running" and v.get("type") == "lxc"), + None, + ) + + if not running_lxc: + runner.record( + "LXC exec (skipped: no running LXC container)", + True, + "Need a running LXC container to test lxc exec", + ) + return + + vmid = running_lxc["vmid"] + print(f" Using CT {vmid} ({running_lxc.get('name', '?')}) for LXC exec tests") + + result = call_tool(tools, "proxmox_run", node="pve1", vmid=vmid, + command="echo LXC-test", timeout=30) + # The Proxmox API does not expose an exec endpoint for LXC containers; the + # tool must return an actionable error pointing to ssh_run + pct exec. + runner.record( + f"proxmox_run on CT {vmid} returns LXC-not-supported guidance", + isinstance(result, dict) + and result.get("status") == "error" + and "pct exec" in result.get("error", ""), + f"Result: {result}", + result, + ) + + +def test_proxmox_vm_lifecycle(runner: TestRunner, tools: dict, test_vmid: int | None) -> None: + runner.section("Proxmox VM Lifecycle") + + if not test_vmid: + runner.record( + "VM lifecycle tests (skipped: no --test-vmid provided)", + True, + "Pass --test-vmid to run start/stop/clone/config tests on a sacrificial VM", + ) + return + + print(f" Using VMID {test_vmid} for lifecycle tests") + + # T13: Read VM config + result = call_tool(tools, "proxmox_vm_config", node="pve1", vmid=test_vmid) + runner.record( + f"proxmox_vm_config (read) for VMID {test_vmid}", + isinstance(result, dict) and ("config" in result or "error" in result), + "Expected config or error", + result, + ) + + if isinstance(result, dict) and "error" in result: + runner.record( + "VM lifecycle tests aborted: test VM not found", + False, + f"VMID {test_vmid} not found. Create it first or use a different --test-vmid.", + result, + ) + return + + vm_type = result.get("type", "qemu") + + # T14: Stop (if running) + status_result = call_tool(tools, "proxmox_vm_status", node="pve1", vmid=test_vmid) + if isinstance(status_result, dict) and status_result.get("status") == "running": + result = call_tool(tools, "proxmox_vm_stop", node="pve1", vmid=test_vmid) + runner.record( + f"proxmox_vm_stop VMID {test_vmid}", + isinstance(result, dict) and "upid" in result, + f"Expected UPID, got: {result}", + result, + ) + print(" Waiting 10s for VM to stop...") + time.sleep(10) + + # T15: Start + result = call_tool(tools, "proxmox_vm_start", node="pve1", vmid=test_vmid) + runner.record( + f"proxmox_vm_start VMID {test_vmid}", + isinstance(result, dict) and ("upid" in result or "error" in result), + f"Result: {result}", + result, + ) + if isinstance(result, dict) and "upid" in result: + print(" Waiting 10s for VM to start...") + time.sleep(10) + + # T16: Verify it's running + status_result = call_tool(tools, "proxmox_vm_status", node="pve1", vmid=test_vmid) + runner.record( + f"VMID {test_vmid} is running after start", + isinstance(status_result, dict) and status_result.get("status") == "running", + f"Status: {status_result.get('status') if isinstance(status_result, dict) else status_result}", + status_result, + ) + + # T17: Restart + result = call_tool(tools, "proxmox_vm_restart", node="pve1", vmid=test_vmid) + runner.record( + f"proxmox_vm_restart VMID {test_vmid}", + isinstance(result, dict) and ("upid" in result or "error" in result), + f"Result: {result}", + result, + ) + + # T18: Modify config (change description, harmless) + result = call_tool(tools, "proxmox_vm_config", node="pve1", vmid=test_vmid, + updates={"description": "BeaconMCP test VM - safe to delete"}) + runner.record( + f"proxmox_vm_config (update description) VMID {test_vmid}", + isinstance(result, dict) and ("applied" in result or "error" in result), + f"Result: {result}", + result, + ) + + # T19: Clone (to VMID test_vmid+1000) + clone_id = test_vmid + 1000 + result = call_tool(tools, "proxmox_vm_clone", node="pve1", vmid=test_vmid, + newid=clone_id, name="beaconmcp-test-clone") + runner.record( + f"proxmox_vm_clone {test_vmid} -> {clone_id}", + isinstance(result, dict) and ("upid" in result or "error" in result), + f"Result: {result}", + result, + ) + + # Cleanup: delete the clone if it was created + if isinstance(result, dict) and "upid" in result: + print(" Waiting 15s for clone to complete...") + time.sleep(15) + # Stop clone if running, then delete + call_tool(tools, "proxmox_vm_stop", node="pve1", vmid=clone_id, force=True) + time.sleep(5) + from beaconmcp.server import proxmox_client + proxmox_client.delete("pve1", f"nodes/pve1/{vm_type}/{clone_id}") + print(f" Cleaned up clone VMID {clone_id}") + + +def test_ssh(runner: TestRunner, tools: dict) -> None: + runner.section("SSH Module") + + if "ssh_run" not in tools: + runner.record( + "SSH tests (skipped: SSH not configured)", + True, + "Declare at least one host under ssh.hosts[] in beaconmcp.yaml to enable SSH tests", + ) + return + + # Use the first declared SSH host so the tests work regardless of whether + # the user named it 'pve1-ssh', 'vps1', etc. Hardcoded 'pve1' no longer + # resolves after the collision rule was introduced. + from beaconmcp.server import config as _cfg + if not _cfg.ssh or not _cfg.ssh.hosts: + runner.record( + "SSH tests (skipped: no ssh.hosts[] declared)", + True, + "Add an entry under ssh.hosts[] to enable SSH tests", + ) + return + ssh_target = _cfg.ssh.hosts[0].name + + # T20: SSH exec + result = call_tool(tools, "ssh_run", host=ssh_target, command="uptime", timeout=30) + runner.record( + f"ssh_run 'uptime' on {ssh_target}", + isinstance(result, dict) and result.get("status") == "ok" and "load average" in result.get("stdout", ""), + f"Result: {result}", + result, + ) + + # T21: SSH exec -- hostname + result = call_tool(tools, "ssh_run", host=ssh_target, command="hostname", timeout=15) + runner.record( + f"ssh_run 'hostname' on {ssh_target}", + isinstance(result, dict) and result.get("status") == "ok" and len(result.get("stdout", "").strip()) > 0, + f"stdout = '{result.get('stdout', '').strip()}'", + result, + ) + + # T22: SSH exec -- df (disk usage) + result = call_tool(tools, "ssh_run", host=ssh_target, command="df -h /", timeout=15) + runner.record( + f"ssh_run 'df -h /' on {ssh_target}", + isinstance(result, dict) and result.get("status") == "ok", + f"exit_code = {result.get('exit_code')}", + result, + ) + + # T23: SSH exec -- failing command + result = call_tool(tools, "ssh_run", host=ssh_target, + command="cat /nonexistent/file/12345", timeout=15) + runner.record( + "ssh_run on nonexistent file returns non-zero exit code", + isinstance(result, dict) and result.get("exit_code", 0) != 0, + f"exit_code = {result.get('exit_code')}, stderr = {result.get('stderr', '')[:100]}", + result, + ) + + # T24: SSH host resolution -- every target must be declared in ssh.hosts[]. + # A raw IP/hostname not declared there now raises SSHHostResolutionError + # (no more implicit passthrough). Exercise both paths. + from beaconmcp.server import config as _cfg + from beaconmcp.server import ssh_client + from beaconmcp.ssh.client import SSHHostResolutionError + + declared = _cfg.ssh.hosts[0] if _cfg.ssh and _cfg.ssh.hosts else None + if declared is not None: + resolved = ssh_client.resolve_host(declared.name) + runner.record( + f"SSH resolve_host('{declared.name}') -> declared host field", + resolved == declared.host, + f"Resolved to: {resolved}, expected {declared.host}", + ) + + try: + ssh_client.resolve_host("definitely-not-declared-xyzzy") + runner.record( + "SSH resolve_host raises on undeclared identifier", + False, + "Expected SSHHostResolutionError, got silent pass", + ) + except SSHHostResolutionError: + runner.record( + "SSH resolve_host raises on undeclared identifier", + True, + "", + ) + + # T25: SSH async start + poll via unified ssh_run + result = call_tool(tools, "ssh_run", host=ssh_target, + command="sleep 2 && echo ssh-async-done", wait=False) + has_id = isinstance(result, dict) and "exec_id" in result + runner.record( + "ssh_run(wait=False) returns exec_id", + has_id, + f"Result: {result}", + result, + ) + + if has_id: + exec_id = result["exec_id"] + print(f" Polling exec_id={exec_id}...") + deadline = time.time() + 30 + final = None + while time.time() < deadline: + final = call_tool(tools, "ssh_run", exec_id=exec_id) + if isinstance(final, dict) and final.get("status") != "running": + break + time.sleep(2) + runner.record( + "ssh_run(exec_id=...) returns completed", + isinstance(final, dict) and final.get("status") == "ok", + f"Final: {final}", + final, + ) + + # T26: List sessions + result = call_tool(tools, "ssh_list_sessions") + runner.record( + "ssh_list_sessions returns session list", + isinstance(result, dict) and "sessions" in result, + f"Result: {result}", + result, + ) + + +def test_bmc(runner: TestRunner, tools: dict) -> None: + runner.section("BMC Module") + + if "bmc_server_info" not in tools: + runner.record( + "BMC tests (skipped: no BMC device configured)", + True, + "Add at least one entry to bmc.devices[] in beaconmcp.yaml to enable BMC tests.", + ) + return + + # T27: Server info + result = call_tool(tools, "bmc_server_info") + runner.record( + "bmc_server_info returns server details", + isinstance(result, dict) and ("product_name" in result or "fru" in result or "error" in result), + f"Result: {json.dumps(result, default=str)[:200]}", + result, + ) + if isinstance(result, dict) and "error" in result: + runner.record( + "BMC tests aborted: cannot reach BMC", + False, + result["error"], + ) + return + + # T28: Health status + result = call_tool(tools, "bmc_health_status") + runner.record( + "bmc_health_status returns health data", + isinstance(result, dict) and "error" not in result, + f"Result type: {type(result).__name__}, keys: {list(result.keys())[:5] if isinstance(result, dict) else 'N/A'}", + result if isinstance(result, dict) and "error" in result else None, + ) + + # T29: Power status + result = call_tool(tools, "bmc_power_status") + runner.record( + "bmc_power_status returns ON/OFF", + isinstance(result, dict) and "power_status" in result, + f"Result: {result}", + result, + ) + if isinstance(result, dict) and "power_status" in result: + runner.record( + "Server power is ON", + str(result["power_status"]).upper() == "ON", + f"power_status = {result['power_status']}", + ) + + # T30: Event log + result = call_tool(tools, "bmc_get_event_log", limit=10) + runner.record( + "bmc_get_event_log returns events", + isinstance(result, dict) and ("events" in result or "error" in result), + f"Total events: {result.get('total', '?')}", + result if isinstance(result, dict) and "error" in result else None, + ) + + # NOTE: power_on / power_off / power_reset are destructive; skipped in + # automated runs. Test them manually if needed. + runner.record( + "bmc_power_on/off/reset (not tested: destructive)", + True, + "Manual testing required. Use bmc_power_status to verify state first.", + ) + + +def test_mcp_resources(runner: TestRunner) -> None: + runner.section("MCP Resources & Prompts") + + from beaconmcp.server import mcp, config + + # T31: Infrastructure resource + resource_fn = None + for key, res in mcp._resource_manager._resources.items(): + if "infrastructure" in str(key): + resource_fn = res + break + + runner.record( + "beaconmcp://infrastructure resource is registered", + resource_fn is not None, + "Expected a resource matching 'infrastructure'", + ) + + # T32: Infrastructure YAML is loaded + runner.record( + "infrastructure.yaml is loaded into config", + len(config.infrastructure) > 0, + f"Keys: {list(config.infrastructure.keys()) if config.infrastructure else 'empty'}", + ) + + # T33: Prompt is registered + prompts = mcp._prompt_manager._prompts + runner.record( + "beaconmcp_context prompt is registered", + "beaconmcp_context" in prompts, + f"Available prompts: {list(prompts.keys())}", + ) + + # T34: Config validation + runner.record( + "Config has at least 1 PVE node", + len(config.pve_nodes) >= 1, + f"Nodes: {[n.name for n in config.pve_nodes]}", + ) + runner.record( + "PVE1 host matches expected", + config.pve_nodes[0].host == os.environ.get("PVE1_HOST", ""), + f"Host: {config.pve_nodes[0].host}", + ) + + +def test_aggregators(runner: TestRunner, tools: dict) -> None: + runner.section("Aggregators & Field Filtering") + + # A1: cluster_overview returns nodes + vms + storage in one call + result = call_tool(tools, "cluster_overview") + runner.record( + "cluster_overview returns nodes, vms, storage", + isinstance(result, dict) + and "nodes" in result and "vms" in result and "storage" in result, + f"Keys: {list(result.keys()) if isinstance(result, dict) else type(result).__name__}", + result, + ) + + # A2: cluster_overview(include_storage=False) omits storage + result = call_tool(tools, "cluster_overview", include_storage=False) + runner.record( + "cluster_overview(include_storage=False) omits storage", + isinstance(result, dict) and "nodes" in result and "vms" in result + and "storage" not in result, + f"storage key present: {'storage' in result if isinstance(result, dict) else 'N/A'}", + result, + ) + + # A3: cluster_health on first configured node + from beaconmcp.server import config as _cfg + first_node = _cfg.pve_nodes[0].name if _cfg.pve_nodes else "" + if first_node: + result = call_tool(tools, "cluster_health", node=first_node) + runner.record( + f"cluster_health(node='{first_node}') returns metrics + errors", + isinstance(result, dict) and ("cpu_pct" in result or "error" in result), + f"Keys: {list(result.keys()) if isinstance(result, dict) else type(result).__name__}", + result, + ) + + # A4: vm_find wildcard + result = call_tool(tools, "vm_find", pattern="*") + runner.record( + "vm_find('*') returns vms list", + isinstance(result, dict) and "vms" in result and "total" in result, + f"Got {result.get('total') if isinstance(result, dict) else 0} matches", + result, + ) + + # A5: fields filter on proxmox_list_nodes + result = call_tool(tools, "proxmox_list_nodes", fields=["name", "status"]) + ok = ( + isinstance(result, dict) + and isinstance(result.get("nodes"), list) + and all(set(n.keys()).issubset({"name", "status", "error"}) for n in result["nodes"]) + ) + runner.record( + "proxmox_list_nodes(fields=['name','status']) trims response", + ok, + f"Sample: {result['nodes'][:1] if isinstance(result, dict) and result.get('nodes') else 'empty'}", + result, + ) + + +def test_error_handling(runner: TestRunner, tools: dict) -> None: + runner.section("Error Handling") + + # T35: Invalid node + result = call_tool(tools, "proxmox_node_status", node="nonexistent-node") + runner.record( + "Invalid node returns actionable error", + isinstance(result, dict) and "error" in result, + f"Result: {result}", + result, + ) + + # T36: Invalid VMID + result = call_tool(tools, "proxmox_vm_status", node="pve1", vmid=99999) + runner.record( + "Invalid VMID returns error (not crash)", + isinstance(result, dict) and "error" in result, + f"Result: {result}", + result, + ) + + # T37: Invalid exec_id (via unified proxmox_run poll mode) + result = call_tool(tools, "proxmox_run", exec_id="nonexistent") + runner.record( + "Invalid exec_id returns error", + isinstance(result, dict) and "error" in result, + f"Result: {result}", + result, + ) + + if "ssh_run" in tools: + result = call_tool(tools, "ssh_run", exec_id="nonexistent") + runner.record( + "Invalid SSH exec_id returns error", + isinstance(result, dict) and "error" in result, + f"Result: {result}", + result, + ) + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +def main() -> None: + parser = argparse.ArgumentParser(description="BeaconMCP integration tests") + parser.add_argument("--section", choices=["proxmox", "ssh", "bmc", "all"], default="all", + help="Which section to test") + parser.add_argument("--test-vmid", type=int, default=None, + help="VMID of a sacrificial test VM for lifecycle tests (start/stop/clone)") + args = parser.parse_args() + + runner = TestRunner() + tools = get_tools() + + print("\nBeaconMCP Integration Tests") + print(f"Tools registered: {len(tools)}") + print(f"Section: {args.section}") + if args.test_vmid: + print(f"Test VMID: {args.test_vmid}") + + sections = args.section + + if sections in ("proxmox", "all"): + test_proxmox_monitoring(runner, tools) + test_proxmox_system(runner, tools) + test_proxmox_exec(runner, tools) + test_proxmox_exec_lxc(runner, tools) + test_proxmox_vm_lifecycle(runner, tools, args.test_vmid) + test_aggregators(runner, tools) + test_error_handling(runner, tools) + + if sections in ("ssh", "all"): + test_ssh(runner, tools) + + if sections in ("bmc", "all"): + test_bmc(runner, tools) + + if sections == "all": + test_mcp_resources(runner) + + runner.summary() + + # Exit with error code if any test failed + failed = sum(1 for r in runner.results if not r.passed) + sys.exit(1 if failed else 0) + + +if __name__ == "__main__": + main() diff --git a/tests/test_mcp_apps_panel.py b/tests/test_mcp_apps_panel.py new file mode 100644 index 0000000..ee9b699 --- /dev/null +++ b/tests/test_mcp_apps_panel.py @@ -0,0 +1,247 @@ +"""Tests for the MCP Apps panels. + +The value here is the wire format: a host only renders a panel if the tool +carries `_meta.ui.resourceUri` and the resource comes back under the mcp-app +MIME type, so these assert what goes out on the wire rather than the Python +objects behind it. +""" + +from __future__ import annotations + +import sys +from pathlib import Path +from typing import Any + +import pytest +from mcp.server.fastmcp import FastMCP + +sys.path.insert(0, str(Path(__file__).parent.parent / "src")) + +from beaconmcp.proxmox.panel import ( + APP_MIME_TYPE, + CLUSTER_PANEL_URI, + LOGS_PANEL_URI, + VM_PANEL_URI, + _syslog_level, + _vm_snapshot, + register_panel_tools, +) + +_NOT_FOUND = {"error": "Configuration file 'nodes/pve1/qemu/100.conf' does not exist"} + +_QEMU_STATUS = { + "status": "running", + "name": "web-101", + "cpu": 0.1234, + "cpus": 4, + "mem": 2 * 1048576 * 1024, + "maxmem": 4 * 1048576 * 1024, + # QEMU reports 0 unless the guest agent answers. + "disk": 0, + "maxdisk": 32 * 1073741824, + "uptime": 7200, +} + +_PANELS = [ + ("proxmox_vm_panel", VM_PANEL_URI), + ("proxmox_logs_panel", LOGS_PANEL_URI), + ("cluster_overview_interactive", CLUSTER_PANEL_URI), +] + + +class _Client: + """Minimal ProxmoxClient stand-in keyed on the path suffix.""" + + def __init__(self, responses: dict[str, Any] | None = None) -> None: + self.responses = responses or {} + self.calls: list[str] = [] + self.kwargs: list[dict[str, Any]] = [] + self.configured_nodes = ["pve1"] + + def get(self, node: str, path: str, **kwargs: Any) -> Any: + self.calls.append(path) + self.kwargs.append(kwargs) + for suffix, value in self.responses.items(): + if path.endswith(suffix): + return value + return _NOT_FOUND + + +def _server(client: Any) -> FastMCP: + mcp = FastMCP("test") + register_panel_tools(mcp, client) + return mcp + + +async def _tool(mcp: FastMCP, name: str) -> Any: + return next(t for t in await mcp.list_tools() if t.name == name) + + +# --- wire format ------------------------------------------------------------ + + +@pytest.mark.parametrize("name,uri", _PANELS) +async def test_tool_points_at_its_panel(name: str, uri: str) -> None: + tool = await _tool(_server(_Client()), name) + assert tool.meta == {"ui": {"resourceUri": uri, "visibility": ["model", "app"]}} + + +@pytest.mark.parametrize("name,uri", _PANELS) +async def test_advertised_resource_resolves_as_an_app(name: str, uri: str) -> None: + """A resourceUri that 404s on resources/read renders as a blank frame.""" + mcp = _server(_Client()) + listed = next(r for r in await mcp.list_resources() if str(r.uri) == uri) + assert listed.mimeType == APP_MIME_TYPE + + contents = list(await mcp.read_resource(uri)) + assert contents, f"{uri} is advertised but reads back empty" + + +@pytest.mark.parametrize("name,uri", _PANELS) +async def test_shared_runtime_is_spliced_in(name: str, uri: str) -> None: + """The marker must be gone and both halves of the runtime present. + + A panel that ships with the literal still in it has no + bridge and no styles: it loads, does nothing, and never speaks to the host. + """ + html = list(await _server(_Client()).read_resource(uri))[0].content + + assert "" not in html + assert "const MCPApp" in html, "bridge.js missing" + assert "--font-sans" in html, "panel.css missing" + + +async def test_handshake_params_are_flat() -> None: + """ui/initialize takes appInfo / appCapabilities / protocolVersion, flat. + + Nesting them under `capabilities`, or sending `clientInfo` instead of + `appInfo`, fails the host's schema check -- and a rejected handshake is + silent: the host just never replies, so the app never sends `initialized` + and never receives the tool result. The panel then sits on "Loading..." + with nothing in the console to explain it. + """ + html = list(await _server(_Client()).read_resource(VM_PANEL_URI))[0].content + params = html.split('request("ui/initialize", {', 1)[1].split("})", 1)[0] + + assert "appInfo:" in params + assert "appCapabilities:" in params + assert "protocolVersion:" in params + assert "clientInfo" not in params + assert "capabilities: {" not in params + + +# --- VM snapshot ------------------------------------------------------------ + + +def test_qemu_disk_usage_is_null_not_zero() -> None: + """0 from QEMU means "no guest agent", which must not render as an empty bar.""" + client = _Client({"qemu/100/status/current": _QEMU_STATUS, "qemu/100/config": {"cores": 4, "memory": 4096}}) + + snap = _vm_snapshot(client, "pve1", 100) + + assert snap["disk_used_gb"] is None + assert snap["disk_max_gb"] == 32.0 + + +def test_qemu_snapshot_maps_status_and_config() -> None: + client = _Client({"qemu/100/status/current": _QEMU_STATUS, "qemu/100/config": {"cores": 4, "memory": 4096}}) + + assert _vm_snapshot(client, "pve1", 100) == { + "node": "pve1", "vmid": 100, "type": "qemu", "name": "web-101", + "status": "running", "cpu_pct": 12.3, "cpus": 4, + "mem_used_mb": 2048, "mem_max_mb": 4096, + "disk_used_gb": None, "disk_max_gb": 32.0, "uptime_h": 2.0, + "cores": 4, "memory_mb": 4096, + } + + +def test_falls_through_to_lxc_when_no_qemu_guest() -> None: + client = _Client({ + "lxc/200/status/current": { + "status": "running", "name": "ct", "cpu": 0, "cpus": 1, + "mem": 0, "maxmem": 536870912, + "disk": 5 * 1073741824, "maxdisk": 10 * 1073741824, "uptime": 0, + }, + "lxc/200/config": {"cores": 1, "memory": 512}, + }) + + snap = _vm_snapshot(client, "pve1", 200) + + assert snap["type"] == "lxc" + assert snap["disk_used_gb"] == 5.0 + + +def test_missing_guest_reports_an_error() -> None: + assert "not found" in _vm_snapshot(_Client(), "pve1", 999)["error"] + + +def test_unreadable_config_still_renders_live_state() -> None: + """A config read can fail on its own; the panel should degrade, not blank out.""" + client = _Client({ + "qemu/100/status/current": _QEMU_STATUS, + "qemu/100/config": {"error": "connection refused"}, + }) + + snap = _vm_snapshot(client, "pve1", 100) + + assert snap["status"] == "running" + assert snap["cores"] is None + + +def test_transport_error_is_not_mistaken_for_a_missing_guest() -> None: + """Only a "does not exist" error means "try the other guest type".""" + client = _Client({"qemu/100/status/current": {"error": "connection refused"}}) + + snap = _vm_snapshot(client, "pve1", 100) + + assert snap["error"] == "connection refused" + assert not any("lxc" in call for call in client.calls) + + +# --- logs panel ------------------------------------------------------------- + + +@pytest.mark.parametrize( + "line,expected", + [ + ("kernel: EXT4-fs error (device sda1)", "error"), + ("corosync: connection refused to node pve2", "error"), + ("pvedaemon: deprecated option 'foo'", "warn"), + ("systemd[1]: Started Daily apt upgrade.", "info"), + ], +) +def test_syslog_lines_are_classified(line: str, expected: str) -> None: + assert _syslog_level(line) == expected + + +async def test_logs_panel_labels_syslog_lines() -> None: + client = _Client({"syslog": [{"t": "kernel: I/O error"}, {"t": "systemd: Started thing"}]}) + tool = await _server(client).call_tool("proxmox_logs_panel", {"node": "pve1"}) + + entries = tool[1]["entries"] + assert [e["level"] for e in entries] == ["error", "info"] + assert entries[0]["text"] == "kernel: I/O error" + + +async def test_logs_panel_flags_failed_tasks() -> None: + """Proxmox writes "OK" on success; anything else is the failure message.""" + client = _Client({"tasks": [ + {"upid": "UPID:1", "type": "vzdump", "status": "OK", "user": "root@pam"}, + {"upid": "UPID:2", "type": "qmstart", "status": "unable to start VM", "user": "root@pam"}, + {"upid": "UPID:3", "type": "qmigrate", "status": None, "user": "root@pam"}, + ]}) + + result = await _server(client).call_tool( + "proxmox_logs_panel", {"node": "pve1", "source": "tasks"}, + ) + + # A running task has no status yet and must not be painted as a failure. + assert [e["level"] for e in result[1]["entries"]] == ["info", "error", "info"] + + +async def test_logs_panel_caps_the_line_count() -> None: + """The Proxmox API rejects anything above 500, and the panel's box has no max.""" + client = _Client({"syslog": []}) + await _server(client).call_tool("proxmox_logs_panel", {"node": "pve1", "limit": 9000}) + + assert client.kwargs[0]["limit"] == 500 diff --git a/tests/test_metrics.py b/tests/test_metrics.py new file mode 100644 index 0000000..dfcfebe --- /dev/null +++ b/tests/test_metrics.py @@ -0,0 +1,50 @@ +"""Tests for the in-process Prometheus-format metrics.""" + +from __future__ import annotations + +from beaconmcp.metrics import Counter, Histogram, Registry + + +def test_counter_increments_with_labels() -> None: + c = Counter("foo_total", "help") + c.inc(tool="ssh_run", status="ok") + c.inc(tool="ssh_run", status="ok") + c.inc(tool="ssh_run", status="err") + out = c.render() + assert 'foo_total{status="ok",tool="ssh_run"} 2' in out + assert 'foo_total{status="err",tool="ssh_run"} 1' in out + + +def test_histogram_buckets_and_sum() -> None: + h = Histogram("lat_ms", "help", buckets_ms=(10, 100)) + h.observe(5, tool="x") + h.observe(50, tool="x") + h.observe(200, tool="x") + out = h.render() + # 5ms: fits bucket <=10, <=100, +Inf + # 50ms: fits <=100, +Inf + # 200ms: only +Inf + assert 'lat_ms_bucket{le="10",tool="x"} 1' in out + assert 'lat_ms_bucket{le="100",tool="x"} 2' in out + assert 'lat_ms_bucket{le="+Inf",tool="x"} 3' in out + assert "lat_ms_sum" in out + assert 'lat_ms_count{tool="x"} 3' in out + + +def test_registry_concats() -> None: + r = Registry() + r.register(Counter("a_total", "a")) + r.register(Counter("b_total", "b")) + out = r.render() + assert "# TYPE a_total counter" in out + assert "# TYPE b_total counter" in out + # No empty tail. + assert out.endswith("\n") + + +def test_histogram_time_context_manager() -> None: + h = Histogram("t_ms", "help") + with h.time(tool="y"): + pass + out = h.render() + assert 't_ms_count{tool="y"} 1' in out diff --git a/tests/test_passkeys.py b/tests/test_passkeys.py new file mode 100644 index 0000000..06279ac --- /dev/null +++ b/tests/test_passkeys.py @@ -0,0 +1,808 @@ +"""Passkey (WebAuthn) tests. + +Drives the real ceremonies end to end against a software authenticator +built here from ``cryptography`` primitives: a fake that only returned +canned dicts would prove nothing, since every interesting failure mode +(wrong origin, replayed challenge, foreign credential) lives inside the +signature verification. + +Run with:: + + pytest tests/test_passkeys.py -v +""" + +from __future__ import annotations + +import hashlib +import json +import os +import struct +import sys +import time +from pathlib import Path + +import pytest +from starlette.applications import Starlette +from starlette.testclient import TestClient + +sys.path.insert(0, str(Path(__file__).parent.parent / "src")) + +pytest.importorskip("webauthn", reason="passkey support is an optional extra") +cbor2 = pytest.importorskip("cbor2") + +from cryptography.hazmat.primitives import hashes # noqa: E402 +from cryptography.hazmat.primitives.asymmetric import ec # noqa: E402 +from cryptography.hazmat.primitives.asymmetric.utils import ( # noqa: E402 + decode_dss_signature, + encode_dss_signature, +) + +from beaconmcp.auth import TotpResult # noqa: E402 +from beaconmcp.dashboard.app import ( # noqa: E402 + DashboardDeps, + SESSION_COOKIE, + build_dashboard_routes, +) +from beaconmcp.dashboard.csrf import CSRF_COOKIE # noqa: E402 +from beaconmcp.dashboard.db import Database # noqa: E402 +from beaconmcp.dashboard.passkeys import ( # noqa: E402 + ChallengeStore, + PasskeyError, + PasskeyService, + PasskeyStore, + b64url_decode, + b64url_encode, + default_label, +) +from beaconmcp.dashboard.session import SessionStore # noqa: E402 + + +ORIGIN = "https://beacon.example" +RP_ID = "beacon.example" + + +# --------------------------------------------------------------------------- +# Software authenticator +# --------------------------------------------------------------------------- + +class SoftwareAuthenticator: + """A minimal ES256 authenticator: enough to produce valid ceremonies.""" + + AAGUID = b"\x00" * 16 + + def __init__(self) -> None: + self.key = ec.generate_private_key(ec.SECP256R1()) + self.credential_id = os.urandom(32) + self.sign_count = 0 + + # --- helpers --------------------------------------------------------- + + def _cose_key(self) -> bytes: + numbers = self.key.public_key().public_numbers() + return cbor2.dumps({ + 1: 2, # kty: EC2 + 3: -7, # alg: ES256 + -1: 1, # crv: P-256 + -2: numbers.x.to_bytes(32, "big"), + -3: numbers.y.to_bytes(32, "big"), + }) + + @staticmethod + def _client_data(kind: str, challenge: str, origin: str) -> bytes: + return json.dumps( + { + "type": kind, + "challenge": challenge, + "origin": origin, + "crossOrigin": False, + }, + separators=(",", ":"), + ).encode("utf-8") + + def _auth_data(self, rp_id: str, flags: int, attested: bool) -> bytes: + data = hashlib.sha256(rp_id.encode("utf-8")).digest() + data += bytes([flags]) + data += struct.pack(">I", self.sign_count) + if attested: + key = self._cose_key() + data += self.AAGUID + data += struct.pack(">H", len(self.credential_id)) + data += self.credential_id + data += key + return data + + # --- ceremonies ------------------------------------------------------ + + def create(self, options: dict, *, origin: str = ORIGIN, rp_id: str = RP_ID) -> dict: + client_data = self._client_data( + "webauthn.create", options["challenge"], origin, + ) + # UP | UV | AT + auth_data = self._auth_data(rp_id, 0x01 | 0x04 | 0x40, attested=True) + attestation = cbor2.dumps({ + "fmt": "none", "attStmt": {}, "authData": auth_data, + }) + return { + "id": b64url_encode(self.credential_id), + "rawId": b64url_encode(self.credential_id), + "type": "public-key", + "clientExtensionResults": {}, + "response": { + "clientDataJSON": b64url_encode(client_data), + "attestationObject": b64url_encode(attestation), + "transports": ["internal"], + }, + } + + def get(self, options: dict, *, origin: str = ORIGIN, rp_id: str = RP_ID) -> dict: + self.sign_count += 1 + client_data = self._client_data("webauthn.get", options["challenge"], origin) + auth_data = self._auth_data(rp_id, 0x01 | 0x04, attested=False) + payload = auth_data + hashlib.sha256(client_data).digest() + raw = self.key.sign(payload, ec.ECDSA(hashes.SHA256())) + # Re-encode so the DER is canonical whatever the backend produced. + r, s = decode_dss_signature(raw) + signature = encode_dss_signature(r, s) + return { + "id": b64url_encode(self.credential_id), + "rawId": b64url_encode(self.credential_id), + "type": "public-key", + "clientExtensionResults": {}, + "response": { + "clientDataJSON": b64url_encode(client_data), + "authenticatorData": b64url_encode(auth_data), + "signature": b64url_encode(signature), + "userHandle": None, + }, + } + + +class FakeRequest: + """Just the surface :mod:`beaconmcp.dashboard.passkeys` reads.""" + + def __init__(self, host: str = RP_ID, scheme: str = "https") -> None: + self.headers = {"host": host, "x-forwarded-proto": scheme} + + class _URL: + netloc = host + + self.url = _URL() + self.url.scheme = scheme # type: ignore[attr-defined] + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + +@pytest.fixture() +def store(tmp_path) -> PasskeyStore: + return PasskeyStore(Database(tmp_path / "dashboard.db")) + + +@pytest.fixture() +def service(store) -> PasskeyService: + return PasskeyService(store) + + +@pytest.fixture() +def enrolled(service): + """A client with one registered passkey, plus its authenticator.""" + request = FakeRequest() + auth = SoftwareAuthenticator() + options, state = service.registration_options( + request, client_id="beaconmcp_test", client_name="Test Client", + ) + record = service.verify_registration( + request, state=state, credential=auth.create(options), + ) + return auth, record + + +# --------------------------------------------------------------------------- +# RP identity +# --------------------------------------------------------------------------- + +def test_rp_id_strips_port_and_scheme(): + from beaconmcp.dashboard.passkeys import origin_for, rp_id_for + + req = FakeRequest(host="beacon.example:8420") + assert rp_id_for(req) == "beacon.example" + assert origin_for(req) == "https://beacon.example:8420" + + +def test_rp_id_handles_ipv6_literal(): + from beaconmcp.dashboard.passkeys import rp_id_for + + assert rp_id_for(FakeRequest(host="[::1]:8420")) == "::1" + + +def test_secure_context_requires_https_or_loopback(): + from beaconmcp.dashboard.passkeys import is_secure_context + + assert is_secure_context(FakeRequest(host="beacon.example", scheme="https")) + assert is_secure_context(FakeRequest(host="localhost:8420", scheme="http")) + assert not is_secure_context(FakeRequest(host="192.168.1.5:8420", scheme="http")) + + +def test_forwarded_host_wins_over_host(): + from beaconmcp.dashboard.passkeys import rp_id_for + + req = FakeRequest(host="internal:8420") + req.headers["x-forwarded-host"] = "public.example" + assert rp_id_for(req) == "public.example" + + +def test_default_label_from_user_agent(): + assert default_label("Mozilla/5.0 (iPhone; CPU iPhone OS 17_0)") == "iPhone / iPad" + assert default_label("Mozilla/5.0 (Windows NT 10.0)") == "Windows" + assert default_label("") == "Passkey" + + +# --------------------------------------------------------------------------- +# Challenge store +# --------------------------------------------------------------------------- + +def test_challenge_is_single_use(): + challenges = ChallengeStore() + state = challenges.issue( + purpose="authenticate", challenge=b"abc", client_id="c1", + ) + assert challenges.consume(state, "authenticate") is not None + assert challenges.consume(state, "authenticate") is None + + +def test_challenge_purpose_must_match(): + challenges = ChallengeStore() + state = challenges.issue(purpose="register", challenge=b"abc", client_id="c1") + assert challenges.consume(state, "authenticate") is None + + +def test_challenge_expires(): + challenges = ChallengeStore(ttl_seconds=-1) + state = challenges.issue(purpose="register", challenge=b"abc", client_id="c1") + assert challenges.consume(state, "register") is None + + +# --------------------------------------------------------------------------- +# Registration +# --------------------------------------------------------------------------- + +def test_registration_round_trip(service, store): + request = FakeRequest() + auth = SoftwareAuthenticator() + options, state = service.registration_options( + request, client_id="beaconmcp_test", client_name="Test Client", + ) + assert options["rp"]["id"] == RP_ID + assert b64url_decode(options["user"]["id"]) == b"beaconmcp_test" + + record = service.verify_registration( + request, state=state, credential=auth.create(options), label="My laptop", + ) + assert record.client_id == "beaconmcp_test" + assert record.label == "My laptop" + assert record.transports == ["internal"] + assert store.count_for_client("beaconmcp_test") == 1 + + +def test_registration_rejects_wrong_origin(service): + request = FakeRequest() + auth = SoftwareAuthenticator() + options, state = service.registration_options( + request, client_id="beaconmcp_test", client_name="Test Client", + ) + with pytest.raises(PasskeyError): + service.verify_registration( + request, + state=state, + credential=auth.create(options, origin="https://evil.example"), + ) + + +def test_registration_state_bound_to_session(service): + request = FakeRequest() + auth = SoftwareAuthenticator() + options, state = service.registration_options( + request, client_id="beaconmcp_test", client_name="Test", + session_id="session-a", + ) + with pytest.raises(PasskeyError, match="another session"): + service.verify_registration( + request, state=state, credential=auth.create(options), + session_id="session-b", + ) + + +def test_registration_excludes_known_credentials(service, enrolled): + _, record = enrolled + options, _ = service.registration_options( + FakeRequest(), client_id="beaconmcp_test", client_name="Test", + ) + assert [c["id"] for c in options["excludeCredentials"]] == [record.credential_id] + + +def test_registration_label_defaults_from_user_agent(service): + request = FakeRequest() + request.headers["user-agent"] = "Mozilla/5.0 (Macintosh; Intel Mac OS X 14_0)" + auth = SoftwareAuthenticator() + options, state = service.registration_options( + request, client_id="beaconmcp_test", client_name="Test", + ) + record = service.verify_registration( + request, state=state, credential=auth.create(options), + ) + assert record.label == "Mac" + + +# --------------------------------------------------------------------------- +# Authentication +# --------------------------------------------------------------------------- + +def test_authentication_round_trip(service, enrolled, store): + auth, record = enrolled + request = FakeRequest() + options, state = service.authentication_options( + request, client_id="beaconmcp_test", + ) + assert [c["id"] for c in options["allowCredentials"]] == [record.credential_id] + + verified = service.verify_authentication( + request, state=state, credential=auth.get(options), + ) + assert verified.client_id == "beaconmcp_test" + assert store.get(record.credential_id).last_used_at is not None + + +def test_authentication_without_credentials_is_refused(service): + with pytest.raises(PasskeyError, match="No passkey"): + service.authentication_options(FakeRequest(), client_id="nobody") + + +def test_authentication_rejects_replayed_challenge(service, enrolled): + auth, _ = enrolled + request = FakeRequest() + options, state = service.authentication_options( + request, client_id="beaconmcp_test", + ) + assertion = auth.get(options) + service.verify_authentication(request, state=state, credential=assertion) + with pytest.raises(PasskeyError, match="expired"): + service.verify_authentication(request, state=state, credential=assertion) + + +def test_authentication_rejects_foreign_credential(service, enrolled): + """A challenge minted for client A must not accept client B's passkey.""" + auth_a, _ = enrolled + request = FakeRequest() + + other = SoftwareAuthenticator() + options, state = service.registration_options( + request, client_id="beaconmcp_other", client_name="Other", + ) + service.verify_registration( + request, state=state, credential=other.create(options), + ) + + options, state = service.authentication_options( + request, client_id="beaconmcp_other", + ) + # Sign the *other* client's challenge with the first client's key. + with pytest.raises(PasskeyError, match="different client"): + service.verify_authentication( + request, state=state, credential=auth_a.get(options), + ) + + +def test_authentication_rejects_wrong_rp_id(service, enrolled): + auth, _ = enrolled + request = FakeRequest() + options, state = service.authentication_options( + request, client_id="beaconmcp_test", + ) + with pytest.raises(PasskeyError): + service.verify_authentication( + request, state=state, + credential=auth.get(options, rp_id="evil.example"), + ) + + +def test_sign_counter_regression_is_rejected(service, enrolled, store): + auth, record = enrolled + request = FakeRequest() + options, state = service.authentication_options( + request, client_id="beaconmcp_test", + ) + service.verify_authentication(request, state=state, credential=auth.get(options)) + + # Roll the authenticator's counter back, as a cloned key would. + auth.sign_count = 0 + options, state = service.authentication_options( + request, client_id="beaconmcp_test", + ) + with pytest.raises(PasskeyError, match="sign count"): + service.verify_authentication( + request, state=state, credential=auth.get(options), + ) + + +# --------------------------------------------------------------------------- +# Storage +# --------------------------------------------------------------------------- + +def test_delete_is_scoped_to_owner(store, service, enrolled): + _, record = enrolled + assert store.delete(record.credential_id, "someone_else") is False + assert store.delete(record.credential_id, "beaconmcp_test") is True + assert store.count_for_client("beaconmcp_test") == 0 + + +def test_service_without_store_reports_unavailable(): + service = PasskeyService(None) + assert service.available is False + assert "database" in service.unavailable_reason + with pytest.raises(PasskeyError): + service.authentication_options(FakeRequest(), client_id="x") + + +def test_working_service_has_no_unavailable_reason(service): + assert service.available is True + assert service.unavailable_reason is None + + +def test_missing_library_is_reported_with_a_fix(monkeypatch, store): + """The banner/doctor message must name the package and the command.""" + import beaconmcp.dashboard.passkeys as mod + + monkeypatch.setattr(mod, "_webauthn", None) + assert mod.webauthn_installed() is False + reason = PasskeyService(store).unavailable_reason + assert "webauthn" in reason and "pip install" in reason + + +# --------------------------------------------------------------------------- +# Dashboard routes +# --------------------------------------------------------------------------- + +class FakeClientStore: + def __init__(self): + self.clients = { + "beaconmcp_test": { + "secret": "sk_test", "name": "Test Client", "totp": "123456", + } + } + + def verify(self, client_id, secret): + c = self.clients.get(client_id) + return bool(c and c["secret"] == secret) + + def check_totp(self, client_id, code): + c = self.clients.get(client_id) + return TotpResult.OK if (c and c["totp"] == code) else TotpResult.INVALID + + def get_name(self, client_id): + c = self.clients.get(client_id) + return c["name"] if c else None + + +class FakeTokenStore: + def __init__(self): + self._tokens: dict[str, str] = {} + self._n = 0 + + def issue(self, client_id, name=None): + self._n += 1 + token = f"bearer_{self._n}" + self._tokens[token] = client_id + return token, 86400 + + def validate(self, token): + return self._tokens.get(token) + + def revoke(self, token): + self._tokens.pop(token, None) + return True + + def list_named(self, client_id): + return [] + + +@pytest.fixture() +def web(tmp_path, monkeypatch): + monkeypatch.setenv("BEACONMCP_DASHBOARD_DB", str(tmp_path / "dashboard.db")) + db = Database(tmp_path / "dashboard.db") + passkey_store = PasskeyStore(db) + deps = DashboardDeps( + database=db, + session_store=SessionStore(db, key=os.urandom(32)), + client_store=FakeClientStore(), + token_store=FakeTokenStore(), + totp_locked=lambda cid: False, + totp_record_failure=lambda cid: None, + totp_record_success=lambda cid: None, + passkeys=PasskeyService(passkey_store), + ) + app = Starlette(routes=build_dashboard_routes(deps)) + client = TestClient( + app, follow_redirects=False, base_url=ORIGIN, + headers={"x-forwarded-proto": "https"}, + ) + return client, deps, passkey_store + + +def _csrf(client) -> str: + client.get("/app/login") + return client.cookies.get(CSRF_COOKIE) + + +def _json_headers(token: str) -> dict[str, str]: + return {"X-CSRF-Token": token, "X-BeaconMCP-Mode": "json"} + + +def _login(client) -> str: + """Sign in with TOTP and return the *rotated* CSRF token.""" + token = _csrf(client) + res = client.post( + "/app/login", + data={ + "csrf_token": token, "client_id": "beaconmcp_test", + "client_secret": "sk_test", "totp": "123456", + }, + headers=_json_headers(token), + ) + assert res.status_code == 200, res.text + return res.json()["csrf_token"] + + +def test_login_json_mode_returns_expiry_and_passkey_state(web): + client, _, _ = web + token = _csrf(client) + res = client.post( + "/app/login", + data={ + "csrf_token": token, "client_id": "beaconmcp_test", + "client_secret": "sk_test", "totp": "123456", + }, + headers=_json_headers(token), + ) + assert res.status_code == 200 + body = res.json() + assert body["ok"] is True + assert body["client_name"] == "Test Client" + assert body["bearer_expires_at"] > time.time() + assert body["session_expires_at"] > body["bearer_expires_at"] + assert body["passkeys_enabled"] is True + assert body["passkey_count"] == 0 + assert client.cookies.get(SESSION_COOKIE) + + +def test_login_form_mode_still_redirects(web): + client, _, _ = web + token = _csrf(client) + res = client.post( + "/app/login", + data={ + "csrf_token": token, "client_id": "beaconmcp_test", + "client_secret": "sk_test", "totp": "123456", + }, + ) + assert res.status_code == 303 + assert res.headers["location"] == "/app/tokens" + + +def test_login_json_mode_reports_errors_as_json(web): + client, _, _ = web + token = _csrf(client) + res = client.post( + "/app/login", + data={ + "csrf_token": token, "client_id": "beaconmcp_test", + "client_secret": "sk_test", "totp": "000000", + }, + headers=_json_headers(token), + ) + assert res.status_code == 401 + assert res.json()["ok"] is False + assert "2FA" in res.json()["error"] + + +def _register_passkey(client, token) -> SoftwareAuthenticator: + auth = SoftwareAuthenticator() + res = client.post( + "/app/api/passkeys/register/options", json={}, headers=_json_headers(token), + ) + assert res.status_code == 200, res.text + body = res.json() + res = client.post( + "/app/api/passkeys/register/verify", + json={"state": body["state"], "credential": auth.create(body["options"])}, + headers=_json_headers(token), + ) + assert res.status_code == 201, res.text + return auth + + +def test_register_then_sign_in_with_passkey(web): + client, _, passkey_store = web + token = _login(client) + auth = _register_passkey(client, token) + assert passkey_store.count_for_client("beaconmcp_test") == 1 + + listed = client.get("/app/api/passkeys", headers=_json_headers(token)).json() + assert len(listed["passkeys"]) == 1 + + # Fresh browser: credentials + passkey, no TOTP anywhere. + fresh = TestClient( + client.app, follow_redirects=False, base_url=ORIGIN, + headers={"x-forwarded-proto": "https"}, + ) + token2 = _csrf(fresh) + res = fresh.post( + "/app/api/passkeys/auth/options", + json={"client_id": "beaconmcp_test", "client_secret": "sk_test"}, + headers=_json_headers(token2), + ) + assert res.status_code == 200, res.text + body = res.json() + res = fresh.post( + "/app/api/passkeys/auth/verify", + json={ + "client_id": "beaconmcp_test", + "client_secret": "sk_test", + "state": body["state"], + "credential": auth.get(body["options"]), + }, + headers=_json_headers(token2), + ) + assert res.status_code == 200, res.text + assert res.json()["ok"] is True + assert res.json()["passkey_count"] == 1 + assert fresh.cookies.get(SESSION_COOKIE) + + +def test_passkey_auth_requires_valid_client_secret(web): + client, _, _ = web + token = _csrf(client) + res = client.post( + "/app/api/passkeys/auth/options", + json={"client_id": "beaconmcp_test", "client_secret": "wrong"}, + headers=_json_headers(token), + ) + assert res.status_code == 401 + assert res.json()["error"] == "Invalid credentials." + + +def test_passkey_auth_verify_rechecks_the_secret(web): + """The state token alone must never be enough to mint a session.""" + client, _, _ = web + token = _login(client) + auth = _register_passkey(client, token) + body = client.post( + "/app/api/passkeys/auth/options", + json={"client_id": "beaconmcp_test", "client_secret": "sk_test"}, + headers=_json_headers(token), + ).json() + res = client.post( + "/app/api/passkeys/auth/verify", + json={ + "client_id": "beaconmcp_test", + "client_secret": "wrong", + "state": body["state"], + "credential": auth.get(body["options"]), + }, + headers=_json_headers(token), + ) + assert res.status_code == 401 + + +def test_passkey_endpoints_require_csrf(web): + client, _, _ = web + _csrf(client) + res = client.post( + "/app/api/passkeys/auth/options", + json={"client_id": "beaconmcp_test", "client_secret": "sk_test"}, + ) + assert res.status_code == 403 + + +def test_passkey_registration_requires_a_session(web): + client, _, _ = web + token = _csrf(client) + res = client.post( + "/app/api/passkeys/register/options", json={}, headers=_json_headers(token), + ) + assert res.status_code == 401 + + +def test_passkey_delete_scoped_to_session_client(web): + client, _, passkey_store = web + token = _login(client) + _register_passkey(client, token) + credential_id = passkey_store.list_for_client("beaconmcp_test")[0].credential_id + res = client.post( + "/app/api/passkeys/delete", + json={"credential_id": credential_id}, + headers=_json_headers(token), + ) + assert res.status_code == 200 + assert passkey_store.count_for_client("beaconmcp_test") == 0 + + +def test_tokens_page_lists_and_removes_passkeys(web): + client, _, passkey_store = web + token = _login(client) + _register_passkey(client, token) + + page = client.get("/app/tokens").text + assert "Passkeys" in page + credential_id = passkey_store.list_for_client("beaconmcp_test")[0].credential_id + assert credential_id in page + assert "never used" in page + + res = client.post( + "/app/passkeys/remove", + data={"csrf_token": token, "credential_id": credential_id}, + ) + assert res.status_code == 303 + assert res.headers["location"] == "/app/tokens" + assert passkey_store.count_for_client("beaconmcp_test") == 0 + + +def test_passkeys_remove_requires_csrf(web): + client, _, passkey_store = web + token = _login(client) + _register_passkey(client, token) + credential_id = passkey_store.list_for_client("beaconmcp_test")[0].credential_id + res = client.post( + "/app/passkeys/remove", + data={"csrf_token": "wrong", "credential_id": credential_id}, + ) + assert res.status_code == 403 + assert passkey_store.count_for_client("beaconmcp_test") == 1 + + +def test_routes_report_unavailable_without_a_service(tmp_path, monkeypatch): + monkeypatch.setenv("BEACONMCP_DASHBOARD_DB", str(tmp_path / "d.db")) + db = Database(tmp_path / "d.db") + deps = DashboardDeps( + database=db, + session_store=SessionStore(db, key=os.urandom(32)), + client_store=FakeClientStore(), + token_store=FakeTokenStore(), + totp_locked=lambda cid: False, + totp_record_failure=lambda cid: None, + totp_record_success=lambda cid: None, + passkeys=None, + ) + client = TestClient( + Starlette(routes=build_dashboard_routes(deps)), + follow_redirects=False, base_url=ORIGIN, + ) + client.get("/app/login") + token = client.cookies.get(CSRF_COOKIE) + res = client.post( + "/app/api/passkeys/auth/options", + json={"client_id": "beaconmcp_test", "client_secret": "sk_test"}, + headers=_json_headers(token), + ) + assert res.status_code == 503 + + +def test_login_page_hides_passkeys_when_unavailable(tmp_path, monkeypatch): + monkeypatch.setenv("BEACONMCP_DASHBOARD_DB", str(tmp_path / "d2.db")) + db = Database(tmp_path / "d2.db") + deps = DashboardDeps( + database=db, + session_store=SessionStore(db, key=os.urandom(32)), + client_store=FakeClientStore(), + token_store=FakeTokenStore(), + totp_locked=lambda cid: False, + totp_record_failure=lambda cid: None, + totp_record_success=lambda cid: None, + passkeys=None, + ) + client = TestClient( + Starlette(routes=build_dashboard_routes(deps)), base_url=ORIGIN, + ) + body = client.get("/app/login").text + assert 'data-passkeys-enabled="false"' in body + assert "passkey-login-btn" not in body diff --git a/tests/test_proxmox_client.py b/tests/test_proxmox_client.py new file mode 100644 index 0000000..f8bf122 --- /dev/null +++ b/tests/test_proxmox_client.py @@ -0,0 +1,182 @@ +"""Tests for :class:`ProxmoxClient`'s connection cache. + +A ProxmoxAPI wraps a requests.Session, which is not thread-safe. Sync Proxmox +tools now run on a worker-thread pool (see ``server._metric_tool``), so the +cache is per-thread: every thread gets its own connection object, and no two +threads are ever handed the same one. +""" + +from __future__ import annotations + +import sys +import threading +from pathlib import Path + +import pytest +from requests.exceptions import ConnectionError as RequestsConnectionError + +sys.path.insert(0, str(Path(__file__).parent.parent / "src")) + +import beaconmcp.proxmox.client as client_mod +from beaconmcp.proxmox.client import NodeNotFoundError, ProxmoxClient + + +class _FakeNode: + def __init__(self, name: str) -> None: + self.name = name + self.host = f"{name}.example.com" + self.token_id = "root@pam!mytoken" + self.token_secret = "secret" + + +class _FakeConfig: + """Minimal stand-in exposing only what ProxmoxClient touches.""" + + def __init__(self, node_names: list[str]) -> None: + self.pve_nodes = [_FakeNode(n) for n in node_names] + self.verify_ssl = False + + def get_node(self, name: str) -> _FakeNode | None: + for n in self.pve_nodes: + if n.name == name: + return n + return None + + +@pytest.fixture() +def fake_proxmox_api(monkeypatch: pytest.MonkeyPatch) -> list[object]: + """Replace ``ProxmoxAPI`` with a counting dummy. + + Returns the list of constructed dummies so a test can assert how many + connections were built. + """ + built: list[object] = [] + build_lock = threading.Lock() + + class _DummyAPI: + def __init__(self, host: str, **kwargs: object) -> None: + self.host = host + self.kwargs = kwargs + with build_lock: + built.append(self) + + monkeypatch.setattr(client_mod, "ProxmoxAPI", _DummyAPI) + return built + + +def test_get_connection_caches_within_a_thread(fake_proxmox_api: list[object]) -> None: + c = ProxmoxClient(_FakeConfig(["pve1"])) + first = c._get_connection("pve1") + second = c._get_connection("pve1") + assert first is second + assert len(fake_proxmox_api) == 1 + + +def test_get_connection_unknown_node_raises(fake_proxmox_api: list[object]) -> None: + c = ProxmoxClient(_FakeConfig(["pve1"])) + with pytest.raises(NodeNotFoundError): + c._get_connection("nope") + + +def test_each_thread_gets_its_own_connection(fake_proxmox_api: list[object]) -> None: + """The point of the thread-local cache: no two threads share a ProxmoxAPI, + and therefore never share a requests.Session.""" + c = ProxmoxClient(_FakeConfig(["pve1"])) + + n_threads = 20 + barrier = threading.Barrier(n_threads) + results: list[object] = [] + errors: list[Exception] = [] + lock = threading.Lock() + + def worker() -> None: + barrier.wait() + try: + conn = c._get_connection("pve1") + # Within one thread the cache must still hold. + assert c._get_connection("pve1") is conn + with lock: + results.append(conn) + except Exception as e: + with lock: + errors.append(e) + + threads = [threading.Thread(target=worker) for _ in range(n_threads)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert not errors, f"unexpected errors: {errors}" + assert len({id(r) for r in results}) == n_threads + assert len(fake_proxmox_api) == n_threads + + +def test_concurrent_get_connection_many_nodes(fake_proxmox_api: list[object]) -> None: + """Threads hammering several distinct nodes always get the connection for + the node they asked for.""" + node_names = [f"pve{i}" for i in range(8)] + c = ProxmoxClient(_FakeConfig(node_names)) + + barrier = threading.Barrier(len(node_names) * 4) + errors: list[Exception] = [] + lock = threading.Lock() + + def worker(name: str) -> None: + barrier.wait() + try: + for _ in range(5): + assert c._get_connection(name).host == f"{name}.example.com" + except Exception as e: + with lock: + errors.append(e) + + threads = [ + threading.Thread(target=worker, args=(name,)) + for name in node_names + for _ in range(4) + ] + for t in threads: + t.start() + for t in threads: + t.join() + + assert not errors, f"unexpected errors: {errors}" + + +class _DeadEndpoint: + """Stands in for ``conn.nodes``; fails the way a dropped socket does.""" + + def get(self, **_kwargs: object) -> None: + raise RequestsConnectionError("socket died") + + +def test_transient_error_evicts_only_the_calling_thread( + fake_proxmox_api: list[object], +) -> None: + """A retry drops the failing thread's socket without yanking the + connection out from under another thread that may be mid-request on it.""" + c = ProxmoxClient(_FakeConfig(["pve1"])) + primed = threading.Event() + main_done = threading.Event() + seen: list[object] = [] + + def worker() -> None: + seen.append(c._get_connection("pve1")) + primed.set() + main_done.wait(timeout=10) + seen.append(c._get_connection("pve1")) + + t = threading.Thread(target=worker) + t.start() + assert primed.wait(timeout=10) + + c._get_connection("pve1").nodes = _DeadEndpoint() # type: ignore[attr-defined] + assert "error" in c.api_call("pve1", "get", "nodes") + + main_done.set() + t.join(timeout=10) + + # The worker's connection survived the main thread's eviction. + assert len(seen) == 2 + assert seen[1] is seen[0] diff --git a/tests/test_proxmox_helpers.py b/tests/test_proxmox_helpers.py new file mode 100644 index 0000000..12fa947 --- /dev/null +++ b/tests/test_proxmox_helpers.py @@ -0,0 +1,122 @@ +"""Unit tests for proxmox tool helpers: backup-type detection and the +guest-agent file-size pre-flight.""" + +from __future__ import annotations + +import base64 + +from beaconmcp.proxmox.system import _qemu_file_size +from beaconmcp.proxmox.vms import _detect_backup_type + + +class FakeClient: + """Minimal ProxmoxClient stand-in: canned post/get responses.""" + + def __init__(self, post_resp=None, get_resps=None): + self.post_resp = post_resp + self.get_resps = list(get_resps or []) + self.post_kwargs: dict = {} + self.get_paths: list[str] = [] + + def post(self, node, path, **kwargs): + self.post_kwargs = kwargs + return self.post_resp + + def get(self, node, path, **kwargs): + self.get_paths.append(path) + return self.get_resps.pop(0) if self.get_resps else {} + + +# --- _detect_backup_type --------------------------------------------------- + + +def test_backup_type_vzdump_names_no_api_call() -> None: + c = FakeClient() + assert _detect_backup_type( + c, "n", "local:backup/vzdump-qemu-100-2026_01_01-00_00_00.vma.zst" + ) == "qemu" + assert _detect_backup_type( + c, "n", "local:backup/vzdump-lxc-200-2026_01_01-00_00_00.tar.zst" + ) == "lxc" + assert c.get_paths == [] + + +def test_backup_type_pbs_namespaces() -> None: + c = FakeClient() + assert _detect_backup_type(c, "n", "pbs:backup/vm/100/2026-01-01T00:00:00Z") == "qemu" + assert _detect_backup_type(c, "n", "pbs:backup/ct/200/2026-01-01T00:00:00Z") == "lxc" + # A *namespace* literally named "vm" must not shadow the real type segment. + assert _detect_backup_type( + c, "n", "pbs:backup/ns/vm/ct/200/2026-01-01T00:00:00Z" + ) == "lxc" + assert c.get_paths == [] + + +def test_backup_type_content_lookup_fallback() -> None: + c = FakeClient(get_resps=[[{"volid": "st:backup/renamed.bin", "subtype": "lxc"}]]) + assert _detect_backup_type(c, "n", "st:backup/renamed.bin") == "lxc" + assert c.get_paths == ["nodes/n/storage/st/content"] + + +def test_backup_type_content_lookup_format_fallback() -> None: + c = FakeClient(get_resps=[[{"volid": "st:backup/renamed.bin", "format": "vma.zst"}]]) + assert _detect_backup_type(c, "n", "st:backup/renamed.bin") == "qemu" + + +def test_backup_type_tolerates_garbage_content() -> None: + # Non-dict items in the listing must be skipped, not crash. + c = FakeClient(get_resps=[ + ["garbage", {"volid": "st:backup/x.bin", "subtype": "qemu"}], + ]) + assert _detect_backup_type(c, "n", "st:backup/x.bin") == "qemu" + + +def test_backup_type_undeterminable() -> None: + assert _detect_backup_type(FakeClient(get_resps=[{"error": "boom"}]), + "n", "st:backup/renamed.bin") is None + # No storage prefix -> no lookup possible. + assert _detect_backup_type(FakeClient(), "n", "renamed.bin") is None + + +# --- _qemu_file_size ------------------------------------------------------- + + +def _b64(s: str) -> str: + return base64.b64encode(s.encode()).decode() + + +def test_file_size_nominal_base64() -> None: + c = FakeClient( + post_resp={"pid": 7}, + get_resps=[{ + "exited": 1, "exitcode": 0, + "out-data": _b64("2048576\n"), "out-data-encoding": "base64", + }], + ) + assert _qemu_file_size(c, "n", 100, "/var/log/big.log") == 2048576 + # argv form: no shell, hostile paths stay inert. + assert c.post_kwargs["command"] == ["stat", "-c", "%s", "/var/log/big.log"] + + +def test_file_size_plain_output() -> None: + c = FakeClient(post_resp={"pid": 7}, + get_resps=[{"exited": 1, "exitcode": 0, "out-data": "512\n"}]) + assert _qemu_file_size(c, "n", 100, "/x") == 512 + + +def test_file_size_none_when_stat_missing() -> None: + # exitcode 127: no `stat` binary in the guest -> caller falls back. + c = FakeClient(post_resp={"pid": 7}, + get_resps=[{"exited": 1, "exitcode": 127, "err-data": "nope"}]) + assert _qemu_file_size(c, "n", 100, "/x") is None + + +def test_file_size_none_when_agent_unavailable() -> None: + c = FakeClient(post_resp={"error": "agent not running"}) + assert _qemu_file_size(c, "n", 100, "/x") is None + + +def test_file_size_none_on_missing_output() -> None: + c = FakeClient(post_resp={"pid": 7}, + get_resps=[{"exited": 1, "exitcode": 0}]) + assert _qemu_file_size(c, "n", 100, "/x") is None diff --git a/tests/test_ratelimit.py b/tests/test_ratelimit.py new file mode 100644 index 0000000..41daf84 --- /dev/null +++ b/tests/test_ratelimit.py @@ -0,0 +1,211 @@ +"""Tests for the in-memory sliding-window rate limiter.""" + +from __future__ import annotations + +import time + +from beaconmcp.ratelimit import RateLimiter, client_ip, forwarded_host + + +def test_allows_up_to_limit_then_blocks() -> None: + rl = RateLimiter(limit=3, window_seconds=60.0) + assert rl.check("1.2.3.4") is True + assert rl.check("1.2.3.4") is True + assert rl.check("1.2.3.4") is True + assert rl.check("1.2.3.4") is False + + +def test_keys_are_independent() -> None: + rl = RateLimiter(limit=2, window_seconds=60.0) + assert rl.check("a") is True + assert rl.check("a") is True + assert rl.check("a") is False + # Different key: fresh budget. + assert rl.check("b") is True + + +def test_window_expiry_frees_slots() -> None: + rl = RateLimiter(limit=2, window_seconds=0.05) + assert rl.check("k") is True + assert rl.check("k") is True + assert rl.check("k") is False + time.sleep(0.08) + # Old events aged out. + assert rl.check("k") is True + + +def test_retry_after_nonzero_when_blocked() -> None: + rl = RateLimiter(limit=1, window_seconds=60.0) + assert rl.check("x") is True + assert rl.check("x") is False + assert rl.retry_after("x") > 0 + + +def test_gc_reclaims_stale_buckets() -> None: + rl = RateLimiter(limit=1, window_seconds=0.01) + for i in range(1100): + assert rl.check(f"k{i}") is True + time.sleep(0.03) + # Any new check past 1024 buckets triggers stale-bucket collection. + assert rl.check("fresh") is True + assert len(rl._buckets) == 1 + + +def test_client_ip_uses_rightmost_untrusted_hop() -> None: + class _H: + def __init__(self, fwd: str | None) -> None: + self._fwd = fwd + + def get(self, k: str) -> str | None: + if k.lower() == "x-forwarded-for": + return self._fwd + return None + + class _Client: + host = "10.0.0.1" + + class _Req: + def __init__(self, fwd: str | None, *, peer: str = "10.0.0.1") -> None: + self.headers = _H(fwd) + c = _Client() + c.host = peer + self.client = c + + # Trusted direct proxy + spoofed left-most value: + # proxy appends the real client to XFF, so we must not return the spoof. + assert ( + client_ip( + _Req("198.51.100.66, 203.0.113.7"), + trusted_proxies=("10.0.0.1",), + ) + == "203.0.113.7" + ) + + # CIDR rules are accepted for trusted proxies. + assert ( + client_ip( + _Req("203.0.113.9", peer="10.1.2.3"), + trusted_proxies=("10.0.0.0/8",), + ) + == "203.0.113.9" + ) + + # Untrusted direct peer -> ignore XFF entirely. + assert ( + client_ip( + _Req("203.0.113.7", peer="192.0.2.8"), + trusted_proxies=("127.0.0.1",), + ) + == "192.0.2.8" + ) + + # Direct peer with no trust config -> use peer IP. + assert client_ip(_Req(None, peer="203.0.113.10"), trusted_proxies=()) == "203.0.113.10" + + +def test_forwarded_host_only_trusts_declared_proxy() -> None: + class _H: + def __init__(self, host: str | None, xfh: str | None) -> None: + self._host = host + self._xfh = xfh + + def get(self, k: str) -> str | None: + k = k.lower() + if k == "host": + return self._host + if k == "x-forwarded-host": + return self._xfh + return None + + class _Client: + host = "10.0.0.1" + + class _Req: + def __init__( + self, host: str | None, xfh: str | None, *, peer: str = "10.0.0.1", + ) -> None: + self.headers = _H(host, xfh) + c = _Client() + c.host = peer + self.client = c + + # No trusted proxies configured -> X-Forwarded-Host is ignored, even when + # the peer looks internal. The request's own Host header wins. + assert ( + forwarded_host(_Req("real.example", "evil.attacker"), trusted_proxies=()) + == "real.example" + ) + + # Trusted direct proxy -> the forwarded host is believed. + assert ( + forwarded_host( + _Req("internal:8420", "public.example", peer="10.0.0.1"), + trusted_proxies=("10.0.0.1",), + ) + == "public.example" + ) + + # A spoofed X-Forwarded-Host from an UNtrusted peer is dropped; Host wins. + assert ( + forwarded_host( + _Req("real.example", "evil.attacker", peer="192.0.2.8"), + trusted_proxies=("10.0.0.0/8",), + ) + == "real.example" + ) + + # Proxy chain: a client-supplied prefix must not win. A proxy that appends + # its own value puts it last, so the last entry is returned -- symmetric + # with client_ip's right-to-left walk. + assert ( + forwarded_host( + _Req("internal", "evil.attacker, public.example", peer="10.0.0.1"), + trusted_proxies=("10.0.0.1",), + ) + == "public.example" + ) + + # Nothing usable -> default. + assert forwarded_host(_Req(None, None), trusted_proxies=()) == "localhost" + + +def test_forwarded_host_through_a_real_starlette_request() -> None: + # Guards against the hand-built _Req above drifting from runtime: build an + # actual Starlette Request from an ASGI scope and confirm the trusted-proxy + # branch opens. This is the shape the app sees once uvicorn is told not to + # rewrite scope["client"] (proxy_headers=False), so request.client.host is + # the real TCP peer -- the proxy -- not the X-Forwarded-For client. + from starlette.requests import Request + + def _req(host: str, xfh: str | None, peer: str) -> Request: + headers = [(b"host", host.encode())] + if xfh is not None: + headers.append((b"x-forwarded-host", xfh.encode())) + scope = { + "type": "http", + "method": "GET", + "path": "/", + "headers": headers, + "client": (peer, 44444), + "scheme": "http", + "server": ("app", 80), + } + return Request(scope) + + # Proxy peer is trusted -> the forwarded host is believed. + assert ( + forwarded_host( + _req("127.0.0.1:8420", "beacon.example.com", "127.0.0.1"), + trusted_proxies=("127.0.0.1",), + ) + == "beacon.example.com" + ) + + # Direct (untrusted) peer -> forwarded host dropped, own Host wins. + assert ( + forwarded_host( + _req("real.example", "evil.attacker", "203.0.113.5"), + trusted_proxies=("127.0.0.1",), + ) + == "real.example" + ) diff --git a/tests/test_security_hardening.py b/tests/test_security_hardening.py new file mode 100644 index 0000000..cd40a4b --- /dev/null +++ b/tests/test_security_hardening.py @@ -0,0 +1,310 @@ +"""Regression tests for the security-audit hardening pass. + +Each test pins one specific weakness that was found and fixed, so a future +refactor that reopens it fails here rather than in production. +""" + +from __future__ import annotations + +import os +import stat +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).parent.parent / "src")) + +from beaconmcp import audit +from beaconmcp.auth import is_trusted_redirect_uri +from beaconmcp.config import Config, ConfigError +from beaconmcp.dashboard.chat import _tool_call_requires_confirmation +from beaconmcp.dashboard.db import Database +from beaconmcp.proxmox.client import ProxmoxClient + + +# --- redirect_uri: loopback userinfo evasion -------------------------------- + + +@pytest.mark.parametrize( + "uri", + [ + # Real host is attacker.example; the string merely *starts with* + # a trusted loopback prefix. + "http://localhost:1@attacker.example/cb", + "http://127.0.0.1:8080@evil.example/callback", + "http://localhost@evil.example/cb", + "http://[::1]:80@evil.example/cb", + # Plain-HTTP host that is simply not loopback. + "http://evil.example/cb", + # An allowed HTTPS origin must not be reachable over plaintext http. + "http://assistant.ai/cb", + ], +) +def test_loopback_prefix_evasion_rejected(uri: str) -> None: + assert not is_trusted_redirect_uri(uri, ["https://assistant.ai"]), uri + + +@pytest.mark.parametrize( + "uri", + [ + "http://localhost:54321/callback", + "http://localhost/callback", + "http://127.0.0.1:3000/oauth/cb", + "http://[::1]:8080/cb", + "vscode://ms-vscode.remote/callback", + ], +) +def test_genuine_loopback_and_scheme_callbacks_still_accepted(uri: str) -> None: + assert is_trusted_redirect_uri(uri), uri + + +# --- Proxmox API path-segment injection ------------------------------------- + + +class _StubConfig: + pve_nodes: list = [] + verify_ssl = False + + def get_node(self, name): # pragma: no cover - never reached + raise AssertionError("connection must not be attempted") + + +@pytest.mark.parametrize( + "path", + [ + # A snapname of "../../../access/users" re-targets the request. + "nodes/pve1/qemu/100/snapshot/../../../access/users/rollback", + "nodes/pve1/storage/../../access/domains/content", + "nodes/pve1/qemu/100/snapshot//rollback", + "nodes/pve1/qemu/100/snapshot/_store/rollback", + ], +) +def test_traversal_path_segments_rejected_before_any_request(path: str) -> None: + client = ProxmoxClient(_StubConfig()) # type: ignore[arg-type] + result = client.get("pve1", path) + assert isinstance(result, dict) + assert "illegal Proxmox API path segment" in result["error"] + + +def test_legitimate_paths_are_not_rejected() -> None: + """The guard must not fire on the shapes real tools build.""" + from beaconmcp.proxmox.client import _split_api_path + + for path in ( + "nodes", + "version", + "nodes/pve-1/qemu/100/status/current", + "nodes/pve1/lxc/200/snapshot/pre-upgrade_2024/rollback", + "nodes/pve1/storage/local-lvm/content", + "nodes/pve1/storage/pbs.backup/content", + "nodes/pve1/qemu/100/agent/exec-status", + "nodes/pve1/vzdump", + ): + assert _split_api_path(path) + + +# --- chat: dangerous-tool confirmation gate --------------------------------- + + +@pytest.mark.parametrize( + "name,args", + [ + # Writing a guest file is code execution in one hop + # (~/.ssh/authorized_keys, /etc/cron.d/...). + ("proxmox_write_file", {"node": "pve1", "vmid": 100, "path": "/x", "content": "y"}), + ("proxmox_upload_file", {"source": "a", "dest": "/b"}), + ("proxmox_run", {"node": "pve1", "vmid": 100, "command": "id"}), + ("ssh_run", {"host": "pve1", "command": "id"}), + # Destructive / irreversible. + ("vm_bulk_action", {"vmids": [1, 2], "action": "stop"}), + ("proxmox_vm_stop", {"node": "pve1", "vmid": 100}), + ("proxmox_snapshot_rollback", {"node": "pve1", "vmid": 100, "snapname": "s"}), + ("proxmox_snapshot_delete", {"node": "pve1", "vmid": 100, "snapname": "s"}), + ("proxmox_backup_restore", {"node": "pve1", "vmid": 100, "archive": "a"}), + ("bmc_power_off", {"device_id": "rack1"}), + ("bmc_power_reset", {"device_id": "rack1"}), + # Read-or-write tool, in its writing shape. + ("proxmox_vm_config", {"node": "pve1", "vmid": 100, "updates": {"memory": 1}}), + ], +) +def test_dangerous_tools_require_confirmation(name: str, args: dict) -> None: + assert _tool_call_requires_confirmation(name, args), name + + +@pytest.mark.parametrize( + "name,args", + [ + # Pure reads. + ("proxmox_list_vms", {}), + ("proxmox_read_file", {"node": "pve1", "vmid": 100, "path": "/etc/hosts"}), + ("cluster_overview", {}), + ("bmc_power_status", {"device_id": "rack1"}), + # Read shape of the read-or-write tool. + ("proxmox_vm_config", {"node": "pve1", "vmid": 100}), + ("proxmox_vm_config", {"node": "pve1", "vmid": 100, "updates": None}), + # Polling an already-approved exec session is read-only. + ("proxmox_run", {"exec_id": "abc123"}), + ("ssh_run", {"exec_id": "abc123"}), + # dry_run tools only describe what they would do. + ("proxmox_snapshot_delete", {"node": "pve1", "vmid": 1, "snapname": "s", "dry_run": True}), + ], +) +def test_read_only_calls_do_not_require_confirmation(name: str, args: dict) -> None: + assert not _tool_call_requires_confirmation(name, args), name + + +@pytest.mark.parametrize( + "name,args", + [ + ("ssh_run", {"host": "pve1", "command": "rm -rf /", "dry_run": True}), + ("proxmox_run", {"node": "pve1", "vmid": 100, "command": "id", "dry_run": True}), + ("proxmox_write_file", { + "node": "pve1", "vmid": 100, + "path": "/root/.ssh/authorized_keys", "content": "ssh-rsa ...", + "dry_run": True, + }), + ("vm_bulk_action", {"vmids": [1, 2], "action": "stop", "dry_run": True}), + ("bmc_power_off", {"device_id": "rack1", "dry_run": True}), + ("proxmox_backup_restore", { + "node": "pve1", "vmid": 100, "archive": "a", "dry_run": True, + }), + ], +) +def test_dry_run_does_not_bypass_gate_on_tools_lacking_it(name: str, args: dict) -> None: + """A stray ``dry_run`` must not talk a tool past the modal. + + None of these tools declare ``dry_run``, and FastMCP validates + arguments with a plain pydantic model, so the extra key is dropped + server-side and the call runs for real. Trusting the argument would + make the gate bypassable by adding one word to an injected + instruction. + """ + assert _tool_call_requires_confirmation(name, args), name + + +def test_dry_run_aware_tools_actually_declare_dry_run() -> None: + """Every entry in the allow-list must really implement ``dry_run``. + + Guards the list against drifting as tools are renamed or reworked: + an entry whose tool has lost the parameter would silently become a + bypass again. + """ + import inspect + + from beaconmcp.dashboard.chat import _DRY_RUN_AWARE + from beaconmcp.proxmox.vms import register_vm_tools + + registered: dict[str, object] = {} + + class _Recorder: + def tool(self, *_a: object, **_kw: object): + def deco(fn): + registered[fn.__name__] = fn + return fn + return deco + + register_vm_tools(_Recorder(), object()) # type: ignore[arg-type] + + assert _DRY_RUN_AWARE <= registered.keys() + for name in _DRY_RUN_AWARE: + assert "dry_run" in inspect.signature(registered[name]).parameters, name + + +# --- BMC verify_tls actually reaches the backend ---------------------------- + + +def _bmc_yaml(tmp_path: Path, verify_tls: str) -> Path: + path = tmp_path / "beaconmcp.yaml" + path.write_text( + "version: 1\n" + "bmc:\n" + " devices:\n" + " - id: rack1\n" + " type: redfish\n" + " host: 10.0.0.5\n" + " user: root\n" + " password: pw\n" + f" verify_tls: {verify_tls}\n" + ) + return path + + +def test_verify_tls_is_parsed_from_yaml(tmp_path: Path) -> None: + cfg = Config.load(config_path=_bmc_yaml(tmp_path, "true")) + assert cfg.bmc_devices[0].verify_tls is True + + cfg = Config.load(config_path=_bmc_yaml(tmp_path, "false")) + assert cfg.bmc_devices[0].verify_tls is False + + +def test_verify_tls_reaches_the_redfish_backend(tmp_path: Path) -> None: + from beaconmcp.bmc import build_registry + + cfg = Config.load(config_path=_bmc_yaml(tmp_path, "true")) + backend = build_registry(cfg)["rack1"] + assert backend._verify is True # type: ignore[attr-defined] + + +def test_verify_tls_defaults_to_false_when_absent(tmp_path: Path) -> None: + path = tmp_path / "beaconmcp.yaml" + path.write_text( + "version: 1\n" + "bmc:\n" + " devices:\n" + " - id: rack1\n" + " type: redfish\n" + " host: 10.0.0.5\n" + " user: root\n" + " password: pw\n" + ) + cfg = Config.load(config_path=path) + assert cfg.bmc_devices[0].verify_tls is False + + +# --- dashboard.db must not be world-readable -------------------------------- + + +@pytest.mark.skipif( + sys.platform == "win32", + reason="Windows os.chmod only toggles the read-only bit; POSIX modes are meaningless there", +) +def test_dashboard_db_is_owner_only(tmp_path: Path) -> None: + db_file = tmp_path / "dashboard.db" + Database(db_file) + mode = stat.S_IMODE(os.stat(db_file).st_mode) + assert mode & 0o077 == 0, f"dashboard.db is group/world accessible: {mode:o}" + for sidecar in ( + db_file.with_name(db_file.name + "-wal"), + db_file.with_name(db_file.name + "-shm"), + ): + if sidecar.exists(): + side_mode = stat.S_IMODE(os.stat(sidecar).st_mode) + assert side_mode & 0o077 == 0, f"{sidecar.name}: {side_mode:o}" + + +# --- audit redaction -------------------------------------------------------- + + +def test_audit_redacts_session_and_oauth_material() -> None: + redacted = audit._redact( + { + "session_id": "s3cr3t-session", + "code_verifier": "verifier", + "totp_secret": "JBSWY3DPEHPK3PXP", + "nested": {"access_token": "tok", "client_id": "beaconmcp_1"}, + "client_id": "beaconmcp_1", + } + ) + assert redacted["session_id"] == "***" + assert redacted["code_verifier"] == "***" + assert redacted["totp_secret"] == "***" + assert redacted["nested"]["access_token"] == "***" + # Non-secret identifiers stay readable -- the log is useless otherwise. + assert redacted["client_id"] == "beaconmcp_1" + assert redacted["nested"]["client_id"] == "beaconmcp_1" + + +def test_config_error_is_importable() -> None: + """Guard against the ConfigError import above going stale.""" + assert issubclass(ConfigError, Exception) diff --git a/tests/test_snapshots.py b/tests/test_snapshots.py new file mode 100644 index 0000000..53bb787 --- /dev/null +++ b/tests/test_snapshots.py @@ -0,0 +1,45 @@ +from unittest.mock import MagicMock +from beaconmcp.proxmox.vms import register_vm_tools + +def test_snapshot_tools(): + mcp_mock = MagicMock() + # Capture the registered tools + tools = {} + + def tool_decorator(*args, **kwargs): + def wrapper(func): + tools[func.__name__] = func + return func + return wrapper + + mcp_mock.tool = tool_decorator + + client_mock = MagicMock() + # Mock _detect_vm_type behavior by returning a dummy status + client_mock.get.side_effect = lambda node, path: {"status": "running"} if "status/current" in path else [{"name": "snap1"}] + + # Mock post/delete for snapshot actions + client_mock.post.return_value = "UPID:node:123:snapshot_create" + client_mock.delete.return_value = "UPID:node:123:snapshot_delete" + + register_vm_tools(mcp_mock, client_mock) + + # Test proxmox_snapshot_list + assert "proxmox_snapshot_list" in tools + res = tools["proxmox_snapshot_list"]("pve1", 100) + assert "snapshots" in res + + # Test proxmox_snapshot_create with vmstate on running QEMU + assert "proxmox_snapshot_create" in tools + res = tools["proxmox_snapshot_create"]("pve1", 100, "snap1", "Test", True) + assert res.get("action") == "snapshot_create" + + # Test proxmox_snapshot_rollback + assert "proxmox_snapshot_rollback" in tools + res = tools["proxmox_snapshot_rollback"]("pve1", 100, "snap1") + assert res.get("action") == "snapshot_rollback" + + # Test proxmox_snapshot_delete + assert "proxmox_snapshot_delete" in tools + res = tools["proxmox_snapshot_delete"]("pve1", 100, "snap1") + assert res.get("action") == "snapshot_delete" diff --git a/tests/test_token_persistence.py b/tests/test_token_persistence.py new file mode 100644 index 0000000..25aa635 --- /dev/null +++ b/tests/test_token_persistence.py @@ -0,0 +1,121 @@ +"""Unit tests for TokenStore named-token persistence (SQLite).""" + +from __future__ import annotations + +import stat +from pathlib import Path + +import pytest + +from beaconmcp.auth import TokenCapExceeded, TokenStore + + +def test_named_token_survives_restart(tmp_path: Path) -> None: + db = tmp_path / "tokens.db" + store = TokenStore(db_path=db) + token, _ = store.issue("client-a", name="laptop") + + reborn = TokenStore(db_path=db) + assert reborn.validate(token) == "client-a" + + +def test_internal_bearer_does_not_persist(tmp_path: Path) -> None: + db = tmp_path / "tokens.db" + store = TokenStore(db_path=db) + session, _ = store.issue("client-a") # unnamed: dashboard session bearer + + reborn = TokenStore(db_path=db) + assert reborn.validate(session) is None + + +def test_db_file_is_owner_only(tmp_path: Path) -> None: + db = tmp_path / "tokens.db" + TokenStore(db_path=db) + assert stat.S_IMODE(db.stat().st_mode) == 0o600 + + +def test_revoked_token_not_resurrected_by_restart(tmp_path: Path) -> None: + db = tmp_path / "tokens.db" + store = TokenStore(db_path=db) + token, _ = store.issue("client-a", name="laptop") + store.revoke(token) + # Still valid in-process during the grace window... + assert store.validate(token) == "client-a" + # ...but a restart inside that window must not bring it back. + reborn = TokenStore(db_path=db) + assert reborn.validate(token) is None + + +def test_expired_rows_dropped_on_load(tmp_path: Path) -> None: + db = tmp_path / "tokens.db" + store = TokenStore(db_path=db) + token, _ = store.issue("client-a", name="laptop") + # Force expiry, then simulate the process being down past the TTL. + store._tokens[token].expires_at = 1.0 + store._persist(store._tokens[token]) + + reborn = TokenStore(db_path=db) + assert reborn.validate(token) is None + assert reborn._db is not None + rows = reborn._db.execute("SELECT COUNT(*) FROM named_tokens").fetchone() + assert rows[0] == 0 + + +def test_named_cap_enforced_after_reload(tmp_path: Path) -> None: + db = tmp_path / "tokens.db" + store = TokenStore(db_path=db) + for i in range(TokenStore.NAMED_TOKEN_CAP): + store.issue("client-b", name=f"t{i}") + + reborn = TokenStore(db_path=db) + with pytest.raises(TokenCapExceeded): + reborn.issue("client-b", name="one-too-many") + + +def test_named_token_uses_configured_ttl(tmp_path: Path) -> None: + db = tmp_path / "tokens.db" + store = TokenStore(db_path=db, named_token_ttl=7200) + _, expires_in = store.issue("client-a", name="laptop") + assert expires_in == 7200 + + +def test_internal_bearer_keeps_24h_ttl_regardless(tmp_path: Path) -> None: + db = tmp_path / "tokens.db" + store = TokenStore(db_path=db, named_token_ttl=7200) + _, expires_in = store.issue("client-a") # unnamed: session bearer + assert expires_in == TokenStore.TOKEN_TTL + + +def test_named_ttl_defaults_to_30_days(tmp_path: Path) -> None: + store = TokenStore(db_path=tmp_path / "tokens.db") + _, expires_in = store.issue("client-a", name="laptop") + assert expires_in == TokenStore.NAMED_TOKEN_TTL == 3600 * 24 * 30 + + +def test_named_ttl_zero_means_no_expiry(tmp_path: Path) -> None: + db = tmp_path / "tokens.db" + store = TokenStore(db_path=db, named_token_ttl=0) + token, expires_in = store.issue("client-a", name="forever") + assert expires_in == 0 + assert store.validate(token) == "client-a" + # Internal session bearers are unaffected by the named TTL. + _, session_ttl = store.issue("client-a") + assert session_ttl == TokenStore.TOKEN_TTL + + # The infinite expiry must survive the SQLite round-trip. + reborn = TokenStore(db_path=db, named_token_ttl=0) + assert reborn.validate(token) == "client-a" + # And revocation must still kill it. + assert reborn.revoke(token) is True + final = TokenStore(db_path=db, named_token_ttl=0) + assert final.validate(token) is None + + +def test_unwritable_db_degrades_to_memory_only(tmp_path: Path) -> None: + target = tmp_path / "blocked" + target.write_text("not a directory") + store = TokenStore(db_path=target / "tokens.db") + assert store._db is None + # The store keeps working in-memory. + token, _ = store.issue("client-c", name="laptop") + assert store.validate(token) == "client-c" diff --git a/tests/test_token_store.py b/tests/test_token_store.py new file mode 100644 index 0000000..096e5a4 --- /dev/null +++ b/tests/test_token_store.py @@ -0,0 +1,134 @@ +"""Tests for :class:`TokenStore` behavior and thread-safety. + +The single-threaded contract (issue / validate / revoke / named-token cap) +must be unchanged by the locking work, and a concurrent issue+validate +storm must not raise ``RuntimeError: dictionary changed size during +iteration`` -- which is exactly what the unlocked version did once tools +started running on a worker-thread pool. +""" + +from __future__ import annotations + +import sys +import threading +import time +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).parent.parent / "src")) + +from beaconmcp.auth import TokenCapExceeded, TokenStore + + +def test_issue_then_validate_returns_client_id() -> None: + store = TokenStore() + token, ttl = store.issue("client_a") + assert ttl == TokenStore.TOKEN_TTL + assert store.validate(token) == "client_a" + + +def test_validate_unknown_token_is_none() -> None: + assert TokenStore().validate("nope") is None + + +def test_expired_token_is_rejected_and_dropped() -> None: + store = TokenStore() + token, _ = store.issue("client_a") + # Force expiry in the past. + store._tokens[token].expires_at = time.time() - 1 + assert store.validate(token) is None + assert token not in store._tokens + + +def test_revoke_applies_grace_then_rejects() -> None: + store = TokenStore() + token, _ = store.issue("client_a") + assert store.revoke(token) is True + # Within the grace window it is still valid. + assert store.validate(token) == "client_a" + # After the grace deadline it is rejected. + store._tokens[token].expires_at = time.time() - 1 + assert store.validate(token) is None + + +def test_revoke_unknown_token_returns_false() -> None: + assert TokenStore().revoke("nope") is False + + +def test_named_token_cap_enforced() -> None: + store = TokenStore() + for i in range(TokenStore.NAMED_TOKEN_CAP): + store.issue("client_a", name=f"tok{i}") + assert store.count_named("client_a") == TokenStore.NAMED_TOKEN_CAP + with pytest.raises(TokenCapExceeded): + store.issue("client_a", name="one-too-many") + # Unnamed (dashboard-session) tokens are not capped. + store.issue("client_a") # must not raise + + +def test_revoke_named_by_prefix() -> None: + store = TokenStore() + token, _ = store.issue("client_a", name="laptop") + assert store.revoke_named(token[:12], "client_a") is True + # Revoked named token starts its grace countdown. + store._tokens[token].expires_at = time.time() - 1 + assert store.validate(token) is None + + +def test_revoke_named_wrong_client_rejected() -> None: + store = TokenStore() + token, _ = store.issue("client_a", name="laptop") + assert store.revoke_named(token[:12], "client_b") is False + + +def test_concurrent_issue_and_validate_no_corruption() -> None: + """Issue + validate + list from many threads at once. + + Without locking the in-memory dict, the iterating readers (count_named, + list_named, validate's _cleanup) would intermittently raise + "dictionary changed size during iteration". This must complete cleanly. + """ + store = TokenStore() + errors: list[Exception] = [] + lock = threading.Lock() + barrier = threading.Barrier(30) + issued: list[str] = [] + issued_lock = threading.Lock() + + def issuer(i: int) -> None: + barrier.wait() + try: + for j in range(50): + token, _ = store.issue(f"client_{i % 5}") + with issued_lock: + issued.append(token) + except Exception as e: # pragma: no cover - failure path + with lock: + errors.append(e) + + def reader(i: int) -> None: + barrier.wait() + try: + for _ in range(50): + store.count_named(f"client_{i % 5}") + store.list_named(f"client_{i % 5}") + with issued_lock: + sample = issued[-1] if issued else None + if sample: + store.validate(sample) + except Exception as e: # pragma: no cover - failure path + with lock: + errors.append(e) + + threads = [threading.Thread(target=issuer, args=(i,)) for i in range(15)] + threads += [threading.Thread(target=reader, args=(i,)) for i in range(15)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert not errors, f"unexpected errors: {errors[:3]}" + # Every issued token is retrievable. + assert len(issued) == 15 * 50 + assert store.validate(issued[0]) is not None diff --git a/tests/test_totp_replay.py b/tests/test_totp_replay.py new file mode 100644 index 0000000..ac6b80f --- /dev/null +++ b/tests/test_totp_replay.py @@ -0,0 +1,116 @@ +"""TOTP replay-protection tests for :meth:`ClientStore.check_totp`. + +A 6-digit TOTP code is valid for its whole 30s step (plus drift), so without +bookkeeping the same code could be redeemed twice. ``check_totp`` records +the last accepted timestep per SEED OWNER and rejects any non-newer code. +For delegated (DCR) clients the key is the owner, so two derived clients +cannot each spend the same code. + +A replay reports ``REPLAY`` rather than ``INVALID`` so callers can avoid +charging it to the 5-strike lockout -- see ``test_dashboard_totp_replay.py``. +""" + +from __future__ import annotations + +import sys +import time +from pathlib import Path + +import pyotp +import pytest + +sys.path.insert(0, str(Path(__file__).parent.parent / "src")) + +from beaconmcp.auth import ClientStore, TotpResult + + +@pytest.fixture() +def clients(tmp_path: Path) -> ClientStore: + return ClientStore(tmp_path / "clients.json") + + +def test_first_use_accepts_then_replay_rejected(clients: ClientStore) -> None: + client_id, _, seed = clients.create("human") + code = pyotp.TOTP(seed).now() + assert clients.verify_totp(client_id, code) is True + # Same code, same step -> replay, must be rejected. + assert clients.verify_totp(client_id, code) is False + # And again. + assert clients.verify_totp(client_id, code) is False + + +def test_fresh_code_in_later_step_accepted(clients: ClientStore) -> None: + client_id, _, seed = clients.create("human") + totp = pyotp.TOTP(seed) + now = time.time() + code_now = totp.at(now) + code_next = totp.at(now + totp.interval) + + assert clients.verify_totp(client_id, code_now) is True + # A code from the following 30s step is strictly newer -> accepted, + # even though the current one was just spent. + assert clients.verify_totp(client_id, code_next) is True + # Replaying either of the now-consumed codes still fails. + assert clients.verify_totp(client_id, code_now) is False + assert clients.verify_totp(client_id, code_next) is False + + +def test_wrong_code_does_not_advance_step(clients: ClientStore) -> None: + """A rejected (wrong) code must not poison the replay counter.""" + client_id, _, seed = clients.create("human") + assert clients.verify_totp(client_id, "000000") is False + # A subsequent valid current code is still accepted. + code = pyotp.TOTP(seed).now() + assert clients.verify_totp(client_id, code) is True + + +def test_replay_is_per_owner_not_per_client(clients: ClientStore) -> None: + """Two derived clients delegating to one owner share the replay window: + a code spent through one cannot be replayed through the other.""" + owner_id, _, owner_seed = clients.create("human") + d1, _ = clients.create_dynamic( + owner_client_id=owner_id, name="d1", registration_source="chatgpt:s1", + ) + d2, _ = clients.create_dynamic( + owner_client_id=owner_id, name="d2", registration_source="chatgpt:s2", + ) + code = pyotp.TOTP(owner_seed).now() + # First derived client spends the code. + assert clients.verify_totp(d1, code) is True + # Second derived client cannot replay the SAME owner code. + assert clients.verify_totp(d2, code) is False + # The owner itself cannot replay it either. + assert clients.verify_totp(owner_id, code) is False + + +def test_distinct_owners_have_independent_windows( + clients: ClientStore, +) -> None: + """Replay state is keyed by owner: burning A's window leaves B's intact.""" + a_id, _, a_seed = clients.create("owner_a") + b_id, _, b_seed = clients.create("owner_b") + now = time.time() + + # A spends the current step; B has not been touched yet. + assert clients.verify_totp(a_id, pyotp.TOTP(a_seed).at(now)) is True + assert clients.verify_totp(a_id, pyotp.TOTP(a_seed).at(now)) is False + # B's *current* code is still spendable despite A's window being burnt. + assert clients.verify_totp(b_id, pyotp.TOTP(b_seed).at(now)) is True + # And B's next step is unaffected by anything A did. + step = pyotp.TOTP(b_seed).interval + assert clients.verify_totp(b_id, pyotp.TOTP(b_seed).at(now + step)) is True + + +def test_replay_is_reported_distinctly_from_a_wrong_code( + clients: ClientStore, +) -> None: + """The dashboard needs to tell "already used" apart from "wrong", because + only the latter may count towards the lockout.""" + client_id, _, seed = clients.create("human") + code = pyotp.TOTP(seed).now() + + assert clients.check_totp(client_id, code) is TotpResult.OK + assert clients.check_totp(client_id, code) is TotpResult.REPLAY + assert clients.check_totp(client_id, "000000") is TotpResult.INVALID + assert clients.check_totp(client_id, "abc") is TotpResult.INVALID + assert clients.check_totp("no-such-client", code) is TotpResult.INVALID diff --git a/tests/test_trusted_redirect.py b/tests/test_trusted_redirect.py new file mode 100644 index 0000000..ee4eec9 --- /dev/null +++ b/tests/test_trusted_redirect.py @@ -0,0 +1,132 @@ +"""Tests for :func:`beaconmcp.auth.is_trusted_redirect_uri`. + +HTTPS redirect trust is sourced from ``server.allowed_origins`` so operator +config drives both CORS and OAuth callback validation. Desktop/CLI callback +forms (custom URI schemes + loopback) are intentionally built-in. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).parent.parent / "src")) + +from beaconmcp.auth import is_trusted_redirect_uri + + +# --- Accepted redirects ------------------------------------------------------ + + +@pytest.mark.parametrize( + "uri", + [ + # Custom OS URI schemes + "vscode://ms-vscode.remote/callback", + "vscode-insiders://ms-vscode.remote/callback", + "cursor://mcp/callback", + # Loopback (Codex, Gemini CLI, Mistral Vibe, OpenCode, mcp-remote) + "http://localhost:54321/callback", + "http://localhost/callback", + "http://127.0.0.1:3000/oauth/cb", + "http://127.0.0.1/cb", + "http://[::1]:8080/cb", + ], +) +def test_builtin_non_origin_redirects_accepted(uri: str) -> None: + assert is_trusted_redirect_uri(uri), f"expected trusted: {uri}" + + +@pytest.mark.parametrize( + "uri", + [ + "https://claude.ai/api/mcp/auth_callback", + "https://assistant.ai/api/organizations/xyz/mcp/callback", + "https://chatgpt.com/connector_platform_oauth/callback", + "https://chat.openai.com/oauth/callback", + "https://platform.openai.com/oauth/callback", + "https://gemini.google.com/oauth/cb", + "https://chat.mistral.ai/connectors/oauth/callback", + "https://vscode.dev/oauth/cb", + "https://github.dev/oauth/cb", + "https://cursor.com/mcp/oauth/callback", + ], +) +def test_https_redirects_follow_allowed_origins(uri: str) -> None: + allowed_origins = [ + "https://claude.ai", + "https://assistant.ai", + "https://chatgpt.com", + "https://chat.openai.com", + "https://platform.openai.com", + "https://gemini.google.com", + "https://chat.mistral.ai", + "https://vscode.dev", + "https://github.dev", + "https://cursor.com", + ] + assert is_trusted_redirect_uri(uri, allowed_origins), f"expected trusted: {uri}" + + +# --- Rejected redirects ------------------------------------------------------ + + +@pytest.mark.parametrize( + "uri", + [ + # Obvious attacker-controlled domains + "https://evil.example.com/cb", + "https://attacker.xyz/callback", + # Typo-squats of real origins + "https://assistant-ai.com/cb", # hyphenated fake + "https://chat.mistral.ai.evil.com/cb", # subdomain confusion + "https://chatgpt.co/cb", # TLD typo + # Valid-looking but not-whitelisted domains + "https://mail.google.com/oauth/cb", + "https://accounts.google.com/oauth/cb", + "https://claude.ai.evil.com/callback", + # Non-HTTP(S) schemes we don't trust + "ftp://assistant.ai/cb", + "file:///etc/passwd", + "javascript:alert(1)", + "data:text/html,