from __future__ import annotations from datetime import datetime from enum import Enum from passlib.context import CryptContext from pydantic import BaseModel, validator from sqlalchemy.exc import IntegrityError from sqlalchemy.orm import Session from . import models class CertificateBase(BaseModel): expiry: datetime class CertificateCreate(CertificateBase): owner_name: str dn_id: int class Certificate(CertificateBase): id: int class Config: orm_mode = True class UserCapability(Enum): admin = "admin" @classmethod def from_value(cls, value) -> UserCapability: if isinstance(value, cls): # use simple value return value elif isinstance(value, models.UserCapability): # create from db return cls(value.capability) else: # create from string representation return cls(str(value)) class UserBase(BaseModel): name: str class UserCreate(UserBase): password: str class User(UserBase): certificates: list[Certificate] capabilities: list[UserCapability] class Config: orm_mode = True @validator("capabilities", pre=True) @classmethod def unify_capabilities( cls, value: list[models.UserCapability | UserCapability | str] ) -> list[UserCapability]: return [ UserCapability.from_value(capability) for capability in value ] @classmethod def from_db( cls, db: Session, name: str, ) -> User | None: if (db_user := models.User.load(db, name)) is None: return None return cls.from_orm(db_user) @classmethod def login( cls, db: Session, name: str, password: str, crypt_context: CryptContext, ) -> User | None: if (db_user := models.User.load(db, name)) is None: # inexistent user, fake doing password verification crypt_context.dummy_verify() return None if not crypt_context.verify(password, db_user.password): # password hash mismatch return None return cls.from_orm(db_user) @classmethod def create( cls, db: Session, user: UserCreate, crypt_context: CryptContext, ) -> User | None: 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: pass def add_capabilities( self, db: Session, capabilities: list[UserCapability], ) -> None: for capability in capabilities: if capability not in self.capabilities: db.add(models.UserCapability( user_name=self.name, capability=capability.value, )) db.commit() 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): id: int certificates: list[Certificate] class Config: orm_mode = True