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

from __future__ import annotations

import datetime
from typing import Optional

from peewee import (BooleanField, CharField, DateTimeField, ForeignKeyField,
                    Model)

from ..config import CRYPT_CONTEXT, DB


class BaseModel(Model):
    class Meta:
        database = DB


class User(BaseModel):
    name = CharField(unique=True)
    password = CharField()

    @classmethod
    def get_by_name(cls, name: str) -> Optional[User]:
        query = (cls.select()
                 .where(cls.name == name)
                 .prefetch(UserCapability))

        if query:
            return query[0]

    def verify(self, password: str) -> bool:
        if CRYPT_CONTEXT.verify(password, self.password):
            return True

        return False


class UserCapability(BaseModel):
    user = ForeignKeyField(User, backref="capabilities")
    capability = CharField()


class DistinguishedName(BaseModel):
    cn_only = BooleanField(default=True)
    common_name = CharField()
    email = CharField()
    organizational_unit = CharField()
    organization = CharField()
    city = CharField()
    state = CharField()
    country = CharField(max_length=2, default="DE")

    class Meta:
        # Certs are unique
        indexes = (
            (('common_name', 'email', 'organizational_unit',
             'organization', 'city', 'state', 'country'), True),
        )


class Certificate(BaseModel):
    owner = ForeignKeyField(User, backref="certs")
    distinguished_name = ForeignKeyField(DistinguishedName)
    expiry = DateTimeField(default=datetime.datetime.now)