pilotwings/backend/pilotwings.py

from os import path
from typing import Annotated

from docker import errors, from_env
from docker.models.containers import Container
from dotenv import load_dotenv
from fastapi import Depends, FastAPI, HTTPException
from fastapi.responses import PlainTextResponse
from fastapi.security import HTTPBasicCredentials
from starlette import status
from uvicorn import run

from .security import AuthStaticFiles, check_auth, security
from .types import ContainerRequest, SerializedContainer
from .utils import serialize_container

load_dotenv()
client = from_env()
app = FastAPI(dependencies=[Depends(check_auth)])


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(
                all=True, filters={"label": ["engine=pilotwings"]}
            )
        ]

    return [
        serialize_container(container)
        for container in client.containers.list(
            all=True,
            filters={"label": ["engine=pilotwings", f"owner={credentials.username}"]},
        )
    ]


@app.get("/api/container/{container_name}")
def get_container(
    container_name: str,
    credentials: Annotated[HTTPBasicCredentials, Depends(security)],
) -> SerializedContainer:
    return serialize_container(select_container(container_name, credentials))


@app.post("/api/container/{container_name}")
def create_or_update_container(
    container_name: str,
    request_body: ContainerRequest,
    credentials: Annotated[HTTPBasicCredentials, Depends(security)],
) -> SerializedContainer:
    owner = None
    networks = client.networks.list(names=["pilotwings"])

    if not networks:
        client.networks.create("pilotwings")

    client.images.pull(request_body.image)

    try:
        container = select_container(container_name, credentials)
        owner = container.labels.get("owner")
        delete_container(container_name, credentials)
    except HTTPException:
        pass

    container = client.containers.run(
        request_body.image,
        detach=True,
        environment=request_body.environment,
        labels={"engine": "pilotwings", "owner": owner or credentials.username},
        name=container_name,
        network="pilotwings",
        restart_policy={"Name": "always"},
    )

    client.images.prune({"dangling": True})

    return serialize_container(container)


@app.post("/api/container/{container_name}/pull")
def pull_container(
    container_name: str,
    credentials: Annotated[HTTPBasicCredentials, Depends(security)],
) -> SerializedContainer:
    container = select_container(container_name, credentials)

    if not container.image:
        raise HTTPException(status_code=status.HTTP_410_GONE)

    request_body = ContainerRequest(
        image=container.image.tags[0], environment=container.attrs["Config"]["Env"]
    )

    return create_or_update_container(container_name, request_body, credentials)


@app.post("/api/container/{container_name}/restart")
def restart_container(
    container_name: str,
    credentials: Annotated[HTTPBasicCredentials, Depends(security)],
) -> SerializedContainer:
    container = select_container(container_name, credentials)
    container.restart()
    container.reload()
    return serialize_container(container)


@app.post("/api/container/{container_name}/start")
def start_container(
    container_name: str,
    credentials: Annotated[HTTPBasicCredentials, Depends(security)],
) -> SerializedContainer:
    container = select_container(container_name, credentials)
    container.start()
    container.reload()
    return serialize_container(container)


@app.post("/api/container/{container_name}/stop")
def stop_container(
    container_name: str,
    credentials: Annotated[HTTPBasicCredentials, Depends(security)],
) -> SerializedContainer:
    container = select_container(container_name, credentials)
    container.stop()
    container.reload()
    return serialize_container(container)


@app.get("/api/container/{container_name}/logs", response_class=PlainTextResponse)
def get_container_logs(
    container_name: str,
    credentials: Annotated[HTTPBasicCredentials, Depends(security)],
) -> str:
    container = select_container(container_name, credentials)
    return container.logs().decode()


@app.delete("/api/container/{container_name}")
def delete_container(
    container_name: str,
    credentials: Annotated[HTTPBasicCredentials, Depends(security)],
) -> None:
    container = select_container(container_name, credentials)
    container.stop()
    container.remove(v=True, force=True)
    client.images.prune({"dangling": True})


app.mount(
    "/",
    AuthStaticFiles(
        directory=f"{path.dirname(path.realpath(__file__))}/dist", html=True
    ),
    name="static",
)


def launch() -> None:
    run(app, host="0.0.0.0")