From 04a5798258068b180726d909a3b5112a3c0a3416 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B6rn-Michael=20Miehe?=
 <40151420+ldericher@users.noreply.github.com>
Date: Mon, 28 Mar 2022 00:48:44 +0000
Subject: [PATCH] capabilities rework

---
 api/kiwi_vpn_api/db_new/capabilities.py |  8 ++++--
 api/kiwi_vpn_api/db_new/user.py         | 37 ++++++++-----------------
 2 files changed, 17 insertions(+), 28 deletions(-)

diff --git a/api/kiwi_vpn_api/db_new/capabilities.py b/api/kiwi_vpn_api/db_new/capabilities.py
index 7d92064..9ed8aba 100644
--- a/api/kiwi_vpn_api/db_new/capabilities.py
+++ b/api/kiwi_vpn_api/db_new/capabilities.py
@@ -13,9 +13,11 @@ class Capability(Enum):
     issue = "issue"
     renew = "renew"
 
+    def __repr__(self) -> str:
+        return self.value
+
 
 class UserCapabilityBase(SQLModel):
-    user_name: str = Field(primary_key=True, foreign_key="user.name")
     capability_name: str = Field(primary_key=True)
 
     @property
@@ -27,6 +29,8 @@ class UserCapabilityBase(SQLModel):
 
 
 class UserCapability(UserCapabilityBase, table=True):
+    user_name: str = Field(primary_key=True, foreign_key="user.name")
+
     user: "User" = Relationship(
-        back_populates="capabilities"
+        back_populates="capabilities",
     )
diff --git a/api/kiwi_vpn_api/db_new/user.py b/api/kiwi_vpn_api/db_new/user.py
index 7d0ec57..cb2fe59 100644
--- a/api/kiwi_vpn_api/db_new/user.py
+++ b/api/kiwi_vpn_api/db_new/user.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any, Sequence
+from typing import Any
 
 from pydantic import root_validator
 from sqlalchemy.exc import IntegrityError
@@ -109,34 +109,19 @@ class User(UserBase, table=True):
             db.delete(self)
             db.commit()
 
-    def extend_capabilities(self, capabilities: Sequence[Capability]) -> None:
-        """
-        Extend this user's capabilities
-        """
+    def get_capabilities(self) -> set[Capability]:
+        return set(
+            capability._
+            for capability in self.capabilities
+        )
 
-        for capability in capabilities:
-            user_capability = UserCapability(
+    def set_capabilities(self, capabilities: set[Capability]) -> None:
+        self.capabilities = [
+            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
+            ) for capability in capabilities
+        ]
 
 
 class UserCreate(UserBase):