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" 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 | str] ) -> list[UserCapability]: return [ UserCapability(str(capability)) for capability in value ] @classmethod def from_db( cls, db: Session, name: str, ) -> User | None: user = (db .query(models.User) .filter(models.User.name == name) .first()) if user is None: return None return cls.from_orm(user) @classmethod def login( cls, db: Session, name: str, password: str, crypt_context: CryptContext, ) -> User | None: user = (db .query(models.User) .filter(models.User.name == name) .first()) if user is None: # inexistent user, fake doing password verification crypt_context.dummy_verify() return None if not crypt_context.verify(password, user.password): # password hash mismatch return None return cls.from_orm(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=[models.UserCapability(capability="admin")], ) db.add(user) db.commit() db.refresh(user) return cls.from_orm(user) except IntegrityError: pass 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