| 1 | import time |
| 2 | import docker |
| 3 | import atexit |
| 4 | from typing import Optional |
| 5 | from helpers.files import get_abs_path |
| 6 | from helpers.errors import format_error |
| 7 | from helpers.print_style import PrintStyle |
| 8 | from helpers.log import Log |
| 9 | |
| 10 | class DockerContainerManager: |
| 11 | 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): |
| 12 | self.logger = logger |
| 13 | self.image = image |
| 14 | self.name = name |
| 15 | self.ports = ports |
| 16 | self.volumes = volumes |
| 17 | self.init_docker() |
| 18 | |
| 19 | def init_docker(self): |
| 20 | self.client = None |
| 21 | while not self.client: |
| 22 | try: |
| 23 | self.client = docker.from_env() |
| 24 | self.container = None |
| 25 | except Exception as e: |
| 26 | err = format_error(e) |
| 27 | if ("ConnectionRefusedError(61," in err or "Error while fetching server API version" in err): |
| 28 | PrintStyle.hint("Connection to Docker failed. Is docker or Docker Desktop running?") # hint for user |
| 29 | if self.logger:self.logger.log(type="hint", content="Connection to Docker failed. Is docker or Docker Desktop running?") |
| 30 | PrintStyle.error(err) |
| 31 | if self.logger:self.logger.log(type="error", content=err) |
| 32 | time.sleep(5) # try again in 5 seconds |
| 33 | else: raise |
| 34 | return self.client |
| 35 | |
| 36 | def cleanup_container(self) -> None: |
| 37 | if self.container: |
| 38 | try: |
| 39 | self.container.stop() |
| 40 | self.container.remove() |
| 41 | PrintStyle.standard(f"Stopped and removed the container: {self.container.id}") |
| 42 | if self.logger: self.logger.log(type="info", content=f"Stopped and removed the container: {self.container.id}") |
| 43 | except Exception as e: |
| 44 | PrintStyle.error(f"Failed to stop and remove the container: {e}") |
| 45 | if self.logger: self.logger.log(type="error", content=f"Failed to stop and remove the container: {e}") |
| 46 | |
| 47 | def get_image_containers(self): |
| 48 | if not self.client: self.client = self.init_docker() |
| 49 | containers = self.client.containers.list(all=True, filters={"ancestor": self.image}) |
| 50 | infos = [] |
| 51 | for container in containers: |
| 52 | infos.append({ |
| 53 | "id": container.id, |
| 54 | "name": container.name, |
| 55 | "status": container.status, |
| 56 | "image": container.image, |
| 57 | "ports": container.ports, |
| 58 | "web_port": (container.ports.get("80/tcp") or [{}])[0].get("HostPort"), |
| 59 | "ssh_port": (container.ports.get("22/tcp") or [{}])[0].get("HostPort"), |
| 60 | # "volumes": container.volumes, |
| 61 | # "data_folder": container.volumes["/a0"], |
| 62 | }) |
| 63 | return infos |
| 64 | |
| 65 | def start_container(self) -> None: |
| 66 | if not self.client: self.client = self.init_docker() |
| 67 | existing_container = None |
| 68 | for container in self.client.containers.list(all=True): |
| 69 | if container.name == self.name: |
| 70 | existing_container = container |
| 71 | break |
| 72 | |
| 73 | if existing_container: |
| 74 | if existing_container.status != 'running': |
| 75 | PrintStyle.standard(f"Starting existing container: {self.name} for safe code execution...") |
| 76 | if self.logger: self.logger.log(type="info", content=f"Starting existing container: {self.name} for safe code execution...") |
| 77 | |
| 78 | existing_container.start() |
| 79 | self.container = existing_container |
| 80 | time.sleep(2) # this helps to get SSH ready |
| 81 | |
| 82 | else: |
| 83 | self.container = existing_container |
| 84 | # PrintStyle.standard(f"Container with name '{self.name}' is already running with ID: {existing_container.id}") |
| 85 | else: |
| 86 | PrintStyle.standard(f"Initializing docker container {self.name} for safe code execution...") |
| 87 | if self.logger: self.logger.log(type="info", content=f"Initializing docker container {self.name} for safe code execution...") |
| 88 | |
| 89 | self.container = self.client.containers.run( |
| 90 | self.image, |
| 91 | detach=True, |
| 92 | ports=self.ports, # type: ignore |
| 93 | name=self.name, |
| 94 | volumes=self.volumes, # type: ignore |
| 95 | ) |
| 96 | # atexit.register(self.cleanup_container) |
| 97 | PrintStyle.standard(f"Started container with ID: {self.container.id}") |
| 98 | if self.logger: self.logger.log(type="info", content=f"Started container with ID: {self.container.id}") |
| 99 | time.sleep(5) # this helps to get SSH ready |