From 32f9573e79228b549a14425bc635bd8feb9420f1 Mon Sep 17 00:00:00 2001
From: Yan <yans@yancomm.net>
Date: Tue, 5 Nov 2024 02:22:39 -0700
Subject: [PATCH] remove verify-module in favor of --module for pwnshop verify

---
 Dockerfile          |  1 +
 pwnshop/__main__.py | 76 ++++++++++++++++++++++++++-------------------
 tests/test.sh       |  2 +-
 3 files changed, 46 insertions(+), 33 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index b7c838a..d7effe7 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -27,6 +27,7 @@ RUN apt-get update && xargs apt-get install --no-install-recommends -yqq <<EOF &
         python3-pip
         python3-yaml
         qemu-system
+        strace
         sudo
         unzip
         wget
diff --git a/pwnshop/__main__.py b/pwnshop/__main__.py
index 2279c78..ffb206f 100644
--- a/pwnshop/__main__.py
+++ b/pwnshop/__main__.py
@@ -25,13 +25,25 @@ def challenge_class(challenge):
 def with_challenges(f):
     @functools.wraps(f)
     def f_with_challenge(args):
-        challenges = [
-            challenge_class(challenge=c)(seed=args.seed, walkthrough=args.walkthrough)
-            for c in args.challenges
-        ] if args.challenges else [
-            c(seed=args.seed, walkthrough=args.walkthrough)
-            for c in pwnshop.ALL_CHALLENGES.values()
-        ]
+        challenges = [ ]
+        for m in (args.module or [ ]) if "module" in args else [ ]:
+            challenges += [
+                c(seed=args.seed, walkthrough=args.walkthrough)
+                for c in pwnshop.MODULE_LEVELS[m]
+            ]
+
+        if args.challenges:
+            challenges += [
+                challenge_class(challenge=c)(seed=args.seed, walkthrough=args.walkthrough)
+                for c in args.challenges
+            ]
+
+        if not challenges:
+            challenges += [
+                c(seed=args.seed, walkthrough=args.walkthrough)
+                for c in pwnshop.ALL_CHALLENGES.values()
+            ]
+
         return f(args, challenges)
     return f_with_challenge
 
@@ -121,9 +133,6 @@ def verify_many(args, challenges):
 def handle_verify(args, challenges):
     return verify_many(args, challenges)
 
-def handle_verify_module(args):
-    return verify_many(args, pwnshop.MODULE_LEVELS[args.module])
-
 def handle_apply(args):
     y = yaml.safe_load(open(args.yaml))
     for c in y['challenges']:
@@ -180,7 +189,6 @@ def main():
     command_render = commands.add_parser("render", help="render the source code of a challenge")
     command_build = commands.add_parser("build", help="build the binary code of a challenge")
     command_verify = commands.add_parser("verify", help="verify the functionality of a challenge")
-    command_verify_module = commands.add_parser("verify-module", help="verify the functionality of all challenges in a module")
     command_apply = commands.add_parser("apply", help="parse a yaml and generate the defined challenges")
 
     command_render.add_argument(
@@ -215,26 +223,32 @@ def main():
         help="Path to the yaml.",
     )
 
-    for c in [ command_verify, command_verify_module ]:
-        c.add_argument(
-            "-t",
-            "--strace",
-            action="store_true",
-            help="print out strace information during verification",
-        )
+    command_verify.add_argument(
+        "-t",
+        "--strace",
+        action="store_true",
+        help="print out strace information during verification",
+    )
 
-        c.add_argument(
-            "-f",
-            "--flag",
-            help="change the flag to be verified against",
-        )
+    command_verify.add_argument(
+        "-f",
+        "--flag",
+        help="change the flag to be verified against",
+    )
 
-        c.add_argument(
-            "-T",
-            "--timeout",
-            type=int,
-            help="set a per-challenge timeout (in seconds) for verification",
-        )
+    command_verify.add_argument(
+        "-T",
+        "--timeout",
+        type=int,
+        help="set a per-challenge timeout (in seconds) for verification",
+    )
+
+    command_verify.add_argument(
+        "-m",
+        "--module",
+        help="Verify all the challenges in a module. Can be specified multiple times.",
+        action="append"
+    )
 
 
     # where to write
@@ -275,9 +289,7 @@ def main():
     for subparser in [ command_render, command_build, command_verify ]:
         subparser.add_argument("challenges", help="the challenges, as multiple ChallengeClassName or ModuleName:level_number. Default: all challenges.", nargs="*")
 
-    command_verify_module.add_argument("module", help="the module to verify")
-
-    parser.epilog = f"""Commands usage:\n\t{command_render.format_usage()}\t{command_build.format_usage()}\t{command_verify.format_usage()}\t{command_verify_module.format_usage()}"""
+    parser.epilog = f"""Commands usage:\n\t{command_render.format_usage()}\t{command_build.format_usage()}\t{command_verify.format_usage()}"""
 
     args = parser.parse_args()
     pwn.context.log_level = "ERROR"
diff --git a/tests/test.sh b/tests/test.sh
index 39be05f..acc9225 100755
--- a/tests/test.sh
+++ b/tests/test.sh
@@ -12,7 +12,7 @@ file /tmp/shell_example | grep ELF
 strings <(pwnshop build ShellOptimized) | grep -- "-O3"
 pwnshop verify ShellExample
 
-( pwnshop verify-module example_module || true ) > /tmp/out
+( pwnshop verify -m example_module || true ) > /tmp/out
 cat /tmp/out | grep "SUCCEEDED: ShellExample"
 cat /tmp/out | grep "SUCCEEDED: ShellOptimized"
 cat /tmp/out | grep "FAILED: ShellBadVerifier"