Patched code sample

import secrets
import webbrowser
from http.server import BaseHTTPRequestHandler, TCPServer
from typing import Any, Dict, List, Optional, Tuple
from urllib.parse import parse_qs, urlparse

import requests

# This code is based on the patched version 4.2.0 of sigstore-python.
# It includes the surrounding classes and definitions from sigstore/_internal/oidc/oauth.py
# to provide a complete, contextual demonstration of the fix.
# The `_OIDCProvider` class is a minimal mock for demonstration purposes.


class OIDCError(Exception):
    pass


class _OIDCProvider:
    """A mock OIDC provider for demonstration."""
    auth_endpoint = "https://oauth.example.com/auth"
    token_endpoint = "https://oauth.example.com/token"
    code_challenge = "mock_challenge"


class _RedirectHandler(BaseHTTPRequestHandler):
    """
    A simple handler for the OAuth redirect.
    """

    def do_GET(self) -> None:
        server: "_RedirectTCPServer"
        server = self.server  # type: ignore

        server.query_params = parse_qs(urlparse(self.path).query)

        self.send_response(200)
        self.send_header("Content-Type", "text/plain; charset=utf-8")
        self.end_headers()
        self.wfile.write(
            b"sigstore-python: success! You may now close this page."
        )


class _RedirectTCPServer(TCPServer):
    """
    A simple TCP server for the OAuth redirect.
    """

    def __init__(
        self,
        server_address: Tuple[str, int],
        RequestHandlerClass: Any,
    ) -> None:
        super().__init__(server_address, RequestHandlerClass)
        self.query_params: Optional[Dict[str, List[str]]] = None


class _OAuthSession:
    """
    Represents an OAuth 2.0 session.
    """

    def __init__(self, client_id: str, client_secret: str, provider: _OIDCProvider):
        self._client_id = client_id
        self._client_secret = client_secret
        self._provider = provider

        # The `state` parameter is used to mitigate CSRF.
        #
        # See: https://datatracker.ietf.org/doc/html/rfc6749#section-10.12
        self._state = secrets.token_urlsafe()

    def _get_identity_token(self, code: str, redirect_uri: str) -> str:
        # In a real implementation, this function would exchange the
        # authorization code for an identity token.
        return "dummy-identity-token"

    def run(self) -> str:
        """
        Runs the OAuth 2.0 flow.
        """
        # Start a local server to handle the OAuth redirect.
        with _RedirectTCPServer(
            ("127.0.0.1", 0), _RedirectHandler
        ) as httpd:
            host, port = httpd.server_address
            redirect_uri = f"http://{host}:{port}"

            params = {
                "response_type": "code",
                "client_id": self._client_id,
                "redirect_uri": redirect_uri,
                "scope": "openid email",
                "state": self._state,
                "code_challenge": self._provider.code_challenge,
                "code_challenge_method": "S256",
            }

            auth_endpoint = self._provider.auth_endpoint
            req = requests.Request("GET", auth_endpoint, params=params)
            url = req.prepare().url
            assert url is not None

            # Open the authorization URL in the user's browser.
            try:
                webbrowser.open(url)
                print(f"Your browser has been opened to visit:\n\n    {url}\n")
            except webbrowser.Error:
                print(f"Please visit the following URL in your browser:\n\n    {url}\n")

            # The server will only handle one request.
            httpd.handle_request()

            # We expect the query parameters to be populated by the handler.
            if not httpd.query_params:
                raise OIDCError("did not receive query parameters from redirect")

            # THE FIX IS HERE:
            # Check that the state is the same, to prevent CSRF.
            # In the vulnerable version, this block was missing.
            #
            # See: https://datatracker.ietf.org/doc/html/rfc6749#section-10.12
            response_state_list = httpd.query_params.get("state")
            if response_state_list is None:
                raise OIDCError("did not receive 'state' in redirect")

            # `parse_qs` returns a list of values for each key.
            # The `state` parameter is not supposed to be repeated, so we
            # expect a list of one item.
            if response_state_list != [self._state]:
                raise OIDCError(
                    f"mismatched 'state' in redirect: "
                    f"expected {self._state!r}, got {response_state_list!r}"
                )

            code_list = httpd.query_params.get("code")
            if code_list is None:
                raise OIDCError("did not receive 'code' in redirect")
            code = code_list[0]

        # Now that we have the authorization code, we can exchange it for an
        # identity token.
        return self._get_identity_token(code=code, redirect_uri=redirect_uri)