# WebSocket Unauthenticated Endpoint Probe — Implementation Plan

> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.

**Goal:** Extend `ExposureScanCheck` to detect WebSocket endpoints that accept unauthenticated connections, producing a `high`-severity finding per path.

**Architecture:** After the existing TCP connect-scan loop, a new `_ws_probe()` method attempts a bare WebSocket upgrade (no credentials) at each of three default paths against every TCP-open port. A `101 Switching Protocols` response means the endpoint is unauthenticated — a `Finding(severity="high")` is appended. All other outcomes (connection refused, timeout, HTTP 4xx/5xx, protocol error) are silently ignored.

**Tech Stack:** Python 3.12, `asyncio`, `websockets>=12.0` (already a prod dependency), `pytest-asyncio`

---

## File Map

| File | Change |
|------|--------|
| `src/agentsec/audit/checks/exposure_scan.py` | Add `DEFAULT_WS_PATHS`; add `ws_paths`/`ws_probe_enabled` constructor params; refactor `run()` to collect `open_ports` first; add `_ws_probe()` |
| `tests/audit/checks/test_exposure_scan.py` | Add 4 new test cases + shared WS server fixture |

---

### Task 1: Tests for unauthenticated WS finding

**Files:**
- Modify: `tests/audit/checks/test_exposure_scan.py`

- [ ] **Step 1: Add WS server fixture and first failing test**

Append to `tests/audit/checks/test_exposure_scan.py`:

```python
import asyncio
from contextlib import asynccontextmanager

import websockets
import websockets.server


@asynccontextmanager
async def _ws_server(host: str = "127.0.0.1"):
    """Start a no-auth WS server on a random port; yield the port."""
    async def _handler(ws):
        await ws.wait_closed()

    async with websockets.server.serve(_handler, host, 0) as server:
        port = server.sockets[0].getsockname()[1]
        yield port


@pytest.mark.asyncio
async def test_ws_probe_unauthenticated_high_finding():
    async with _ws_server() as port:
        check = ExposureScanCheck(
            host="127.0.0.1",
            ports=[port],
            allowlist=[],
            ws_paths=["/ws"],
        )
        ctx = CheckContext(
            executor=None, redactor=Redactor(), ioc_dir=None,
            profile=LINUX, status=CheckStatus(),
        )
        findings = await check.run(ctx)

    ws_findings = [f for f in findings if f.evidence.get("transport") == "ws"]
    assert len(ws_findings) == 1
    assert ws_findings[0].severity == "high"
    assert ws_findings[0].evidence["path"] == "/ws"
    assert ws_findings[0].evidence["port"] == port
```

- [ ] **Step 2: Run test to confirm it fails**

```bash
cd /Users/roger/Work/AgentSec && .venv/bin/pytest tests/audit/checks/test_exposure_scan.py::test_ws_probe_unauthenticated_high_finding -v
```

Expected: `FAILED` — `TypeError` or `unexpected keyword argument 'ws_paths'`

---

### Task 2: Tests for no-finding cases

**Files:**
- Modify: `tests/audit/checks/test_exposure_scan.py`

- [ ] **Step 1: Add three more failing tests**

Append after the previous test:

```python
@pytest.mark.asyncio
async def test_ws_probe_no_finding_when_port_closed():
    check = ExposureScanCheck(
        host="127.0.0.1",
        ports=[1],          # privileged port, not bound
        allowlist=[],
        ws_paths=["/ws"],
    )
    ctx = CheckContext(
        executor=None, redactor=Redactor(), ioc_dir=None,
        profile=LINUX, status=CheckStatus(),
    )
    findings = await check.run(ctx)
    ws_findings = [f for f in findings if f.evidence.get("transport") == "ws"]
    assert ws_findings == []


@pytest.mark.asyncio
async def test_ws_probe_disabled():
    async with _ws_server() as port:
        check = ExposureScanCheck(
            host="127.0.0.1",
            ports=[port],
            allowlist=[],
            ws_paths=["/ws"],
            ws_probe_enabled=False,
        )
        ctx = CheckContext(
            executor=None, redactor=Redactor(), ioc_dir=None,
            profile=LINUX, status=CheckStatus(),
        )
        findings = await check.run(ctx)

    ws_findings = [f for f in findings if f.evidence.get("transport") == "ws"]
    assert ws_findings == []
    # TCP finding still present
    tcp_findings = [f for f in findings if f.evidence.get("transport") == "tcp"]
    assert len(tcp_findings) == 1


@pytest.mark.asyncio
async def test_ws_probe_no_finding_when_tcp_only():
    """Plain TCP socket (non-WS) — upgrade is refused, no WS finding."""
    s = _listen_on(0)   # reuse existing helper from the same file
    port = s.getsockname()[1]
    try:
        check = ExposureScanCheck(
            host="127.0.0.1",
            ports=[port],
            allowlist=[port],   # suppress TCP finding; we only care about WS
            ws_paths=["/ws"],
        )
        ctx = CheckContext(
            executor=None, redactor=Redactor(), ioc_dir=None,
            profile=LINUX, status=CheckStatus(),
        )
        findings = await check.run(ctx)
    finally:
        s.close()

    ws_findings = [f for f in findings if f.evidence.get("transport") == "ws"]
    assert ws_findings == []
```

- [ ] **Step 2: Run all four new tests to confirm they all fail**

```bash
.venv/bin/pytest tests/audit/checks/test_exposure_scan.py -k "ws_probe" -v
```

Expected: all 4 `FAILED` (constructor doesn't accept `ws_paths` yet)

---

### Task 3: Implement WS probe in `ExposureScanCheck`

**Files:**
- Modify: `src/agentsec/audit/checks/exposure_scan.py`

- [ ] **Step 1: Replace the entire file with the updated implementation**

```python
"""exposure_scan check (spec §3.2 #3).

TCP connect-scan from the evaluator side against `host` for an explicit list
of ports. Each LISTEN that is not on the allowlist becomes one Finding.

WebSocket probe: after TCP scan, each open port is probed with a bare WS
upgrade (no credentials) at each path in `ws_paths`. A 101 response means
the endpoint is unauthenticated — Finding(severity="high"). Per-port timeout
`timeout_sec` applies to both TCP connect and WS handshake.

`exposure_scan` does **not** use SSH; it observes the target's exposed
network surface from wherever the evaluator is running. Authorization is
already enforced at the orchestrator level by `Authorization.scope`
containing `server-audit` (which implies external network observation of
the consenting target).
"""

from __future__ import annotations

import asyncio
import socket
from typing import Iterable

from ..findings import Finding
from .base import Check, CheckContext

DEFAULT_PORTS: tuple[int, ...] = (18789, 18790, 8443, 80, 443, 22)
DEFAULT_ALLOWLIST: tuple[int, ...] = (18789, 22)
DEFAULT_WS_PATHS: tuple[str, ...] = ("/ws", "/v1/ws", "/")


class ExposureScanCheck(Check):
    id = "exposure_scan"

    def __init__(
        self,
        *,
        host: str,
        ports: Iterable[int] = DEFAULT_PORTS,
        allowlist: Iterable[int] = DEFAULT_ALLOWLIST,
        timeout_sec: float = 0.5,
        ws_paths: Iterable[str] = DEFAULT_WS_PATHS,
        ws_probe_enabled: bool = True,
    ):
        super().__init__()
        self._host = host
        self._ports = list(ports)
        self._allow = set(allowlist)
        self._timeout = timeout_sec
        self._ws_paths = list(ws_paths)
        self._ws_probe_enabled = ws_probe_enabled

    async def run(self, ctx: CheckContext) -> list[Finding]:
        open_ports = [p for p in self._ports if self._is_listening(p)]
        findings: list[Finding] = []

        for port in open_ports:
            if port in self._allow:
                continue
            findings.append(Finding(
                check=self.id,
                severity="medium",
                title=f"Unexpected listening port: {port}/tcp on {self._host}",
                description=(
                    f"Port {port}/tcp on {self._host} accepted a TCP connection "
                    f"from the evaluator and is not in the allowlist {sorted(self._allow)}."
                ),
                evidence={"host": self._host, "port": port, "transport": "tcp"},
                remediation=(
                    f"If port {port} is intentionally exposed, add it to the "
                    f"`exposure_scan.allowlist`. Otherwise close the port or "
                    f"firewall it from the evaluator's network position."
                ),
                threat_refs=["TI-OPENCLAW-EXPOSURE"],
            ))

        if self._ws_probe_enabled:
            for port in open_ports:
                for path in await self._ws_probe(port):
                    findings.append(Finding(
                        check=self.id,
                        severity="high",
                        title=f"Unauthenticated WebSocket endpoint: {self._host}:{port}{path}",
                        description=(
                            f"A WebSocket upgrade at {path} on {self._host}:{port} "
                            f"was accepted without credentials from the evaluator's "
                            f"network position."
                        ),
                        evidence={"host": self._host, "port": port, "path": path, "transport": "ws"},
                        remediation=(
                            "Require authentication (e.g. token query-parameter or "
                            "Authorization header) before accepting the WebSocket upgrade handshake."
                        ),
                        threat_refs=["TI-OPENCLAW-EXPOSURE"],
                    ))

        return findings

    def _is_listening(self, port: int) -> bool:
        try:
            with socket.create_connection((self._host, port), timeout=self._timeout):
                return True
        except (OSError, socket.timeout):
            return False

    async def _ws_probe(self, port: int) -> list[str]:
        """Return paths where unauthenticated WS upgrade was accepted (101)."""
        import websockets
        import websockets.exceptions

        accepted: list[str] = []
        for path in self._ws_paths:
            url = f"ws://{self._host}:{port}{path}"
            try:
                async with asyncio.timeout(self._timeout):
                    async with websockets.connect(url) as _ws:
                        accepted.append(path)
            except (OSError, TimeoutError, websockets.exceptions.WebSocketException):
                pass
        return accepted
```

- [ ] **Step 2: Run all four new tests — expect PASS**

```bash
.venv/bin/pytest tests/audit/checks/test_exposure_scan.py -k "ws_probe" -v
```

Expected output (all pass):
```
PASSED tests/audit/checks/test_exposure_scan.py::test_ws_probe_unauthenticated_high_finding
PASSED tests/audit/checks/test_exposure_scan.py::test_ws_probe_no_finding_when_port_closed
PASSED tests/audit/checks/test_exposure_scan.py::test_ws_probe_disabled
PASSED tests/audit/checks/test_exposure_scan.py::test_ws_probe_no_finding_when_tcp_only
```

- [ ] **Step 3: Run the full existing test file to verify no regressions**

```bash
.venv/bin/pytest tests/audit/checks/test_exposure_scan.py -v
```

Expected: all 7 tests pass (3 original + 4 new)

- [ ] **Step 4: Commit**

```bash
git add src/agentsec/audit/checks/exposure_scan.py tests/audit/checks/test_exposure_scan.py
git commit -m "feat(audit): extend exposure_scan with unauthenticated WS endpoint probe

Adds _ws_probe() to ExposureScanCheck: after TCP scan, each open port is
probed with a bare WS upgrade at /ws, /v1/ws, / (configurable). A 101
response — no credentials supplied — produces a high-severity finding.
ws_probe_enabled=False disables the probe without affecting TCP findings.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>"
```

---

### Task 4: Full test suite regression check

**Files:** (read-only run)

- [ ] **Step 1: Run the complete test suite**

```bash
.venv/bin/pytest tests/ -x -q
```

Expected: all tests pass, no failures. If any unrelated test fails, investigate before proceeding.

- [ ] **Step 2: Run ruff lint**

```bash
.venv/bin/ruff check src/ --fix
```

Expected: no errors. If auto-fixed, stage the changes:
```bash
git add -p src/ && git commit -m "chore: ruff auto-fix"
```

---

### Task 5: Update ROADMAP and CHANGELOG

**Files:**
- Modify: `ROADMAP.md`
- Modify: `CHANGELOG.md`

- [ ] **Step 1: Mark ROADMAP item done**

In `ROADMAP.md`, find:
```markdown
- [ ] WebSocket 端点探测（`exposure_scan` 扩展）
```
Replace with:
```markdown
- [x] WebSocket 端点探测（`exposure_scan` 扩展）— v0.5.1，2026-05-04
```

- [ ] **Step 2: Add CHANGELOG entry**

Open `CHANGELOG.md`. The file already has `## [Unreleased]` → `### Added`. Append one bullet to that existing list:

```markdown
- `ExposureScanCheck` now probes open ports for unauthenticated WebSocket endpoints.
  A bare WS upgrade at `/ws`, `/v1/ws`, `/` (configurable via `ws_paths`) that returns
  `101 Switching Protocols` produces a `high`-severity finding. The probe can be disabled
  with `ws_probe_enabled=False`. No findings are generated for auth-gated endpoints.
```

- [ ] **Step 3: Commit**

```bash
git add ROADMAP.md CHANGELOG.md
git commit -m "chore: mark WS exposure probe complete in ROADMAP; update CHANGELOG

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>"
```
