285 lines
5.7 KiB
Python
285 lines
5.7 KiB
Python
"""
|
|
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
|