support verification images

example_module/__init__.py

@@ -38,6 +38,7 @@ def verify(self, **kwargs):
             shellcode = pwnlib.asm.asm(
                 pwnlib.shellcraft.open("/flag") + pwnlib.shellcraft.sendfile(1, 3, 0, 1024) + pwnlib.shellcraft.exit(0)
             )
+            process.clean()
             process.write(shellcode)
             assert self.flag in process.readall()
 
@@ -61,6 +62,15 @@ class Shell1604(ShellExample):
     The same example, built using Ubuntu 16.04.
     """
 
-    BUILD_IMAGE = "pwncollege/pwnshop-builder:ubuntu1604"
+    BUILD_IMAGE = "ubuntu:16.04"
     APT_DEPENDENCIES = [ "libcapstone-dev" ]
     PIN_LIBRARIES = True
+
+class Shell1604InVitu(ShellExample):
+    """
+    The same example, built using Ubuntu 16.04 and verified in the same container, without the need for library pinning.
+    """
+
+    BUILD_IMAGE = "ubuntu:16.04"
+    VERIFY_IMAGE = "ubuntu:16.04"
+    APT_DEPENDENCIES = [ "libcapstone-dev" ]

pwnshop/__main__.py

@@ -1,15 +1,19 @@
 import traceback
 import functools
 import argparse
+import logging
 import pwnshop
 import pwnlib.context
+import pwnlib.log
 import signal
 import random
 import yaml
 import glob
 import sys
 import os
 
+pwnlib.log.install_default_handler()
+
 def challenge_class(challenge):
     if ":" not in challenge:
         assert challenge in pwnshop.ALL_CHALLENGES, "Unknown challenge specified!"

pwnshop/challenge.py

@@ -220,7 +220,6 @@ def build(self):
 
         cmd = self.build_compiler_cmd()
 
-        self.libraries = None
         if self._build_container:
             ret, out = self._build_container.exec_run(cmd)
             if ret != 0:
@@ -231,6 +230,7 @@ def build(self):
             self._build_container.exec_run(f'chown {os.getuid()}:{os.getgid()} {self.bin_path}')
             self.libraries = self.pin_libraries() if self.PIN_LIBRARIES else []
         else:
+            self.libraries = None
             subprocess.check_output(cmd)
 
         with open(self.bin_path, 'rb') as f:
@@ -253,12 +253,21 @@ def run_challenge(
         strace=False,
         **kwargs,
     ):
-        environment_ctx = None
-        if flag_symlink:
-            os.symlink("/flag", f"{flag_symlink}")
+        if not self._verify_container:
+            self._verify_container = self._create_container(self.VERIFY_IMAGE)
 
-        with open("/flag", "rb") as f:
-            assert f.read() == self.flag
+        if flag_symlink:
+            self.run_sh(f"ln -s /flag {flag_symlink}")
+
+        if self._verify_container:
+            p = self.run_sh(f"tee /flag")
+            p.send(self.flag)
+            p.stdin.close()
+            assert p.proc.wait() == 0
+            assert self.run_sh("cat /flag").readall() == self.flag
+        else:
+            with open("/flag", "rb") as f:
+                assert f.read() == self.flag
 
         if argv is None:
             argv = [self.bin_path]
@@ -273,18 +282,22 @@ def run_challenge(
         if not self.binary:
             self.build()
 
+        if self._verify_container:
+            argv = f"docker exec -i {self._verify_container.name}".split() + argv
+
         with pwnlib.tubes.process.process(
             argv, **kwargs
         ) as process:
             if close_stdin:
                 process.stdin.close()
-            try:
-                yield process
-            finally:
-                if environment_ctx:
-                    environment_ctx.__exit__(*sys.exc_info())
+            yield process
 
     def run_sh(self, command, **kwargs):
+        if not self._verify_container:
+            self._verify_container = self._create_container(self.VERIFY_IMAGE)
+        if self._verify_container:
+            command = f"docker exec -i {self._verify_container.name} {command}"
+
         return pwnlib.tubes.process.process(command, shell=True, **kwargs)
 
     def _create_container(self, image=None):
@@ -304,7 +317,7 @@ def _create_container(self, image=None):
             volumes = {self.work_dir : {'bind': self.work_dir, 'mode': 'rw'}}
         )
         ret, out = container.exec_run(
-            f'/bin/bash -c "apt-get update && apt-get install -y gcc patchelf {" ".join(self.APT_DEPENDENCIES)} && mkdir -p /tmp/pwnshop"'
+            f'/bin/bash -c "apt-get update && apt-get install -y gcc patchelf {" ".join(self.APT_DEPENDENCIES)}"'
         )
 
         if ret != 0:

tests/test.sh

@@ -11,17 +11,12 @@ file /tmp/shell_example | grep ELF
 strings <(pwnshop build ShellOptimized) | grep -- "-O3"
 pwnshop verify ShellExample
 
-( pwnshop verify -m example_module || true ) | tee /tmp/out
-cat /tmp/out | grep "SUCCEEDED: ShellExample"
-cat /tmp/out | grep "SUCCEEDED: ShellOptimized"
-cat /tmp/out | grep "FAILED: ShellBadVerifier"
-cat /tmp/out | grep "SUCCEEDED: Shell1604"
-
 ( pwnshop verify || true ) | tee /tmp/out
 cat /tmp/out | grep "SUCCEEDED: ShellExample"
 cat /tmp/out | grep "SUCCEEDED: ShellOptimized"
 cat /tmp/out | grep "FAILED: ShellBadVerifier"
 cat /tmp/out | grep "SUCCEEDED: Shell1604"
+cat /tmp/out | grep "SUCCEEDED: Shell1604InVitu"
 
 pwnshop apply ../example_deploy/pwnshop.yml
 SOURCES=( ../example_deploy/*/*/*.c )