Connection.use()

api/kiwi_vpn_api/db/connection.py

@@ -6,6 +6,19 @@ from sqlalchemy.orm import Session, sessionmaker
 from .models import ORMBaseModel
 
 
+class SessionManager:
+    __session: Session
+
+    def __init__(self, session: Session) -> None:
+        self.__session = session
+
+    def __enter__(self) -> Session:
+        return self.__session
+
+    def __exit__(self, *args) -> None:
+        self.__session.close()
+
+
 class Connection:
     engine: Engine | None = None
     session_local: sessionmaker | None = None
@@ -18,6 +31,13 @@ class Connection:
         )
         ORMBaseModel.metadata.create_all(bind=engine)
 
+    @classmethod
+    def use(cls) -> SessionManager | None:
+        if cls.session_local is None:
+            return None
+
+        return SessionManager(cls.session_local())
+
     @classmethod
     async def get(cls) -> Generator[Session | None, None, None]:
         if cls.session_local is None:

api/kiwi_vpn_api/db/models.py

@@ -1,9 +1,11 @@
+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 relationship
+from sqlalchemy.orm import Session, relationship
 
 ORMBaseModel = declarative_base()
 
@@ -17,6 +19,13 @@ class User(ORMBaseModel):
     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"
@@ -29,9 +38,6 @@ class UserCapability(ORMBaseModel):
     )
     capability = Column(String, primary_key=True)
 
-    def __str__(self) -> str:
-        return self.capability
-
 
 class DistinguishedName(ORMBaseModel):
     __tablename__ = "distinguished_names"

api/kiwi_vpn_api/db/schemas.py

@@ -30,8 +30,19 @@ class Certificate(CertificateBase):
 class UserCapability(Enum):
     admin = "admin"
 
-    def __str__(self) -> str:
-        return self._value_
+    @classmethod
+    def from_value(cls, value) -> UserCapability:
+        if isinstance(value, cls):
+            # use simple value
+            return value
+
+        elif isinstance(value, models.UserCapability):
+            # create from db
+            return cls(value.capability)
+
+        else:
+            # create from string representation
+            return cls(str(value))
 
 
 class UserBase(BaseModel):
@@ -53,10 +64,10 @@ class User(UserBase):
     @classmethod
     def unify_capabilities(
         cls,
-        value: list[models.UserCapability | str]
+        value: list[models.UserCapability | UserCapability | str]
     ) -> list[UserCapability]:
         return [
-            UserCapability(str(capability))
+            UserCapability.from_value(capability)
             for capability in value
         ]
 
@@ -66,15 +77,10 @@ class User(UserBase):
         db: Session,
         name: str,
     ) -> User | None:
-        user = (db
-                .query(models.User)
-                .filter(models.User.name == name)
-                .first())
-
-        if user is None:
+        if (db_user := models.User.load(db, name)) is None:
             return None
 
-        return cls.from_orm(user)
+        return cls.from_orm(db_user)
 
     @classmethod
     def login(
@@ -84,21 +90,16 @@ class User(UserBase):
         password: str,
         crypt_context: CryptContext,
     ) -> User | None:
-        user = (db
-                .query(models.User)
-                .filter(models.User.name == name)
-                .first())
-
-        if user is None:
+        if (db_user := models.User.load(db, name)) is None:
             # inexistent user, fake doing password verification
             crypt_context.dummy_verify()
             return None
 
-        if not crypt_context.verify(password, user.password):
+        if not crypt_context.verify(password, db_user.password):
             # password hash mismatch
             return None
 
-        return cls.from_orm(user)
+        return cls.from_orm(db_user)
 
     @classmethod
     def create(
@@ -127,10 +128,16 @@ class User(UserBase):
         self,
         db: Session,
         capabilities: list[UserCapability],
-    ) -> bool:
-        # TODO
+    ) -> None:
+        for capability in capabilities:
+            if capability not in self.capabilities:
+                cap = models.UserCapability(
+                    user_name=self.name,
+                    capability=capability.value,
+                )
+                db.add(cap)
 
-        return True
+        db.commit()
 
 
 class DistinguishedNameBase(BaseModel):

api/kiwi_vpn_api/main.py

@@ -39,7 +39,7 @@ async def on_startup() -> None:
         Connection.connect(await current_config.db.db_engine)
 
         # some testing
-        async for db in Connection.get():
+        with Connection.use() as db:
             print(schemas.User.from_db(db, "admin"))
             print(schemas.User.from_db(db, "nonexistent"))
 

api/kiwi_vpn_api/routers/admin.py

@@ -16,7 +16,7 @@ router = APIRouter(prefix="/admin")
 )
 async def install(
     config: Config,
-    user: schemas.UserCreate,
+    admin_user: schemas.UserCreate,
     current_config: Config | None = Depends(Config.load),
 ):
     if current_config is not None:
@@ -25,10 +25,10 @@ async def install(
     await config.save()
     Connection.connect(await config.db.db_engine)
 
-    async for db in Connection.get():
+    with Connection.use() as db:
         admin_user = schemas.User.create(
             db=db,
-            user=user,
+            user=admin_user,
             crypt_context=await config.crypto.crypt_context,
         )