gatekeeper/app/model/models.py

from datetime import UTC, datetime
from typing import Literal

from pydantic import model_validator
from sqlalchemy import DateTime, TypeDecorator
from sqlmodel import Field, Relationship, SQLModel


class Base(SQLModel):
    pass


#### Extra Column Types
class UTCDateTime(TypeDecorator):
    """Modified datetime for timezones
    Input time gets read as given tz or local tz, then converted to UTC
    Output everything as UTC
    """

    impl = DateTime(timezone=True)
    LOCAL_TIMEZONE = datetime.now(UTC).astimezone().tzinfo
    cache_ok = True

    def process_bind_param(self, value: datetime | None, dialect) -> datetime | None:
        if value is None:
            return None
        if value.tzinfo is None:
            value = value.astimezone(self.LOCAL_TIMEZONE)
        return value.astimezone(UTC)

    def process_result_value(self, value: datetime | None, dialect) -> datetime | None:
        if value is None:
            return None
        if value.tzinfo is None:
            return value.replace(tzinfo=UTC)
        return value.astimezone(UTC)


#### User
class UserBase(Base):
    name: str = Field(index=True, unique=True)
    email: str | None = None
    is_admin: bool = False


class UserResponse(UserBase):
    id: int


class UserCreate(UserBase):
    password: str


class UserDB(UserBase, table=True):
    id: int | None = Field(default=None, primary_key=True)
    passwordhash: str


class UserUpdate(Base):
    name: str | None = None
    email: str | None = None
    is_admin: bool | None = None
    password: str | None = None


#### Special
class AaGroupLink(Base, table=True):
    group_id: int | None = Field(
        default=None, foreign_key="groupdb.id", primary_key=True
    )
    accessauth_id: int | None = Field(
        default=None, foreign_key="accessauthorizationdb.id", primary_key=True
    )


#### Token
class Token(Base):
    access_token: str
    token_type: str


class TokenData(Base):
    username: str | None = None


#### Group
class GroupBase(Base):
    name: str = Field(index=True, unique=True)


class GroupCreate(GroupBase):
    pass


class GroupDB(GroupBase, table=True):
    id: int | None = Field(default=None, primary_key=True)
    cards: list["Card"] = Relationship(back_populates="group")
    accessauths: list["AccessAuthorizationDB"] = Relationship(
        back_populates="groups", link_model=AaGroupLink
    )


class GroupResponse(GroupBase):
    id: int
    cards: list["Card"] | None
    accessauths: list["AccessAuthorizationDB"] | None


#### AccessAuthorization
class AccessAuthorizationBase(Base):
    name: str = Field(index=True)
    type: Literal["oneshot", "somefuturespec"]
    is_active: bool


class AccessAuthorizationDB(AccessAuthorizationBase, table=True):
    id: int | None = Field(default=None, primary_key=True)
    type: str
    groups: list["GroupDB"] = Relationship(
        back_populates="accessauths", link_model=AaGroupLink
    )
    oneshot: "OneShotAccess" = Relationship(
        back_populates="accessauth", cascade_delete=True
    )


class OneShotAccessBase(Base):
    uses: int = 1
    ends_at: datetime = Field(sa_type=UTCDateTime)


class OneShotAccess(OneShotAccessBase, table=True):
    id: int | None = Field(default=None, primary_key=True)
    accessauth_id: int = Field(default=None, foreign_key="accessauthorizationdb.id")
    accessauth: AccessAuthorizationDB = Relationship(back_populates="oneshot")


class AccessAuthorizationCreate(AccessAuthorizationBase):
    oneshot: OneShotAccessBase | None = None

    @model_validator(mode="after")
    def check_type(self):
        if self.type == "oneshot":
            if not self.oneshot:
                raise ValueError("oneshot auths require a oneshot object")
        elif self.type == "somefuturespec":
            raise ValueError("somefuturespec is not jet implemented. Please do not use")
        return self


class AccessAuthorizationResponse(AccessAuthorizationBase):
    id: int
    oneshot: OneShotAccessBase | None = None
    groups: list["GroupDB"]


class AccessAuthorizationUpdate(Base):
    name: str | None = None
    type: Literal["oneshot", "somefuturespec"] | None = None
    is_active: bool | None = None
    oneshot: OneShotAccessBase | None = None


#### Card
class Card(Base, table=True):
    id: int | None = Field(default=None, primary_key=True)
    key: str
    card_serial: str
    enabled: bool = True
    name: str = Field(unique=True, max_length=32)
    group_id: int | None = Field(default=None, foreign_key="groupdb.id")
    group: GroupDB | None = Relationship(back_populates="cards")


class CardCreate(Base):
    name: str = Field(unique=True, max_length=32)
    enabled: bool = True
    group_id: int


class CardUpdate(Base):
    name: str | None = None
    enabled: bool | None = None
    group_id: int | None = None