cereal/src/services/auth.py

100 lines
3 KiB
Python

from datetime import datetime, UTC, timedelta
import abc
from src.config.email import EmailAddress
from src.config.parse import assert_key_of_type, ParseError
from typing import Any, TypeVar
from dataclasses import dataclass, asdict
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
import jwt
from src.services.users.data import User
from src.config.auth import Auth as AuthConfig
JWT = str
T = TypeVar('T', bound='Claim')
@dataclass
class Claim(abc.ABC):
sub: int
exp: datetime
iat: datetime
@classmethod
@abc.abstractmethod
def from_dict(cls, claims: dict[str, Any]) -> T:
pass
@classmethod
@abc.abstractmethod
def from_user(cls, user: User, expire_in_secs: int) -> T:
pass
class ExpiredTokenError(Exception):
def __init__(self):
super().__init__('token was expired')
@dataclass
class BaseClaims(Claim):
@classmethod
def from_user(cls, user: User, expire_in_secs: int) -> 'BaseClaims':
now = datetime.now(tz=UTC)
expires_at = now + timedelta(seconds=expire_in_secs)
return BaseClaims(sub=user.id, exp=expires_at, iat=now)
@classmethod
def from_dict(cls, claims: dict[str, Any]) -> 'BaseClaims':
assert_key_of_type(claims, 'sub', int)
assert_key_of_type(claims, 'exp', int)
assert_key_of_type(claims, 'iat', int)
exp = datetime.fromtimestamp(claims['exp'], UTC)
now = datetime.now(UTC)
if exp < now:
raise ExpiredTokenError()
return BaseClaims(
sub=claims['sub'],
exp=exp,
iat=datetime.fromtimestamp(claims['iat'], UTC),
)
@dataclass
class EmailConfirmationClaim(BaseClaims):
email: EmailAddress
@classmethod
def from_user(cls, user: User, expire_in_secs) -> 'EmailConfirmationClaim':
base = BaseClaims.from_user(user, expire_in_secs)
return EmailConfirmationClaim(**asdict(base), email=user.email)
@classmethod
def from_dict(cls, claims: dict[str, Any]) -> 'EmailConfirmationClaim':
base = BaseClaims.from_dict(claims)
assert_key_of_type(claims, 'email', str)
return EmailConfirmationClaim(**asdict(base), email=claims['email'])
class AuthService:
def __init__(self, config: AuthConfig):
self._private_key = Ed25519PrivateKey.from_private_bytes(config.ed25519_private_key.expose_secret())
self._public_key = self._private_key.public_key()
def mint_claim_from_user(self, claim: type[Claim], user: User, expires_in_secs=600) -> JWT:
claims = claim.from_user(user, expires_in_secs)
return jwt.encode(asdict(claims), self._private_key, algorithm='EdDSA')
def validate_token[T: Claim](self, claim: type[T], token: JWT) -> T | None:
try:
claims = jwt.decode(token, key=self._public_key, algorithms=['EdDSA'])
except jwt.InvalidTokenError:
return None
try:
return claim.from_dict(claims)
except ParseError:
return None