""" 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, Field, validator from sqlalchemy.exc import IntegrityError from sqlalchemy.orm import Session from . import models ########## # table: user_capabilities ########## class UserCapability(Enum): admin = "admin" login = "login" issue = "issue" renew = "renew" def __repr__(self) -> str: return self.value @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)) @property def model(self) -> models.UserCapability: return models.UserCapability( capability=self.value, ) ########## # table: users ########## class UserBase(BaseModel): name: str country: str state: str city: str organization: str organizational_unit: str email: str capabilities: list[UserCapability] = [] class UserCreate(UserBase): password: str class User(UserBase): devices: list[Device] = Field( default=[], repr=False ) 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. """ db_user = models.User(name=name) db.refresh(db_user) 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: db_user = models.User( name=user.name, password=crypt_context.hash(user.password), capabilities=[ capability.model for capability in user.capabilities ], ) db.add(db_user) db.commit() db.refresh(db_user) return cls.from_orm(db_user) except IntegrityError: # user already existed pass def is_admin(self) -> bool: return UserCapability.admin in self.capabilities def authenticate( self, db: Session, password: str, crypt_context: CryptContext, ) -> User | None: """ Authenticate with name/password against users in database. """ db_user = models.User(name=self.name) db.refresh(db_user) if db_user 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. """ db_user = models.User(name=self.name) db.refresh(db_user) for capability in db_user.capabilities: db.delete(capability) db_user.capabilities = [ capability.model for capability in self.capabilities ] db.commit() def delete( self, db: Session, ) -> bool: """ Delete this user from the database. """ db_user = models.User(name=self.name) db.refresh(db_user) if db_user is None: # nonexistent user return False db.delete(db_user) db.commit() return True ########## # table: devices ########## class DeviceBase(BaseModel): name: str type: str expiry: datetime class DeviceCreate(DeviceBase): owner_name: str class Device(DeviceBase): class Config: orm_mode = True @classmethod def create( cls, db: Session, device: DeviceCreate, ) -> Device | None: """ Create a new device in the database. """ try: db_device = models.Device( owner_name=device.owner_name, name=device.name, type=device.type, expiry=device.expiry, ) db.add(db_device) db.commit() db.refresh(db_device) return cls.from_orm(db_device) except IntegrityError: # device already existed pass def delete( self, db: Session, ) -> bool: """ Delete this device from the database. """ db_device = models.Device( # owner_name= name=self.name, ) db.refresh(db_device) if db_device is None: # nonexistent device return False db.delete(db_device) db.commit() return True