chore: refactor api endpoints (#31)

Author: msterhuj

fob_api/auth/__init__.py

@@ -66,3 +66,13 @@ def basic_auth_validator(username: str, password: str) -> User:
         if not user or not password_context.verify(password, user.password):
             return False
         return user
+
+def is_admin(user: User):
+    """Check if the user is an admin"""
+    if not user.is_admin:
+        raise HTTPException(status_code=403, detail="Not enough permissions")
+
+def is_admin_or_self(user: User, username: str):
+    """Check if the user is an admin or the user itself"""
+    if not user.is_admin and user.username != username:
+        raise HTTPException(status_code=403, detail="Not enough permissions")

fob_api/models/api/openstack.py

@@ -6,6 +6,7 @@ class OpenStackProjectCreate(BaseModel):
 
 class OpenStackProject(OpenStackProjectCreate):
     id: int
+    type: str
 
 class OpenStackUserPassword(BaseModel):
     username: str

fob_api/routes/devices.py

@@ -16,6 +16,14 @@
 
 template = Jinja2Templates(directory="templates")
 
+MAX_ALLOWED_DEVICES = 5
+
+def count_devices_for_user(username: str) -> int:
+    return len(headscale_driver.node.list(username=username))
+
+def can_add_device(username: str) -> bool:
+    return count_devices_for_user(username) <= MAX_ALLOWED_DEVICES
+
 @router.get("/register/{mkey}", tags=["vpn"], response_class=HTMLResponse)
 def register_device_get(request: Request, mkey: str):
     """
@@ -34,7 +42,6 @@ async def register_device_post(request: Request, mkey: str):
     """
     Handle device registration for headscale require basic auth
     """
-    # todo add limit to max allowed devices
     form_data = await request.form()
     username = form_data.get("username")
     password = form_data.get("password")
@@ -63,6 +70,13 @@ async def register_device_post(request: Request, mkey: str):
 
     headscale_tasks.get_or_create_user(username)
 
+    if not can_add_device(username):
+        return template.TemplateResponse(
+            request=request,
+            name="register_device.html.j2",
+            context={"mkey": mkey, "error": f"Max allowed devices reached ({MAX_ALLOWED_DEVICES})"}
+        )
+
     try:
         headscale_driver.node.register(username, mkey)
     except Exception as e:
@@ -79,22 +93,27 @@ async def register_device_post(request: Request, mkey: str):
     )
 
 @router.get("/{username}", tags=["vpn"])
-def list_devices(user: Annotated[User, Depends(auth.get_current_user)], username: str):
+def list_devices(
+        username: str,
+        user: Annotated[User, Depends(auth.get_current_user)]
+    ):
     """
     List all devices
     """
-    if user.username != username and not user.is_admin:
-        raise HTTPException(status_code=403, detail="You are not an admin")
+    auth.is_admin_or_self(user, username)
     nodes = headscale_driver.node.list(username=username)
     return [ApiDeviceResponse(**node.__dict__) for node in nodes]
 
 @router.delete("/{username}/{name}", tags=["vpn"], response_model=DeviceDeleteResponse)
-def delete_device(user: Annotated[User, Depends(auth.get_current_user)], username: str, name: str):
+def delete_device(
+        username: str,
+        name: str,
+        user: Annotated[User, Depends(auth.get_current_user)]
+    ):
     """
     Delete a device
     """
-    if user.username != username and not user.is_admin:
-        raise HTTPException(status_code=403, detail="You are not an admin")
+    auth.is_admin_or_self(user, username)
     user_nodes: List[Node] = headscale_driver.node.list(username=username)
     for node in user_nodes:
         if node.givenName == name:
@@ -107,10 +126,11 @@ def generate_preauth_key(user: Annotated[User, Depends(auth.get_current_user)],
     """
     Generate a preauth key active for 5 minutes
     """
-    # todo add limit to max allowed devices
-    if user.username != username:
-        raise HTTPException(status_code=403, detail="You are not allowed to generate preauth key for other users")
+    auth.is_admin_or_self(user, username)
     headscale_tasks.get_or_create_user(username=username)
+    
+    if not can_add_device(username):
+        raise HTTPException(status_code=400, detail=f"Max allowed devices reached ({MAX_ALLOWED_DEVICES})")
+    
     pre_auth_key: PreAuthKey = headscale_driver.preauthkey.create(username=username, expiration=datetime.now() + timedelta(minutes=5))
-    print(pre_auth_key.__dict__)
     return DevicePreAuthKeyResponse(**pre_auth_key.__dict__)

fob_api/routes/headscale.py

@@ -1,11 +1,10 @@
 from typing import Annotated, List
-from datetime import datetime, timedelta
 
 from fastapi import APIRouter, Depends, HTTPException, Request
 from sqlmodel import Session, select
 from sqlalchemy.exc import IntegrityError
 
-from fob_api import auth, engine
+from fob_api import auth, get_session
 from fob_api.models.database import User, HeadScalePolicyACL, HeadScalePolicyHost
 from fob_api.models.api import HeadScalePolicyAcl as HeadScalePolicyAclAPI
 from fob_api.models.api import HeadScalePolicyAclCreate as HeadScalePolicyAclCreateAPI
@@ -16,14 +15,15 @@
 router = APIRouter(prefix="/headscale")
 
 @router.get("/acls/", tags=["vpn"])
-def list(user: Annotated[User, Depends(auth.get_current_user)]) -> List[HeadScalePolicyAclAPI]:
+def list(
+        user: Annotated[User, Depends(auth.get_current_user)],
+        session: Session = Depends(get_session),
+    ) -> List[HeadScalePolicyAclAPI]:
     """
     Return list of HeadScale ACLs in Policy
     """
-    if not user.is_admin:
-        raise HTTPException(status_code=403, detail="Not an admin")
-    with Session(engine) as session:
-        acls = session.exec(select(HeadScalePolicyACL)).all()
+    auth.is_admin(user)
+    acls = session.exec(select(HeadScalePolicyACL)).all()
     return [HeadScalePolicyAclAPI(
         id=acl.id,
         action=acl.action,
@@ -34,114 +34,110 @@ def list(user: Annotated[User, Depends(auth.get_current_user)]) -> List[HeadScal
 
 @router.post("/acls/", tags=["vpn"])
 def create(
-    acl: HeadScalePolicyAclCreateAPI,
-    user: Annotated[User, Depends(auth.get_current_user)],
-) -> HeadScalePolicyAclAPI | None:
+        acl: HeadScalePolicyAclCreateAPI,
+        user: Annotated[User, Depends(auth.get_current_user)],
+        session: Session = Depends(get_session),
+    ) -> HeadScalePolicyAclAPI | None:
     """
     Create a new HeadScale ACL in Policy
     """
-    if not user.is_admin:
-        raise HTTPException(status_code=403, detail="Not an admin")
-    with Session(engine) as session:
-        new_acl = HeadScalePolicyACL(**acl.model_dump())
-        session.add(new_acl)
-        session.commit()
-        session.refresh(new_acl)
+    auth.is_admin(user)
+    new_acl = HeadScalePolicyACL(**acl.model_dump())
+    session.add(new_acl)
+    session.commit()
+    session.refresh(new_acl)
 
-        try:
-            update_headscale_policy()
-        except Exception as e:
-            session.delete(new_acl)
-            session.commit()
-            raise HTTPException(status_code=400, detail=f"Failed to apply new policy: {e}")
-        return HeadScalePolicyAclAPI(
-            id=new_acl.id,
-            action=new_acl.action,
-            src=new_acl.src.split(","),
-            dst=new_acl.dst.split(","),
-            proto=new_acl.proto,
-        )
+    try:
+        update_headscale_policy()
+    except Exception as e:
+        session.delete(new_acl)
+        session.commit()
+        raise HTTPException(status_code=400, detail=f"Failed to apply new policy: {e}")
+    return HeadScalePolicyAclAPI(
+        id=new_acl.id,
+        action=new_acl.action,
+        src=new_acl.src.split(","),
+        dst=new_acl.dst.split(","),
+        proto=new_acl.proto,
+    )
 
 @router.delete("/acls/{acl_id}/", tags=["vpn"])
 def delete(
-    acl_id: int,
-    user: Annotated[User, Depends(auth.get_current_user)],
-) -> None:
+        acl_id: int,
+        user: Annotated[User, Depends(auth.get_current_user)],
+        session: Session = Depends(get_session),
+    ) -> None:
     """
     Delete a HeadScale ACL from Policy
     """
-    if not user.is_admin:
-        raise HTTPException(status_code=403, detail="Not an admin")
-    with Session(engine) as session:
-        acl = session.get(HeadScalePolicyACL, acl_id)
-        if not acl:
-            raise HTTPException(status_code=404, detail="ACL not found")
-        session.delete(acl)
+    auth.is_admin(user)
+    acl = session.get(HeadScalePolicyACL, acl_id)
+    if not acl:
+        raise HTTPException(status_code=404, detail="ACL not found")
+    session.delete(acl)
+    session.commit()
+    try:
+        update_headscale_policy()
+    except Exception as e:
+        session.add(acl)
         session.commit()
-        try:
-            update_headscale_policy()
-        except Exception as e:
-            session.add(acl)
-            session.commit()
-            raise HTTPException(status_code=400, detail=f"Failed to apply new policy: {e}")
+        raise HTTPException(status_code=400, detail=f"Failed to apply new policy: {e}")
 
 @router.get("/host/", tags=["vpn"])
-def list_hosts(user: Annotated[User, Depends(auth.get_current_user)]) -> List[HeadScalePolicyHostAPI]:
+def list_hosts(
+        user: Annotated[User, Depends(auth.get_current_user)],
+        session: Session = Depends(get_session),
+    ) -> List[HeadScalePolicyHostAPI]:
     """
     Return list of HeadScale hosts
     """
-    if not user.is_admin:
-        raise HTTPException(status_code=403, detail="Not an admin")
-    with Session(engine) as session:
-        headscale_policy_host = session.exec(select(HeadScalePolicyHost)).all()
+    auth.is_admin(user)
+    headscale_policy_host = session.exec(select(HeadScalePolicyHost)).all()
     return headscale_policy_host
 
 @router.post("/host/", tags=["vpn"])
 def create_host(
-    host: HeadScalePolicyHostCreateAPI,
-    user: Annotated[User, Depends(auth.get_current_user)],
-) -> HeadScalePolicyHostAPI | None:
+        host: HeadScalePolicyHostCreateAPI,
+        user: Annotated[User, Depends(auth.get_current_user)],
+        session: Session = Depends(get_session),
+    ) -> HeadScalePolicyHostAPI | None:
     """
     Create a new HeadScale host
     """
-    if not user.is_admin:
-        raise HTTPException(status_code=403, detail="Not an admin")
-    with Session(engine) as session:
-
-        new_host = HeadScalePolicyHost(**host.model_dump())
-        session.add(new_host)
-        try:
-            session.commit()
-        except IntegrityError:
-            raise HTTPException(status_code=400, detail="This host binding already exists")
-        session.refresh(new_host)
-        try:
-            update_headscale_policy()
-        except Exception as e:
-            session.delete(new_host)
-            session.commit()
-            raise HTTPException(status_code=400, detail=f"Failed to apply new policy: {e}")
-        return new_host
+    auth.is_admin(user)
+    new_host = HeadScalePolicyHost(**host.model_dump())
+    session.add(new_host)
+    try:
+        session.commit()
+    except IntegrityError:
+        raise HTTPException(status_code=400, detail="This host binding already exists")
+    session.refresh(new_host)
+    try:
+        update_headscale_policy()
+    except Exception as e:
+        session.delete(new_host)
+        session.commit()
+        raise HTTPException(status_code=400, detail=f"Failed to apply new policy: {e}")
+    return new_host
 
 @router.delete("/host/{host_id}/", tags=["vpn"])
 def delete_host(
-    host_id: int,
-    user: Annotated[User, Depends(auth.get_current_user)],
-) -> None:
+        host_id: int,
+        user: Annotated[User, Depends(auth.get_current_user)],
+        session: Session = Depends(get_session),
+    ) -> None:
     """
     Delete a HeadScale host
     """
-    if not user.is_admin:
-        raise HTTPException(status_code=403, detail="Not an admin")
-    with Session(engine) as session:
-        host = session.get(HeadScalePolicyHost, host_id)
-        if not host:
-            raise HTTPException(status_code=404, detail="Host not found")
-        session.delete(host)
+    auth.is_admin(user)
+    host = session.get(HeadScalePolicyHost, host_id)
+    if not host:
+        raise HTTPException(status_code=404, detail="Host not found")
+    session.delete(host)
+    session.commit()
+    try:
+        update_headscale_policy()
+    except Exception as e:
+        session.add(host)
         session.commit()
-        try:
-            update_headscale_policy()
-        except Exception as e:
-            session.add(host)
-            session.commit()
-            raise HTTPException(status_code=400, detail=f"Failed to apply new policy: {e}")
+        raise HTTPException(status_code=400, detail=f"Failed to apply new policy: {e}")