kiwi-vpn/api/kiwi_vpn_api/db/schemas.py
2022-03-20 01:37:27 +00:00

162 lines
3.4 KiB
Python

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