keywright-server 0.0.1__py3-none-any.whl

keywright_server/web/app.py

@@ -0,0 +1,154 @@
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import logging
+from collections.abc import AsyncIterator
+
+from fastapi import FastAPI
+from fastapi.responses import JSONResponse
+from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser
+from mcp.server.auth.provider import AccessToken
+from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
+from mcp.server.transport_security import TransportSecuritySettings
+from starlette.routing import Route
+
+from .. import __version__, approvals, llm_proxy, mail, nudges, tool_access
+from ..config import settings
+from ..db import init_db, session
+from ..mcp_server import build_server, resolve_agent
+from ..notify.manager import manager
+from ..settings_store import warm
+from . import (
+    admin,
+    discovery_ui,
+    friction_ui,
+    inventory_ui,
+    knowledge_ui,
+    mail_ui,
+    mcp_oauth_ui,
+    network_map_ui,
+    oidc,
+    usage_ui,
+)
+from .auth import password_is_set, set_password
+from .security import BrowserSecurityHeaders
+
+log = logging.getLogger("keywright")
+
+
+class MCPAuthASGI:
+    """Rejects MCP traffic that doesn't carry a valid agent bearer token, then hands off to the SDK."""
+
+    def __init__(self, session_manager: StreamableHTTPSessionManager):
+        self.sm = session_manager
+
+    async def __call__(self, scope, receive, send):
+        if scope["type"] != "http":
+            await send({"type": "http.response.start", "status": 404, "headers": []})
+            await send({"type": "http.response.body", "body": b""})
+            return
+        headers = {k.decode().lower(): v.decode() for k, v in scope.get("headers", [])}
+        async with session() as db:
+            agent = await resolve_agent(db, headers)
+        if agent is None:
+            # A JSON-RPC error body: MCP SDK clients surface its message instead of a generic
+            # "server returned an error", so an agent learns it is the token that is wrong.
+            resp = JSONResponse({"jsonrpc": "2.0", "id": None,
+                                 "error": {"code": -32001, "message": "unauthorized: Authorization: Bearer <agent token>",
+                                           "data": {"error": "unauthorized"}}},
+                                status_code=401, headers={"WWW-Authenticate": "Bearer"})
+            await resp(scope, receive, send)
+            return
+        # Give the SDK a principal so GET/POST/DELETE sessions cannot be shared across agents/tokens.
+        scope = dict(scope)
+        scope["user"] = AuthenticatedUser(AccessToken(
+            token=agent.token_hash, client_id=agent.id + ":" + agent.token_hash, scopes=[],
+        ))
+        await self.sm.handle_request(scope, receive, send)
+
+
+async def _janitor() -> None:
+    while True:
+        await asyncio.sleep(60)
+        try:
+            async with session() as db:
+                n = await approvals.expire_stale(db)
+                if n:
+                    log.info("expired %d stale requests", n)
+                n = await nudges.expire_stale(db)
+                if n:
+                    log.info("cleared %d nudges nobody answered", n)
+        except Exception:
+            log.exception("janitor failed")
+
+
+TOOL_WATCH_S = 3600
+
+
+async def _tool_watch() -> None:
+    """Hourly: what each MCP server offers, against the baseline the owner accepted."""
+    await asyncio.sleep(30)
+    while True:
+        try:
+            async with session() as db:
+                await tool_access.watch_once(db)
+        except Exception:
+            log.exception("tool watch failed")
+        await asyncio.sleep(TOOL_WATCH_S)
+
+
+def create_app() -> FastAPI:
+    server = build_server()
+    sm = StreamableHTTPSessionManager(
+        app=server,
+        security_settings=TransportSecuritySettings(enable_dns_rebinding_protection=False),
+    )
+
+    @contextlib.asynccontextmanager
+    async def lifespan(app: FastAPI) -> AsyncIterator[None]:
+        logging.basicConfig(level=settings.log_level, format="%(asctime)s %(levelname)s %(name)s: %(message)s")
+        await init_db()
+        async with session() as db:
+            await warm(db)
+            if settings.auth_mode != "oidc" and settings.admin_password and not password_is_set():
+                await set_password(db, settings.admin_password)
+                log.info("admin password seeded from KW_ADMIN_PASSWORD")
+        await manager.reload()
+        janitor = asyncio.create_task(_janitor(), name="janitor")
+        watch = asyncio.create_task(_tool_watch(), name="tool-watch")
+        mail_worker = asyncio.create_task(mail.worker(), name="mail-inbox")
+        async with sm.run():
+            log.info("Keywright listening; MCP at /mcp, admin at /ui")
+            try:
+                yield
+            finally:
+                janitor.cancel()
+                watch.cancel()
+                mail_worker.cancel()
+                with contextlib.suppress(asyncio.CancelledError):
+                    await mail_worker
+                await manager.stop()
+
+    app = FastAPI(title="Keywright", version=__version__, lifespan=lifespan, docs_url=None, redoc_url=None)
+    oidc.configure(app)
+    app.add_middleware(BrowserSecurityHeaders)
+    app.include_router(admin.router)
+    app.include_router(oidc.router)
+    app.include_router(mcp_oauth_ui.router)
+    app.include_router(inventory_ui.router)
+    app.include_router(knowledge_ui.router)
+    app.include_router(network_map_ui.router)
+    app.include_router(discovery_ui.router)
+    app.include_router(usage_ui.router)
+    app.include_router(friction_ui.router)
+    app.include_router(mail_ui.router)
+    app.include_router(admin.ws_router)
+    app.include_router(llm_proxy.router)
+    app.router.routes.append(Route("/mcp", endpoint=MCPAuthASGI(sm), methods=["GET", "POST", "DELETE"]))
+
+    @app.get("/healthz")
+    async def healthz():
+        return {"ok": True, "notifier": manager.driver_id, "version": __version__}
+
+    return app

keywright_server/web/auth.py

@@ -0,0 +1,144 @@
+"""Admin authentication: local password or an OIDC browser session."""
+
+from __future__ import annotations
+
+import hashlib
+import hmac
+import secrets
+import time
+from collections import OrderedDict, deque
+
+from fastapi import HTTPException, Request
+from fastapi.responses import RedirectResponse
+from itsdangerous import BadSignature, TimestampSigner
+
+from ..config import settings
+from ..settings_store import cached, put
+from .security import safe_redirect
+
+COOKIE = "kw_admin"
+MAX_AGE = 30 * 86400
+_signer: TimestampSigner | None = None
+
+
+def signer() -> TimestampSigner:
+    global _signer
+    if _signer is None:
+        _signer = TimestampSigner(settings.secret_key(), salt="admin-session")
+    return _signer
+
+
+def hash_password(pw: str, salt: str | None = None) -> str:
+    salt = salt or secrets.token_hex(16)
+    dk = hashlib.pbkdf2_hmac("sha256", pw.encode(), salt.encode(), 200_000)
+    return f"{salt}${dk.hex()}"
+
+
+def verify_password(pw: str, stored: str) -> bool:
+    try:
+        salt, _ = stored.split("$", 1)
+    except ValueError:
+        return False
+    return hmac.compare_digest(hash_password(pw, salt), stored)
+
+
+def password_is_set() -> bool:
+    return bool((cached("admin") or {}).get("password_hash"))
+
+
+async def set_password(db, pw: str) -> None:
+    await put(db, "admin", {"password_hash": hash_password(pw)})
+
+
+def check_password(pw: str) -> bool:
+    if len(pw) > 1024:
+        return False
+    stored = (cached("admin") or {}).get("password_hash", "")
+    return bool(stored) and verify_password(pw, stored)
+
+
+def _session_payload() -> bytes:
+    # A password reset invalidates previously issued cookies, including across restarts.
+    stored = (cached("admin") or {}).get("password_hash", "")
+    return b"admin:" + hashlib.sha256(stored.encode()).hexdigest().encode()
+
+
+def is_logged_in(request: Request) -> bool:
+    # Legacy password cookie validation stays separate from server-side OIDC sessions.
+    if settings.auth_mode == "oidc":
+        return False
+    raw = request.cookies.get(COOKIE)
+    if not raw or not password_is_set():
+        return False
+    try:
+        return hmac.compare_digest(signer().unsign(raw, max_age=MAX_AGE), _session_payload())
+    except BadSignature:
+        return False
+
+
+def oidc_enabled() -> bool:
+    return settings.auth_mode in ("both", "oidc")
+
+
+async def authenticated(request) -> bool:
+    from .oidc import session_identity
+
+    identity = await session_identity(request) if oidc_enabled() else None
+    if identity is not None:
+        request.state.admin_identity = identity
+        return True
+    if is_logged_in(request):
+        request.state.admin_identity = {"method": "password"}
+        return True
+    return False
+
+
+def login_response(url: str) -> RedirectResponse:
+    resp = RedirectResponse(safe_redirect(url), status_code=303)
+    resp.set_cookie(COOKIE, signer().sign(_session_payload()).decode(), max_age=MAX_AGE, httponly=True, samesite="lax")
+    return resp
+
+
+def logout_response(url: str) -> RedirectResponse:
+    resp = RedirectResponse(safe_redirect(url), status_code=303)
+    resp.delete_cookie(COOKIE)
+    return resp
+
+
+class RequireAdmin:
+    """FastAPI dependency: redirect to setup/login when needed."""
+
+    async def __call__(self, request: Request) -> None:
+        if not password_is_set() and not oidc_enabled():
+            raise HTTPException(status_code=303, headers={"Location": "/ui/setup"})
+        if not await authenticated(request):
+            raise HTTPException(status_code=303, headers={"Location": f"/ui/login?next={request.url.path}"})
+
+
+require_admin = RequireAdmin()
+
+
+# Bounded per-process throttles. Requests never trust client-supplied forwarding headers here.
+_login_attempts: OrderedDict[str, deque[float]] = OrderedDict()
+_global_attempts: deque[float] = deque()
+
+
+def limit_login(request: Request) -> str:
+    address = request.client.host if request.client else "unknown"
+    at = time.monotonic()
+    attempts = _login_attempts.setdefault(address, deque())
+    _login_attempts.move_to_end(address)
+    while len(_login_attempts) > 4096:
+        _login_attempts.popitem(last=False)
+    for queue in (attempts, _global_attempts):
+        while queue and queue[0] <= at - 60:
+            queue.popleft()
+    if len(attempts) >= 8 or len(_global_attempts) >= 100:
+        raise HTTPException(429, "Too many login attempts; try again in a minute", headers={"Retry-After": "60"})
+    attempts.append(at)
+    _global_attempts.append(at)
+    return address
+
+
+def login_succeeded(address: str) -> None:
+    _login_attempts.pop(address, None)

keywright_server/web/discovery_ui.py

@@ -0,0 +1,112 @@
+"""Admin-only discovery previews; all source reads are audited."""
+from datetime import timedelta
+from urllib.parse import urlencode
+
+from fastapi import APIRouter, Depends, Request
+from sqlalchemy import select
+from sqlalchemy.exc import IntegrityError
+
+from .. import discovery, inventory, reconciliation
+from ..db import session
+from ..models import DiscoveryLink, DiscoveryObservation, DiscoveryRun, now
+from . import auth
+from .admin import redirect, render
+from .security import require_same_origin
+
+router = APIRouter(dependencies=[Depends(require_same_origin), Depends(auth.require_admin)])
+
+
+@router.get('/ui/inventory/discovery')
+async def discovery_page(request: Request, source: str = '', q: str = '', page: int = 1, state: str = ''):
+    async with session() as db:
+        data = await inventory.snapshot(db)
+        runs = list((await db.scalars(select(DiscoveryRun).order_by(DiscoveryRun.started_at.desc()))).all())
+        observations = list((await db.scalars(select(DiscoveryObservation))).all())
+        links = {r.observation_id: r for r in (await db.scalars(select(DiscoveryLink))).all()}
+    latest, succeeded = {}, {}
+    for run in runs:
+        key = (run.resource_id, run.provider)
+        latest.setdefault(key, run)
+        if run.status == 'succeeded':
+            succeeded.setdefault(key, run)
+    by_identity = {(r.resource_id, r.provider, r.external_id): r for r in observations}
+    def review_state(record):
+        link = links.get(record.id)
+        return ('changed' if link.reviewed_data != record.data else 'reviewed') if link else 'new'
+    filtered = [r for r in observations if (not source or r.resource_id == source) and q.casefold() in
+                ' '.join([r.data['name'], r.data['mac'], *r.data['addresses'], r.data['site']]).casefold()
+                and (state not in ('new', 'changed', 'reviewed') or review_state(r) == state)]
+    filtered.sort(key=lambda r: (r.data['name'].casefold(), r.resource_id, r.external_id))
+    total_pages = max(1, (len(filtered) + 99) // 100)
+    page = min(max(1, page), total_pages)
+    rows = []
+    for record in filtered[(page - 1) * 100:page * 100]:
+        link = links.get(record.id)
+        baseline = link.reviewed_data if link else record.previous
+        key = (record.resource_id, record.provider)
+        success = succeeded.get(key)
+        parent = by_identity.get((*key, record.data['parent']))
+        attachment = by_identity.get((*key, record.data['attachment']))
+        labels = {'name': 'Name', 'kind': 'Device type', 'site': 'Site', 'mac': 'MAC address',
+                  'addresses': 'IP addresses', 'medium': 'Connection', 'attachment': 'Connected to',
+                  'port': 'Port / SSID', 'parent': 'Hosted on', 'status': 'Source state',
+                  'device_type': 'Network device type', 'model': 'Model', 'uplink_note': 'Uplink details'}
+        def display(field, value, source_key=key):
+            if field in ('parent', 'attachment') and value:
+                target = by_identity.get((*source_key, value))
+                return target.data['name'] if target else value
+            if field == 'medium':
+                return inventory.MEDIA.get(value, 'Unknown')
+            if field == 'kind':
+                return inventory.KINDS.get(value, 'Unknown')
+            return (', '.join(value) if isinstance(value, list) else value) or 'Not reported'
+        rows.append({'record': record, 'matches': discovery.matches(record, data), 'link': link,
+                     'review_state': ('Needs review' if link.reviewed_data != record.data else 'Reviewed') if link else 'New',
+                     'parent': parent.data['name'] if parent else record.data['parent'],
+                     'attachment': attachment.data['name'] if attachment else record.data['attachment'],
+                     'missing': success is not None and record.run_id != success.id,
+                     'stale': record.last_seen < now() - timedelta(hours=24),
+                     'changes': [{'label': labels[k], 'before': display(k, baseline.get(k)),
+                                  'after': display(k, v)} for k, v in record.data.items()
+                                 if baseline is not None and baseline.get(k) != v]})
+    return render(request, 'inventory_discovery.html', **data, rows=rows, latest=list(latest.values()),
+                  kinds=inventory.KINDS, media=inventory.MEDIA, providers=discovery.PROVIDERS,
+                  source=source, q=q, state=state, page=page, total_pages=total_pages, total_count=len(filtered),
+                  previous_query=urlencode({'source': source, 'q': q, 'page': page - 1, 'state': state}),
+                  next_query=urlencode({'source': source, 'q': q, 'page': page + 1, 'state': state}),
+                  interrupted_before=now() - timedelta(minutes=5))
+
+
+@router.post('/ui/inventory/discovery/refresh')
+async def discovery_refresh(request: Request):
+    form = await request.form()
+    try:
+        resource_id = inventory.value(form, 'resource_id', required=True)
+        provider = inventory.value(form, 'provider', required=True)
+        run = await discovery.refresh(resource_id, provider)
+    except (discovery.DiscoveryError, inventory.InventoryError) as exc:
+        return redirect('/ui/inventory/discovery', str(exc))
+    message = f'Refreshed {run.count} observations. Inventory records are unchanged.' if run.status == 'succeeded' else run.error
+    return redirect('/ui/inventory/discovery', message)
+
+
+@router.get('/ui/inventory/discovery/{observation_id}/review')
+async def observation_review(request: Request, observation_id: str, target_id: str = '', interface_id: str = ''):
+    try:
+        async with session() as db:
+            ctx = await reconciliation.review(db, observation_id, target_id, interface_id)
+    except (discovery.DiscoveryError, inventory.InventoryError) as exc:
+        return redirect('/ui/inventory/discovery', str(exc))
+    return render(request, 'inventory_discovery_review.html', **ctx)
+
+
+@router.post('/ui/inventory/discovery/{observation_id}/review')
+async def observation_apply(request: Request, observation_id: str):
+    form = await request.form()
+    try:
+        asset_id = await reconciliation.apply(observation_id, form)
+    except (discovery.DiscoveryError, inventory.InventoryError, IntegrityError) as exc:
+        message = str(exc) if not isinstance(exc, IntegrityError) else 'A conflicting record was saved. Review again.'
+        query = urlencode({'target_id': form.get('target_id', ''), 'interface_id': form.get('interface_id', '')})
+        return redirect(f'/ui/inventory/discovery/{observation_id}/review?{query}', message)
+    return redirect(f'/ui/inventory/assets/{asset_id}', 'Discovery review saved. Only the selected inventory changes were applied.')

keywright_server/web/friction_ui.py

@@ -0,0 +1,92 @@
+"""Admin report and explicit review/save flow for constrained friction suggestions."""
+from datetime import UTC, date, datetime, time, timedelta
+from urllib.parse import urlencode
+
+from fastapi import APIRouter, Depends, Form, Request
+from itsdangerous import BadSignature, SignatureExpired, URLSafeTimedSerializer
+from sqlalchemy import text
+
+from .. import friction, policy_preview
+from ..config import settings
+from ..db import session
+from ..models import now
+from . import auth
+from .admin import redirect, render
+from .security import require_same_origin
+
+router = APIRouter(dependencies=[Depends(require_same_origin), Depends(auth.require_admin)])
+
+
+def dates(start, end):
+    today = now().date()
+    first = date.fromisoformat(start) if start else today - timedelta(days=29)
+    last = date.fromisoformat(end) if end else today
+    if first > last or last > today or (last - first).days >= 90:
+        raise ValueError("Choose 1 to 90 days, ending no later than today.")
+    return (datetime.combine(first, time.min, UTC),
+            datetime.combine(last + timedelta(days=1), time.min, UTC))
+
+
+def signer():
+    return URLSafeTimedSerializer(settings.secret_key(), salt="friction-rule-review")
+
+
+@router.get("/ui/friction")
+async def friction_page(request: Request, start: str = "", end: str = "", agent: str = ""):
+    error = None
+    try:
+        since, until = dates(start, end)
+    except (ValueError, OverflowError):
+        since, until = dates("", "")
+        error = "Choose a valid range of 1 to 90 days, ending no later than today. Showing the last 30 days."
+    filters = {"start": since.date().isoformat(), "end": (until - timedelta(days=1)).date().isoformat(),
+               "agent": agent}
+    async with session() as db:
+        report = await friction.report(db, since, until, agent_id=agent)
+    report.pop("proposals")  # Raw exact constraints are never passed to the report template.
+    return render(request, "friction.html", **report, filters=filters, error=error,
+                  review_query=urlencode(filters))
+
+
+@router.get("/ui/friction/review/{rid}")
+async def review(request: Request, rid: str, start: str = "", end: str = "", agent: str = ""):
+    try:
+        since, until = dates(start, end)
+    except (ValueError, OverflowError):
+        return redirect("/ui/friction", "Choose a valid date range before reviewing a suggestion.")
+    filters = {"start": since.date().isoformat(), "end": (until - timedelta(days=1)).date().isoformat(),
+               "agent": agent}
+    async with session() as db:
+        report = await friction.report(db, since, until, agent_id=agent)
+        proposal = report["proposals"].get(rid)
+        if proposal is None:
+            return redirect("/ui/friction?" + urlencode(filters), "This suggestion is no longer eligible. Refresh the report.")
+        rules = proposal["rules"]
+        preview = await policy_preview.preview(db, rules[0], additional=tuple(rules[1:]))
+    token = signer().dumps({"request": rid, "fingerprint": proposal["fingerprint"], "filters": filters})
+    public = {k: proposal[k] for k in ("request_id", "agent", "resource", "tool", "support", "examples")}
+    return render(request, "friction_review.html", proposal=public, preview=preview, token=token,
+                  rule_count=len(rules), back="/ui/friction?" + urlencode(filters))
+
+
+@router.post("/ui/friction/review/{rid}")
+async def save(rid: str, token: str = Form(...)):
+    try:
+        reviewed = signer().loads(token, max_age=600)
+        if reviewed["request"] != rid:
+            raise ValueError
+        filters = reviewed["filters"]
+        since, until = dates(filters["start"], filters["end"])
+    except (BadSignature, SignatureExpired, ValueError, TypeError, KeyError, OverflowError):
+        return redirect("/ui/friction", "Review expired or invalid. Preview the suggestion again before saving.")
+    async with session() as db:
+        # Serialize SQLite writers while revalidating support, targets, hazards and standing policy.
+        if db.bind.dialect.name == "sqlite":
+            await db.execute(text("BEGIN IMMEDIATE"))
+        report = await friction.report(db, since, until, agent_id=filters["agent"])
+        proposal = report["proposals"].get(rid)
+        if proposal is None or proposal["fingerprint"] != reviewed["fingerprint"]:
+            return redirect("/ui/friction?" + urlencode(filters),
+                            "Evidence or access settings changed. Review a fresh preview before saving.")
+        added = await friction.save_proposal(db, proposal)
+    return redirect("/ui/policies", f"Added {added} exact-operation rule(s) for {proposal['agent']}.")