""" Pydantic representation of database contents. """ from __future__ import annotations from datetime import datetime from enum import Enum from typing import Any from passlib.context import CryptContext from pydantic import BaseModel, validator from sqlalchemy.exc import IntegrityError from sqlalchemy.orm import Session from . import models ########## # table: distinguished_names ########## class DistinguishedNameBase(BaseModel): cn_only: bool country: str state: str city: str organization: str organizational_unit: str email: str common_name: str class DistinguishedNameCreate(DistinguishedNameBase): pass class DistinguishedName(DistinguishedNameBase): class Config: orm_mode = True ########## # table: certificates ########## class CertificateBase(BaseModel): expiry: datetime class CertificateCreate(CertificateBase): pass class Certificate(CertificateBase): distinguished_name: DistinguishedName class Config: orm_mode = True ########## # table: user_capabilities ########## class UserCapability(Enum): admin = "admin" @classmethod def from_value(cls, value) -> UserCapability: """ Create UserCapability from various formats """ if isinstance(value, cls): # value is already a UserCapability, use that return value elif isinstance(value, models.UserCapability): # create from db format return cls(value.capability) else: # create from string representation return cls(str(value)) ########## # table: users ########## class UserBase(BaseModel): name: str class UserCreate(UserBase): password: str class User(UserBase): capabilities: list[UserCapability] = [] distinguished_names: list[DistinguishedName] = [] certificates: list[Certificate] = [] class Config: orm_mode = True @validator("capabilities", pre=True) @classmethod def unify_capabilities(cls, value: list[Any]) -> list[UserCapability]: """ Import the capabilities from various formats """ return [ UserCapability.from_value(capability) for capability in value ] @classmethod def from_db( cls, db: Session, name: str, ) -> User | None: """ Load user from database by name. """ if (db_user := models.User.load(db, name)) is None: return None return cls.from_orm(db_user) @classmethod def create( cls, db: Session, user: UserCreate, crypt_context: CryptContext, ) -> User | None: """ Create a new user in the database. """ try: user = models.User( name=user.name, password=crypt_context.hash(user.password), capabilities=[], ) db.add(user) db.commit() db.refresh(user) return cls.from_orm(user) except IntegrityError: # user already existed pass def authenticate( self, db: Session, password: str, crypt_context: CryptContext, ) -> User | None: """ Authenticate with name/password against users in database. """ if (db_user := models.User.load(db, self.name)) is None: # nonexistent user, fake doing password verification crypt_context.dummy_verify() return False if not crypt_context.verify(password, db_user.password): # password hash mismatch return False self.from_orm(db_user) return True def update( self, db: Session, ) -> None: """ Update this user in the database. """ old_dbuser = models.User.load(db, self.name) old_user = self.from_orm(old_dbuser) for capability in self.capabilities: if capability not in old_user.capabilities: old_dbuser.capabilities.append( models.UserCapability(capability=capability.value) ) db.commit()