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

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)

    capabilities = relationship("UserCapability", lazy="joined")
    certificates = relationship("Certificate", lazy="joined")

    @classmethod
    def load(cls, db: Session, name: str) -> User | None:
        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)

    cn_only = Column(Boolean, default=True)
    country = Column(String(2))
    state = Column(String)
    city = Column(String)
    organization = Column(String)
    organizational_unit = Column(String)
    email = Column(String)
    common_name = Column(String)

    certificates = relationship("Certificate", lazy="joined")

    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"))
    expiry = Column(DateTime, default=datetime.datetime.now)