OAuth Authorization Code flow from Python desktop

Introduction

A few days ago, I found myself requiring a way to contact an application secured with OAuth 2, with these restrictions:

• I needed to use user credentials (no Client Credentials flow),
• the source code had to be Python,
• the code had to be runnable from VS Code or a Jupyter Notebook

It's quite easy to find tutorials to achieve the above using C#, through .NET's Kestrel server. But either my Google ninja skills failed or this is not common in Python - regardless, I had to do it from scratch. I'm not proficient in Python in any sense of the word, so it's likely that it can be improved significantly.

Building the solution

In this tutorial, we are going to be building the following flow:

1. We start an HTTP server that will handle the OAuth redirection.
2. We then build the Authorization URI so that the user can sign in.
3. Next, we open the website (and browser, if needed), for the user to enter their credentials, which happens at the OAuth server.
4. The OAuth server then redirects the user back to our HTTP server, which contains the PKCE Code we need for the second part of the authorization flow. At this point, the HTTP server can be shut down since it's no longer needed.
5. Finally, we build the Token request body and call the OAuth server to exchange the Code for:
   • An Access Token,
   • A Refresh Token (if the offline scope is used),
   • An ID Token (if the openid scope is used and the server supports OIDC)

Custom HTTP server

The following simple code represents an HTTP server that allows us to parse the response back

from the OAuth server:
from http.server import BaseHTTPRequestHandler, HTTPServer
from urllib import parse

class OAuthHttpServer(HTTPServer):
    def __init__(self, *args, **kwargs):
        HTTPServer.__init__(self, *args, **kwargs)
        self.authorization_code = ""


class OAuthHttpHandler(BaseHTTPRequestHandler):
    def do_GET(self):
        self.send_response(200)
        self.send_header("Content-Type", "text/html")
        self.end_headers()
        self.wfile.write("<script type=\"application/javascript\">window.close();</script>".encode("UTF-8"))
        
        parsed = parse.urlparse(self.path)
        qs = parse.parse_qs(parsed.query)
        
        self.server.authorization_code = qs["code"][0]

You will notice that we're returning a simple script to close the window/tab automatically since there is nothing for users to do with it.

Performing the login

The login process is done using the custom server implemented above.

from typing import Any
import webbrowser
import requests
from oauthlib.oauth2 import WebApplicationClient
        
def login(config: dict[str, Any]) -> str:
    # 1
    with OAuthHttpServer(('', config["port"]), OAuthHttpHandler) as httpd:
        client = WebApplicationClient(config["client_id"])
        
        # 2
        code_verifier, code_challenge = generate_code()

        auth_uri = client.prepare_request_uri(config["auth_uri"], redirect_uri=config["redirect_uri"], 
            scope=config["scopes"], state="test_doesnotmatter", code_challenge= code_challenge, code_challenge_method = "S256" )

        # 3
        webbrowser.open_new(auth_uri)

        # 4
        httpd.handle_request()

        auth_code = httpd.authorization_code

        data = {
            "code": auth_code,
            "client_id": config["client_id"],
            "grant_type": "authorization_code",
            "scopes": config["scopes"],
            "redirect_uri": config["redirect_uri"],
            "code_verifier": code_verifier
        }

        # 5
        response = requests.post(config["token_uri"], data=data, verify=False)

        access_token = response.json()["access_token"]
        clear_output()

        print("Logged in successfully")

        return access_token

If we look at the numbered comments, we can see the following:

1. Starts the server.
2. Obtains an PKCE code.
3. Opens the generated Authentication URI in the browser.
4. Handles the request back from the OAuth server.
5. Exchanges the PKCE code for a token (as explained above)

Generating the PKCE code

We need two parts here: 1) the code_challenge that is sent with the authentication request, and 2) the code_verifier which is used for verifying the originator together with the code when requesting the token. The code_challenge is a hashed version of the code_verifier and which hashing method is supported depends on your OAuth server, but a common one is SHA-256. This can be generated with the following:

def generate_code() -> tuple[str, str]:
    rand = random.SystemRandom()
    code_verifier = ''.join(rand.choices(string.ascii_letters + string.digits, k=128))

    code_sha_256 = hashlib.sha256(code_verifier.encode('utf-8')).digest()
    b64 = base64.urlsafe_b64encode(code_sha_256)
    code_challenge = b64.decode('utf-8').replace('=', '')

    return (code_verifier, code_challenge)

Getting a token

The final part of this is generating the configuration needed for the login method:

config = {
    "port": "port_for_listening",
    "client_id": "{your_client_id}",
    "redirect_uri": f"http://localhost:{same_as_port}",
    "auth_uri": "{authentication_uri}",
    "token_uri": "{token_uri}",
    "scopes": [ "openid", "profile", "any_other_scope" ]
}

access_token = login(config)
headers = { "Authorization": "Bearer " + access_token }

The URIs for authentication and token must be retrieved from the documentation of your OAuth server. For OIDC servers, it can be found in the /.well-known/openid-configuration endpoint.

Demo

You can try out this code in the notebook I have uploaded to my GitHub samples repository. Notice that you will need to either run the sample code for "OAuth PKCE flow for ASP.NET Core with Swagger" (if you have .NET installed) or update the URLs to your systems otherwise.