helpers/docker.py

import time
import docker
import atexit
from typing import Optional
from helpers.files import get_abs_path
from helpers.errors import format_error
from helpers.print_style import PrintStyle
from helpers.log import Log

class DockerContainerManager:
    def __init__(self, image: str, name: str, ports: Optional[dict[str, int]] = None, volumes: Optional[dict[str, dict[str, str]]] = None,logger: Log|None=None):
        self.logger = logger
        self.image = image
        self.name = name
        self.ports = ports
        self.volumes = volumes
        self.init_docker()
                
    def init_docker(self):
        self.client = None
        while not self.client:
            try:
                self.client = docker.from_env()
                self.container = None
            except Exception as e:
                err = format_error(e)
                if ("ConnectionRefusedError(61," in err or "Error while fetching server API version" in err):
                    PrintStyle.hint("Connection to Docker failed. Is docker or Docker Desktop running?") # hint for user
                    if self.logger:self.logger.log(type="hint", content="Connection to Docker failed. Is docker or Docker Desktop running?")
                    PrintStyle.error(err)
                    if self.logger:self.logger.log(type="error", content=err)
                    time.sleep(5) # try again in 5 seconds
                else: raise
        return self.client
                            
    def cleanup_container(self) -> None:
        if self.container:
            try:
                self.container.stop()
                self.container.remove()
                PrintStyle.standard(f"Stopped and removed the container: {self.container.id}")
                if self.logger: self.logger.log(type="info", content=f"Stopped and removed the container: {self.container.id}")
            except Exception as e:
                PrintStyle.error(f"Failed to stop and remove the container: {e}")
                if self.logger: self.logger.log(type="error", content=f"Failed to stop and remove the container: {e}")

    def get_image_containers(self):
        if not self.client: self.client = self.init_docker()
        containers = self.client.containers.list(all=True, filters={"ancestor": self.image})
        infos = []
        for container in containers:
            infos.append({                
                "id": container.id,
                "name": container.name,
                "status": container.status,
                "image": container.image,
                "ports": container.ports,
                "web_port": (container.ports.get("80/tcp") or [{}])[0].get("HostPort"),
                "ssh_port": (container.ports.get("22/tcp") or [{}])[0].get("HostPort"),
                # "volumes": container.volumes,
                # "data_folder": container.volumes["/a0"],
            })
        return infos

    def start_container(self) -> None:
        if not self.client: self.client = self.init_docker()
        existing_container = None
        for container in self.client.containers.list(all=True):
            if container.name == self.name:
                existing_container = container
                break

        if existing_container:
            if existing_container.status != 'running':
                PrintStyle.standard(f"Starting existing container: {self.name} for safe code execution...")
                if self.logger: self.logger.log(type="info", content=f"Starting existing container: {self.name} for safe code execution...")
                
                existing_container.start()
                self.container = existing_container
                time.sleep(2) # this helps to get SSH ready
                
            else:
                self.container = existing_container
                # PrintStyle.standard(f"Container with name '{self.name}' is already running with ID: {existing_container.id}")
        else:
            PrintStyle.standard(f"Initializing docker container {self.name} for safe code execution...")
            if self.logger: self.logger.log(type="info", content=f"Initializing docker container {self.name} for safe code execution...")

            self.container = self.client.containers.run(
                self.image,
                detach=True,
                ports=self.ports, # type: ignore
                name=self.name,
                volumes=self.volumes, # type: ignore
            ) 
            # atexit.register(self.cleanup_container)
            PrintStyle.standard(f"Started container with ID: {self.container.id}")
            if self.logger: self.logger.log(type="info", content=f"Started container with ID: {self.container.id}")
            time.sleep(5) # this helps to get SSH ready