pwnshop: add challenge.yml support for privileged containers

Author: zardus

tools/pwnshop/pyproject.toml

@@ -9,6 +9,7 @@ dependencies = [
   "click",
   "jinja2",
   "pyastyle",
+  "pyyaml",
   "rich",
 ]
 

tools/pwnshop/src/pwnshop/commands/run.py

@@ -32,11 +32,12 @@
 @click.argument("command", nargs=-1, default=("/bin/bash",))
 def run_command(challenge_path, user, volumes, command):
     """Run interactive shell for a challenge."""
+    config = lib.load_challenge_config(challenge_path)
     try:
         image_id = lib.build_challenge(challenge_path)
     except RuntimeError as error:
         raise click.ClickException(str(error)) from error
     resolved_volumes = [path.resolve() for path in volumes]
     logger.info("running %s as uid=%d, command=%s", challenge_path, user, list(command))
-    with lib.run_challenge(image_id, volumes=resolved_volumes) as (container, flag):
+    with lib.run_challenge(image_id, volumes=resolved_volumes, privileged=config["privileged"]) as (container, flag):
         subprocess.run(["docker", "exec", f"--user={user}", "-it", container, *command])

tools/pwnshop/src/pwnshop/commands/test.py

@@ -64,6 +64,7 @@ def test_challenge(challenge_path):
         challenge_path = pathlib.Path(challenge_path)
         rendered = None
         try:
+            config = lib.load_challenge_config(challenge_path)
             rendered = lib.render_challenge(challenge_path)
             image_id = lib.build_challenge(challenge_path)
             tests = sorted(rendered.rglob("test*/test_*"))
@@ -75,7 +76,7 @@ def test_challenge(challenge_path):
             for test in tests:
                 test_name = test.relative_to(rendered)
                 logger.debug("running test %s in %s", test_name, challenge_path)
-                with lib.run_challenge(image_id, volumes=[test]) as (container, _):
+                with lib.run_challenge(image_id, volumes=[test], privileged=config["privileged"]) as (container, _):
                     try:
                         run = subprocess.run(
                             ["docker", "exec", "--user=1000:1000", container, f"{test}"],

tools/pwnshop/src/pwnshop/lib/__init__.py

@@ -8,16 +8,32 @@
 import shutil
 import subprocess
 import tempfile
-from typing import Iterable, Iterator, List, Optional, Sequence
+from typing import Any, Dict, Iterable, Iterator, List, Optional, Sequence
 
 import black
 import jinja2
 import pyastyle
+import yaml
 
 logger = logging.getLogger(__name__)
 
 CHALLENGE_SEED = int(os.environ.get("CHALLENGE_SEED", "0"))
 
+CHALLENGE_CONFIG_DEFAULTS: Dict[str, Any] = {
+    "privileged": False,
+}
+
+
+def load_challenge_config(challenge_path: pathlib.Path) -> Dict[str, Any]:
+    config = dict(CHALLENGE_CONFIG_DEFAULTS)
+    config_file = challenge_path / "challenge.yml"
+    if config_file.is_file():
+        logger.debug("loading challenge config from %s", config_file)
+        with open(config_file) as f:
+            user_config = yaml.safe_load(f) or {}
+        config.update(user_config)
+    return config
+
 
 class _NoSelfExtendLoader(jinja2.FileSystemLoader):
     """FileSystemLoader that prevents templates from extending/including themselves.
@@ -136,7 +152,7 @@ def ignore_git_crypt(current, names):
 
 @contextlib.contextmanager
 def run_challenge(
-    challenge_image: str, *, volumes: Optional[Sequence[pathlib.Path]] = None
+    challenge_image: str, *, volumes: Optional[Sequence[pathlib.Path]] = None, privileged: bool = False
 ) -> Iterator[tuple[str, str]]:
     flag = "pwn.college{" + base64.b64encode(os.urandom(32)).decode() + "}"
     env_options = []
@@ -149,6 +165,8 @@ def run_challenge(
     logger.info("starting container for image %s", challenge_image)
     if volumes:
         logger.debug("mounting volumes: %s", volumes)
+    if privileged:
+        logger.debug("running container in privileged mode")
     container = (
         subprocess.check_output(
             [
@@ -159,8 +177,7 @@ def run_challenge(
                 "--detach",
                 "--init",
                 "--user=0:0",
-                "--device=/dev/kvm",
-                "--cap-add=SYS_PTRACE",
+                *(["--privileged"] if privileged else ["--device=/dev/kvm", "--cap-add=SYS_PTRACE"]),
                 *env_options,
                 *[f"--volume={volume}:{volume}:ro" for volume in (volumes or [])],
                 challenge_image,