From f6ba1ac4761297ea5ccf333a5339aea0c1cf697d Mon Sep 17 00:00:00 2001
From: Michel Roux <xefir@crystalyx.net>
Date: Mon, 4 Nov 2024 11:04:25 +0100
Subject: [PATCH] fix: :zap: improve creation of containers

---
 backend/pilotwings.py | 101 ++++++++++++++++++------------------------
 1 file changed, 44 insertions(+), 57 deletions(-)

diff --git a/backend/pilotwings.py b/backend/pilotwings.py
index 8bb7109..d5f0ecb 100644
--- a/backend/pilotwings.py
+++ b/backend/pilotwings.py
@@ -52,6 +52,7 @@ class SerializedContainer(BaseModel):
     labels: dict[str, str]
     status: str
     health: str
+    engine: str | None
     owner: str | None
     environment: list[str]
 
@@ -64,24 +65,46 @@ def serialize_container(container: Container) -> SerializedContainer:
         labels=container.labels,
         status=container.status,
         health=container.health,
+        engine=container.labels.get("engine"),
         owner=container.labels.get("owner"),
         environment=container.attrs["Config"]["Env"],
     )
 
 
+def select_container(
+    container_name: str, credentials: Annotated[HTTPBasicCredentials, Depends(security)]
+) -> Container:
+    try:
+        container = client.containers.get(container_name)
+    except errors.APIError:
+        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND)
+
+    if (
+        credentials.username != "admin"
+        and container.labels.get("engine") != "pilotwings"
+        and container.labels.get("owner") != credentials.username
+    ):
+        raise HTTPException(status_code=status.HTTP_403_FORBIDDEN)
+
+    return container
+
+
 @app.get("/api/containers")
 def get_containers(
     credentials: Annotated[HTTPBasicCredentials, Depends(security)],
 ) -> list[SerializedContainer]:
     if credentials.username == "admin":
         return [
-            serialize_container(container) for container in client.containers.list()
+            serialize_container(container)
+            for container in client.containers.list(
+                filters={"label": ["engine=pilotwings"]}
+            )
         ]
 
     return [
         serialize_container(container)
         for container in client.containers.list(
-            filters={"label": f"owner={credentials.username}"}
+            filters={"label": ["engine=pilotwings", f"owner={credentials.username}"]},
         )
     ]
 
@@ -91,76 +114,40 @@ def get_container(
     container_name: str,
     credentials: Annotated[HTTPBasicCredentials, Depends(security)],
 ) -> SerializedContainer:
-    try:
-        container = client.containers.get(container_name)
-    except errors.APIError:
-        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND)
-
-    if (
-        credentials.username != "admin"
-        and f"owner={credentials.username}" not in container.labels
-    ):
-        raise HTTPException(status_code=status.HTTP_403_FORBIDDEN)
-
-    return serialize_container(container)
+    return serialize_container(select_container(container_name, credentials))
 
 
-class UpdateContainerRequest(BaseModel):
+class ContainerRequest(BaseModel):
+    image: str
     environment: dict[str, str]
 
 
-class CreateContainerRequest(UpdateContainerRequest):
-    image: str
-
-
 @app.post("/api/container/{container_name}")
-def create_container(
+def create_or_update_container(
     container_name: str,
-    request_body: CreateContainerRequest,
+    request_body: ContainerRequest,
     credentials: Annotated[HTTPBasicCredentials, Depends(security)],
 ) -> SerializedContainer:
+    networks = client.networks.list(names=["pilotwings"])
+
+    if not networks:
+        client.networks.create("pilotwings")
+
+    try:
+        container = select_container(container_name, credentials)
+        container.stop()
+        container.remove(v=True, force=True)
+    except errors.APIError:
+        pass
+
     return serialize_container(
         client.containers.run(
             request_body.image,
             detach=True,
             environment=request_body.environment,
-            labels={"owner": credentials.username},
-            name=container_name,
-            restart_policy={"Name": "always"},
-        )
-    )
-
-
-@app.put("/api/container/{container_name}")
-def update_container(
-    container_name: str,
-    request_body: UpdateContainerRequest,
-    credentials: Annotated[HTTPBasicCredentials, Depends(security)],
-) -> SerializedContainer:
-    try:
-        container = client.containers.get(container_name)
-    except errors.APIError:
-        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND)
-
-    if (
-        credentials.username != "admin"
-        and f"owner={credentials.username}" not in container.labels
-    ):
-        raise HTTPException(status_code=status.HTTP_403_FORBIDDEN)
-
-    if not container.image:
-        raise HTTPException(status_code=status.HTTP_410_GONE)
-
-    container.stop()
-    container.remove(v=True, force=True)
-
-    return serialize_container(
-        client.containers.run(
-            container.image.tags[0],
-            detach=True,
-            environment=request_body.environment,
-            labels={"owner": credentials.username},
+            labels={"engine": "pilotwings", "owner": credentials.username},
             name=container_name,
+            network="pilotwings",
             restart_policy={"Name": "always"},
         )
     )