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

"""
SQLAlchemy representation of database contents.
"""

from __future__ import annotations

import datetime

from sqlalchemy import (Boolean, Column, DateTime, ForeignKey, Integer, String,
                        UniqueConstraint)
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import Session, relationship

ORMBaseModel = declarative_base()


class User(ORMBaseModel):
    __tablename__ = "users"

    name = Column(String, primary_key=True, index=True)
    password = Column(String, nullable=False)

    capabilities: list[UserCapability] = relationship(
        "UserCapability", lazy="joined", cascade="all, delete-orphan"
    )
    certificates: list[Certificate] = relationship(
        "Certificate", lazy="select",  back_populates="owner"
    )
    distinguished_names: list[DistinguishedName] = relationship(
        "DistinguishedName", lazy="select", back_populates="owner"
    )

    @classmethod
    def load(cls, db: Session, name: str) -> User | None:
        """
        Load user from database by name.
        """

        return (db
                .query(User)
                .filter(User.name == name)
                .first())


class UserCapability(ORMBaseModel):
    __tablename__ = "user_capabilities"

    user_name = Column(
        String,
        ForeignKey("users.name"),
        primary_key=True,
        index=True,
    )
    capability = Column(String, primary_key=True)


class DistinguishedName(ORMBaseModel):
    __tablename__ = "distinguished_names"

    id = Column(Integer, primary_key=True, autoincrement=True)

    owner_name = Column(String, ForeignKey("users.name"))
    cn_only = Column(Boolean, default=True, nullable=False)
    country = Column(String(2))
    state = Column(String)
    city = Column(String)
    organization = Column(String)
    organizational_unit = Column(String)
    email = Column(String)
    common_name = Column(String, nullable=False)

    owner: User = relationship(
        "User", lazy="joined", back_populates="distinguished_names"
    )

    UniqueConstraint(
        country,
        state,
        city,
        organization,
        organizational_unit,
        email,
        common_name,
    )


class Certificate(ORMBaseModel):
    __tablename__ = "certificates"

    id = Column(Integer, primary_key=True, autoincrement=True)

    owner_name = Column(String, ForeignKey("users.name"))
    dn_id = Column(
        Integer,
        ForeignKey("distinguished_names.id"),
        nullable=False,
    )
    expiry = Column(DateTime, default=datetime.datetime.now)

    distinguished_name: DistinguishedName = relationship(
        "DistinguishedName", lazy="joined"
    )

    owner: User = relationship(
        "User", lazy="joined", back_populates="certificates"
    )