Personal Cookie Handler script with Redis

cookie_handler.py

"""Package to assist cookie creation and management"""

from typing import Any, Dict, Union
import aioredis
from uuid import uuid4

REDIS_URL = "" # Your Redis instance's URL goes here


redis = aioredis.from_url(REDIS_URL)

class CookieHandler:
    """Class to handle cookies"""

    async def create_cookie(self, data: Dict, ttl: int):
        cookie_key = str(uuid4())
        await redis.hmset(name=cookie_key, mapping=data)
        await redis.expire(cookie_key, ttl) # If you want to make the expire optional, make this line only execute conditionally, with the ttl param to this function being an optional one
        return cookie_key

    async def retrieve_cookie(self, cookie: str) -> Union[Dict, None]:
        data = await redis.hgetall(cookie)
        if data is None:
            return None

        cleaned = {}
        for key, value in data.items():
            cleaned[key.decode()] = value.decode()

        return cleaned

example.py

"""An example script to display how the Cookie Handler can be used to securely store any python dictionary.
I made the Cookie Handler Script to help with managing state during multi-provider oauth (First Microsoft Teams, then Discord). Maintaining state during OAuth is good practice, and can prevent serious security vulnerabilities.

In this example, we have two routes, "/add_cookie" and "/show_cookie", with the former redirecting to the latter. When a user visits the add cookie route, a random dictionary is generated with their IP Address and a uuid4(), which is securely stored by the cookie_handler, a uuid is returned which can be set as the cookie's value. The user is then redirected to the show cookie route, where the cookie is fetched and displayed to the user.
"""

from fastapi import FastAPI, Request, Cookie
from fastapi.responses import RedirectResponse
from cookie_handler import CookieHandler
from uuid import uuid4

ch = CookieHandler()

@app.get("/add_cookie")
async def add_cookie(request: Request):
  resp = RedirectResponse("/show_cookie")
  cookie_key = await ch.create_cookie(data={"ip_address": request.client.host, "password": str(uuid4)}, ttl=30)
  # You can securely store any data here, the key expires after 30 seconds. I personally use the cookie_handler to store state during oauth, especially when the user has to undergo multiple steps, it adds vulnerability to your flow, which can be alleviated by maintaining state.
  resp.set_cookie(key="session", value=cookie_key)
  return resp

@app.get("/show_cookie")
async def show_cookie(request: Request, session: str = Cookie(default=None)):
  return ch.retrieve_cookie(session)

if __name__ == "__main__":
  import uvicorn
  uvicorn.run(app, host="127.0.0.1", port=8080)