pushmetrics-cli 0.1.0__py3-none-any.whl

pm/auth/session.py

@@ -0,0 +1,440 @@
+"""PMSession — authenticated requests.Session wrapper with silent recovery."""
+from __future__ import annotations
+
+import http.cookiejar as cookiejar
+import ipaddress
+import logging
+import os
+import time
+
+from urllib.parse import urlsplit
+
+import click
+import requests
+
+from pm.auth import device_flow, exchange, jwt_cache, store
+from pm.auth import token as token_auth
+from pm.config import Profile
+from pm.log import log_request, log_response
+
+logger = logging.getLogger(__name__)
+
+
+# pm-consolidation follow-up — when the prefect agent invokes the pm
+# subprocess, it sets PM_AGENT_UID so the API can scope block.list /
+# block.get to the agent's assigned blocks (see
+# superset/pushmetrics/security/agent_scope.py for the backend half).
+# Read once at session construction; tagged onto every request below.
+_AGENT_UID_HEADER = "X-Agent-Uid"
+
+
+# Anything with whitespace or control characters must never reach requests
+# (InvalidHeader echoes the value). One pattern, shared with the JWT cache.
+_JWT_SHAPE_RE = jwt_cache.JWT_SHAPE_RE
+
+JWT_SHAPE_HINT = (
+    "PM_API_JWT is malformed: it must be a single token with no spaces or "
+    "line breaks."
+)
+
+
+class InsecureApiUrlError(click.ClickException):
+    """A bearer credential would travel in cleartext to a non-loopback host."""
+
+
+# Harness-only opt-in: the prefect pod reaches the app over the in-cluster
+# plain-http Service (TLS terminates at the ingress), so cli_adapter.run_pm
+# sets this for the delegation subprocess. It is honoured ONLY in delegation
+# mode (PM_API_JWT); a PAT is never sent in cleartext, whatever the env says.
+ALLOW_PLAIN_HTTP_ENV_VAR = "PM_ALLOW_PLAIN_HTTP"
+
+
+def _plain_http_opted_in() -> bool:
+    return (os.environ.get(ALLOW_PLAIN_HTTP_ENV_VAR) or "").strip() == "1"
+
+
+def _refuse_bearer_over_plain_http(api_url: str, *, is_delegation: bool = False) -> None:
+    """Token and delegation mode attach `Authorization: Bearer` to the very
+    first request, before any server-side 301 to https could help. Loopback
+    http is the local and agent profiles' legitimate case (security S6);
+    in-cluster http is the harness's, behind an explicit opt-in."""
+    parts = urlsplit(api_url)
+    if parts.scheme != "http" or _is_loopback_host(parts.hostname or ""):
+        return
+    if is_delegation and _plain_http_opted_in():
+        logger.debug("plain http to %s allowed: delegation mode with %s=1", parts.hostname, ALLOW_PLAIN_HTTP_ENV_VAR)
+        return
+    raise InsecureApiUrlError(
+        f"Refusing to send a bearer credential over plain http to "
+        f"{parts.hostname}. Use https:// in --api-url, PM_API_URL, or "
+        f"`pm config set <profile> api_url https://...`."
+    )
+
+
+# A 502/503/504 on a read is almost always one bad moment at the ingress
+# (production: one page of a 14-page `pm report list` 502'd, the next run of
+# the same pages passed). GETs are safe to repeat, so retry them a few times
+# before failing the whole command. Writes are never retried.
+_RETRY_STATUSES = frozenset({502, 503, 504})
+_RETRY_BACKOFF_SECONDS = (0.5, 1.0, 2.0)
+_sleep = time.sleep  # patched in tests
+
+
+class PMConnectionError(click.ClickException):
+    """The API could not be reached. The transport error's own text is never
+    shown: it can embed the request, headers included."""
+
+
+class NotAuthenticatedError(click.ClickException):
+    def __init__(
+        self,
+        message: str = (
+            "Not logged in. Run `pm login --token` with a personal access "
+            "token, or set PM_API_TOKEN."
+        ),
+    ):
+        super().__init__(message)
+
+
+def _is_loopback_host(host: str) -> bool:
+    host = (host or "").lower().rstrip(".")
+    if host == "localhost":
+        return True
+    try:
+        return ipaddress.ip_address(host).is_loopback
+    except ValueError:
+        return False
+
+
+def _strip_secure_for_local_http(jar, api_url: str) -> None:
+    """Clear the Secure flag on cookies bound for a loopback http backend.
+
+    Browsers treat http://localhost and http://127.0.0.1 as trustworthy
+    origins, so a backend running with SESSION_COOKIE_SECURE=True (prod
+    parity) still works in the local web app. Python's cookie handling has no
+    such carve-out — and requests rebuilds the jar with the default policy on
+    every send (Session.prepare_request -> merge_cookies), so a custom policy
+    can't override it either. Without this, the Secure ``superset_session``
+    cookie is silently dropped and every CLI request against a local backend
+    goes out anonymous, 401ing on the first protected endpoint. Non-loopback
+    http targets keep browser semantics: Secure cookies stay unsent.
+
+    Only cookies whose own domain is loopback are touched: the jar is written
+    back to the credential store by ``persist()``, and a cookie for any other
+    host must never be de-secured durably. http.cookiejar stores unqualified
+    request hosts with a ``.local`` suffix (``localhost.local``), so that
+    suffix is peeled before the loopback check.
+    """
+    parts = urlsplit(api_url)
+    if parts.scheme != "http" or not _is_loopback_host(parts.hostname or ""):
+        return
+    for cookie in jar:
+        domain = (cookie.domain or "").lstrip(".")
+        if domain.endswith(".local"):
+            domain = domain[: -len(".local")]
+        if _is_loopback_host(domain):
+            cookie.secure = False
+
+
+class PMSession:
+    """Authenticated session against the PushMetrics API.
+
+    Three auth modes:
+
+    - **Cookie mode**: the Auth0 device-flow session cookie (carries
+      workspace_id for Postgres RLS). On 401, a single recovery is attempted via
+      the stored Auth0 refresh token. Internal-team interactive use; the
+      lowest-priority mode (used only when no PAT / delegation JWT is present).
+    - **Token mode** (``token=...``): a raw Personal Access Token is exchanged for
+      a short-lived workspace-scoped JWT, sent as ``Authorization: Bearer``. No
+      cookie jar and no Auth0 path. On 401 the PAT is re-exchanged once.
+    - **Delegation mode** (``static_jwt=...``): a pre-minted workspace-scoped
+      JWT (the agent harness sets ``PM_API_JWT``) used as-is. No exchange and
+      no refresh authority - a 401 is final.
+
+    ``from_profile`` resolution order: ``PM_API_JWT`` > PAT (``PM_API_TOKEN``
+    env or stored ``pm login --token`` token) > stored cookie credentials.
+    """
+
+    def __init__(
+        self,
+        profile: Profile,
+        creds: store.StoredCredentials | None = None,
+        *,
+        token: str | None = None,
+        static_jwt: str | None = None,
+    ):
+        self.profile = profile
+        self.creds = creds
+        self._session = requests.Session()
+        # Token-mode state. _token is the raw PAT; _jwt is the cached
+        # short-lived exchange result, refreshed shortly before expiry.
+        self._token = token
+        self._jwt: str | None = None
+        self._jwt_expires_at = 0.0
+        # Static-bearer mode (agent delegation): a pre-minted workspace-scoped
+        # JWT used as-is. No exchange, no refresh authority - a 401 is final.
+        # Normalize once: a blank/whitespace static_jwt is treated as unset so
+        # it can't send an empty ``Authorization: Bearer`` nor suppress the
+        # cookie jar / other auth modes.
+        self._static_jwt = (static_jwt or "").strip() or None
+        if token is not None or self._static_jwt is not None:
+            _refuse_bearer_over_plain_http(
+                profile.api_url, is_delegation=self._static_jwt is not None
+            )
+        if self._static_jwt is not None:
+            # Delegation invariant: the pre-minted JWT is the WHOLE session.
+            # Freeze the jar so a Flask-re-issued ``superset_session`` (or any
+            # other Set-Cookie) is never stored and never rides along on the
+            # next request of a multi-request command — auth and workspace
+            # scope must come from the Bearer header alone. Same immutability
+            # idea as the env-cookie _RejectServerCookiesPolicy this branch
+            # removed with the cookie delegation path.
+            self._session.cookies.set_policy(
+                cookiejar.DefaultCookiePolicy(allowed_domains=[])
+            )
+        if token is None and self._static_jwt is None:
+            if creds is None:
+                raise ValueError(
+                    "PMSession needs credentials: pass `creds` (cookie mode), "
+                    "`token` (bearer mode), or `static_jwt` (delegation mode)."
+                )
+            jar = creds.cookies()
+            _strip_secure_for_local_http(jar, profile.api_url)
+            self._session.cookies = jar
+        # Read PM_AGENT_UID once at construction. Empty string treated
+        # same as unset so misconfigured callers don't accidentally
+        # send "X-Agent-Uid: " to the API.
+        self._agent_uid = (os.environ.get("PM_AGENT_UID") or "").strip() or None
+
+    @classmethod
+    def from_profile(cls, profile: Profile) -> "PMSession":
+        # Delegation JWT wins outright: PM_API_JWT is set only by the agent
+        # harness (cli_adapter), where the pre-minted credential IS the
+        # session and nothing may fall through to stored user credentials.
+        static_jwt = (os.environ.get("PM_API_JWT") or "").strip()
+        if static_jwt:
+            if not _JWT_SHAPE_RE.fullmatch(static_jwt):
+                raise NotAuthenticatedError(JWT_SHAPE_HINT)
+            return cls(profile, static_jwt=static_jwt)
+        # Bearer PAT takes priority over the cookie jar — it's the explicit
+        # programmatic-auth signal (env var or `pm login --token`).
+        token = token_auth.load_token(profile.name)
+        if token:
+            return cls(profile, token=token)
+        creds = store.load(profile.name)
+        if creds is None:
+            raise NotAuthenticatedError()
+        return cls(profile, creds)
+
+    def request(self, method: str, path: str, **kwargs) -> requests.Response:
+        url = f"{self.profile.api_url}{path}"
+        kwargs.setdefault("timeout", 60)
+        headers = dict(kwargs.get("headers") or {})
+        # Forward agent identity to the API for backend scope filters.
+        if self._agent_uid:
+            headers.setdefault(_AGENT_UID_HEADER, self._agent_uid)
+        # Static-bearer mode: attach the pre-minted delegation JWT as-is.
+        if self._static_jwt is not None:
+            headers["Authorization"] = f"Bearer {self._static_jwt}"
+        # Token mode: attach a fresh Bearer JWT (exchanged from the PAT).
+        elif self._token is not None:
+            headers["Authorization"] = f"Bearer {self._ensure_jwt()}"
+        if headers:
+            kwargs["headers"] = headers
+        # Cookie mode: any earlier response in this process may have re-issued
+        # the session cookie with Secure (SESSION_COOKIE_SECURE=True backend),
+        # which would silently drop it from this send — re-normalize each time.
+        if self._token is None and self._static_jwt is None:
+            _strip_secure_for_local_http(self._session.cookies, self.profile.api_url)
+        resp = self._send(method, url, kwargs)
+        # Recover at most once per request (request-local, not session-wide):
+        # a JWT can expire mid-session, so a later request must still be able to
+        # re-exchange. There's no retry loop, so this can't spin.
+        if resp.status_code == 401:
+            logger.debug("Got 401; attempting a single auth recovery.")
+            if self._token is not None:
+                # The rejected JWT must not be presented again by the next
+                # process either: forget it before minting a replacement.
+                self._forget_jwt()
+            if self._recover():
+                # Re-stamp the Bearer header with the freshly minted JWT.
+                if self._token is not None:
+                    kwargs.setdefault("headers", headers)
+                    kwargs["headers"]["Authorization"] = f"Bearer {self._jwt}"
+                resp = self._send(method, url, kwargs)
+            if resp.status_code == 401:
+                # CONTRACT: these messages carry "auth failed: 401" so the
+                # prefect cli_adapter's stderr status scrape
+                # (_API_STATUS_RE, "failed:\s*(\d{3})") classifies them as
+                # an auth denial - report tools must NOT fall back to the
+                # direct-DB path on an expired/invalid credential.
+                if self._static_jwt is not None:
+                    raise NotAuthenticatedError(
+                        "API auth failed: 401 - delegation credential expired "
+                        "or invalid (PM_API_JWT). The delegation cannot be "
+                        "refreshed; a new one must be minted."
+                    )
+                if self._token is not None:
+                    # The replacement JWT was rejected too: never present it
+                    # again from the cache.
+                    self._forget_jwt()
+                    raise NotAuthenticatedError(
+                        "API auth failed: 401 - bearer token rejected "
+                        "(expired, revoked, or invalid). Check PM_API_TOKEN "
+                        "or run `pm login --token`."
+                    )
+                raise NotAuthenticatedError(
+                    "API auth failed: 401 - session expired. "
+                    "Run `pm login` again."
+                )
+        return resp
+
+    def _send(self, method: str, url: str, kwargs: dict) -> requests.Response:
+        # Assumes every GET is idempotent. The only state-changing GETs the CLI
+        # sends, /workspace/<pk>/switch and /logout/, are idempotent (verified by security).
+        retryable = method.upper() == "GET"
+        for delay in (*_RETRY_BACKOFF_SECONDS, None):
+            started = log_request(logger, method, url)
+            try:
+                resp = self._session.request(method, url, **kwargs)
+            except requests.RequestException as exc:
+                raise PMConnectionError(
+                    token_auth.unreachable_message(self.profile.api_url, exc)
+                )
+            log_response(logger, resp, started)
+            if not retryable or resp.status_code not in _RETRY_STATUSES or delay is None:
+                return resp
+            logger.debug("got %d; retrying in %.1fs", resp.status_code, delay)
+            _sleep(delay)
+        return resp
+
+    def get(self, path: str, **kwargs) -> requests.Response:
+        return self.request("GET", path, **kwargs)
+
+    def post(self, path: str, **kwargs) -> requests.Response:
+        return self.request("POST", path, **kwargs)
+
+    def put(self, path: str, **kwargs) -> requests.Response:
+        return self.request("PUT", path, **kwargs)
+
+    def delete(self, path: str, **kwargs) -> requests.Response:
+        return self.request("DELETE", path, **kwargs)
+
+    def persist(self) -> None:
+        """Write the current cookiejar back to the credential store.
+
+        No-op in token / delegation mode (there is no cookie jar to persist;
+        the PAT is managed via `pm login --token` / `PM_API_TOKEN`, and the
+        delegation JWT via `PM_API_JWT`).
+        """
+        if self._token is not None or self._static_jwt is not None:
+            return
+        self.creds.set_cookies(self._session.cookies)
+        store.save(self.creds)
+
+    def _ensure_jwt(self) -> str:
+        """Return a valid Bearer JWT: the in-process one, else the one cached
+        by an earlier process for this profile/host/PAT, else a fresh
+        exchange."""
+        if self._jwt is None:
+            self._adopt_cached_jwt()
+        if self._jwt is None or time.monotonic() >= self._jwt_expires_at:
+            self._exchange()
+        return self._jwt
+
+    def _adopt_cached_jwt(self) -> None:
+        cached = jwt_cache.load(self.profile.name, self.profile.api_url, self._token)
+        if cached is None:
+            return
+        self._jwt = cached.jwt
+        self._jwt_expires_at = time.monotonic() + cached.seconds_left()
+
+    def _forget_jwt(self) -> None:
+        self._jwt = None
+        self._jwt_expires_at = 0.0
+        jwt_cache.clear(self.profile.name)
+
+    def _exchange(self) -> None:
+        result = token_auth.exchange_pat_for_jwt(
+            self._session, self.profile.api_url, self._token
+        )
+        self._jwt = result.access_token
+        # Refresh before the server-side expiry to avoid racing it
+        # mid-request (same margin as the cross-process cache), but never
+        # past the server TTL: if expires_in is tiny, cache_for clamps to 0
+        # so the next request re-exchanges rather than reusing an
+        # already-expired JWT.
+        cache_for = max(0, result.expires_in - jwt_cache.EXPIRY_MARGIN_SECONDS)
+        self._jwt_expires_at = time.monotonic() + cache_for
+        jwt_cache.save(
+            self.profile.name,
+            self.profile.api_url,
+            self._token,
+            result.access_token,
+            result.expires_in,
+        )
+
+    def _recover(self) -> bool:
+        """Single-shot auth recovery after a 401: re-exchange the PAT in token
+        mode, or silently re-login via Auth0 in cookie mode."""
+        if self._static_jwt is not None:
+            # A pre-minted delegation JWT has no refresh authority: the PAT
+            # or session that could mint a fresh one lives with the caller.
+            return False
+        if self._token is not None:
+            try:
+                self._exchange()
+                return True
+            except Exception as exc:
+                logger.debug("PAT re-exchange failed: %s", exc)
+                return False
+        return self._silent_relogin()
+
+    def _silent_relogin(self) -> bool:
+        if not self.creds.auth0_refresh_token or not self.profile.auth0_client_id:
+            return False
+        try:
+            token = device_flow.refresh_auth0_token(
+                domain=self.profile.auth0_domain,
+                client_id=self.profile.auth0_client_id,
+                refresh_token=self.creds.auth0_refresh_token,
+            )
+            self._session.cookies.clear()
+            result = exchange.exchange_for_session(
+                self._session, self.profile.api_url, token.access_token
+            )
+            # The re-issued session cookie carries Secure again; re-normalize
+            # so the retry and the workspace restore below aren't anonymous.
+            _strip_secure_for_local_http(self._session.cookies, self.profile.api_url)
+            self.creds.pm_access_token = (
+                result.pm_access_token or self.creds.pm_access_token
+            )
+            self.creds.pm_refresh_token = (
+                result.pm_refresh_token or self.creds.pm_refresh_token
+            )
+            if token.refresh_token:
+                self.creds.auth0_refresh_token = token.refresh_token
+
+            if self.creds.last_workspace_uid:
+                ws_resp = self._session.get(
+                    f"{self.profile.api_url}/api/v1/workspace/", timeout=15
+                )
+                if ws_resp.status_code == 200:
+                    body = ws_resp.json()
+                    items = body.get("result", []) if isinstance(body, dict) else body
+                    for ws in items:
+                        if ws.get("uid") == self.creds.last_workspace_uid:
+                            pk = ws.get("id")
+                            if pk:
+                                self._session.get(
+                                    f"{self.profile.api_url}/api/v1/workspace/{pk}/switch",
+                                    timeout=15,
+                                )
+                            break
+
+            self.persist()
+            return True
+        except Exception as exc:
+            logger.debug("Silent re-login failed: %s", exc)
+            return False

pm/auth/store.py

@@ -0,0 +1,160 @@
+"""Credential persistence.
+
+Primary backend: OS keyring (Windows Credential Manager / macOS Keychain /
+Linux Secret Service). Fallback: ~/.pm/credentials.json (chmod 0600).
+
+Stores the Auth0 device-flow login's artifacts (session cookiejar + refresh
+token) for interactive/internal use. The cookiejar carries workspace_id for
+RLS; the Auth0 refresh token silently re-establishes the session when cookies
+expire. Programmatic auth (PATs via ``PM_API_TOKEN``, delegation JWTs via
+``PM_API_JWT``) does not touch this store - see ``pm/auth/token.py`` and
+``pm/auth/session.py``.
+"""
+from __future__ import annotations
+
+import base64
+import json
+import logging
+from dataclasses import asdict, dataclass
+from typing import List, Optional
+
+from requests.cookies import RequestsCookieJar, create_cookie
+
+from pm.auth.secret_files import read_private_json, write_private_json
+from pm.config import CONFIG_DIR, keyring_account
+from pm.log import timed
+
+try:
+    import keyring
+    _KEYRING_AVAILABLE = True
+except Exception:
+    _KEYRING_AVAILABLE = False
+
+logger = logging.getLogger(__name__)
+
+KEYRING_SERVICE = "pushmetrics-cli"
+FALLBACK_FILE = CONFIG_DIR / "credentials.json"
+
+
+# Cookie serialization: JSON only. ``pickle.loads`` on stored bytes is
+# RCE-by-design — we encode cookies as a list of plain dicts and base64 the
+# resulting JSON.
+def encode_cookies_json(jar: RequestsCookieJar) -> str:
+    """Serialize a RequestsCookieJar to a base64-encoded JSON list of cookie dicts."""
+    items: List[dict] = []
+    for c in jar:
+        items.append({
+            "name": c.name,
+            "value": c.value,
+            "domain": c.domain or "",
+            "path": c.path or "/",
+            "secure": bool(c.secure),
+            # `expires` is either an int (unix epoch) or None.
+            "expires": c.expires,
+        })
+    return base64.b64encode(json.dumps(items).encode("utf-8")).decode("ascii")
+
+
+def decode_cookies_json(b64: str) -> RequestsCookieJar:
+    """Inverse of ``encode_cookies_json``. Raises on malformed input."""
+    raw = base64.b64decode(b64.encode("ascii"))
+    items = json.loads(raw.decode("utf-8"))
+    if not isinstance(items, list):
+        raise ValueError("expected list of cookie dicts")
+    jar = RequestsCookieJar()
+    for d in items:
+        if not isinstance(d, dict) or "name" not in d or "value" not in d:
+            raise ValueError("cookie dict missing name/value")
+        jar.set_cookie(create_cookie(
+            name=d["name"],
+            value=d["value"],
+            domain=d.get("domain") or "",
+            path=d.get("path") or "/",
+            secure=bool(d.get("secure", False)),
+            expires=d.get("expires"),
+        ))
+    return jar
+
+
+@dataclass
+class StoredCredentials:
+    profile: str
+    api_url: str
+    cookies_b64: str = ""
+    auth0_refresh_token: Optional[str] = None
+    pm_access_token: Optional[str] = None
+    pm_refresh_token: Optional[str] = None
+    last_workspace_uid: Optional[str] = None
+
+    def cookies(self) -> RequestsCookieJar:
+        if not self.cookies_b64:
+            return RequestsCookieJar()
+        # JSON only. A record that does not decode (including the pickle
+        # format internal builds wrote before 2026-07) reads as "not logged
+        # in"; pickle.loads on stored bytes is never acceptable in a shipped
+        # credential store (security S5).
+        try:
+            return decode_cookies_json(self.cookies_b64)
+        except Exception:
+            logger.debug("stored cookies for profile %s are unreadable; ignoring", self.profile)
+            return RequestsCookieJar()
+
+    def set_cookies(self, jar: RequestsCookieJar) -> None:
+        self.cookies_b64 = encode_cookies_json(jar)
+
+
+def _fallback_load_all() -> dict:
+    return read_private_json(FALLBACK_FILE)
+
+
+def _fallback_save_all(data: dict) -> None:
+    # The file holds the Auth0 refresh token and session cookies: it must be
+    # 0600 from creation, never write-then-chmod (security report S2).
+    write_private_json(FALLBACK_FILE, data)
+
+
+def load(profile: str) -> Optional[StoredCredentials]:
+    if _KEYRING_AVAILABLE:
+        try:
+            with timed(logger, "keyring.load"):
+                raw = keyring.get_password(KEYRING_SERVICE, keyring_account(profile))
+            if raw:
+                return StoredCredentials(**json.loads(raw))
+        except Exception:
+            pass
+    data = _fallback_load_all().get(profile)
+    if not data:
+        return None
+    return StoredCredentials(**data)
+
+
+def _to_storable_dict(creds: StoredCredentials) -> dict:
+    """Serialize creds to a dict suitable for keyring/fallback storage."""
+    return dict(asdict(creds))
+
+
+def save(creds: StoredCredentials) -> None:
+    storable = _to_storable_dict(creds)
+    payload = json.dumps(storable)
+    if _KEYRING_AVAILABLE:
+        try:
+            with timed(logger, "keyring.save"):
+                keyring.set_password(KEYRING_SERVICE, keyring_account(creds.profile), payload)
+            return
+        except Exception:
+            pass
+    data = _fallback_load_all()
+    data[creds.profile] = storable
+    _fallback_save_all(data)
+
+
+def clear(profile: str) -> None:
+    if _KEYRING_AVAILABLE:
+        try:
+            keyring.delete_password(KEYRING_SERVICE, keyring_account(profile))
+        except Exception:
+            pass
+    data = _fallback_load_all()
+    if profile in data:
+        del data[profile]
+        _fallback_save_all(data)