Auth module¶

def apply_system_certs() -> None:
    """
    Patch the Python SSL module to use the operating system's certificate store
    and propagate the configuration to child processes.

    Calls :func:`truststore.inject_into_ssl`, which replaces the default
    ``ssl.create_default_context`` factory so that all subsequent TLS
    connections (including those made by ``requests``) trust the OS certificate
    store instead of the bundled ``certifi`` CA bundle.

    Also sets ``UV_SYSTEM_CERTS=true`` in the process environment so that
    ``uv`` child processes also trust the OS certificate store.

    This is a global, process-wide side effect.  Calling it multiple times is
    harmless (it is idempotent).

    Raises
    ------
    ImportError
        If the ``truststore`` package is not installed.
    """
    import truststore

    truststore.inject_into_ssl()

    # Propagate to uv child processes — uv respects this env var natively.
    if "UV_SYSTEM_CERTS" not in os.environ:
        os.environ["UV_SYSTEM_CERTS"] = "true"

def resolve_system_certs(profile: str | None = None) -> bool:
    """
    Apply system certificates if enabled via environment variable or profile config.

    Resolution order:
    1. ``NEXTMV_SYSTEM_CERTS`` env var (``1``, ``true``, ``yes`` → enabled).
    2. ``system_certs`` flag on the given profile (or the default profile).

    Returns ``True`` if system certificates were applied, ``False`` otherwise.

    Parameters
    ----------
    profile : str | None
        The profile name.  If ``None``, the default (top-level) profile is used.

    Returns
    -------
    bool
        Whether system certificates were applied.
    """
    from nextmv.config import get_system_certs, load_config

    env_val = os.environ.get("NEXTMV_SYSTEM_CERTS", "").strip().lower()
    if env_val in ("1", "true", "yes"):
        apply_system_certs()
        return True

    config = load_config()
    if get_system_certs(config, profile):
        apply_system_certs()
        return True

    return False

def token_dir(session: str) -> Path:
    """
    Returns the directory that holds token files for *session*.

    Parameters
    ----------
    session : str
        The auth session name.  The reserved value ``"default"`` (case-
        insensitive) maps to ``~/.nextmv/auth/default/``.  Obtain the correct
        session name for a profile via
        :func:`nextmv.config.get_auth_session`.

    Returns
    -------
    Path
        The directory path ``~/.nextmv/auth/<session>/``.

    Raises
    ------
    ValueError
        If *session* is empty or would escape the auth directory (e.g. path
        traversal).
    """
    name = session.strip() if session else DEFAULT_AUTH_SESSION
    if not name:
        name = DEFAULT_AUTH_SESSION
    path = (AUTH_DIR / name).resolve()
    if not path.is_relative_to(AUTH_DIR.resolve()):
        raise ValueError(f"Invalid session name {name!r}: must not escape the auth directory.")
    return path

def load_tokens(session: str) -> dict[str, Any] | None:
    """
    Load stored tokens for *session* from disk.

    Parameters
    ----------
    session : str
        The auth session name.  Resolve this from a profile via
        :func:`nextmv.config.get_auth_session`.

    Returns
    -------
    dict[str, Any] | None
        The token dict, or ``None`` if no token file exists.
    """
    path = _token_path(session)
    if not path.exists():
        return None
    with path.open() as fh:
        return json.load(fh)

def save_tokens(session: str, tokens: dict[str, Any]) -> None:
    """
    Persist *tokens* for *session* to disk.

    Creates any missing parent directories with mode 0o700.

    Parameters
    ----------
    session : str
        The auth session name.  Resolve this from a profile via
        :func:`nextmv.config.get_auth_session`.
    tokens : dict[str, Any]
        The token dict to persist.  Must contain at least ``access_token``.
    """
    path = _token_path(session)
    path.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
    # Restrict read/write access to the owner only.  Use os.open to set the
    # mode atomically when the file is created, avoiding a window where the
    # file exists with broader permissions.  Best-effort: silently ignored on
    # filesystems or platforms that don't support POSIX permissions (e.g. Windows).
    try:
        fd = os.open(str(path), os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
        with os.fdopen(fd, "w") as fh:
            json.dump(tokens, fh, indent=2)
    except (OSError, PermissionError):
        # Fallback: write without restrictive mode.
        with path.open("w") as fh:
            json.dump(tokens, fh, indent=2)

def delete_tokens(session: str) -> None:
    """
    Delete stored tokens for *session* from disk.

    Silently does nothing if no token file exists.

    Parameters
    ----------
    session : str
        The auth session name.  Resolve this from a profile via
        :func:`nextmv.config.get_auth_session`.
    """
    path = _token_path(session)
    try:
        path.unlink()
    except FileNotFoundError:
        pass

def is_invalid_grant_error(exc: Exception) -> bool:
    """
    Return ``True`` if *exc* represents an OAuth2 ``invalid_grant`` error.

    Checks for an HTTP 400 response with ``{"error": "invalid_grant"}`` in the
    body, per `RFC 6749 §5.2`_.

    .. _RFC 6749 §5.2: https://datatracker.ietf.org/doc/html/rfc6749#section-5.2
    """
    if not isinstance(exc, requests.HTTPError) or exc.response is None:
        return False
    if exc.response.status_code != 400:
        return False
    try:
        body = exc.response.json()
        return body.get("error") == "invalid_grant"
    except Exception:
        return False

def is_token_expired(tokens: dict[str, Any]) -> bool:
    """
    Return ``True`` when the stored access token is expired (or will expire within the
    next 30 seconds), ``False`` otherwise.

    If ``expires_at`` is absent the token is treated as *not* expired so that tokens
    without an explicit expiry still work.

    Parameters
    ----------
    tokens : dict[str, Any]
        The token dict loaded from disk.

    Returns
    -------
    bool
    """
    expires_at_str: str | None = tokens.get("expires_at")
    if not expires_at_str:
        return False
    try:
        expires_at = datetime.fromisoformat(expires_at_str)
        # Ensure timezone-aware comparison.
        if expires_at.tzinfo is None:
            expires_at = expires_at.replace(tzinfo=timezone.utc)
        now = datetime.now(tz=timezone.utc)
        # Consider expired if within 30-second buffer.
        return (expires_at - now).total_seconds() < 30
    except ValueError:
        # Unparseable expiry - treat as expired so the token gets refreshed
        # rather than sending a potentially corrupt token to the API.
        return True

def refresh_tokens(
    refresh_token: str,
    token_endpoint: str | None = None,
    client_id: str | None = None,
    oidc_discovery_url: str | None = None,
) -> dict[str, Any]:
    """
    Use a refresh token to obtain a new access token.

    Parameters
    ----------
    refresh_token : str
        A valid refresh token previously obtained via :func:`run_pkce_flow`.
    token_endpoint : str | None
        The token endpoint URL.  If ``None``, the OIDC discovery document is fetched to
        resolve it.
    client_id : str | None
        The OAuth2 client ID.  When ``None`` the module-level :data:`CLIENT_ID`
        constant is used.
    oidc_discovery_url : str | None
        The OIDC discovery document URL.  Used only when *token_endpoint* is
        ``None``.  When ``None`` the module-level :data:`OIDC_DISCOVERY_URL` is
        used.

    Returns
    -------
    dict[str, Any]
        The refreshed token response body augmented with an ``expires_at`` field
        (ISO-8601 UTC string).

    Raises
    ------
    requests.HTTPError
        If the token endpoint returns a non-2xx response.
    """
    if token_endpoint is None:
        _, _, token_endpoint = _discover_endpoints(oidc_discovery_url)

    cid = client_id or CLIENT_ID
    payload = {
        "grant_type": "refresh_token",
        "client_id": cid,
        "refresh_token": refresh_token,
    }
    resp = requests.post(
        token_endpoint,
        data=payload,
        headers={"Content-Type": "application/x-www-form-urlencoded"},
        timeout=30,
    )
    try:
        resp.raise_for_status()
    except requests.HTTPError as exc:
        raise requests.HTTPError(f"Token refresh failed ({resp.status_code}): {resp.text}") from exc

    tokens: dict[str, Any] = resp.json()
    _validate_token_response(tokens)
    expires_in: int = tokens.get("expires_in", 3600)
    expires_at = datetime.now(tz=timezone.utc) + timedelta(seconds=expires_in)
    tokens["expires_at"] = expires_at.isoformat()
    # The Cognito refresh response does not include a new refresh_token;
    # preserve the existing one.
    if "refresh_token" not in tokens:
        tokens["refresh_token"] = refresh_token
    return tokens

def run_pkce_flow(
    oidc_discovery_url: str | None = None,
    client_id: str | None = None,
    force: bool = False,
    identity_provider: str | None = None,
) -> dict[str, Any]:
    """
    Execute the full PKCE authorization-code flow.

    This function:

    1. Resolves the OIDC endpoints.
    2. Generates a PKCE pair.
    3. Listens for the redirect callback on the fixed port ``CALLBACK_PORT`` (``56734``).
    4. Opens the system browser at the authorization URL (or the logout URL when
       ``force=True``, which clears the identity provider session and chains
       into a fresh login).
    5. Waits for the redirect callback (up to ``_BROWSER_TIMEOUT`` seconds).
    6. Exchanges the authorization code for tokens.

    Parameters
    ----------
    oidc_discovery_url : str | None
        The OIDC discovery document URL for the identity provider backing this
        profile's endpoint.  When ``None`` the module-level
        :data:`OIDC_DISCOVERY_URL` constant is used (production endpoint).
    client_id : str | None
        The OAuth2 client ID for the identity provider.  When ``None`` the
        module-level :data:`CLIENT_ID` constant is used.
    force : bool
        When ``True``, opens the identity provider's logout endpoint first,
        which clears any active browser session, and chains all authorization
        parameters onto it so that the provider immediately presents the login
        page.  Defaults to ``False``.
    identity_provider : str | None
        When set, adds ``identity_provider=<value>`` to the authorization URL
        so that Cognito delegates to a third-party SSO provider.  Use
        :func:`fetch_sso_provider` to resolve the identifier from a domain.

    Returns
    -------
    dict[str, Any]
        A token dict containing at minimum ``access_token``, ``token_type``,
        and ``expires_at``.  Also contains ``refresh_token`` and ``id_token``
        when the provider issues them.

    Raises
    ------
    TimeoutError
        If the user does not complete the browser flow within the timeout.
    RuntimeError
        If the provider returns an error.
    requests.HTTPError
        If the token exchange request fails.
    """
    auth_endpoint, token_endpoint, logout_endpoint = _discover_endpoints(oidc_discovery_url)
    cid = client_id or CLIENT_ID
    code_verifier, code_challenge = _generate_pkce_pair()
    redirect_uri = f"http://127.0.0.1:{CALLBACK_PORT}"
    # Generate a random state value to prevent login CSRF / authorization-code injection.
    # The callback is validated against this value before the code is exchanged for
    # tokens.
    state = secrets.token_urlsafe(16)

    auth_params = {
        "response_type": "code",
        "client_id": cid,
        "redirect_uri": redirect_uri,
        "scope": SCOPES,
        "code_challenge_method": "S256",
        "code_challenge": code_challenge,
        "state": state,
    }
    if identity_provider:
        auth_params["identity_provider"] = identity_provider

    if force:
        # Use the logout endpoint to clear the Cognito session, then chain the
        # full authorization request onto it.  Cognito will log the user out
        # and immediately redirect to the login page with all auth params intact.
        # The redirect_uri must be registered as an Allowed Callback URL (same
        # requirement as the normal authorization flow - no additional sign-out
        # URL registration is needed).
        logout_params = {
            "client_id": cid,
            **auth_params,
        }
        open_url = logout_endpoint + "?" + urllib.parse.urlencode(logout_params)
    else:
        open_url = auth_endpoint + "?" + urllib.parse.urlencode(auth_params)

    # Start the callback listener in a background thread so we can open the
    # browser on the main thread without blocking.
    callback_result: dict[str, str] = {}
    exc_holder: list[Exception] = []

    def _listen() -> None:
        try:
            result = _wait_for_callback(CALLBACK_PORT)
            callback_result.update(result)
        except Exception as exc:
            exc_holder.append(exc)

    listener = threading.Thread(target=_listen, daemon=True)
    listener.start()

    webbrowser.open(open_url)

    listener.join(timeout=_BROWSER_TIMEOUT + 5)

    if exc_holder:
        raise exc_holder[0]

    # Verify the state before trusting the code - protects against login CSRF where a
    # malicious local page hits our callback with an attacker-supplied code.
    returned_state = callback_result.get("state")
    if returned_state != state:
        raise RuntimeError(
            "OAuth2 state mismatch: the callback state does not match the expected value. "
            "This may indicate a login CSRF attempt. Please try running `nextmv auth login` again."
        )

    code = callback_result.get("code")
    if not code:
        raise RuntimeError("No authorization code received. Please try running `nextmv auth login` again.")

    return _exchange_code_for_tokens(token_endpoint, code, code_verifier, redirect_uri, cid)

def fetch_organizations(access_token: str, endpoint: str) -> list[dict[str, Any]]:
    """
    Fetch the list of organizations (teams) the authenticated user belongs to.

    Calls ``GET https://<endpoint>/v1/internal/me/organization`` and returns
    the response body as a list of organization dicts.

    Each dict contains at minimum:

    - ``id`` (str) - the team UUID; use this for the ``nextmv-account`` header.
    - ``name`` (str) - the human-readable team name; show this to the user.
    - ``role`` (str) - the user's role in the team.
    - ``pending_invite`` (bool) - whether the user has a pending invite.

    Parameters
    ----------
    access_token : str
        A valid access token (or id_token) for the authenticated user.
    endpoint : str
        The API endpoint hostname, e.g. ``"api.cloud.nextmv.io"``.  A leading
        ``https://`` scheme is accepted and stripped.  Plain ``http://`` is
        rejected to prevent sending tokens over an unencrypted connection.

    Returns
    -------
    list[dict[str, Any]]
        The list of organization objects returned by the API.

    Raises
    ------
    ValueError
        If *endpoint* uses ``http://``.
    requests.HTTPError
        If the API returns a non-2xx response.
    """
    if endpoint.startswith("http://"):
        raise ValueError(
            f"Refusing to send tokens over plain HTTP for endpoint {endpoint!r}. Use https:// or a bare hostname."
        )
    # Strip https:// if present so we always build a consistent URL.
    bare = _strip_scheme(endpoint)
    url = f"https://{bare}/v1/internal/me/organization"
    resp = requests.get(
        url,
        headers={"Authorization": f"Bearer {access_token}", "Content-Type": "application/json"},
        timeout=15,
    )
    resp.raise_for_status()
    return resp.json()

def fetch_sso_provider(domain: str, endpoint: str) -> str | None:
    """
    Check whether a third-party SSO provider is enabled for *domain*.

    Calls ``GET https://<endpoint>/v1/enterprise/sso/domain?domain=<domain>``
    (no authentication required).  Returns the ``domain_identifier`` if SSO is
    enabled, or ``None`` otherwise.

    Parameters
    ----------
    domain : str
        The email domain to check, e.g. ``"nextmv.io"``.
    endpoint : str
        The API endpoint hostname, e.g. ``"api.cloud.nextmv.io"``.

    Returns
    -------
    str | None
        The ``domain_identifier`` when SSO is enabled, ``None`` otherwise.
    """
    if endpoint.startswith("http://"):
        raise ValueError(
            f"Refusing to connect over plain HTTP for endpoint {endpoint!r}. Use https:// or a bare hostname."
        )
    bare = _strip_scheme(endpoint)
    url = f"https://{bare}/v1/enterprise/sso/domain?domain={domain}"
    resp = requests.get(url, timeout=10)
    resp.raise_for_status()
    body: dict[str, str] = resp.json()
    if body.get("enabled") and body.get("domain_identifier"):
        return body["domain_identifier"]
    return None