96 lines
2.9 KiB
Python
96 lines
2.9 KiB
Python
from __future__ import annotations
|
|
|
|
import base64
|
|
import hmac
|
|
import json
|
|
import os
|
|
import time
|
|
from dataclasses import dataclass
|
|
from hashlib import sha256
|
|
from typing import Any, Dict, Tuple
|
|
|
|
|
|
def _b64url_encode(data: bytes) -> str:
|
|
return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii")
|
|
|
|
|
|
def _b64url_decode(data: str) -> bytes:
|
|
padding = "=" * (-len(data) % 4)
|
|
return base64.urlsafe_b64decode(data + padding)
|
|
|
|
|
|
@dataclass
|
|
class JwtConfig:
|
|
issuer: str
|
|
audience: str
|
|
access_secret: bytes
|
|
refresh_secret: bytes
|
|
access_ttl_seconds: int
|
|
refresh_ttl_seconds: int
|
|
|
|
|
|
def _sign(secret: bytes, msg: bytes) -> str:
|
|
sig = hmac.new(secret, msg, sha256).digest()
|
|
return _b64url_encode(sig)
|
|
|
|
|
|
def _encode_header() -> str:
|
|
header = {"alg": "HS256", "typ": "JWT"}
|
|
return _b64url_encode(json.dumps(header, separators=(",", ":")).encode("utf-8"))
|
|
|
|
|
|
def _encode_payload(claims: Dict[str, Any]) -> str:
|
|
return _b64url_encode(json.dumps(claims, separators=(",", ":")).encode("utf-8"))
|
|
|
|
|
|
def create_jwt(
|
|
config: JwtConfig, subject: str, kind: str = "access", extra: Dict[str, Any] | None = None
|
|
) -> str:
|
|
now = int(time.time())
|
|
ttl = config.access_ttl_seconds if kind == "access" else config.refresh_ttl_seconds
|
|
secret = config.access_secret if kind == "access" else config.refresh_secret
|
|
claims: Dict[str, Any] = {
|
|
"iss": config.issuer,
|
|
"aud": config.audience,
|
|
"sub": subject,
|
|
"iat": now,
|
|
"exp": now + ttl,
|
|
"typ": kind,
|
|
}
|
|
if extra:
|
|
claims.update(extra)
|
|
header = _encode_header()
|
|
payload = _encode_payload(claims)
|
|
signing_input = f"{header}.{payload}".encode("ascii")
|
|
signature = _sign(secret, signing_input)
|
|
return f"{header}.{payload}.{signature}"
|
|
|
|
|
|
def verify_jwt(
|
|
config: JwtConfig, token: str, expected_kind: str = "access"
|
|
) -> Tuple[Dict[str, Any], Dict[str, Any]]:
|
|
try:
|
|
header_b64, payload_b64, sig = token.split(".")
|
|
except ValueError:
|
|
raise ValueError("Invalid token format")
|
|
signing_input = f"{header_b64}.{payload_b64}".encode("ascii")
|
|
header = json.loads(_b64url_decode(header_b64))
|
|
if header.get("alg") != "HS256" or header.get("typ") != "JWT":
|
|
raise ValueError("Unsupported JWT header")
|
|
payload = json.loads(_b64url_decode(payload_b64))
|
|
kind = payload.get("typ")
|
|
secret = config.access_secret if kind == "access" else config.refresh_secret
|
|
if not hmac.compare_digest(sig, _sign(secret, signing_input)):
|
|
raise ValueError("Invalid signature")
|
|
now = int(time.time())
|
|
if payload.get("iss") != config.issuer or payload.get("aud") != config.audience:
|
|
raise ValueError("Invalid claims")
|
|
if kind != expected_kind:
|
|
raise ValueError("Invalid token type")
|
|
if int(payload.get("exp", 0)) < now:
|
|
raise ValueError("Token expired")
|
|
return header, payload
|
|
|
|
|
|
def random_secret(n: int = 32) -> bytes:
|
|
return os.urandom(n)
|