From e2f916debcd669d38213ad7fdbbbaa65013e61ec Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B6rn-Michael=20Miehe?=
 <40151420+ldericher@users.noreply.github.com>
Date: Sun, 27 Mar 2022 13:47:38 +0000
Subject: [PATCH] Capabilities

---
 api/kiwi_vpn_api/db_new/capabilities.py | 32 ++++++++++++++++++
 api/kiwi_vpn_api/db_new/user.py         | 44 +++++++++++++++++++++++--
 api/kiwi_vpn_api/main.py                | 10 +++++-
 3 files changed, 83 insertions(+), 3 deletions(-)
 create mode 100644 api/kiwi_vpn_api/db_new/capabilities.py

diff --git a/api/kiwi_vpn_api/db_new/capabilities.py b/api/kiwi_vpn_api/db_new/capabilities.py
new file mode 100644
index 0000000..7d92064
--- /dev/null
+++ b/api/kiwi_vpn_api/db_new/capabilities.py
@@ -0,0 +1,32 @@
+from enum import Enum
+from typing import TYPE_CHECKING
+
+from sqlmodel import Field, Relationship, SQLModel
+
+if TYPE_CHECKING:
+    from .user import User
+
+
+class Capability(Enum):
+    admin = "admin"
+    login = "login"
+    issue = "issue"
+    renew = "renew"
+
+
+class UserCapabilityBase(SQLModel):
+    user_name: str = Field(primary_key=True, foreign_key="user.name")
+    capability_name: str = Field(primary_key=True)
+
+    @property
+    def _(self) -> Capability:
+        return Capability(self.capability_name)
+
+    def __repr__(self) -> str:
+        return self.capability_name
+
+
+class UserCapability(UserCapabilityBase, table=True):
+    user: "User" = Relationship(
+        back_populates="capabilities"
+    )
diff --git a/api/kiwi_vpn_api/db_new/user.py b/api/kiwi_vpn_api/db_new/user.py
index 45235e6..3d3eea7 100644
--- a/api/kiwi_vpn_api/db_new/user.py
+++ b/api/kiwi_vpn_api/db_new/user.py
@@ -1,12 +1,13 @@
 from __future__ import annotations
 
-from typing import Any
+from typing import Any, Sequence
 
 from pydantic import root_validator
 from sqlalchemy.exc import IntegrityError
-from sqlmodel import Field, SQLModel
+from sqlmodel import Field, Relationship, SQLModel
 
 from ..config import Config
+from .capabilities import Capability, UserCapability
 from .connection import Connection
 
 
@@ -24,6 +25,14 @@ class UserBase(SQLModel):
 class User(UserBase, table=True):
     password: str
 
+    capabilities: list[UserCapability] = Relationship(
+        back_populates="user",
+        sa_relationship_kwargs={
+            "lazy": "joined",
+            "cascade": "all, delete-orphan",
+        },
+    )
+
     @classmethod
     def create(cls, **kwargs) -> User | None:
         """
@@ -94,6 +103,37 @@ class User(UserBase, table=True):
         with Connection.session as db:
             db.delete(self)
             db.commit()
+
+    def extend_capabilities(self, capabilities: Sequence[Capability]) -> None:
+        """
+        Extend this user's capabilities
+        """
+
+        for capability in capabilities:
+            user_capability = UserCapability(
+                user_name=self.name,
+                capability_name=capability.value,
+            )
+
+            if user_capability not in self.capabilities:
+                self.capabilities.append(user_capability)
+
+    def remove_capabilities(self, capabilities: Sequence[Capability]) -> None:
+        """
+        Remove from this user's capabilities
+        """
+
+        for capability in capabilities:
+            try:
+                self.capabilities.remove(UserCapability(
+                    user_name=self.name,
+                    capability_name=capability.value,
+                ))
+
+            except ValueError:
+                pass
+
+
 class UserCreate(UserBase):
     password: str | None = Field(default=None)
     password_clear: str | None = Field(default=None)
diff --git a/api/kiwi_vpn_api/main.py b/api/kiwi_vpn_api/main.py
index 335d058..074e5b5 100755
--- a/api/kiwi_vpn_api/main.py
+++ b/api/kiwi_vpn_api/main.py
@@ -15,7 +15,7 @@ from fastapi import FastAPI
 from .config import Config, Settings
 from .db import Connection
 from .db.schemata import User
-from .db_new import connection, user
+from .db_new import capabilities, connection, user
 from .routers import main_router
 
 settings = Settings.get()
@@ -65,6 +65,14 @@ async def on_startup() -> None:
 
         uwe = user.User.authenticate("Uwe", "ulf")
 
+        uwe.extend_capabilities([capabilities.Capability.admin])
+        uwe.update()
+        print(uwe)
+
+        uwe.remove_capabilities([capabilities.Capability.admin])
+        uwe.update()
+        print(uwe)
+
 
 def main() -> None:
     uvicorn.run(