pkce_examples.py

import argparse
import base64
import hashlib
import secrets


def generate_code_verifier() -> str:
    """
    Generate a code verifier for PKCE (Proof Key for Code Exchange).

    Returns:
        str: A base64 URL-safe encoded string that serves as the code verifier.
    """

    return (
        base64.urlsafe_b64encode(secrets.token_bytes(64)).rstrip(b"=").decode("utf-8")
    )


def generate_code_challenge(code_verifier: str) -> str:
    """
    Generate a code challenge from the code verifier using SHA-256.

    Args:
        code_verifier (str): The code verifier string.

    Returns:
        str: A base64 URL-safe encoded string that serves as the code challenge.
    """

    if not code_verifier:
        raise ValueError("Code verifier must not be empty.")

    code_challenge: bytes = hashlib.sha256(code_verifier.encode("utf-8")).digest()

    return base64.urlsafe_b64encode(code_challenge).rstrip(b"=").decode("utf-8")


def generate_state() -> str:
    """
    Generate a random state string for CSRF protection.

    Returns:
        str: A base64 URL-safe encoded string that serves as the state.
    """

    return (
        base64.urlsafe_b64encode(secrets.token_bytes(16)).rstrip(b"=").decode("utf-8")
    )


def main():
    parser: argparse.ArgumentParser = argparse.ArgumentParser(
        description="PKCE helper utilities: generate code verifier, code challenge, and state."
    )
    parser.add_argument(
        "--generate",
        choices=["verifier", "challenge", "state", "all"],
        required=True,
        help="What to generate: verifier, challenge, or state",
    )
    parser.add_argument(
        "--verifier", help="The code verifier to use for generating a code challenge"
    )

    args: argparse.Namespace = parser.parse_args()

    if args.generate == "verifier":
        verifier = generate_code_verifier()
        print(f"verifier={verifier}")

    elif args.generate == "challenge":
        if not args.verifier:
            print("No verifier provided. Use --verifier to specify the code verifier.")
            print("Generating a new verifier instead.")
            args.verifier = generate_code_verifier()

        challenge = generate_code_challenge(args.verifier)

        print(f"verifier={args.verifier}")
        print(f"challenge={challenge}")

    elif args.generate == "state":
        state = generate_state()
        print(f"state={state}")

    elif args.generate == "all":
        verifier = generate_code_verifier()
        challenge = generate_code_challenge(verifier)
        state = generate_state()

        print(f"verifier={verifier}")
        print(f"challenge={challenge}")
        print(f"state={state}")


if __name__ == "__main__":
    main()