kiwi-vpn/api/kiwi_vpn_api/db/schemas.py

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