OAuth2 Token Fetcher and Refresher

OAuth2 Token Fetcher and Refresher

Problem Statement

Modern APIs (Google, GitHub, Salesforce, Slack, internal microservices) almost universally require OAuth2 access tokens rather than static API keys. These tokens expire — typically in 1 hour — which means any long-running script, cron job, or CI pipeline will eventually fail with a 401 Unauthorized unless it knows how to refresh tokens. Handling this manually is error-prone: developers often hardcode tokens, forget to persist refresh tokens, or race against expiration windows.

This script solves that problem by providing a reusable, CLI-driven OAuth2 client that can fetch a fresh access token (via client_credentials or refresh_token grant) and cache it to disk with a safety margin. Subsequent invocations reuse the cached token until it nears expiry, at which point the script transparently refreshes it. The result is a drop-in token provider that any downstream tool can call.

How It Works

The script uses authlib's OAuth2Session to perform standards-compliant token exchanges against the configured --token-url, while requests is used for lightweight validation calls (e.g., verifying the token works). Tokens are persisted to a JSON file in a configurable cache directory (~/.cache/oauth2_client/tokens.json by default) keyed by client ID, so multiple services can coexist. The cache stores access_token, refresh_token, expires_at (absolute epoch time), and token_type.

The core flow is: (1) load cached token, (2) check if it's still valid with a 60-second safety buffer, (3) if expired, either use the refresh_token grant or fall back to client_credentials, (4) write the new token to cache atomically (write to temp file, then os.replace). The CLI exits with status 0 on success, printing the access token to stdout so it can be consumed by shell scripts like TOKEN=$(python3 oauth2_client.py ...).

Full Script

#!/usr/bin/env python3
"""
OAuth2 Token Fetcher and Refresher

Fetches and caches OAuth2 access tokens using the client_credentials or
refresh_token grant. Reuses cached tokens until they near expiry, then
transparently refreshes them. Prints the access token to stdout on success.

Usage:
    python3 oauth2_client.py --client-id ID --client-secret SECRET \
        --token-url https://auth.example.com/token
"""

from __future__ import annotations

import argparse
import json
import logging
import os
import sys
import tempfile
import time
from dataclasses import dataclass, asdict
from pathlib import Path
from typing import Optional

import requests
from authlib.integrations.requests_client import OAuth2Session
from authlib.common.errors import AuthlibBaseError

# ---------- Configuration (env-var driven with sane defaults) ----------
DEFAULT_CACHE_DIR = Path(
    os.environ.get("OAUTH2_CACHE_DIR", Path.home() / ".cache" / "oauth2_client")
)
DEFAULT_SCOPE = os.environ.get("OAUTH2_SCOPE", "")
DEFAULT_TIMEOUT = float(os.environ.get("OAUTH2_HTTP_TIMEOUT", "15"))
# Refresh if the token expires within this many seconds (safety margin)
EXPIRY_SKEW_SECONDS = int(os.environ.get("OAUTH2_EXPIRY_SKEW", "60"))

logger = logging.getLogger("oauth2_client")


# ---------- Data model ----------
@dataclass
class TokenRecord:
    """A persisted OAuth2 token with absolute expiry time."""

    access_token: str
    token_type: str = "Bearer"
    refresh_token: Optional[str] = None
    expires_at: float = 0.0  # epoch seconds; 0 means "unknown"

    @property
    def is_expired(self) -> bool:
        if self.expires_at <= 0:
            return True
        return time.time() >= (self.expires_at - EXPIRY_SKEW_SECONDS)

    @classmethod
    def from_response(cls, payload: dict) -> "TokenRecord":
        expires_in = payload.get("expires_in")
        expires_at = time.time() + float(expires_in) if expires_in else 0.0
        return cls(
            access_token=payload["access_token"],
            token_type=payload.get("token_type", "Bearer"),
            refresh_token=payload.get("refresh_token"),
            expires_at=expires_at,
        )


# ---------- Cache I/O ----------
class TokenCache:
    """Simple JSON-backed token cache keyed by client_id."""

    def __init__(self, path: Path):
        self.path = path
        self.path.parent.mkdir(parents=True, exist_ok=True)

    def _read_all(self) -> dict:
        if not self.path.exists():
            return {}
        try:
            with self.path.open("r", encoding="utf-8") as fh:
                return json.load(fh)
        except (json.JSONDecodeError, OSError) as exc:
            logger.warning("Cache unreadable (%s); starting fresh", exc)
            return {}

    def load(self, client_id: str) -> Optional[TokenRecord]:
        raw = self._read_all().get(client_id)
        if not raw:
            return None
        try:
            return TokenRecord(**raw)
        except TypeError:
            logger.warning("Cache entry malformed for %s; ignoring", client_id)
            return None

    def save(self, client_id: str, record: TokenRecord) -> None:
        data = self._read_all()
        data[client_id] = asdict(record)
        # Atomic write: temp file in same dir, then os.replace
        tmp_fd, tmp_path = tempfile.mkstemp(
            dir=str(self.path.parent), prefix=".tokens-", suffix=".json"
        )
        try:
            with os.fdopen(tmp_fd, "w", encoding="utf-8") as fh:
                json.dump(data, fh, indent=2)
            os.chmod(tmp_path, 0o600)  # tokens are secrets
            os.replace(tmp_path, self.path)
        except Exception:
            Path(tmp_path).unlink(missing_ok=True)
            raise


# ---------- OAuth2 client ----------
class OAuth2TokenClient:
    """Fetches and refreshes OAuth2 tokens, using a disk cache."""

    def __init__(
        self,
        client_id: str,
        client_secret: str,
        token_url: str,
        scope: str = "",
        cache: Optional[TokenCache] = None,
        timeout: float = DEFAULT_TIMEOUT,
    ):
        self.client_id = client_id
        self.client_secret = client_secret
        self.token_url = token_url
        self.scope = scope
        self.timeout = timeout
        self.cache = cache or TokenCache(DEFAULT_CACHE_DIR / "tokens.json")

    def _new_session(self) -> OAuth2Session:
        return OAuth2Session(
            client_id=self.client_id,
            client_secret=self.client_secret,
            scope=self.scope,
            token_endpoint=self.token_url,
        )

    def _fetch_client_credentials(self) -> TokenRecord:
        logger.info("Requesting new token via client_credentials grant")
        session = self._new_session()
        token = session.fetch_token(
            self.token_url,
            grant_type="client_credentials",
            timeout=self.timeout,
        )
        return TokenRecord.from_response(token)

    def _refresh(self, refresh_token: str) -> TokenRecord:
        logger.info("Refreshing token via refresh_token grant")
        session = self._new_session()
        token = session.refresh_token(
            self.token_url,
            refresh_token=refresh_token,
            timeout=self.timeout,
        )
        record = TokenRecord.from_response(token)
        # Some providers omit refresh_token on refresh; preserve the old one
        if not record.refresh_token:
            record.refresh_token = refresh_token
        return record

    def get_token(self, force_refresh: bool = False) -> TokenRecord:
        """Return a valid token, refreshing or fetching as needed."""
        cached = None if force_refresh else self.cache.load(self.client_id)
        if cached and not cached.is_expired:
            logger.debug("Using cached token (expires in %.0fs)",
                         cached.expires_at - time.time())
            return cached

        try:
            if cached and cached.refresh_token:
                record = self._refresh(cached.refresh_token)
            else:
                record = self._fetch_client_credentials()
        except AuthlibBaseError as exc:
            # Authlib wraps OAuth error responses; surface them clearly
            raise RuntimeError(f"OAuth2 token endpoint rejected request: {exc}") from exc
        except requests.RequestException as exc:
            raise RuntimeError(f"Network error contacting {self.token_url}: {exc}") from exc

        self.cache.save(self.client_id, record)
        return record


# ---------- CLI ----------
def parse_args(argv: Optional[list[str]] = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Fetch or refresh an OAuth2 access token and print it.",
    )
    parser.add_argument("--client-id", required=True,
                        default=os.environ.get("OAUTH2_CLIENT_ID"),
                        help="OAuth2 client ID (or $OAUTH2_CLIENT_ID)")
    parser.add_argument("--client-secret", required=True,
                        default=os.environ.get("OAUTH2_CLIENT_SECRET"),
                        help="OAuth2 client secret (or $OAUTH2_CLIENT_SECRET)")
    parser.add_argument("--token-url", required=True,
                        default=os.environ.get("OAUTH2_TOKEN_URL"),
                        help="Token endpoint URL (or $OAUTH2_TOKEN_URL)")
    parser.add_argument("--scope", default=DEFAULT_SCOPE,
                        help="Space-separated scopes (or $OAUTH2_SCOPE)")
    parser.add_argument("--cache-file", type=Path,
                        default=DEFAULT_CACHE_DIR / "tokens.json",
                        help="Path to token cache file")
    parser.add_argument("--force-refresh", action="store_true",
                        help="Ignore cached token and fetch a new one")
    parser.add_argument("--json", action="store_true",
                        help="Print the full token record as JSON")
    parser.add_argument("-v", "--verbose", action="store_true",
                        help="Enable debug logging")
    return parser.parse_args(argv)


def main(argv: Optional[list[str]] = None) -> int:
    args = parse_args(argv)
    logging.basicConfig(
        level=logging.DEBUG if args.verbose else logging.INFO,
        format="%(asctime)s %(levelname)s %(name)s: %(message)s",
        stream=sys.stderr,  # keep stdout clean for the token itself
    )

    client = OAuth2TokenClient(
        client_id=args.client_id,
        client_secret=args.client_secret,
        token_url=args.token_url,
        scope=args.scope,
        cache=TokenCache(args.cache_file),
    )

    try:
        record = client.get_token(force_refresh=args.force_refresh)
    except RuntimeError as exc:
        logger.error("%s", exc)
        return 1

    if args.json:
        print(json.dumps(asdict(record)))
    else:
        print(record.access_token)
    return 0


if __name__ == "__main__":
    sys.exit(main())

Configuration

Usage Examples

1. Basic client_credentials fetch (first run — cache miss):

$ python3 oauth2_client.py \
    --client-id svc-reporting \
    --client-secret "$SECRET" \
    --token-url https://auth.example.com/token \
    --scope "reports:read"
2025-01-15 10:22:01 INFO oauth2_client: Requesting new token via client_credentials grant
eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiJzdmMtcmVwb3J0aW5nIi...

2. Second run within the token's lifetime — cache hit, no network call:

$ python3 oauth2_client.py --client-id svc-reporting \
    --client-secret "$SECRET" --token-url https://auth.example.com/token
eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiJzdmMtcmVwb3J0aW5nIi...

3. Force a fresh token and capture it in a shell variable:

$ TOKEN=$(python3 oauth2_client.py --client-id svc-reporting \
    --client-secret "$SECRET" --token-url https://auth.example.com/token \
    --force-refresh 2>/dev/null)
$ curl -H "Authorization: Bearer $TOKEN" https://api.example.com/v1/reports
{"reports": [...]}

Error Handling

• Missing required arguments — argparse prints a usage message and exits with code 2.
• Network failures / timeouts — requests.RequestException is caught and re-raised as a RuntimeError with the target URL, so the user knows which endpoint is unreachable.
• OAuth error responses (e.g. invalid_client, invalid_grant) — AuthlibBaseError is caught and surfaced with the provider's message. Exit code 1.
• Corrupt or unreadable cache — logged as a warning, then treated as a cache miss so the script self-heals on the next successful fetch.
• Malformed cache entries — TokenRecord(**raw) is guarded by a TypeError check; bad entries are ignored.
• Partial writes — tokens are written to a temp file and os.replaced atomically, so a crash mid-write never corrupts the cache. File mode is 0600.
• Token leakage in logs — logging goes to stderr; only the token itself (or JSON record) is printed to stdout, keeping shell pipelines clean.

Extending Ideas

• Add a --validate-url option that performs a lightweight GET against the API with the new token and exits non-zero if the token is rejected — useful as a pre-flight check in CI.
• Support the authorization code flow by adding --auth-url, --redirect-uri, and a local callback server (using http.server) so interactive logins can seed the cache once, after which refresh tokens take over.
• Encrypt the cache at rest with cryptography.fernet, keyed by an env var like OAUTH2_CACHE_KEY, for environments where the cache directory is on shared storage.
• Wrap as a systemd timer or Kubernetes CronJob that refreshes tokens into a shared Secret/ConfigMap, so downstream workloads never need client credentials at all.