diff --git a/docs/configuration.md b/docs/configuration.md index 5f13f31..ad8b7d3 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -41,7 +41,7 @@ Tailscale IP, a VPN address, a bastion. | `server.named_token_ttl` | Lifetime of named API tokens, in seconds. Default `2592000` (30 days); `0` means never expires, revoke-only. OAuth access tokens and dashboard session bearers keep their fixed 24 h TTL either way. Env override: `BEACONMCP_NAMED_TOKEN_TTL`. | | `server.refresh_token_ttl` | Idle lifetime of OAuth refresh tokens (issued to `authorization_code` clients such as claude.ai, ChatGPT or Le Chat), in seconds. Default and maximum `86400` (24 h); `0` is rejected. Every refresh rotates the token and restarts this clock, but no sign-in outlives 24 h: every refresh token of a sign-in expires 24 h after the authorization, whatever this value says, and the client goes back through the TOTP page. A shorter value only adds an idle timeout inside that day. Env override: `BEACONMCP_REFRESH_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` / `proxmox_upload_link` / `proxmox_fetch_url`. Default 500. | +| `server.transfers_max_mb` | Size cap for `proxmox_upload_file` / `proxmox_download_file` / `proxmox_upload_link` / `proxmox_fetch_url`. Default 500. Behind Cloudflare, an upload link request body is also capped by Cloudflare (100 MB on the free plan, 413 from Cloudflare above that); `proxmox_fetch_url` is not affected. | ## Proxmox diff --git a/docs/tools.md b/docs/tools.md index 1aaaf51..1d5bcf9 100644 --- a/docs/tools.md +++ b/docs/tools.md @@ -123,7 +123,7 @@ The token in the URL is the credential, so the route needs no bearer. A link is node, VMID and path chosen when it was minted, works once, expires after `expires_in` seconds (default 600, at most 3600) and is forgotten on restart. The base of the URL is `features.dashboard.public_url` when set, otherwise the address the MCP request came in on. -Behind Cloudflare, uploads are limited to 100 MB per request on the free plan. +Behind Cloudflare, uploads are limited to 100 MB per request on the free plan. The link needs a running guest: minting one for a stopped VM/CT is refused, and an upload to a guest stopped since then answers 409. ## SSH (2) diff --git a/docs/updates.md b/docs/updates.md index ad8279c..19ad283 100644 --- a/docs/updates.md +++ b/docs/updates.md @@ -36,7 +36,7 @@ Detection is not a guess about how you *should* have installed it: | Detected | What you are told | |----------|-------------------| -| git checkout | `cd ` → `git pull --ff-only` (release channel: `git fetch --tags origin` → `git merge --ff-only `) → `/bin/pip install -e . -c constraints.txt` (the real venv path, when there is one) → `systemctl restart beaconmcp` if a unit file exists | +| git checkout | `cd ` → `git pull --ff-only` (release channel: `git fetch --tags origin` → `git merge --ff-only `; the tool itself mirrors the tags into `refs/beaconmcp/tags/`, so a re-pointed upstream tag never blocks it) → `/bin/pip install -e . -c constraints.txt` (the real venv path, when there is one) → `systemctl restart beaconmcp` if a unit file exists | | container | `docker compose pull` → `docker compose up -d` | | pip distribution | `pip install --upgrade 'beaconmcp @ git+https://github.com/Showdown76py/BeaconMCP.git'` | | unknown | Re-run `deploy/install.sh` from a checkout | @@ -79,6 +79,8 @@ In order, stopping at the first failure: *actual* configuration, imports every module, builds the HTTP app and answers `/health`, against scratch copies of the token store and dashboard database. 7. **Restart**: `systemctl restart`, deferred a few seconds so the response reaches you first. + Only when the server actually runs under the unit: a `beaconmcp serve` started by hand is told + to restart itself instead. Step 6 is a hard gate before the restart. If the new revision cannot start with your config (a setting was renamed, a new one is required, an import is broken) the checkout is reset to exactly @@ -92,7 +94,9 @@ breaks at runtime). So the restart itself is done by a watchdog, started as its with `systemd-run` (a plain child would be killed with the service's cgroup). It: 1. restarts the service; -2. waits up to 90 seconds for `/health` to answer, and to still answer 10 seconds later; +2. waits up to 90 seconds for `/health` to answer, then requires it to keep answering for 30 + seconds with the same main PID and restart count (a process that crashes and is restarted by + systemd every few seconds can answer every single probe); 3. otherwise runs `git reset --hard `, reinstalls the previous pinned dependencies and restarts again. @@ -126,8 +130,13 @@ features: - `require_signed: true` makes `git verify-tag` (release) or `git verify-commit` (branch) a hard gate. Git checks the signature against the keys the service user trusts: import the maintainers' GPG keys, or for SSH signatures set `gpg.ssh.allowedSignersFile` in that user's git config. + **Every** public key in that user's keyring is accepted, trusted or not, so keep it to the + maintainers' keys (or point `GNUPGHOME` at a dedicated keyring in the unit). A release tag must + also carry its own name: a signed `v2.1.0-rc1` republished as `v2.1.0` is refused. On the branch channel this needs a signed tip commit; merges made through GitHub's web UI are - signed by GitHub's key, which you then have to trust. + signed by GitHub's key, which you then have to trust. That key signs every web-UI commit of every + GitHub user, so on the branch channel the gate only proves the commit went through github.com: + it does not stop anyone who can push or merge to the branch. Only one update runs at a time. The dashboard button and the MCP tool reach the same code, and two `git pull` / `pip install` runs in one checkout would fight over `index.lock`. A second caller is diff --git a/src/beaconmcp/__main__.py b/src/beaconmcp/__main__.py index 1f12d46..49e9d5c 100644 --- a/src/beaconmcp/__main__.py +++ b/src/beaconmcp/__main__.py @@ -1212,6 +1212,10 @@ def _run_http(mcp, host: str, port: int, *, serve: bool = True): ) from .ratelimit import RateLimiter, TotpLockout, client_ip, forwarded_host from .server import config, transport_security + # The same list the SDK enforces for CORS, including its built-in default + # when beaconmcp.yaml leaves server.allowed_origins empty: OAuth + # redirects to claude.ai must keep working out of the box. + redirect_origins = list(transport_security.allowed_origins or []) from .metrics import REGISTRY, auth_events, http_requests @@ -1343,9 +1347,24 @@ async def dispatch(self, request: Request, call_next): # client_id alone, anyone who knew a client_id could lock its owner out # of every 2FA path. totp_lockout = TotpLockout(per_source=5, per_client=10, lockout_seconds=300) - totp_locked = totp_lockout.locked - totp_record_failure = totp_lockout.record_failure - totp_record_success = totp_lockout.record_success + + # A derived (DCR) client checks codes against its owner's seed, so the + # lockout is keyed on the seed owner: otherwise every derived client_id + # would get its own guess budget against the same seed. + def _seed_owner(client_id: str) -> str: + client = client_store.get(client_id) + if client is not None and client.owner_client_id: + return client.owner_client_id + return client_id + + def totp_locked(client_id: str, ip: str | None = None) -> bool: + return totp_lockout.locked(_seed_owner(client_id), ip) + + def totp_record_failure(client_id: str, ip: str | None = None) -> None: + totp_lockout.record_failure(_seed_owner(client_id), ip) + + def totp_record_success(client_id: str, ip: str | None = None) -> None: + totp_lockout.record_success(_seed_owner(client_id), ip) def _ip(request: Request) -> str: return client_ip(request, tuple(config.server.trusted_proxies)) @@ -1409,7 +1428,7 @@ def _validate_authorize_params( # allowlist. Prevents authorization-code exfiltration via a # typo-squat or attacker-controlled client that somehow got a # valid client_id. - if not auth.is_trusted_redirect_uri(redirect_uri, config.server.allowed_origins): + if not auth.is_trusted_redirect_uri(redirect_uri, redirect_origins): return {}, JSONResponse( {"error": "invalid_request", "error_description": ( @@ -2265,7 +2284,7 @@ async def dcr_register(request: Request) -> Response: status_code=400, ) bad = [u for u in redirect_uris_raw - if not auth.is_trusted_redirect_uri(u, config.server.allowed_origins)] + if not auth.is_trusted_redirect_uri(u, redirect_origins)] if bad: return JSONResponse( {"error": "invalid_redirect_uri", diff --git a/src/beaconmcp/auth.py b/src/beaconmcp/auth.py index de2f604..35ea266 100644 --- a/src/beaconmcp/auth.py +++ b/src/beaconmcp/auth.py @@ -992,6 +992,17 @@ def refresh(self, refresh_token: str, client_id: str) -> OAuthGrant: rt = self._refresh_tokens.get(_hash_secret(refresh_token)) if rt is None or rt.client_id != client_id: raise RefreshTokenError("invalid") + # Reuse is checked before the token's own idle expiry: a + # rotated-out token replayed after its idle TTL is still a + # stolen copy, and must still take the family down. + if rt.used and now < rt.family_expires_at: + if ( + rt.replay is not None + and now - rt.used_at <= self.REFRESH_REUSE_GRACE + ): + return rt.replay + self.revoke_family(rt.family_id) + raise RefreshTokenError("reused", family_id=rt.family_id) if now > rt.expires_at or now >= rt.family_expires_at: # The whole sign-in is over with the family end; a token # that idled past its own TTL only takes itself out. @@ -1000,16 +1011,12 @@ def refresh(self, refresh_token: str, client_id: str) -> OAuthGrant: else: self._drop_refresh(rt.token_hash) raise RefreshTokenError("expired", family_id=rt.family_id) - if rt.used: - if ( - rt.replay is not None - and now - rt.used_at <= self.REFRESH_REUSE_GRACE - ): - return rt.replay - self.revoke_family(rt.family_id) - raise RefreshTokenError("reused", family_id=rt.family_id) rt.used = True rt.used_at = now + # Keep the rotated-out token until the family ends (not its own + # idle expiry, shorter with a custom refresh_token_ttl) so a + # replay is still recognised as reuse. + rt.expires_at = rt.family_expires_at self._persist_refresh(rt) grant = self._mint_oauth(client_id, rt.family_id, rt.family_expires_at) rt.replay = grant diff --git a/src/beaconmcp/dashboard/mcp_bridge.py b/src/beaconmcp/dashboard/mcp_bridge.py index 0ed9a77..f212d20 100644 --- a/src/beaconmcp/dashboard/mcp_bridge.py +++ b/src/beaconmcp/dashboard/mcp_bridge.py @@ -25,7 +25,6 @@ import httpx2 from mcp.client.session import ClientSession from mcp.client.streamable_http import streamable_http_client -from pydantic import AnyUrl # Extension identifier and content type from the MCP Apps spec (SEP-1865). UI_EXTENSION_ID = "io.modelcontextprotocol/ui" @@ -144,7 +143,8 @@ async def read_ui_resource(session: ClientSession, uri: str) -> str: if not uri.startswith("ui://"): raise UiResourceError("not a ui:// resource") try: - result = await session.read_resource(AnyUrl(uri)) + # mcp 2.x takes a plain str: an AnyUrl fails request validation. + result = await session.read_resource(uri) except Exception as exc: # noqa: BLE001 raise UiResourceError(str(exc)) from exc diff --git a/src/beaconmcp/proxmox/client.py b/src/beaconmcp/proxmox/client.py index d4f3f23..24f973a 100644 --- a/src/beaconmcp/proxmox/client.py +++ b/src/beaconmcp/proxmox/client.py @@ -10,7 +10,8 @@ from proxmoxer import ProxmoxAPI from requests.adapters import HTTPAdapter -from requests.exceptions import ConnectionError, Timeout +from requests.exceptions import ConnectionError, ConnectTimeout, Timeout +from urllib3.exceptions import ConnectTimeoutError, NewConnectionError from ..config import Config @@ -98,6 +99,14 @@ def _trust_ca_bundle(conn: ProxmoxAPI, ca_path: str) -> None: session.mount("https://", _CABundleAdapter(ca_path)) +def _never_sent(exc: Exception) -> bool: + """True when the TCP connection itself failed, so no request went out.""" + if isinstance(exc, ConnectTimeout): + return True + reason = getattr(exc.args[0], "reason", None) if exc.args else None + return isinstance(reason, (NewConnectionError, ConnectTimeoutError)) + + class ProxmoxClient: """Manages connections to one or more Proxmox VE nodes via API tokens.""" @@ -166,6 +175,17 @@ def api_call(self, node_name: str, method: str, path: str, **kwargs: Any) -> Any raise except (ConnectionError, Timeout) as e: last_exc = e + # A GET is safe to send twice. Anything else is retried only + # when the request provably never reached the node: a read + # timeout on vm_clone or snapshot_create may mean the task is + # already running, and a second POST would start it again. + if method != "get" and not _never_sent(e): + self._cache().pop(node_name, None) + return { + "error": f"Node '{node_name}' did not answer in time: {e}. " + "The request may still have been applied; check " + "proxmox_get_tasks before retrying." + } # 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. diff --git a/src/beaconmcp/proxmox/system.py b/src/beaconmcp/proxmox/system.py index 02200c5..00e4c99 100644 --- a/src/beaconmcp/proxmox/system.py +++ b/src/beaconmcp/proxmox/system.py @@ -13,6 +13,7 @@ from typing import Any from urllib.parse import urlparse +import httpcore import httpx from mcp.server.mcpserver import MCPServer @@ -40,26 +41,57 @@ class ExecSession: def _prune_exec_sessions() -> None: + # Runs on executor threads while other threads add sessions: snapshot the + # items first, or the dict can change size under the iteration. now = time.time() stale = [ eid - for eid, s in _exec_sessions.items() + for eid, s in list(_exec_sessions.items()) if s.status != "running" and now - s.started_at > _EXEC_SESSION_TTL ] for eid in stale: - del _exec_sessions[eid] + _exec_sessions.pop(eid, None) -def _detect_vm_type(client: ProxmoxClient, node: str, vmid: int) -> str | None: +def _detect_guest(client: ProxmoxClient, node: str, vmid: int) -> tuple[str, str] | None: + """``(vm_type, status)`` of a guest, e.g. ``("lxc", "running")``.""" 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 vm_type, str(data["status"]) return None +def _detect_vm_type(client: ProxmoxClient, node: str, vmid: int) -> str | None: + guest = _detect_guest(client, node, vmid) + return guest[0] if guest else None + + +async def _running_guest( + client: ProxmoxClient, node: str, vmid: int, +) -> tuple[str | None, dict[str, Any] | None]: + """Resolve a push target: ``(vm_type, None)`` or ``(None, error result)``. + + Files can only be pushed into a running guest (``pct push`` refuses a + stopped CT, a stopped VM has no SSH), so check it up front rather than + after the bytes were transferred. + """ + loop = asyncio.get_running_loop() + guest = await loop.run_in_executor(None, _detect_guest, client, node, vmid) + if guest is None: + return None, {"status": "error", "error": f"VM/CT {vmid} not found on node '{node}'."} + vm_type, status = guest + if status != "running": + return None, { + "status": "error", + "error": f"VM/CT {vmid} is {status}: start it first, files can only be pushed into a running guest.", + "guest_status": status, + } + return vm_type, None + + #: Added to an agent/exec error when the guest OS could not be detected, so a #: Windows guest whose agent blocks guest-get-osinfo is not left guessing. _OS_UNKNOWN_HINT = ( @@ -81,6 +113,11 @@ def _guest_argv( ``cmd.exe /c`` instead. When the agent can't report its OS (old agent, command blocked), assume a POSIX guest, which is by far the common case, and return a hint to attach to an exec error (``None`` when the OS is known). + + On Windows the agent joins the argv into one command line with C runtime + quoting (``"`` becomes ``\\"``), which cmd.exe does not understand: a + command that itself contains double quotes can reach cmd.exe mangled. + Commands without inner double quotes are unaffected. """ info = client.get(node, f"nodes/{node}/qemu/{vmid}/agent/get-osinfo") result = info.get("result") if isinstance(info, dict) else None @@ -131,6 +168,40 @@ def _sha256_file(path: Path) -> str: # Hard cap for the guest-agent file-read path (matches proxmox_write_file). _FILE_READ_MAX_BYTES = 1024 * 1024 +#: PVE caps agent/file-write ``content`` at 60 KiB, and the content goes +#: base64-encoded (4/3 larger): 45 KiB of raw bytes is what fits. +_FILE_WRITE_MAX_BYTES = 45 * 1024 + + +def _agent_read_bytes( + client: ProxmoxClient, node: str, vmid: int, path: str, +) -> bytes | dict[str, Any]: + """Raw bytes of a guest file via agent/file-read, or an error result. + + Reads one byte past the limit so an oversized file is told apart from + one exactly at it. ``decode=0`` keeps the guest agent's base64 as is: + PVE's default decodes it server-side and the bytes then come back as + JSON text, which isn't base64 any more. + """ + api = f"nodes/{node}/qemu/{vmid}/agent/file-read" + count = _FILE_READ_MAX_BYTES + 1 + result = client.get(node, api, file=path, count=count, decode=0) + legacy = False + if isinstance(result, dict) and "error" in result and "decode" in str(result["error"]): + # PVE before the `decode` parameter: content arrives decoded, one + # JSON code point per byte. + result = client.get(node, api, file=path, count=count) + legacy = True + if isinstance(result, dict) and "error" in result: + return {"status": "error", "error": result["error"]} + if not isinstance(result, dict) or not isinstance(result.get("content"), str): + return {"status": "error", "error": f"Unexpected file-read response: {result!r}"[:300]} + try: + if legacy: + return result["content"].encode("latin-1") + return base64.b64decode(result["content"], validate=True) + except (ValueError, UnicodeEncodeError) as e: + return {"status": "error", "error": f"Failed to decode file content: {e}"} def _qemu_file_size(client: ProxmoxClient, node: str, vmid: int, path: str) -> int | None: @@ -192,9 +263,11 @@ async def _refuse_local_target(request: httpx.Request) -> None: Keeps proxmox_fetch_url from reading services that only listen on the BeaconMCP host (its own /metrics, the Proxmox API on 127.0.0.1) or a - cloud metadata endpoint into a guest. The address is resolved here and - again by httpx when it connects, so a DNS answer that changes in between - is not covered; this is a guard rail, not a sandbox. + cloud metadata endpoint into a guest. This early check gives a clear + error before any connection; the one that holds is + :class:`_GuardedNetworkBackend`, which connects to the very address it + checked (a DNS answer that changes between two lookups can't slip a + loopback address past it). """ host = request.url.host try: @@ -213,6 +286,89 @@ async def _refuse_local_target(request: httpx.Request) -> None: ) +class _GuardedNetworkBackend(httpcore.AsyncNetworkBackend): + """Resolve once, refuse local addresses, connect to the checked IP. + + TLS still verifies against the URL's hostname: httpcore takes the SNI + and certificate name from the request origin, not from this address. + Resolution failures fail closed. + """ + + def __init__(self) -> None: + self._inner = httpcore.AnyIOBackend() + + async def connect_tcp( + self, + host: str, + port: int, + timeout: float | None = None, + local_address: str | None = None, + socket_options: Any = None, + ) -> httpcore.AsyncNetworkStream: + try: + ipaddress.ip_address(host) + addrs = [host] + except ValueError: + try: + addrs = await _host_addresses(host) + except OSError as e: + raise httpcore.ConnectError(f"cannot resolve {host}: {e}") from e + if not addrs: + raise httpcore.ConnectError(f"cannot resolve {host}") + local = next((a for a in addrs if _is_local_address(a)), None) + if local is not None: + raise _LocalTargetRefused( + f"Refusing to fetch from {host} ({local}): loopback and link-local " + "addresses are not allowed." + ) + return await self._inner.connect_tcp( + addrs[0], port, timeout=timeout, + local_address=local_address, socket_options=socket_options, + ) + + async def sleep(self, seconds: float) -> None: + await self._inner.sleep(seconds) + + +def _fetch_transport() -> httpx.AsyncBaseTransport: + """Transport for proxmox_fetch_url: every connection goes through the guard. + + Passing a transport also turns off HTTP(S)_PROXY from the environment: + a proxy would resolve the name itself, out of the guard's reach. + """ + transport = httpx.AsyncHTTPTransport() + transport._pool = httpcore.AsyncConnectionPool( + ssl_context=httpx.create_ssl_context(), + network_backend=_GuardedNetworkBackend(), + ) + return transport + + +async def _remove_quietly(ssh_client: Any, host: str, path: str) -> None: + """Best-effort cleanup that also runs when the caller is being cancelled + (a restart cancels in-flight requests after a few seconds).""" + try: + await asyncio.shield(ssh_client.sftp_remove(host, path)) + except (Exception, asyncio.CancelledError): + pass + + +def _checksum_from(sum_res: dict[str, Any]) -> tuple[str | None, str | None]: + """``(sha, error)`` from a ``sha256sum `` run in the guest. + + Only a missing ``sha256sum`` (exit 127) or no output counts as "can't + verify". Any other failure (a directory, no permission, a wrong path) + means the file is not where the caller asked, so it is an error. + """ + code = sum_res.get("exit_code") + if code == 0 and sum_res.get("stdout"): + return sum_res["stdout"].split()[0], None + if code in (0, 127): + return None, None + detail = (sum_res.get("stderr") or sum_res.get("error") or f"exit {code}").strip() + return None, f"Checksum check of the pushed file failed: {detail}" + + async def _push_staged( client: ProxmoxClient, ssh_client: Any, @@ -233,9 +389,9 @@ async def _push_staged( size_bytes = local_path.stat().st_size 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}'."} + vm_type, err = await _running_guest(client, node, vmid) + if err is not None: + return err started = time.monotonic() local_sha = await loop.run_in_executor(None, _sha256_file, local_path) if verify_checksum else None @@ -243,15 +399,19 @@ async def _push_staged( 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) + 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, + ) + finally: + # Also after a partial SFTP (full /tmp) or a cancelled request: + # the node's /tmp must not keep the copy. + await _remove_quietly(ssh_client, node, node_tmp) if push_res.get("exit_code") != 0: return { "status": "error", @@ -264,8 +424,9 @@ async def _push_staged( 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] + remote_sha, sum_err = _checksum_from(sum_res) + if sum_err: + return {"status": "error", "error": sum_err, "bytes": size_bytes} if remote_sha and remote_sha != local_sha: return { "status": "error", @@ -297,6 +458,16 @@ async def _push_staged( "or use proxmox_write_file for files ≤1MB." ), } + # SFTP put into an existing directory would drop the file inside it + # under the staging name (upload-), not at `dest`. + dir_check = await ssh_client.exec_command( + str(vmid), f"test -d {shlex.quote(dest)}", timeout=30, + ) + if dir_check.get("exit_code") == 0: + return { + "status": "error", + "error": f"`dest` {dest!r} is a directory in VM {vmid}: give the full file path.", + } try: await ssh_client.sftp_put(str(vmid), str(local_path), dest) except Exception as e: @@ -306,8 +477,9 @@ async def _push_staged( 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] + remote_sha, sum_err = _checksum_from(sum_res) + if sum_err: + return {"status": "error", "error": sum_err, "bytes": size_bytes} if remote_sha and remote_sha != local_sha: return { "status": "error", @@ -372,31 +544,26 @@ async def proxmox_read_file(node: str, vmid: int, path: str, binary: bool = Fals "limit. Use proxmox_download_file for large files." ), } - result = await loop.run_in_executor( - None, - lambda: client.get(node, f"nodes/{node}/qemu/{vmid}/agent/file-read", file=path), + data = await loop.run_in_executor( + None, _agent_read_bytes, client, node, vmid, path, ) - if isinstance(result, dict) and "error" in result: - return {"status": "error", "error": result["error"]} - # QEMU Guest Agent returns file content base64-encoded + if isinstance(data, dict): + return data + if len(data) > _FILE_READ_MAX_BYTES: + return { + "status": "error", + "error": "File exceeds the 1MB guest-agent read limit. Use proxmox_download_file for large files.", + } + if binary: + return { + "status": "success", "vmid": vmid, "node": node, "path": path, + "content_base64": base64.b64encode(data).decode("ascii"), + } 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}"} + content = data.decode("utf-8") + 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": content} return {"status": "error", "error": "LXC file reading is currently unsupported via API. Please use ssh_run to cat the file."} @@ -413,17 +580,26 @@ async def proxmox_write_file(node: str, vmid: int, path: str, content: str) -> d 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."} - + raw = content.encode("utf-8") + if len(raw) > _FILE_WRITE_MAX_BYTES: + return { + "status": "error", + "error": ( + f"Content is {len(raw)} bytes, over the {_FILE_WRITE_MAX_BYTES} byte " + "guest-agent write limit. Use proxmox_upload_link or proxmox_upload_file." + ), + } + if vm_type == "qemu": - import base64 - encoded = base64.b64encode(content.encode("utf-8")).decode("utf-8") + # encode=0: the content is already base64. With PVE's default + # (encode=1) it would be encoded a second time and the guest + # would get the base64 text instead of the file. + encoded = base64.b64encode(raw).decode("ascii") result = await loop.run_in_executor( None, lambda: client.post( node, f"nodes/{node}/qemu/{vmid}/agent/file-write", - file=path, content=encoded, encode=1, + file=path, content=encoded, encode=0, ), ) if isinstance(result, dict) and "error" in result: @@ -791,10 +967,9 @@ async def proxmox_upload_link( return {"status": "error", "error": "`node` and `vmid` are required to mint a link."} 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}'."} + vm_type, err = await _running_guest(client, node, vmid) + if err is not None: + return err max_bytes = client._config.server.transfers_max_mb * 1024 * 1024 try: @@ -860,9 +1035,9 @@ async def proxmox_fetch_url( if not isinstance(dest, str) or not dest.startswith("/"): return {"status": "error", "error": "`dest` must be an absolute path inside the guest."} # Check the target before pulling a possibly large file for nothing. - loop = asyncio.get_running_loop() - if not await loop.run_in_executor(None, _detect_vm_type, client, node, vmid): - return {"status": "error", "error": f"VM/CT {vmid} not found on node '{node}'."} + _vm_type, err = await _running_guest(client, node, vmid) + if err is not None: + return err max_bytes = client._config.server.transfers_max_mb * 1024 * 1024 staged = _staging_dir(client) / f"fetch-{uuid.uuid4().hex[:12]}" @@ -873,6 +1048,7 @@ async def proxmox_fetch_url( async with httpx.AsyncClient( follow_redirects=True, timeout=min(max(timeout, 10), 1800), event_hooks={"request": [_refuse_local_target]}, + transport=_fetch_transport(), ) as http: async with http.stream("GET", url) as resp: if resp.status_code >= 400: diff --git a/src/beaconmcp/ratelimit.py b/src/beaconmcp/ratelimit.py index 27845de..5c6b612 100644 --- a/src/beaconmcp/ratelimit.py +++ b/src/beaconmcp/ratelimit.py @@ -106,6 +106,11 @@ class TotpLockout: whose cooldown has passed starts again from zero. A success clears the caller's own per-source counter only, so a distributed guesser can't reset the per-client cap by having the owner log in. + + A source that entered a valid code recently is exempt from the + per-client cap (its own per-source cap still applies). Otherwise + anyone with two addresses could spend the per-client budget every + cooldown and keep the owner locked out of every TOTP path for good. """ def __init__( @@ -114,13 +119,17 @@ def __init__( per_source: int = 5, per_client: int = 10, lockout_seconds: float = 300.0, + known_good_seconds: float = 30 * 86400.0, ) -> None: self._per_source = per_source self._per_client = per_client self._lockout = lockout_seconds + self._known_good_ttl = known_good_seconds # key -> (failures, cooldown_until). Keys are ("src", client_id, ip) # and ("client", client_id). self._strikes: dict[tuple[str, ...], tuple[int, float]] = {} + # ("src", client_id, ip) -> until: sources that passed TOTP lately. + self._known_good: dict[tuple[str, ...], float] = {} self._lock = threading.Lock() @staticmethod @@ -137,14 +146,24 @@ def _count_locked(self, key: tuple[str, ...], now: float) -> int: return 0 return count + def _is_known_good(self, key: tuple[str, ...], now: float) -> bool: + until = self._known_good.get(key) + if until is None: + return False + if now >= until: + del self._known_good[key] + return False + return True + def locked(self, client_id: str, ip: str | None = None) -> bool: now = time.time() + src = self._src_key(client_id, ip) with self._lock: - return ( - self._count_locked(("client", client_id), now) >= self._per_client - or self._count_locked(self._src_key(client_id, ip), now) - >= self._per_source - ) + if self._count_locked(src, now) >= self._per_source: + return True + if self._is_known_good(src, now): + return False + return self._count_locked(("client", client_id), now) >= self._per_client def record_failure(self, client_id: str, ip: str | None = None) -> None: now = time.time() @@ -158,8 +177,18 @@ def record_failure(self, client_id: str, ip: str | None = None) -> None: del self._strikes[key] def record_success(self, client_id: str, ip: str | None = None) -> None: + now = time.time() + src = self._src_key(client_id, ip) with self._lock: - self._strikes.pop(self._src_key(client_id, ip), None) + self._strikes.pop(src, None) + self._known_good[src] = now + self._known_good_ttl + if len(self._known_good) > 4096: + for key, until in list(self._known_good.items()): + if now >= until: + del self._known_good[key] + # Still full: forget the oldest, the cap only bounds memory. + while len(self._known_good) > 4096: + del self._known_good[min(self._known_good, key=self._known_good.__getitem__)] def _coerce_ip(value: object) -> str | None: @@ -214,22 +243,25 @@ def _is_trusted_proxy(ip_value: str, trusted_proxies: tuple[str, ...]) -> bool: def _warn_untrusted_forwarding(forwarded_for: str | None, peer: str | None) -> None: - """Say once that a proxy is forwarding but ``trusted_proxies`` is empty. + """Say once that a proxy is forwarding but is not in ``trusted_proxies``. Behind a reverse proxy every request then comes from the proxy's own address, so the per-IP limits (login, authorize, TOTP lockout) become one budget shared by every visitor, and one attacker can spend it for - everyone. + everyone. Also fires with a non-empty list: a cloudflared tunnel + connects from its own LAN address, which ``cloudflare`` (the edge + ranges) does not cover. """ global _forwarding_warned if _forwarding_warned or not forwarded_for: return _forwarding_warned = True logging.getLogger("beaconmcp").warning( - "X-Forwarded-For received from %s but server.trusted_proxies is empty: " + "X-Forwarded-For received from %s, which is not in server.trusted_proxies: " "every client behind that proxy shares one rate-limit and TOTP-lockout " - "budget. Add the proxy's address to server.trusted_proxies.", - peer or "an unknown peer", + "budget. If %s is your reverse proxy or cloudflared tunnel, add it to " + "server.trusted_proxies.", + peer or "an unknown peer", peer or "it", ) @@ -246,9 +278,10 @@ def client_ip(request: object, trusted_proxies: tuple[str, ...] = ()) -> str: direct_ip = _coerce_ip(direct_peer_raw) headers = getattr(request, "headers", None) - if not trusted_proxies and headers is not None and hasattr(headers, "get"): + peer_trusted = bool(direct_ip) and _is_trusted_proxy(direct_ip, trusted_proxies) + if not peer_trusted and headers is not None and hasattr(headers, "get"): _warn_untrusted_forwarding(headers.get("x-forwarded-for"), direct_ip) - if headers is not None and direct_ip and _is_trusted_proxy(direct_ip, trusted_proxies): + if headers is not None and peer_trusted: fwd = headers.get("x-forwarded-for") if hasattr(headers, "get") else None if fwd: chain: list[str] = [ diff --git a/src/beaconmcp/ssh/client.py b/src/beaconmcp/ssh/client.py index 26d5419..6f327e5 100644 --- a/src/beaconmcp/ssh/client.py +++ b/src/beaconmcp/ssh/client.py @@ -115,10 +115,29 @@ def _known_host_pattern(host: str, port: int) -> str: return host if port == 22 else f"[{host}]:{port}" -def _load_known_hosts(path: str) -> asyncssh.SSHKnownHosts: - if os.path.isfile(path): - return asyncssh.read_known_hosts(path) - return asyncssh.import_known_hosts("") +def _pinned_keys(path: str, host: str, port: int) -> list[asyncssh.SSHKey]: + """Keys recorded in the managed file for exactly ``host`` on ``port``. + + asyncssh's own matching retries a non-22 port without the port, so a + ``[host]:2222`` target would be checked against the key of ``host`` on + port 22, which is often a different machine behind the same NAT address. + """ + if not os.path.isfile(path): + return [] + pattern = _known_host_pattern(host, port) + keys: list[asyncssh.SSHKey] = [] + with open(path, encoding="utf-8") as f: + for raw in f: + parts = raw.split() + if len(parts) < 3 or parts[0].startswith(("#", "@")): + continue + if pattern not in parts[0].split(","): + continue + try: + keys.append(asyncssh.import_public_key(f"{parts[1]} {parts[2]}")) + except (asyncssh.KeyImportError, ValueError): + continue + return keys class _TofuClient(asyncssh.SSHClient): @@ -139,7 +158,7 @@ def __init__(self, path: str, host: str, port: int) -> None: def validate_host_public_key(self, host: str, addr: str, port: int, key) -> bool: with _tofu_lock: - pinned = _load_known_hosts(self._path).match(host, addr, port)[0] + pinned = _pinned_keys(self._path, self._host, self._port) if key in pinned: return True if pinned: @@ -176,7 +195,9 @@ async def _connect_to_host( connect_kwargs["known_hosts"] = None elif policy.mode == "tofu": assert policy.path is not None - connect_kwargs["known_hosts"] = _load_known_hosts(policy.path) + connect_kwargs["known_hosts"] = ( + _pinned_keys(policy.path, spec.host, spec.port), [], [], + ) tofu_client = _TofuClient(policy.path, spec.host, spec.port) connect_kwargs["client_factory"] = lambda: tofu_client # "system": omit the kwarg -> asyncssh uses ~/.ssh/known_hosts. diff --git a/src/beaconmcp/transfers.py b/src/beaconmcp/transfers.py index ea1adf0..864353c 100644 --- a/src/beaconmcp/transfers.py +++ b/src/beaconmcp/transfers.py @@ -17,6 +17,7 @@ from __future__ import annotations +import ipaddress import secrets import threading import time @@ -25,6 +26,7 @@ from dataclasses import dataclass, field from pathlib import Path from typing import Any, Awaitable, Callable +from urllib.parse import urlsplit from starlette.requests import Request from starlette.responses import JSONResponse @@ -185,9 +187,28 @@ def register_pusher(pusher: Pusher | None) -> None: def public_base_url(configured: str | None) -> str | None: - """Base URL for upload links: the configured public URL, else the request's.""" - base = (configured or "").strip() or current_public_base.get() - return base.rstrip("/") if base else None + """Base URL for upload links: the configured public URL, else the request's. + + A loopback request base (the dashboard chat's own MCP session runs over + 127.0.0.1) is no use to whoever uploads, so it counts as unknown. + """ + configured = (configured or "").strip() + if configured: + return configured.rstrip("/") + base = current_public_base.get() + if not base or _is_loopback_url(base): + return None + return base.rstrip("/") + + +def _is_loopback_url(url: str) -> bool: + host = (urlsplit(url).hostname or "").lower() + if host == "localhost" or host.endswith(".localhost"): + return True + try: + return ipaddress.ip_address(host).is_loopback + except ValueError: + return False def staged_name(link: UploadLink) -> str: @@ -211,18 +232,23 @@ async def upload_endpoint(request: Request) -> JSONResponse: }, status_code=404, ) - result, status_code = await _receive_and_push(request, link) - _store.finish(link, result) - audit.emit( - "transfer.upload", - link_id=link.link_id, - node=link.node, - vmid=link.vmid, - dest=link.dest, - status=result.get("status"), - bytes=result.get("bytes"), - ip=client_ip(request, trusted_proxies), - ) + # The link is spent: settle it and audit it whatever happens, a + # request cancelled by a restart included. + result, status_code = _error("Upload interrupted (server shutting down)."), 500 + try: + result, status_code = await _receive_and_push(request, link) + finally: + _store.finish(link, result) + audit.emit( + "transfer.upload", + link_id=link.link_id, + node=link.node, + vmid=link.vmid, + dest=link.dest, + status=result.get("status"), + bytes=result.get("bytes"), + ip=client_ip(request, trusted_proxies), + ) return JSONResponse(result, status_code=status_code) return upload_endpoint @@ -272,4 +298,17 @@ async def _receive_and_push(request: Request, link: UploadLink) -> tuple[dict[st result.setdefault("bytes", received) result["link_id"] = link.link_id - return result, 200 if result.get("status") == "success" else 502 + return result, _status_code(result) + + +def _status_code(result: dict[str, Any]) -> int: + """HTTP status for a finished transfer. + + Not 502: Cloudflare swaps an origin 502 for its own "error code: 502" + page, so the uploader would never see why the push failed. + """ + if result.get("status") == "success": + return 200 + if "guest_status" in result: + return 409 # target stopped since the link was minted + return 500 diff --git a/src/beaconmcp/updates.py b/src/beaconmcp/updates.py index 8ae3748..6f95658 100644 --- a/src/beaconmcp/updates.py +++ b/src/beaconmcp/updates.py @@ -56,6 +56,8 @@ #: Release tags the release channel follows. Pre-releases (v2.1.0-rc1) are #: skipped on purpose. _RELEASE_TAG = re.compile(r"^v\d+(?:\.\d+)*$") +#: Where the release channel keeps its copy of the upstream tags. +_TAG_NAMESPACE = "refs/beaconmcp/tags/" #: 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. @@ -70,11 +72,16 @@ _PIP_TIMEOUT = 900 #: ``check-startup`` imports the whole server and builds the HTTP app. _STARTUP_CHECK_TIMEOUT = 180 +#: After the first good /health, the new process must keep answering, with +#: the same PID and restart count, for this long before the update counts as +#: healthy. The unit restarts a crashed process after 5 s, so two bare +#: samples a few seconds apart can both land on different incarnations. +_SETTLE_SECONDS = 30 #: Longest a post-restart watchdog can still act: two health waits (with #: their settle re-check), a pip reinstall and some slack. While the last #: status is a ``pending`` younger than this, a new update is refused, so a #: watchdog that is still waiting cannot roll back a newer update. -_WATCHDOG_WINDOW_SECONDS = 2 * (HEALTH_WAIT_SECONDS + 15) + _PIP_TIMEOUT + 60 +_WATCHDOG_WINDOW_SECONDS = 2 * (HEALTH_WAIT_SECONDS + _SETTLE_SECONDS + 15) + _PIP_TIMEOUT + 60 @dataclass @@ -584,10 +591,20 @@ class _Target: def _resolve_target(root: Path, channel: str) -> tuple[_Target | None, str | None]: """Fetch and name the update target, or explain why there is none.""" if channel == "release": - code, _, err = _git(root, "fetch", "--quiet", "--tags", "origin") + # Mirrored into a private namespace, forced and pruned: a plain + # `fetch --tags` exits 1 for good once upstream re-points a tag the + # checkout already has, and never forgets a deleted one. Forcing is + # safe, the ff-only merge and the signature check still gate the move. + code, _, err = _git( + root, "fetch", "--quiet", "--no-tags", "--prune", "--force", "origin", + f"+refs/tags/v*:{_TAG_NAMESPACE}v*", + ) if code != 0: return None, f"could not reach the remote ({err or 'git fetch failed'})" - code, out, err = _git(root, "tag", "--list", "v*", "--sort=-v:refname") + code, out, err = _git( + root, "for-each-ref", "--sort=-v:refname", "--format=%(refname:strip=3)", + _TAG_NAMESPACE, + ) if code != 0: return None, f"could not list release tags ({err or 'git failed'})" tags = [t for t in out.splitlines() if _RELEASE_TAG.match(t.strip())] @@ -597,7 +614,7 @@ def _resolve_target(root: Path, channel: str) -> tuple[_Target | None, str | Non "features.updates.channel to 'branch' to follow the default branch" ) tag = tags[0].strip() - code, sha, err = _git(root, "rev-parse", f"refs/tags/{tag}^{{commit}}") + code, sha, err = _git(root, "rev-parse", f"{_TAG_NAMESPACE}{tag}^{{commit}}") if code != 0: return None, f"could not resolve {tag} ({err or 'git failed'})" return _Target("release", tag, sha), None @@ -621,7 +638,16 @@ def _verify_signature(root: Path, target: _Target) -> tuple[bool, str]: (``git verify-commit``; GitHub signs the merges it makes). """ if target.kind == "release": - code, out, err = _git(root, "verify-tag", f"refs/tags/{target.name}") + ref = f"{_TAG_NAMESPACE}{target.name}" + # verify-tag checks the signature, not that the signed object is + # this tag: a signed v2.1.0-rc1 republished as v2.1.0 would pass. + _, own_name, _ = _git(root, "for-each-ref", "--format=%(tag)", ref) + if own_name != target.name: + return False, ( + f"the signed tag object is named {own_name or 'nothing'!r}, " + f"not {target.name!r}" + ) + code, out, err = _git(root, "verify-tag", ref) else: code, out, err = _git(root, "verify-commit", target.sha) detail = (err or out or "").strip() @@ -854,15 +880,21 @@ def _startup_check_command(install: Installation, config_path: Path | None) -> l #: (``systemd-run``): ``systemctl restart`` kills every process in the #: service's cgroup, so a child of the server would die with it. All values #: arrive through argv, never interpolated into the script. It restarts the -#: service, waits for the new process to answer /health (and still answer -#: ``settle`` seconds later, so a crash loop is not mistaken for a start), and +#: service, waits for the new process to answer /health (and keep answering +#: for ``settle`` seconds with the same MainPID and NRestarts, so a crash +#: loop is not mistaken for a start), and #: otherwise resets the checkout to the previous revision, #: reinstalls its pinned dependencies and restarts again. The outcome lands in #: STATUS_FILE so the next check (and the operator) can see it. +#: +#: The script must never contain a dollar-brace or a double dollar: PID 1 +#: expands ``${VAR}`` (and collapses ``$$``) in the ExecStart of the +#: transient unit before ``sh`` sees it, so ``${10}`` would arrive empty. _WATCHDOG_SCRIPT = r""" delay=$1 service=$2 url=$3 root=$4 from=$5 to=$6 py=$7 status=$8 wait=$9 -settle=${10} -shift 10 +shift 9 +settle=$1 +shift log() { echo "beaconmcp-update: $*"; } write_status() { printf '{"at": %s, "from": "%s", "to": "%s", "outcome": "%s"}\n' \ @@ -872,17 +904,29 @@ def _startup_check_command(install: Installation, config_path: Path | None) -> l "$py" -c 'import sys, urllib.request as u; u.build_opener(u.ProxyHandler({})).open(sys.argv[1], timeout=3).read()' \ "$url" >/dev/null 2>&1 } +ident() { + systemctl show -p MainPID -p NRestarts --value "$service" 2>/dev/null | tr '\n' ' ' +} +stable() { healthy && [ "$(ident)" = "$first" ]; } wait_healthy() { end=$(( $(date +%s) + wait )) while [ "$(date +%s)" -lt "$end" ]; do if healthy; then - sleep "$settle" - if healthy; then return 0; fi + first=$(ident) + waited=0 + while [ "$waited" -lt "$settle" ] && stable; do + sleep 2 + waited=$((waited + 2)) + done + if [ "$waited" -ge "$settle" ] && stable; then return 0; fi fi sleep 2 done return 1 } +pip_install() { + if command -v timeout >/dev/null 2>&1; then timeout 900 "$@"; else "$@"; fi +} sleep "$delay" systemctl restart "$service" if wait_healthy; then @@ -890,7 +934,7 @@ def _startup_check_command(install: Installation, config_path: Path | None) -> l log "$to is up and healthy" exit 0 fi -log "no healthy answer from $url within ${wait}s; rolling back to $from" +log "no steady healthy answer from $url within $wait seconds; rolling back to $from" if [ "$(git -C "$root" rev-parse HEAD)" != "$to" ]; then write_status superseded log "the checkout is no longer at $to (another update ran); not rolling back" @@ -898,7 +942,7 @@ def _startup_check_command(install: Installation, config_path: Path | None) -> l fi if git -C "$root" reset --hard "$from"; then cd "$root" || exit 1 - if [ -f constraints.txt ]; then "$@" install -e . -c constraints.txt; else "$@" install -e .; fi \ + if [ -f constraints.txt ]; then pip_install "$@" install -e . -c constraints.txt; else pip_install "$@" install -e .; fi \ || log "reinstalling the dependencies of $from FAILED; run pip install by hand" systemctl restart "$service" if wait_healthy; then @@ -927,7 +971,7 @@ class _Watchdog: python: str pip: list[str] wait_seconds: int = HEALTH_WAIT_SECONDS - settle_seconds: int = 10 + settle_seconds: int = _SETTLE_SECONDS def _write_status(root: Path, payload: dict[str, Any]) -> None: @@ -1167,7 +1211,10 @@ def _rollback(reason: str) -> UpdateResult: result.ok = True watched = False - if restart and install.service: + # A unit file on disk does not mean this process runs under it: restarting + # the unit while `beaconmcp serve` runs by hand would leave this process + # answering /health and the watchdog calling the update healthy. + if restart and install.service and install.under_systemd: watchdog = None if _runtime.health_url: watchdog = _Watchdog( @@ -1209,7 +1256,7 @@ def _rollback(reason: str) -> UpdateResult: f"If it does not answer /health within {HEALTH_WAIT_SECONDS}s it is " f"rolled back to {result.from_ref} automatically." ) - elif install.service: + elif install.service and install.under_systemd: bits.append(f"Restart it with: systemctl restart {install.service}") else: bits.append("Restart the server process to run the new code.") diff --git a/tests/test_agent_file_io.py b/tests/test_agent_file_io.py new file mode 100644 index 0000000..c648c83 --- /dev/null +++ b/tests/test_agent_file_io.py @@ -0,0 +1,118 @@ +"""proxmox_read_file / proxmox_write_file against PVE's agent/file-* semantics. + +The fake client mirrors qemu-server's API2/Qemu/Agent.pm: file-read decodes +the guest agent's base64 unless ``decode=0``, file-write base64-encodes +``content`` itself unless ``encode=0``, and ``content`` is capped at 60 KiB. +""" + +from __future__ import annotations + +import base64 +import sys +from pathlib import Path +from typing import Any + +import pytest + +sys.path.insert(0, str(Path(__file__).parent.parent / "src")) + +from beaconmcp.proxmox import system # noqa: E402 + +NODE, VM = "pve1", 100 + + +class FakeMCP: + def __init__(self) -> None: + self.tools: dict[str, Any] = {} + + def tool(self): + def deco(fn): + self.tools[fn.__name__] = fn + return fn + return deco + + +class FakePVE: + def __init__(self, *, has_decode_param: bool = True) -> None: + self.files: dict[str, bytes] = {} + self.has_decode_param = has_decode_param + + def get(self, node: str, path: str, **params: Any) -> Any: + if path == f"nodes/{node}/qemu/{VM}/status/current": + return {"status": "running"} + if path == f"nodes/{node}/qemu/{VM}/agent/file-read": + if "decode" in params and not self.has_decode_param: + return {"error": "400 Parameter verification failed: decode: property is not defined in schema"} + data = self.files[params["file"]][: params.get("count", 16 * 1024 * 1024)] + if params.get("decode", 1): + # Perl hands JSON the decoded bytes, one code point each. + return {"content": data.decode("latin-1"), "bytes-read": len(data)} + return {"content": base64.b64encode(data).decode(), "bytes-read": len(data)} + return {"error": "not found"} + + def post(self, node: str, path: str, **params: Any) -> Any: + if path == f"nodes/{node}/qemu/{VM}/agent/file-write": + if len(params["content"]) > 60 * 1024: + return {"error": "400 Parameter verification failed: content: value may only be 61440 characters long"} + encode = params.get("encode", 1) + buf = base64.b64encode(params["content"].encode()).decode() if encode else params["content"] + self.files[params["file"]] = base64.b64decode(buf) + return {} + return {"error": "not found"} + + def exec_stat(self, path: str) -> int | None: + return len(self.files.get(path, b"")) + + +@pytest.fixture +def pve(monkeypatch: pytest.MonkeyPatch): + fake = FakePVE() + monkeypatch.setattr(system, "_qemu_file_size", lambda c, n, v, p: fake.exec_stat(p)) + mcp = FakeMCP() + system.register_system_tools(mcp, fake, None) + fake.tools = mcp.tools + return fake + + +async def test_write_stores_the_content_not_its_base64(pve) -> None: + res = await pve.tools["proxmox_write_file"](node=NODE, vmid=VM, path="/etc/motd", content="héllo\n") + assert res["status"] == "success", res + assert pve.files["/etc/motd"] == "héllo\n".encode() + + +async def test_write_over_the_agent_cap_points_to_transfers(pve) -> None: + res = await pve.tools["proxmox_write_file"](node=NODE, vmid=VM, path="/x", content="a" * (46 * 1024)) + assert res["status"] == "error" + assert "proxmox_upload_link" in res["error"] + ok = await pve.tools["proxmox_write_file"](node=NODE, vmid=VM, path="/x", content="a" * (45 * 1024)) + assert ok["status"] == "success", ok + + +async def test_read_text_round_trips(pve) -> None: + pve.files["/etc/hostname"] = "web01\n".encode() + res = await pve.tools["proxmox_read_file"](node=NODE, vmid=VM, path="/etc/hostname") + assert res == {"status": "success", "vmid": VM, "node": NODE, "path": "/etc/hostname", "content": "web01\n"} + + +async def test_read_binary_returns_real_base64(pve) -> None: + blob = bytes(range(256)) + pve.files["/bin/blob"] = blob + res = await pve.tools["proxmox_read_file"](node=NODE, vmid=VM, path="/bin/blob", binary=True) + assert base64.b64decode(res["content_base64"]) == blob + text = await pve.tools["proxmox_read_file"](node=NODE, vmid=VM, path="/bin/blob") + assert text["status"] == "error" and "binary=True" in text["error"] + + +async def test_read_works_on_pve_without_the_decode_param(pve) -> None: + pve.has_decode_param = False + pve.files["/etc/os"] = "Ubuntu é\n".encode() + res = await pve.tools["proxmox_read_file"](node=NODE, vmid=VM, path="/etc/os") + assert res["content"] == "Ubuntu é\n" + + +async def test_read_refuses_more_than_the_limit(pve, monkeypatch: pytest.MonkeyPatch) -> None: + # stat unavailable: the post-read guard must still catch it. + monkeypatch.setattr(system, "_qemu_file_size", lambda c, n, v, p: None) + pve.files["/big"] = b"x" * (system._FILE_READ_MAX_BYTES + 10) + res = await pve.tools["proxmox_read_file"](node=NODE, vmid=VM, path="/big") + assert res["status"] == "error" diff --git a/tests/test_dashboard_mcp_apps.py b/tests/test_dashboard_mcp_apps.py index 536e350..af8fe6f 100644 --- a/tests/test_dashboard_mcp_apps.py +++ b/tests/test_dashboard_mcp_apps.py @@ -169,6 +169,9 @@ def __init__(self, contents): self._contents = contents async def read_resource(self, uri): + # Like mcp 2.x's ClientSession: ReadResourceRequestParams.uri is a + # strict str, an AnyUrl is rejected before anything is sent. + assert type(uri) is str, f"read_resource needs a str, got {type(uri).__name__}" return _FakeReadResult(self._contents) diff --git a/tests/test_mcp_http_stack.py b/tests/test_mcp_http_stack.py index f92def0..31604a4 100644 --- a/tests/test_mcp_http_stack.py +++ b/tests/test_mcp_http_stack.py @@ -39,6 +39,12 @@ @pytest.fixture def stack(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> dict[str, Any]: """Build the app exactly as `_cmd_serve` does and return it with a bearer.""" + return _build_stack(tmp_path, monkeypatch) + + +def _build_stack( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, *, yaml_origins: list[str] | None = None, +) -> dict[str, Any]: for key, value in _CI_ENV.items(): monkeypatch.setenv(key, value) monkeypatch.setenv("BEACONMCP_CONFIG", str(ROOT / "beaconmcp.yaml.example")) @@ -50,7 +56,13 @@ def stack(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> dict[str, Any]: from beaconmcp import __main__ as entry from beaconmcp import auth - from beaconmcp.server import mcp + from beaconmcp.server import config, mcp, transport_security + + if yaml_origins is not None: + # What server.py computes when the YAML list is empty: its built-in + # fallback goes to the SDK, the YAML value stays empty. + monkeypatch.setattr(config.server, "allowed_origins", yaml_origins) + monkeypatch.setattr(transport_security, "allowed_origins", ["https://claude.ai"]) captured: dict[str, Any] = {} # `_run_http` does `import uvicorn` locally, so patching the attribute on @@ -67,6 +79,7 @@ def stack(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> dict[str, Any]: # least invasive way to mint a token that the middleware will accept. token, _ttl = auth._active_token_store.issue(client_id) captured["token"] = token + captured["client_id"] = client_id return captured @@ -120,6 +133,25 @@ async def _handshake(c: httpx.AsyncClient, token: str, path: str = "/mcp") -> di return h +async def test_authorize_trusts_claude_without_yaml_origins(tmp_path, monkeypatch) -> None: + # With server.allowed_origins left out of beaconmcp.yaml the SDK falls back + # to its built-in list for CORS; the OAuth redirect check must use that + # same list, or the Claude connector gets a 400 on a default install. + stack = _build_stack(tmp_path, monkeypatch, yaml_origins=[]) + app = stack["app"] + params = { + "response_type": "code", + "client_id": stack["client_id"], + "redirect_uri": "https://claude.ai/api/mcp/auth_callback", + "state": "s", + "code_challenge": "x" * 43, + "code_challenge_method": "S256", + } + async with app.router.lifespan_context(app), _client(app) as c: + resp = await c.get("/oauth/authorize", params=params) + assert resp.status_code == 200, resp.text + + def test_stop_does_not_wait_forever_on_open_requests(stack) -> None: # uvicorn waits for every connection to close by default; one stalled # request then held restarts until systemd SIGKILLed the process at 90 s. diff --git a/tests/test_oauth_refresh.py b/tests/test_oauth_refresh.py index 3f5f6c8..457f62e 100644 --- a/tests/test_oauth_refresh.py +++ b/tests/test_oauth_refresh.py @@ -117,6 +117,27 @@ def test_replayed_refresh_token_revokes_whole_family() -> None: store.refresh(second.refresh_token, "client-a") +def test_reuse_is_detected_after_a_short_refresh_ttl(monkeypatch: pytest.MonkeyPatch) -> None: + # With refresh_token_ttl below the family lifetime, a rotated-out token + # replayed after its own idle TTL must still count as reuse, not as an + # unknown token that leaves the family alive. + now = 1_000_000.0 + _clock(monkeypatch, now) + store = TokenStore(refresh_token_ttl=3600) + first = store.issue_oauth("client-a", refresh=True) + _clock(monkeypatch, now + 600) + second = store.refresh(first.refresh_token, "client-a") + + _clock(monkeypatch, now + 4200) # past R1's own 1 h TTL, family still alive + store._cleanup() + with pytest.raises(RefreshTokenError) as exc: + store.refresh(first.refresh_token, "client-a") + assert exc.value.reason == "reused" + assert store.validate(second.access_token) is None + with pytest.raises(RefreshTokenError): + store.refresh(second.refresh_token, "client-a") + + def test_reuse_does_not_touch_other_families() -> None: store = TokenStore() a = store.issue_oauth("client-a", refresh=True) diff --git a/tests/test_proxmox_client.py b/tests/test_proxmox_client.py index f8bf122..3338310 100644 --- a/tests/test_proxmox_client.py +++ b/tests/test_proxmox_client.py @@ -180,3 +180,48 @@ def worker() -> None: # The worker's connection survived the main thread's eviction. assert len(seen) == 2 assert seen[1] is seen[0] + + +class _Endpoint: + def __init__(self, exc: Exception) -> None: + self.exc = exc + self.calls = 0 + + def _raise(self, **kwargs: object) -> object: + self.calls += 1 + raise self.exc + + get = post = _raise + + +def _client_with(monkeypatch: pytest.MonkeyPatch, endpoint: _Endpoint) -> ProxmoxClient: + c = ProxmoxClient(_FakeConfig(["pve1"])) + monkeypatch.setattr(c, "_get_connection", lambda name: type("C", (), {"nodes": endpoint})()) + monkeypatch.setattr(client_mod, "_RETRY_BACKOFF_SECONDS", 0) + return c + + +def test_post_is_not_resent_after_a_read_timeout(monkeypatch: pytest.MonkeyPatch) -> None: + # The node may already be cloning: a second POST would start it twice. + from requests.exceptions import ReadTimeout + + ep = _Endpoint(ReadTimeout("read timed out")) + result = _client_with(monkeypatch, ep).post("pve1", "nodes") + assert ep.calls == 1 + assert "proxmox_get_tasks" in result["error"] + + +def test_post_is_retried_when_the_connection_never_opened(monkeypatch: pytest.MonkeyPatch) -> None: + from requests.exceptions import ConnectTimeout + + ep = _Endpoint(ConnectTimeout("connect timed out")) + _client_with(monkeypatch, ep).post("pve1", "nodes") + assert ep.calls == 2 + + +def test_get_is_retried_after_a_read_timeout(monkeypatch: pytest.MonkeyPatch) -> None: + from requests.exceptions import ReadTimeout + + ep = _Endpoint(ReadTimeout("read timed out")) + _client_with(monkeypatch, ep).get("pve1", "nodes") + assert ep.calls == 2 diff --git a/tests/test_tls_hostkeys.py b/tests/test_tls_hostkeys.py index 8e4c0e2..c086b9d 100644 --- a/tests/test_tls_hostkeys.py +++ b/tests/test_tls_hostkeys.py @@ -331,3 +331,28 @@ async def test_tofu_parallel_first_contacts_both_accept_the_same_key(tmp_path) - finally: server.close() await server.wait_closed() + + +async def test_tofu_keeps_ports_apart(tmp_path) -> None: + """A key pinned for port 22 says nothing about another port on that address. + + Port forwards on one NAT address usually lead to different machines; + asyncssh's matcher would retry ``[host]:port`` as plain ``host``. + """ + cfg = _load(tmp_path, _ssh()) + server = await _serve(asyncssh.generate_private_key("ssh-ed25519")) + try: + spec = cfg.get_ssh_host("h") + spec.host, spec.port = "127.0.0.1", server.sockets[0].getsockname()[1] + policy = host_key_policy(cfg, spec) + other = asyncssh.generate_private_key("ssh-ed25519").export_public_key("openssh") + kh = Path(policy.path) + kh.parent.mkdir(parents=True, exist_ok=True) + kh.write_text(f"127.0.0.1 {other.decode().strip()}\n") + + conn = await _connect_to_host(spec, policy) + conn.close() + assert kh.read_text().splitlines()[1].startswith(f"[127.0.0.1]:{spec.port} ") + finally: + server.close() + await server.wait_closed() diff --git a/tests/test_totp_lockout.py b/tests/test_totp_lockout.py index d43ad38..44c7f31 100644 --- a/tests/test_totp_lockout.py +++ b/tests/test_totp_lockout.py @@ -108,6 +108,41 @@ def test_success_does_not_reset_the_per_client_cap(clock) -> None: for _ in range(4): guard.record_failure("c1", f"203.0.113.{i}") guard.record_success("c1", OWNER) + # The owner's own source is exempt now, every other one stays locked. + assert not guard.locked("c1", OWNER) + assert guard.locked("c1", "203.0.113.200") + + +def test_known_good_source_survives_a_distributed_lockout(clock) -> None: + """Two attacker IPs at the per-source limit must not lock the owner out + of every TOTP path once the owner has logged in from their address.""" + guard = TotpLockout() + guard.record_success("c1", OWNER) + for ip in ("203.0.113.1", "203.0.113.2"): + for _ in range(5): + guard.record_failure("c1", ip) + assert not guard.locked("c1", OWNER) + assert guard.locked("c1", "203.0.113.3") + + +def test_known_good_source_still_has_its_own_cap(clock) -> None: + guard = TotpLockout() + guard.record_success("c1", OWNER) + for _ in range(5): + guard.record_failure("c1", OWNER) + assert guard.locked("c1", OWNER) + + +def test_known_good_expires(clock) -> None: + guard = TotpLockout(known_good_seconds=60) + guard.record_success("c1", OWNER) + for ip in ("203.0.113.1", "203.0.113.2"): + for _ in range(5): + guard.record_failure("c1", ip) + clock.now += 61 + for ip in ("203.0.113.4", "203.0.113.5"): + for _ in range(5): + guard.record_failure("c1", ip) assert guard.locked("c1", OWNER) @@ -139,11 +174,26 @@ def test_forwarded_header_without_trusted_proxies_warns_once(caplog, monkeypatch with caplog.at_level("WARNING", logger="beaconmcp"): assert ratelimit.client_ip(request) == "10.0.0.2" ratelimit.client_ip(request) - warnings = [r for r in caplog.records if "trusted_proxies is empty" in r.getMessage()] + warnings = [r for r in caplog.records if "not in server.trusted_proxies" in r.getMessage()] assert len(warnings) == 1 assert "10.0.0.2" in warnings[0].getMessage() +def test_forwarded_header_from_an_unlisted_proxy_warns(caplog, monkeypatch) -> None: + """``trusted_proxies: [cloudflare]`` doesn't cover a cloudflared tunnel, + which connects from its own LAN address: say so instead of staying quiet.""" + from types import SimpleNamespace + + monkeypatch.setattr(ratelimit, "_forwarding_warned", False) + request = SimpleNamespace( + client=SimpleNamespace(host="192.168.1.201"), + headers={"x-forwarded-for": ATTACKER}, + ) + with caplog.at_level("WARNING", logger="beaconmcp"): + assert ratelimit.client_ip(request, ("127.0.0.1", "173.245.48.0/20")) == "192.168.1.201" + assert [r for r in caplog.records if "192.168.1.201" in r.getMessage()] + + def test_no_warning_when_the_proxy_is_trusted(caplog, monkeypatch) -> None: from types import SimpleNamespace diff --git a/tests/test_updates.py b/tests/test_updates.py index 629e74f..02cbd8b 100644 --- a/tests/test_updates.py +++ b/tests/test_updates.py @@ -374,6 +374,7 @@ def test_apply_pulls_installs_validates_and_restarts( _advance(upstream, "feat: shiny") install = _install(checkout) install.service = "beaconmcp" + install.under_systemd = True result = updates.apply_update(install=install, restart_delay=3) @@ -609,7 +610,7 @@ def __init__(self, cmd, **kwargs): assert cmd[script_at + 1:] == [ "beaconmcp-update", "5", "beaconmcp", "http://127.0.0.1:8420/health", str(tmp_path), "a" * 40, "b" * 40, "/opt/venv/bin/python", - str(tmp_path / updates.STATUS_FILE), str(updates.HEALTH_WAIT_SECONDS), "10", + str(tmp_path / updates.STATUS_FILE), str(updates.HEALTH_WAIT_SECONDS), str(updates._SETTLE_SECONDS), "/opt/venv/bin/pip", ] @@ -721,6 +722,7 @@ def test_apply_hands_the_watchdog_both_revisions(checkout, upstream, monkeypatch _advance(upstream, "feat: thing") install = _install(checkout) install.service = "beaconmcp" + install.under_systemd = True result = updates.apply_update(install=install) @@ -799,6 +801,43 @@ def test_release_channel_follows_the_newest_version_tag(checkout, upstream, stub assert _head(checkout) == tagged, "must stop at the tag, not the branch tip" +def test_release_channel_survives_a_moved_tag(checkout, upstream, stub_side_effects): + _advance(upstream, "release v1.0.0") + _tag(upstream, "v1.0.0") + _run("git", "fetch", "--quiet", "--tags", "origin", cwd=checkout) + assert updates.check_for_update(install=_install(checkout), settings=RELEASE).release_tag == "v1.0.0" + # Upstream re-tags v1.0.0 to fix the release, then ships v1.1.0. A plain + # `fetch --tags` now exits 1 ("would clobber existing tag") forever. + _advance(upstream, "fix the release") + _run("git", "tag", "-f", "-a", "-m", "v1.0.0", "v1.0.0", cwd=upstream) + _advance(upstream, "release v1.1.0") + _tag(upstream, "v1.1.0") + + info = updates.check_for_update(install=_install(checkout), settings=RELEASE) + assert info.error is None, info.error + assert info.release_tag == "v1.1.0" + + +def test_require_signed_refuses_a_tag_object_under_another_name( + checkout, upstream, stub_side_effects, gpg_signer, +): + # A signed pre-release, republished under a release name. + _advance(upstream, "release v2.0.0-rc1") + _tag(upstream, "v2.0.0-rc1", signed=True) + obj = subprocess.run( + ["git", "rev-parse", "refs/tags/v2.0.0-rc1"], cwd=str(upstream), + capture_output=True, text=True, check=True, + ).stdout.strip() + _run("git", "update-ref", "refs/tags/v2.0.0", obj, cwd=upstream) + before = _head(checkout) + + result = updates.apply_update(install=_install(checkout), settings=SIGNED_RELEASE) + + assert result.ok is False + assert "v2.0.0-rc1" in result.steps[-1].detail + assert _head(checkout) == before + + def test_release_channel_without_tags_says_so(checkout, upstream, stub_side_effects): _advance(upstream, "feat: thing") info = updates.check_for_update(install=_install(checkout), settings=RELEASE) @@ -931,12 +970,21 @@ def _exe(path: Path, body: str) -> Path: return path -def _run_watchdog(tmp_path, checkout, from_sha, to_sha, probe_body): - """Run _WATCHDOG_SCRIPT with fake systemctl, pip and health probe.""" +def _run_watchdog( + tmp_path, checkout, from_sha, to_sha, probe_body, *, settle="0", show_body="echo 1; echo 0\n", +): + """Run _WATCHDOG_SCRIPT with fake systemctl, pip and health probe. + + ``systemctl show`` (MainPID, NRestarts) answers through ``show_body`` and + is left out of the call log; every other systemctl call is logged. + """ bindir = tmp_path / "bin" bindir.mkdir(exist_ok=True) log = tmp_path / "calls.log" - _exe(bindir / "systemctl", f'echo "systemctl $*" >> {log}\n') + _exe( + bindir / "systemctl", + f'if [ "$1" = show ]; then\n{show_body}exit 0\nfi\necho "systemctl $*" >> {log}\n', + ) pip = _exe(bindir / "fake-pip", f'echo "pip $* (in $(pwd))" >> {log}\n') # Stands in for "$py -c ": healthy or not, by exit code. probe = _exe(bindir / "fake-python", probe_body) @@ -944,7 +992,7 @@ def _run_watchdog(tmp_path, checkout, from_sha, to_sha, probe_body): proc = subprocess.run( ["/bin/sh", "-c", updates._WATCHDOG_SCRIPT, "beaconmcp-update", "0", "beaconmcp", "http://127.0.0.1:1/health", str(checkout), - from_sha, to_sha, str(probe), str(status), "3", "0", str(pip)], + from_sha, to_sha, str(probe), str(status), "3", settle, str(pip)], capture_output=True, text=True, timeout=60, env={**os.environ, "PATH": f"{bindir}:{os.environ['PATH']}"}, ) @@ -1015,6 +1063,52 @@ def test_watchdog_does_not_undo_a_newer_update(tmp_path, checkout, upstream): assert status["outcome"] == "superseded" +def test_watchdog_script_survives_systemd_expansion(): + # PID 1 expands ${VAR} and collapses $$ in a transient unit's ExecStart, + # script text included: ${10} once reached sh as an empty settle. + assert "${" not in updates._WATCHDOG_SCRIPT + assert "$$" not in updates._WATCHDOG_SCRIPT + + +def test_watchdog_rolls_back_a_crash_loop_that_answers_between_restarts( + tmp_path, checkout, upstream, +): + from_sha = _head(checkout) + _advance(upstream, "feat: dies a few seconds after each start") + _run("git", "pull", "--ff-only", cwd=checkout) + to_sha = _head(checkout) + + # /health answers every time it is probed, but on the new revision every + # probe meets a new MainPID: systemd keeps restarting a crashing process. + show = ( + f'if [ "$(git -C {checkout} rev-parse HEAD)" = "{to_sha}" ]; then ' + 'date +%s%N; else echo 1; fi; echo 0\n' + ) + proc, _, status = _run_watchdog( + tmp_path, checkout, from_sha, to_sha, "exit 0\n", settle="2", show_body=show, + ) + + assert proc.returncode == 1, proc.stdout + proc.stderr + assert status["outcome"] == "rolled_back" + assert _head(checkout) == from_sha + + +def test_no_restart_when_the_server_runs_outside_its_unit(checkout, upstream, stub_side_effects): + # The unit file exists, but this process was started by hand: restarting + # the unit would leave it answering /health for code it never loaded. + _advance(upstream, "feat: thing") + install = _install(checkout) + install.service = "beaconmcp" + install.under_systemd = False + + result = updates.apply_update(install=install, restart_delay=3) + + assert result.ok is True + assert result.restart_scheduled is False + assert "restart" not in [s.name for s in result.steps] + assert "Restart the server process" in result.message + + def test_check_reports_how_the_last_update_ended(checkout): updates._write_status(checkout, {"outcome": "rolled_back", "from": "a", "to": "b"}) info = updates.check_for_update(install=_install(checkout)) diff --git a/tests/test_upload_links.py b/tests/test_upload_links.py index d8ce996..5e60160 100644 --- a/tests/test_upload_links.py +++ b/tests/test_upload_links.py @@ -27,6 +27,8 @@ NODE = "pve1" CT = 201 +STOPPED_CT = 202 +VM = 300 # a QEMU VM reached over SSH class FakeMCP: @@ -41,7 +43,7 @@ def deco(fn): class FakeProxmox: - """Knows a single LXC container, ``CT``.""" + """Knows a running LXC container, ``CT``, and a stopped one, ``STOPPED_CT``.""" def __init__(self, tmp_path: Path, max_mb: int = 1, public_url: str | None = "https://beacon.example") -> None: self._config = SimpleNamespace( @@ -49,21 +51,38 @@ def __init__(self, tmp_path: Path, max_mb: int = 1, public_url: str | None = "ht features=SimpleNamespace(dashboard=SimpleNamespace(public_url=public_url)), ) + self.stopped = {STOPPED_CT} + def get(self, node: str, path: str, **_: Any) -> dict[str, Any]: - if path == f"nodes/{node}/lxc/{CT}/status/current": + for vmid in (CT, STOPPED_CT): + if path == f"nodes/{node}/lxc/{vmid}/status/current": + return {"status": "stopped" if vmid in self.stopped else "running"} + if path == f"nodes/{node}/qemu/{VM}/status/current": return {"status": "running"} return {"error": "not found"} class FakeSSH: - """Records node uploads and ``pct push`` targets.""" + """Records node uploads and ``pct push`` targets, and SFTP into ``VM``.""" def __init__(self) -> None: self.node_files: dict[str, bytes] = {} self.guest_files: dict[str, bytes] = {} + self.vm_files: dict[str, bytes] = {} + self.vm_dirs: set[str] = {"/opt/app"} + self.sha256sum_exit: int | None = None # force a sha256sum exit code + + def resolve(self, host: str) -> str: + if host != str(VM): + raise KeyError(host) + return "192.0.2.30" async def sftp_put(self, host: str, local_path: str, remote_path: str) -> None: - self.node_files[remote_path] = Path(local_path).read_bytes() + data = Path(local_path).read_bytes() + if host == str(VM): + self.vm_files[remote_path] = data + else: + self.node_files[remote_path] = data async def sftp_remove(self, host: str, remote_path: str) -> None: self.node_files.pop(remote_path, None) @@ -74,8 +93,15 @@ async def exec_command(self, host: str, command: str, timeout: int = 60) -> dict self.guest_files[argv[4]] = self.node_files[argv[3]] return {"exit_code": 0, "stdout": "", "stderr": ""} if argv[:2] == ["pct", "exec"] and argv[4] == "sha256sum": + if self.sha256sum_exit is not None: + return {"exit_code": self.sha256sum_exit, "stdout": "", "stderr": "sha256sum: failed"} data = self.guest_files[argv[5]] return {"exit_code": 0, "stdout": f"{hashlib.sha256(data).hexdigest()} {argv[5]}\n"} + if host == str(VM) and argv[:2] == ["test", "-d"]: + return {"exit_code": 0 if argv[2] in self.vm_dirs else 1, "stdout": ""} + if host == str(VM) and argv[0] == "sha256sum": + data = self.vm_files[argv[1]] + return {"exit_code": 0, "stdout": f"{hashlib.sha256(data).hexdigest()} {argv[1]}\n"} return {"exit_code": 1, "stderr": f"unexpected command: {command}"} @@ -92,7 +118,9 @@ def env(tmp_path: Path): ), ]) http = httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="https://beacon.example") - yield SimpleNamespace(tools=mcp.tools, ssh=ssh, http=http, staging=tmp_path / "staging") + yield SimpleNamespace( + tools=mcp.tools, ssh=ssh, http=http, staging=tmp_path / "staging", proxmox=proxmox, + ) transfers.register_pusher(None) @@ -202,6 +230,117 @@ async def test_multipart_upload_is_refused(env) -> None: assert resp.status_code == 415 +async def test_mint_refuses_a_stopped_guest(env) -> None: + res = await _mint(env, vmid=STOPPED_CT) + assert res["status"] == "error" + assert "stopped" in res["error"] + + +async def test_guest_stopped_after_mint_is_409(env) -> None: + link = await _mint(env) + env.proxmox.stopped.add(CT) + resp = await env.http.put(link["url"], content=b"data") + assert resp.status_code == 409 + assert "stopped" in resp.json()["error"] + assert env.ssh.guest_files == {} and env.ssh.node_files == {} + assert list(env.staging.iterdir()) == [] + + +async def test_failed_push_is_500_not_502(env, monkeypatch: pytest.MonkeyPatch) -> None: + # A 502 from the origin is replaced by Cloudflare's own error page, + # which would hide the reason from the uploader. + async def failing_exec(host: str, command: str, timeout: int = 60) -> dict[str, Any]: + return {"exit_code": 1, "stdout": "", "stderr": "no such directory"} + + monkeypatch.setattr(env.ssh, "exec_command", failing_exec) + link = await _mint(env) + resp = await env.http.put(link["url"], content=b"data") + assert resp.status_code == 500 + assert "no such directory" in resp.json()["error"] + + +async def test_node_copy_is_removed_after_a_failed_sftp(env, monkeypatch: pytest.MonkeyPatch) -> None: + async def partial_put(host: str, local_path: str, remote_path: str) -> None: + env.ssh.node_files[remote_path] = b"partial" + raise OSError("No space left on device") + + monkeypatch.setattr(env.ssh, "sftp_put", partial_put) + link = await _mint(env) + resp = await env.http.put(link["url"], content=b"data") + assert resp.status_code == 500 + assert env.ssh.node_files == {} + + +async def test_cancelled_upload_cleans_up_and_settles_the_link(env, monkeypatch: pytest.MonkeyPatch) -> None: + # A restart cancels in-flight requests (uvicorn's graceful-shutdown cap). + import asyncio + + started = asyncio.Event() + + async def hanging_exec(host: str, command: str, timeout: int = 60) -> dict[str, Any]: + started.set() + await asyncio.sleep(3600) + return {} + + events: list[dict[str, Any]] = [] + monkeypatch.setattr(transfers.audit, "emit", lambda event, **kw: events.append({"event": event, **kw})) + monkeypatch.setattr(env.ssh, "exec_command", hanging_exec) + link = await _mint(env) + task = asyncio.create_task(env.http.put(link["url"], content=b"data")) + await asyncio.wait_for(started.wait(), 5) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + await asyncio.sleep(0) # let the shielded cleanup run + + assert env.ssh.node_files == {} + assert list(env.staging.iterdir()) == [] + status = await env.tools["proxmox_upload_link"](link_id=link["link_id"]) + assert status["status"] == "error" + assert [e["event"] for e in events] == ["transfer.upload"] + + +@pytest.mark.parametrize("code, ok", [(127, True), (1, False)]) +async def test_checksum_failure_is_not_reported_as_unverified_success(env, code: int, ok: bool) -> None: + env.ssh.sha256sum_exit = code + link = await _mint(env) + resp = await env.http.put(link["url"], content=b"data") + body = resp.json() + if ok: # sha256sum missing: success, flagged as unverified + assert body["status"] == "success" and body["checksum_verified"] is False + else: # anything else: the file is not where asked + assert body["status"] == "error" and "Checksum check" in body["error"] + + +async def test_qemu_push_lands_at_dest(env) -> None: + link = await _mint(env, vmid=VM, dest="/opt/app/tool.bin") + resp = await env.http.put(link["url"], content=b"bin") + assert resp.json()["status"] == "success", resp.json() + assert env.ssh.vm_files == {"/opt/app/tool.bin": b"bin"} + + +async def test_qemu_push_refuses_a_directory_dest(env) -> None: + link = await _mint(env, vmid=VM, dest="/opt/app") + resp = await env.http.put(link["url"], content=b"bin") + assert resp.json()["status"] == "error" + assert "directory" in resp.json()["error"] + assert env.ssh.vm_files == {} + + +async def test_loopback_request_base_is_not_used_for_links(tmp_path: Path) -> None: + # The dashboard chat reaches the MCP endpoint over 127.0.0.1. + mcp = FakeMCP() + system.register_system_tools(mcp, FakeProxmox(tmp_path, public_url=None), FakeSSH()) + var = transfers.current_public_base.set("http://127.0.0.1:8420") + try: + res = await mcp.tools["proxmox_upload_link"](node=NODE, vmid=CT, dest="/x") + finally: + transfers.current_public_base.reset(var) + transfers.register_pusher(None) + assert res["url"].startswith(transfers.UPLOAD_PATH_PREFIX) + assert "warning" in res + + async def test_unknown_token_is_404(env) -> None: resp = await env.http.put(transfers.UPLOAD_PATH_PREFIX + "nope", content=b"x") assert resp.status_code == 404 @@ -242,7 +381,7 @@ async def fake_dns(host: str) -> list[str]: real = httpx.AsyncClient monkeypatch.setattr( system.httpx, "AsyncClient", - lambda **kw: real(transport=httpx.MockTransport(handler), **kw), + lambda **kw: real(**{**kw, "transport": httpx.MockTransport(handler)}), ) monkeypatch.setattr(system, "_host_addresses", fake_dns) return SimpleNamespace(files=files, requested=requested) @@ -297,6 +436,15 @@ async def test_fetch_url_allows_private_lan(env, remote) -> None: assert res["status"] == "success", res +async def test_fetch_url_refuses_a_stopped_guest_before_downloading(env, remote) -> None: + res = await env.tools["proxmox_fetch_url"]( + node=NODE, vmid=STOPPED_CT, url="https://files.example/tool.bin", dest="/x", + ) + assert res["status"] == "error" + assert "stopped" in res["error"] + assert remote.requested == [] + + async def test_fetch_url_checks_target_before_downloading(env, remote) -> None: res = await env.tools["proxmox_fetch_url"]( node=NODE, vmid=999, url="https://files.example/tool.bin", dest="/x", @@ -304,3 +452,95 @@ async def test_fetch_url_checks_target_before_downloading(env, remote) -> None: assert res["status"] == "error" assert "999" in res["error"] assert remote.requested == [] + + +# --- the connect-time guard (real sockets) ----------------------------------- + + +@pytest.fixture +def loopback_server(): + """A real HTTP server on 127.0.0.1 holding a "secret" only reachable locally.""" + import http.server + import threading + + class Handler(http.server.BaseHTTPRequestHandler): + def do_GET(self): # noqa: N802 + body = b"LOOPBACK-ONLY-SECRET " + self.headers.get("Host", "").encode() + self.send_response(200) + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, *args): + pass + + server = http.server.ThreadingHTTPServer(("127.0.0.1", 0), Handler) + threading.Thread(target=server.serve_forever, daemon=True).start() + yield server.server_address[1] + server.shutdown() + + +def _dns_sequence(monkeypatch: pytest.MonkeyPatch, answers: list) -> None: + """Each lookup of a ``.test`` name pops the next answer (an exception is + raised). Patched at socket.getaddrinfo, so the guard's lookups and the + connect's own lookup (if any) both see the sequence.""" + import socket + + real = socket.getaddrinfo + + def fake_getaddrinfo(host, port, *args, **kwargs): + name = host.decode() if isinstance(host, bytes) else str(host) # anyio passes IDNA bytes + if not name.endswith(".test"): + return real(host, port, *args, **kwargs) + answer = answers.pop(0) if len(answers) > 1 else answers[0] + if isinstance(answer, Exception): + raise socket.gaierror(str(answer)) + return [ + (socket.AF_INET, socket.SOCK_STREAM, 6, "", (ip, port or 0)) for ip in answer + ] + + monkeypatch.setattr(socket, "getaddrinfo", fake_getaddrinfo) + + +@pytest.mark.parametrize( + "answers", + [ + # DNS rebinding: public for the early check, loopback for the connect. + [["93.184.216.34"], ["127.0.0.1"]], + # Lookup error first: the early check lets it through, the connect + # must not. + [OSError("SERVFAIL"), ["127.0.0.1"]], + ], +) +async def test_fetch_url_connects_only_to_the_checked_address( + env, loopback_server, monkeypatch: pytest.MonkeyPatch, answers, +) -> None: + _dns_sequence(monkeypatch, answers) + res = await env.tools["proxmox_fetch_url"]( + node=NODE, vmid=CT, url=f"http://rebind.test:{loopback_server}/", dest="/x", + ) + assert res["status"] == "error", res + assert env.ssh.guest_files == {} + + +async def test_fetch_url_unresolvable_host_fails_closed(env, monkeypatch: pytest.MonkeyPatch) -> None: + _dns_sequence(monkeypatch, [OSError("NXDOMAIN")]) + res = await env.tools["proxmox_fetch_url"]( + node=NODE, vmid=CT, url="http://nowhere.test/", dest="/x", + ) + assert res["status"] == "error" + assert env.ssh.guest_files == {} + + +async def test_guarded_connect_keeps_the_hostname( + env, loopback_server, monkeypatch: pytest.MonkeyPatch, +) -> None: + # Pretend 127.0.0.1 is a public address to exercise the success path: + # the connection goes to the resolved IP, the Host header keeps the name. + _dns_sequence(monkeypatch, [["127.0.0.1"]]) + monkeypatch.setattr(system, "_is_local_address", lambda addr: False) + res = await env.tools["proxmox_fetch_url"]( + node=NODE, vmid=CT, url=f"http://files.test:{loopback_server}/f", dest="/f", + ) + assert res["status"] == "success", res + assert env.ssh.guest_files["/f"] == f"LOOPBACK-ONLY-SECRET files.test:{loopback_server}".encode()