ahaz_cli 0.0.4__tar.gz

ahaz_cli-0.0.4/PKG-INFO

@@ -0,0 +1,9 @@
+Metadata-Version: 2.4
+Name: ahaz_cli
+Version: 0.0.4
+Summary: Command-line management interface for Ahaz
+Requires-Dist: docker>=7.1.0
+Requires-Dist: rich>=14.2.0
+Requires-Dist: typer>=0.20.0
+Requires-Dist: ahaz_common
+Requires-Dist: pyyaml>=6.0.3

ahaz_cli-0.0.4/ahaz_cli/__main__.py

@@ -0,0 +1,4 @@
+if __name__ == "__main__":
+    from .cli import app
+
+    app()

ahaz_cli-0.0.4/ahaz_cli/ahaz.py

@@ -0,0 +1,128 @@
+import colorsys
+import logging
+import time
+from pathlib import Path
+from typing import Annotated
+
+import docker
+import docker.errors
+import typer
+from rich.status import Status
+
+from .lib.docker import cleanup_env, create_env, log_docker_logs, try_build_image
+from .lib.file import test_for_file
+from .lib.task import deserialise_task, normalise_task_name
+from .templates import copy_example_images, write_task_yaml
+
+log = logging.getLogger(__name__)
+CWD = Path.cwd()
+
+
+def test(
+    build: Annotated[bool, typer.Option("--build", "-b", help="Always build Docker images.")] = False,
+    verbose: Annotated[bool, typer.Option("--verbose", "-v", help="Enable verbose logging.")] = False,
+    up: Annotated[bool, typer.Option("--up", "-u", help="Start the task environment after testing.")] = False,
+) -> None:
+    """
+    Test the task configuration by validating the config file and attempting to build task images.
+
+    Optionally starts the task environment after testing.
+    """
+
+    if verbose:
+        log.setLevel(logging.DEBUG)
+
+    config = "task.yaml"
+    if not test_for_file(config):
+        config = "task.yml"
+        if not test_for_file(config):
+            log.error("No task configuration file found (task.yaml or task.yml)")
+            log.error("Are you sure you are in the task directory?")
+            return
+
+    task = deserialise_task(Path(config).read_text())
+    log.info(f"Loaded task: {task.name}")
+
+    log.info("Testing Docker images for all pods...")
+    client = docker.from_env()
+    for pod in task.pods:
+        with Status(f"Checking image for pod '{pod.name}'...", spinner="dots") as status:
+            image_tag = f"{pod.image.image_name}:{task.version}"
+            if not build:
+                # See if we can find the image locally
+                try:
+                    client.images.get(image_tag)
+                    status.update(f"Image '{image_tag}' found locally.")
+                    log.info(f"Image '{image_tag}' found locally for pod '{pod.name}'.")
+                    continue
+                except docker.errors.ImageNotFound:
+                    log.info(f"Image '{image_tag}' not found locally for pod '{pod.name}', building...")
+            # Build the image
+            status.update(f"Building image '{image_tag}'...")
+            log.info(f"Building image '{image_tag}' for pod '{pod.name}'...")
+            build_args = {arg.name: arg.value for arg in (pod.image.build_args or [])}
+            try:
+                try_build_image(image_tag, pod.image.build_context, build_args, verbose)
+            except Exception as e:
+                log.error(f"Failed to build image '{image_tag}' for pod '{pod.name}': {e}")
+                return
+
+    # Attempt to set up the entire task environment
+    if up:
+        log.info("Setting up the task environment...")
+        containers = create_env(task)
+        log_docker_logs(
+            containers,
+            lambda: cleanup_env(
+                task.name, [pod.name for pod in task.pods], [net.name for net in task.networks]
+            ),
+        )
+    log.info("Task test completed.")
+
+
+def init(
+    name: Annotated[str, typer.Argument(help="Name of the task.")],
+    path: Annotated[Path, typer.Argument(help="Path to the task.", show_default=False)] = CWD,
+) -> None:
+    """
+    Initialize a new task directory with a template configuration and an example image.
+
+    Creates the following structure:\n
+    <out_dir>/\n
+    ├─task.yaml\n
+    └─hello-world/\n
+    ....└──Dockerfile\n
+    """
+
+    log.info(f"Initializing new task '{name}' at '{path}'...")
+    out_dir = path / normalise_task_name(name)
+    if out_dir.exists():
+        log.error(f"Directory '{out_dir}' already exists.")
+        return
+    out_dir.mkdir(parents=True)
+    log.info(f"Writing task template to '{out_dir}'...")
+    write_task_yaml(out_dir, name)
+    log.info(f"Copying example images to '{out_dir}'...")
+    copy_example_images(out_dir)
+
+
+def useless_gradient_function(step: int) -> str:
+    # Calculate hue based on step
+    hue = (step * 3 % 360) / 360.0
+    # Convert hue to RGB
+    red, green, blue = colorsys.hsv_to_rgb(hue, 1.0, 1.0)
+    # Output as ANSI
+    return f"\x1b[38;2;{int(red * 255)};{int(green * 255)};{int(blue * 255)}m"
+
+
+# Load bearing function, do not remove :3
+def epic() -> None:
+    """
+    A very important function.
+    """
+
+    step = 0
+    while True:
+        print(f"\r{useless_gradient_function(step)}Epic function executed successfully.\x1b[0m", end="")
+        step += 1
+        time.sleep(0.1)

ahaz_cli-0.0.4/ahaz_cli/assets/hello-world/Dockerfile

@@ -0,0 +1,3 @@
+FROM alpine:latest
+
+CMD ["echo", "Hello, world!"]

ahaz_cli-0.0.4/ahaz_cli/assets/task.yaml

@@ -0,0 +1,36 @@
+name: "My Very Cool Task" # Name of the task; shown to the user
+version: "1.0.0" # Version of the task; used for image tagging, increment when making changes to already-pushed tasks
+description: "Very cool description UwU" # Description of the task; shown to the user
+score: 150 # Maximum score for the task
+scoring_type: "dynamic" # Scoring type; can be "static" or "dynamic"
+
+# Definition of the pods that will be created for the task
+pods:
+- name: "hello-world" # Name of the pod; used for referencing in networks and env_vars
+  visible_to_user: true # Whether the user can see this pod
+  # Definition of the container image for the pod
+  image:
+    image_name: "hello-world" # Name of the image; used for pulling/building the image
+    build_context: "./hello-world" # Context path for building the image, needs to at least contain a Dockerfile
+  limits_ram: "128Mi" # RAM limit for the pod
+  limits_cpu: 1 # CPU limit for the pod
+  # Testing configuration for the pod; used for `ahaz test` command
+  testing:
+    # Defines ports that will be exposed on the host machine via Docker
+    exposed_ports:
+    - 1337:80
+
+# Network definitions for the task
+networks:
+# Teamnet will always contain the VPN pod - use for user-accessible services
+- name: "teamnet" # Name of the network
+  devices: [ "hello-world" ] # Pods connected to this network
+# Other networks are inaccessible to the players
+- name: "super-secret-internal-network"
+  devices: [ "hello-world" ]
+
+# Environment variables to be set in the pods
+env_vars:
+- pod_name: "hello-world" # Pod to which the environment variable will be set
+  name: "COOL_ENV_VAR" # Name of the environment variable
+  value: "hello-world" # Value of the environment variable

ahaz_cli-0.0.4/ahaz_cli/cli.py

@@ -0,0 +1,28 @@
+import logging
+from pathlib import Path
+
+import rich.logging
+import typer
+
+from .ahaz import epic, init, test
+
+log = logging.getLogger()
+log.addHandler(rich.logging.RichHandler(markup=True))
+log.setLevel(logging.INFO)
+
+SCRIPTS_ROOT = Path(__file__).parent.parent.resolve()
+
+app = typer.Typer(
+    no_args_is_help=True,
+    help="""\
+CLI for interacting with the Ahaz CTF task manager.
+""",
+)
+
+app.command()(test)
+app.command()(epic)
+app.command()(init)
+
+
+if __name__ == "__main__":
+    app()

ahaz_cli-0.0.4/ahaz_cli/lib/docker.py

@@ -0,0 +1,215 @@
+import logging
+import threading
+import time
+from typing import Callable
+
+import docker
+import rich
+import rich.ansi
+import rich.style
+from ahaz_common.task import Pod, Task
+from docker.errors import BuildError, NotFound
+from docker.models.containers import Container
+from rich.status import Status
+
+from .task import normalise_task_name
+
+log = logging.getLogger(__name__)
+
+
+def calculate_string_colour(s: str) -> int:
+    """Calculate a consistent colour code for a given string."""
+    hash_value = 0
+    for char in s:
+        hash_value = (hash_value * 31 + ord(char)) & 0xFFFFFFFF
+    return 16 + (hash_value % 216)  # Use colours from 16 to 231
+
+
+def number_to_hex_colour(n: int) -> str:
+    """Convert a number to a hex colour code."""
+    r = (n >> 16) & 0xFF
+    g = (n >> 8) & 0xFF
+    b = n & 0xFF
+    return f"#{r:02x}{g:02x}{b:02x}"
+
+
+class _ContainerLoggerAdapter(logging.LoggerAdapter):
+    def process(self, msg, kwargs):
+        if not self.extra:
+            self.extra = {}
+        cname: str = self.extra.get("container", "<unknown>")  # type: ignore
+        colour = calculate_string_colour(cname)
+        return f"[bold {number_to_hex_colour(colour)}]{cname}[/] [dim]|[/] {msg}", kwargs
+
+
+def get_container_name(task_name: str, pod_k8s_name: str) -> str:
+    return f"ahaz-{normalise_task_name(task_name)}-{pod_k8s_name}"
+
+
+def get_network_name(task_name: str, network_netname: str) -> str:
+    return f"ahaz-{normalise_task_name(task_name)}-{network_netname}"
+
+
+def try_build_image(image_tag: str, build_context: str, build_args: dict[str, str], verbose: bool) -> None:
+    client = docker.from_env()
+    log.info(f"Building image '{image_tag}'...")
+    try:
+        build_logs = client.api.build(
+            path=build_context,
+            tag=image_tag,
+            buildargs=build_args,
+            decode=True,
+        )
+        print("\x1b[2m")
+        for chunk in build_logs:
+            if "stream" in chunk and verbose:
+                # dim the build output
+                print(f"{chunk['stream']}", end="")
+        print("\x1b[0m")
+    except BuildError as e:
+        raise e
+
+
+def config_units_to_docker_units(config_value: str) -> str:
+    """
+    Convert resource limit strings from config format to Docker-compatible format.
+    E.g., "512Mi" -> "512m", "2Gi" -> "2048m"
+    """
+    if config_value.endswith("Mi"):
+        return config_value[:-2] + "m"
+    elif config_value.endswith("Gi"):
+        gi_value = int(config_value[:-2])
+        return str(gi_value * 1024) + "m"
+    else:
+        raise ValueError(f"Unsupported resource unit in value: {config_value}")
+
+
+def create_env(task: Task) -> list[tuple[Pod, Container]]:
+    client = docker.from_env()
+    try:
+        with Status("Setting up task environment...", spinner="dots") as status:
+            # Set up networks
+            network_guide = {}
+            for net in task.networks:
+                log.info(f"Setting up network '{net.name}'...")
+                status.update(f"Setting up network '{net.name}'...")
+                network = client.networks.create(
+                    name=f"{get_network_name(task.name, net.name)}", check_duplicate=True
+                )
+                for device in net.devices:
+                    network_guide.setdefault(device, []).append(network)
+
+            status.update("Setting up containers...")
+            containers = []
+            for pod in task.pods:
+                status.update(f"Creating container for pod '{pod.name}'...")
+
+                image_tag = f"{pod.image.image_name}:{task.version}"
+                env_vars = {env.name: env.value for env in task.env_vars or [] if env.pod_name == pod.name}
+                testing = getattr(pod, "testing", None)
+                exposed_ports = {
+                    port.split(":")[1]: int(port.split(":")[0])
+                    for port in (testing.exposed_ports if testing and testing.exposed_ports else [])
+                }
+
+                container = client.containers.run(
+                    image=image_tag,
+                    name=f"{get_container_name(task.name, pod.name)}",
+                    detach=True,
+                    tty=True,
+                    stdin_open=True,
+                    environment=env_vars,
+                    mem_limit=config_units_to_docker_units(pod.limits_ram),
+                    cpu_count=pod.limits_cpu,
+                    ports=exposed_ports,
+                    hostname=pod.name,
+                    stream=True,
+                )
+                containers.append((pod, container))
+
+                # Connect to networks
+                for network in network_guide.get(pod.name, []):
+                    status.update(f"Connecting pod '{pod.name}' to network '{network.name}'...")
+                    network.connect(container)
+
+            status.update("Task environment set up successfully.")
+            log.info("Task environment set up successfully.")
+    except Exception as e:
+        # Clean up any created containers and networks
+        log.error(f"Error setting up task environment: {e}")
+        cleanup_env(task.name, [pod.name for pod in task.pods], [net.name for net in task.networks])
+        raise e
+    return containers
+
+
+def cleanup_env(task_name: str, pod_names: list[str], network_names: list[str]):
+    client = docker.from_env()
+    log.info("Cleaning up created containers and networks...")
+    for pod_name in pod_names:
+        container_name = get_container_name(task_name, pod_name)
+        log.info(f"Cleaning up container '{container_name}'...")
+        try:
+            container = client.containers.get(container_name)
+            container.stop()
+            container.remove()
+        except NotFound:
+            log.warning(f"Container '{container_name}' not found during cleanup.")
+            pass  # Container does not exist, nothing to clean up
+    for network_name in network_names:
+        network_name_full = get_network_name(task_name, network_name)
+        log.info(f"Cleaning up network '{network_name_full}'...")
+        try:
+            network = client.networks.get(network_name_full)
+            network.remove()
+        except NotFound:
+            log.warning(f"Network '{network_name_full}' not found during cleanup.")
+            pass  # Network does not exist, nothing to clean up
+    log.info("Docker test environment cleaned up successfully!")
+
+
+def stream_logs(pod: Pod, container: Container, logger: _ContainerLoggerAdapter) -> None:
+    while True:
+        try:
+            buffer = ""
+            for chunk in container.logs(stream=True, tail=0, follow=True):
+                # chunk can be bytes, str, or a (stdout, stderr) tuple
+                if isinstance(chunk, tuple):
+                    chunk = chunk[0] or chunk[1]
+                if isinstance(chunk, bytes):
+                    buffer += chunk.decode("utf-8", errors="replace")
+                else:
+                    buffer += str(chunk)
+                while "\n" in buffer:
+                    line, buffer = buffer.split("\n", 1)
+                    if line.strip():
+                        ansi_line = rich.ansi.AnsiDecoder().decode(line)
+                        for segment in ansi_line:
+                            logger.info(segment.plain)
+        except Exception as e:
+            log.debug(f"Log stream for pod '{pod.name}' ended: {e}")
+
+
+def log_docker_logs(containers: list[tuple[Pod, Container]], exit_callback: Callable) -> None:
+    # Hang until the user interrupts
+    log.info("Task environment is running. Press [bold]Ctrl+C[/bold] to stop.")
+    print()
+    # Create loggers for each container
+    loggers: dict[str, _ContainerLoggerAdapter] = {}
+    for pod, _ in containers:
+        logger = logging.getLogger(pod.name)
+        loggers[pod.name] = _ContainerLoggerAdapter(logger, {"container": pod.name})
+    try:
+        # Fetch logs from containers
+        # TODO: I am fairly certain this allows the threads to log also the container shutdown logs,
+        # but I have not tested this
+        threads = []
+        for pod, container in containers:
+            t = threading.Thread(target=stream_logs, args=(pod, container, loggers[pod.name]), daemon=True)
+            t.start()
+            threads.append(t)
+
+        while True:
+            time.sleep(1)
+    except KeyboardInterrupt:
+        print()
+        exit_callback()

ahaz_cli-0.0.4/ahaz_cli/lib/file.py

@@ -0,0 +1,21 @@
+from pathlib import Path
+
+
+def read_file(file_path: str) -> str:
+    try:
+        with open(file_path, "r", encoding="utf-8") as f:
+            return f.read()
+    except FileNotFoundError as e:
+        raise FileNotFoundError(f"File not found: {file_path}") from e
+
+
+def write_file(file_path: str, content: str) -> None:
+    try:
+        with open(file_path, "w", encoding="utf-8") as f:
+            f.write(content)
+    except Exception as e:
+        raise IOError(f"Error writing to file: {file_path}") from e
+
+
+def test_for_file(file_path: str) -> bool:
+    return Path(file_path).is_file()

ahaz_cli-0.0.4/ahaz_cli/lib/task.py

@@ -0,0 +1,82 @@
+import logging
+
+import yaml
+
+from ahaz_common import Task
+
+log = logging.getLogger()
+
+
+def deserialise_task(task_config: str) -> Task:
+    try:
+        config_dict = yaml.safe_load(task_config)
+    except yaml.YAMLError as e:
+        log.error(f"Error parsing YAML: {e}")
+        raise ValueError("Invalid task configuration") from e
+
+    try:
+        return Task(**config_dict)
+    except TypeError as e:
+        log.error(f"Error constructing Task object: {e}")
+        raise ValueError("Invalid task configuration structure") from e
+
+
+# TODO: Figure out a better way to handle serialisation of nested objects
+# TODO: Is this even necessary?
+def serialise_task(task: Task) -> str:
+    """
+    Serialise a Task into YAML with a fixed field order matching the example:
+    name, version, description, score, scoring_type, pods, networks, env_vars
+    """
+
+    # helper getters to be defensive about missing attributes
+    def _image_dict(pod):
+        img = getattr(pod, "image", None) or {}
+        return {
+            "image_name": getattr(img, "image_name", None),
+            "build_context": getattr(img, "build_context", None),
+        }
+
+    def _testing_dict(pod):
+        testing = getattr(pod, "testing", None)
+        return {"exposed_ports": testing.exposed_ports if testing is not None else []}
+
+    # Build top-level mapping in the desired order (insertion order is preserved)
+    config_dict = {
+        "name": task.name,
+        "version": getattr(task, "version", "1.0.0"),
+        "description": task.description,
+        "score": task.score,
+        "scoring_type": task.scoring_type,
+        "pods": [
+            {
+                "name": pod.name,
+                "image": _image_dict(pod),
+                "limits_ram": getattr(pod, "limits_ram", None),
+                "limits_cpu": getattr(pod, "limits_cpu", None),
+                "visible_to_user": getattr(pod, "visible_to_user", None),
+                # include build for convenience (matches example where build == build_context)
+                "build": getattr(getattr(pod, "image", None), "build_context", None),
+                "testing": _testing_dict(pod),
+            }
+            for pod in task.pods or []
+        ],
+        "networks": [{"name": net.name, "devices": list(net.devices)} for net in task.networks or []],
+        "env_vars": [
+            {
+                # map original fields to the example-style names
+                "pod_name": env.pod_name,
+                "name": env.name,
+                "value": env.value,
+            }
+            for env in task.env_vars or []
+        ],
+    }
+
+    # Preserve insertion order in output
+    return yaml.safe_dump(config_dict, sort_keys=False)
+
+
+def normalise_task_name(name: str) -> str:
+    # Make the task lowercase, replace spaces with hyphens, and remove special characters
+    return name.lower().replace(" ", "-").replace(r"([^a-z0-9-])", "")

ahaz_cli-0.0.4/ahaz_cli/templates.py

@@ -0,0 +1,24 @@
+from pathlib import Path
+
+from .lib.task import deserialise_task, serialise_task
+
+
+def write_task_yaml(file_path: Path, task_name: str) -> None:
+    template_path = Path(__file__).parent / "assets" / "task.yaml"
+    task = deserialise_task(template_path.read_text())
+
+    task.name = task_name
+
+    task_dest = file_path / "task.yaml"
+    task_dest.write_text(serialise_task(task))
+
+
+def copy_example_images(dest_path: Path) -> None:
+    example_image_src = Path(__file__).parent / "assets" / "hello-world"
+    example_image_dest = dest_path / "hello-world"
+    example_image_dest.mkdir(exist_ok=True)
+
+    for item in example_image_src.iterdir():
+        if item.is_file():
+            dest_file = example_image_dest / item.name
+            dest_file.write_text(item.read_text())

ahaz_cli-0.0.4/ahaz_cli.egg-info/PKG-INFO

@@ -0,0 +1,9 @@
+Metadata-Version: 2.4
+Name: ahaz_cli
+Version: 0.0.4
+Summary: Command-line management interface for Ahaz
+Requires-Dist: docker>=7.1.0
+Requires-Dist: rich>=14.2.0
+Requires-Dist: typer>=0.20.0
+Requires-Dist: ahaz_common
+Requires-Dist: pyyaml>=6.0.3

ahaz_cli-0.0.4/ahaz_cli.egg-info/SOURCES.txt

@@ -0,0 +1,17 @@
+pyproject.toml
+ahaz_cli/__init__.py
+ahaz_cli/__main__.py
+ahaz_cli/ahaz.py
+ahaz_cli/cli.py
+ahaz_cli/templates.py
+ahaz_cli.egg-info/PKG-INFO
+ahaz_cli.egg-info/SOURCES.txt
+ahaz_cli.egg-info/dependency_links.txt
+ahaz_cli.egg-info/entry_points.txt
+ahaz_cli.egg-info/requires.txt
+ahaz_cli.egg-info/top_level.txt
+ahaz_cli/assets/task.yaml
+ahaz_cli/assets/hello-world/Dockerfile
+ahaz_cli/lib/docker.py
+ahaz_cli/lib/file.py
+ahaz_cli/lib/task.py

ahaz_cli-0.0.4/ahaz_cli.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

ahaz_cli-0.0.4/ahaz_cli.egg-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+ahaz = ahaz_cli.cli:app

ahaz_cli-0.0.4/ahaz_cli.egg-info/requires.txt

@@ -0,0 +1,5 @@
+docker>=7.1.0
+rich>=14.2.0
+typer>=0.20.0
+ahaz_common
+pyyaml>=6.0.3

ahaz_cli-0.0.4/ahaz_cli.egg-info/top_level.txt

@@ -0,0 +1 @@
+ahaz_cli

ahaz_cli-0.0.4/pyproject.toml

@@ -0,0 +1,26 @@
+[project]
+name = "ahaz_cli"
+version = "0.0.4"
+description = "Command-line management interface for Ahaz"
+dependencies = [
+    "docker>=7.1.0",
+    "rich>=14.2.0",
+    "typer>=0.20.0",
+    "ahaz_common",
+    "pyyaml>=6.0.3",
+]
+
+[tool.uv.sources]
+ahaz_common = { workspace = true }
+
+[project.scripts]
+ahaz = "ahaz_cli.cli:app"
+
+[tool.setuptools]
+include-package-data = true
+
+[tool.setuptools.package-data]
+"ahaz_cli" = ["assets/**"]
+
+[tool.setuptools.packages.find]
+include = ["ahaz_cli*"]

ahaz_cli-0.0.4/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+