From 7d176fad3941393d52d0f9d90e81943f9c1706ff Mon Sep 17 00:00:00 2001 From: ReLoia Date: Mon, 13 Apr 2026 20:33:14 +0200 Subject: [PATCH] feat(docker): add async port change and docker compose reload for services - Introduce async `change_port` to update service port and reload docker compose - Add `run_docker_command` for running docker compose commands asynchronously - Implement `_get_random_port` utility for selecting available ports - Refactor service discovery and port change logic to use new async methods - Update proxy and manager logic to support new port change workflow --- app/integrations/docker/__init__.py | 8 ++++---- app/integrations/docker/compose.py | 31 ++++++++++++++++++++++++++--- app/integrations/docker/editor.py | 21 +++++++++++++------ app/integrations/docker/ports.py | 13 ++++++++++++ app/main.py | 4 ++-- app/repositories/proxy.py | 14 ++++++++++--- app/services/mitmproxy/manager.py | 24 +++++++++++++++------- app/services/mitmproxy/models.py | 5 ----- 8 files changed, 90 insertions(+), 30 deletions(-) delete mode 100644 app/services/mitmproxy/models.py diff --git a/app/integrations/docker/__init__.py b/app/integrations/docker/__init__.py index 92297f2..e12fd8e 100644 --- a/app/integrations/docker/__init__.py +++ b/app/integrations/docker/__init__.py @@ -1,7 +1,7 @@ from integrations.docker.compose import ( - find_docker_compose, - get_all_services, - get_services, + find_docker_compose, + get_all_services_in_dir, + get_services, ) from integrations.docker.editor import change_port from integrations.docker.models import DockerContainer @@ -10,6 +10,6 @@ __all__ = [ "DockerContainer", "find_docker_compose", "get_services", - "get_all_services", + "get_all_services_in_dir", "change_port", ] diff --git a/app/integrations/docker/compose.py b/app/integrations/docker/compose.py index e58e613..4de0fe0 100644 --- a/app/integrations/docker/compose.py +++ b/app/integrations/docker/compose.py @@ -1,3 +1,4 @@ +import asyncio from pathlib import Path import yaml @@ -19,7 +20,7 @@ def _get_visible_folders(directory: Path) -> list[Path]: ] -def find_docker_compose(path: Path) -> tuple[dict, str] | None: +def find_docker_compose(path: Path) -> tuple[dict, Path] | None: if not path or not path.is_dir(): return None @@ -27,7 +28,7 @@ def find_docker_compose(path: Path) -> tuple[dict, str] | None: file = path / name if file.is_file(): content = yaml.safe_load(file.read_text()) or {} - return content, file.name + return content, file return None @@ -61,7 +62,7 @@ def get_services(folder: Path) -> list[DockerContainer] | None: return services -def get_all_services(directory: str | Path) -> list[DockerContainer]: +def get_all_services_in_dir(directory: str | Path) -> list[DockerContainer]: path = Path(directory) if not path.is_dir(): return [] @@ -73,3 +74,27 @@ def get_all_services(directory: str | Path) -> list[DockerContainer]: services.extend(sub_services) return services + + +async def run_docker_command( + file: Path, + command: list[str] +) -> dict | None: + if not command: + return None + + full_command = ["docker", "compose", "-f", f"{file}", *command] + + process = await asyncio.create_subprocess_exec( + *full_command, + cwd=str(file.parent), + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE + ) + + stdout, stderr = await process.communicate() + + if process.returncode == 0: + return {"status": "success", "output": stdout.decode()} + else: + return {"status": "error", "reason": stderr.decode()} diff --git a/app/integrations/docker/editor.py b/app/integrations/docker/editor.py index dbd952a..fd0dd3d 100644 --- a/app/integrations/docker/editor.py +++ b/app/integrations/docker/editor.py @@ -2,7 +2,7 @@ from pathlib import Path import yaml -from integrations.docker.compose import find_docker_compose +from integrations.docker.compose import find_docker_compose, run_docker_command from integrations.docker.ports import _get_port_index, _replace_port_in_mapping @@ -12,9 +12,9 @@ def _backup_path_for(compose_path: Path) -> Path: return compose_path.with_name(f"{compose_path.name}.bak") -def change_port( - folder: Path, +async def change_port( service_name: str, + folder: Path, port_from: int, port_to: int, ) -> bool: @@ -22,7 +22,7 @@ def change_port( if compose_data is None: return False - docker_compose, filename = compose_data + docker_compose, compose_path = compose_data compose_services = docker_compose.get("services") or {} if not compose_services: return False @@ -41,10 +41,19 @@ def change_port( mappings[idx] = _replace_port_in_mapping(mappings[idx], port_from, port_to) - compose_path = folder / filename backup_path = _backup_path_for(compose_path) + original_text = compose_path.read_text() + backup_path.write_text(compose_path.read_text()) compose_path.write_text(yaml.dump(docker_compose, default_flow_style=False, sort_keys=False)) - # TODO: run docker compose commands + result = await run_docker_command( + compose_path, + ["up", "-d", service_name], + ) + + if result is None or result.get("status") != "success": + compose_path.write_text(original_text) + return False + return True diff --git a/app/integrations/docker/ports.py b/app/integrations/docker/ports.py index e5cecbc..ad437b3 100644 --- a/app/integrations/docker/ports.py +++ b/app/integrations/docker/ports.py @@ -1,6 +1,19 @@ +import socket from collections.abc import Mapping +def _get_random_port() -> int: + # os automatically assigns an available port when binding to port 0, so we can use that to find a free port + try: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(('', 0)) + s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + port = s.getsockname()[1] + return port + except Exception as e: + raise RuntimeError(f"Could not find an available port: {e}") + + def _get_service_port(mapping: str | int | Mapping[str, object]) -> int | None: if isinstance(mapping, int): return mapping diff --git a/app/main.py b/app/main.py index fe1ef19..9f258bc 100644 --- a/app/main.py +++ b/app/main.py @@ -11,7 +11,7 @@ from core.config import settings from core.mongodb import close_mongodb, init_mongodb from core.redis import close_redis, init_redis from core.security import hash_password -from integrations.docker import get_all_services +from integrations.docker import get_all_services_in_dir from repositories.auth import ensure_auth_indexes, ensure_single_account from repositories.proxy import ensure_proxy_indexes @@ -28,7 +28,7 @@ async def lifespan(_: FastAPI): await ensure_proxy_indexes() home_dir = str(Path.home()) - app.state.services = get_all_services(home_dir) or [] + app.state.services = get_all_services_in_dir(home_dir) or [] if not app.state.services: warning = ( "\n" + "!" * 60 + "\n" diff --git a/app/repositories/proxy.py b/app/repositories/proxy.py index 8895374..5d00a49 100644 --- a/app/repositories/proxy.py +++ b/app/repositories/proxy.py @@ -8,7 +8,6 @@ from pymongo import ReturnDocument from core.mongodb import get_database from models.proxy import ProxyDocument, ProxyUpdateInput from models.service import ServiceEndpointDocument -from services.mitmproxy.models import ServiceConfig COLLECTION_NAME = "proxies" @@ -38,16 +37,25 @@ async def ensure_proxy_indexes() -> None: async def create_proxy( proxy_port: int, name: str, - service: ServiceConfig, + upstream_port: int, code: str, ) -> ProxyDocument: + """ + Saves a Proxy to the database. + + :param proxy_port: the port the proxy is listening on + :param name: the name of the service + :param upstream_port: the upstream port the proxy is forwarding to + :param code: the code of the service + :return: + """ now = _utcnow() proxy = ProxyDocument( proxy_port=proxy_port, name=name, code=code, service=ServiceEndpointDocument( - port=service["port"], + port=upstream_port, ), created_at=now, updated_at=now, diff --git a/app/services/mitmproxy/manager.py b/app/services/mitmproxy/manager.py index 3393c5f..d331bdb 100644 --- a/app/services/mitmproxy/manager.py +++ b/app/services/mitmproxy/manager.py @@ -1,27 +1,37 @@ -import uuid from pathlib import Path +from integrations import docker from integrations.docker import DockerContainer -from services.mitmproxy.models import ServiceConfig +from integrations.docker.ports import _get_random_port async def create_proxy( port: int, - name: str, - service: ServiceConfig, + service: DockerContainer, code: str, - services: list[DockerContainer] ) -> int: """ Utility function that will: - run ``integrations/docker`` utilities to change the old port of original service to a random one - - write ``code`` into ``/tmp/theplug/{random}.py`` + /- write ``code`` into ``/tmp/theplug/{random}.py`` - start a subprocess via asyncio - return the proxy infos (pid, ) """ - path = Path("/tmp/theplug/{random}.py".format(random=str(uuid.uuid4()))) + path = Path(f"/tmp/theplug/{service["name"]}.py") path.parent.mkdir(parents=True, exist_ok=True) path.write_text(code, encoding="utf-8") + # da qui bisogna: + # scegliere una porta random alla quale spostare il docker + # modificare il docker via integrations.docker.editor.change_port ( service.folder, service.name, port, random_nuova_port ) + random_port = _get_random_port() + docker.editor.change_port( + service["name"], + service["folder"], + port, + random_port, + ) + + return 0 diff --git a/app/services/mitmproxy/models.py b/app/services/mitmproxy/models.py deleted file mode 100644 index 531e1e9..0000000 --- a/app/services/mitmproxy/models.py +++ /dev/null @@ -1,5 +0,0 @@ -from typing import TypedDict - - -class ServiceConfig(TypedDict): - port: int