From edbac0268d55df8e63d30e1f96582763e4d76d53 Mon Sep 17 00:00:00 2001 From: Yan Date: Thu, 21 Dec 2023 12:25:36 -0700 Subject: [PATCH] redesign pwnshop's CLI interface (and remove yaml stuff) --- .github/workflows/test.yml | 6 +- README.md | 17 ++- pwnshop/__main__.py | 282 +++++++++++++------------------------ 3 files changed, 105 insertions(+), 200 deletions(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 4229a2c..9d5ab30 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -21,7 +21,7 @@ jobs: run: | echo "pwncollege{TESTING}" | sudo tee /flag cd .. - pwnshop -I pwnshop/example_module -c ShellExample --src - pwnshop -I pwnshop/example_module -c ShellExample --bin > /tmp/shell_example + pwnshop -I pwnshop/example_module render ShellExample + pwnshop -I pwnshop/example_module build ShellExample > /tmp/shell_example file /tmp/shell_example | grep ELF - pwnshop -I pwnshop/example_module -c ShellExample --verify + pwnshop -I pwnshop/example_module verify ShellExample diff --git a/README.md b/README.md index 5fefe48..33bf830 100644 --- a/README.md +++ b/README.md @@ -17,18 +17,17 @@ Let's generate some things! ```sh -# example challenge in testing mode -pwnshop -I /path/to/example_module --challenge ShellExample --src +# render example challenge source code in testing mode +pwnshop -I /path/to/example_module render ShellExample -# example challenge in teaching mode -pwnshop -I /path/to/example_module --challenge ShellExample --walkthrough --src +# render example challenge source code in teaching mode +pwnshop -I /path/to/example_module render ShellExample --walkthrough -# make sure the example challenge compiles and the reference solution works -pwnshop -I /path/to/example_module --challenge ShellExample --walkthrough --verify - -# generate the example challenge binary -pwnshop -I /path/to/example_module --challenge ShellExample --walkthrough --bin > example_shell +# test the example challenge binary and solution +pwnshop -I /path/to/example_module verify ShellExample --walkthrough +# build the example challenge binary +pwnshop -I /path/to/example_module build ShellExample --walkthrough -O example_shell ``` ## Writing challenges diff --git a/pwnshop/__main__.py b/pwnshop/__main__.py index 58ed5f6..fdff579 100644 --- a/pwnshop/__main__.py +++ b/pwnshop/__main__.py @@ -1,132 +1,106 @@ import argparse import pwnshop import random -import yaml +import glob import sys import pwn import os def challenge_class(challenge=None, module=None, level=None): - if challenge: + if ":" not in challenge: assert challenge in pwnshop.ALL_CHALLENGES, "Unknown challenge specified!" return pwnshop.ALL_CHALLENGES[challenge] - elif level and module: + else: + module, level_src = challenge.split(":") + level = int(level_src) + assert module in pwnshop.MODULE_LEVELS, "Uknown module specified!" assert 0 < level <= len(pwnshop.MODULE_LEVELS[module]), "Invalid level specified." return pwnshop.MODULE_LEVELS[module][level-1] - else: - raise AssertionError("Improper challenge specification (need challenge or module&level).") def make_challenge(challenge=None, module=None, level=None, **kwargs): return challenge_class(challenge=challenge, module=module, level=level)(**kwargs) def main(): - parser = argparse.ArgumentParser(description="pwnshop challenge emitter") + parser = argparse.ArgumentParser(description="pwnshop challenge emitter", formatter_class=argparse.RawDescriptionHelpFormatter) parser.add_argument( "-I", "--import", required=False, - nargs="*", - action="extend", - help="a path to import additional challenges from (either /path/to/module.py or /path/to/package/)", - ) - parser.add_argument( - "-c", - "--challenge", - required=False, - help="the challenge to generate", - ) - parser.add_argument( - "-m", - "--module", - required=False, - help="the module to use", - ) - parser.add_argument( - "-L", - "--level", - required=False, - type=int, - help="the level (of the given module) to generate", - ) - parser.add_argument( - "-s", - "--seed", - required=False, - default=random.randrange(0, 999999), - help="the seed from which to generate the challenge", + help="a path glob to import additional challenges from (either /path/to/module.py or /path/to/package/ or /some/glob/*, but avoid shell-expansion for the latter!)", ) - parser.add_argument( - "-d", - "--debug", - action="store_true", - help="print out debugging information", - ) - parser.add_argument( + commands = parser.add_subparsers(help="the action for pwnshop to perform", required=True, dest="ACTION") + 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.add_argument( "-t", "--strace", action="store_true", - help="print out strace information", + help="print out strace information during verification", ) - parser.add_argument( + + command_render.add_argument( "-l", "--lineno", action="store_true", - help="print out line numbers", + help="render line numbers", ) - parser.add_argument( + + command_verify.add_argument( "-f", "--flag", help="change the flag to be verified against", ) - # where to write - parser.add_argument( - "-O", - "--out", - type=argparse.FileType('w'), - default='-', - help="change the output destination" - ) - parser.add_argument( - "-D", - "--dojo", - help="the dojo to insert the yml spec into" + command_build.add_argument( + "--lpath", + help="Location to store needed library files", ) - group = parser.add_mutually_exclusive_group(required=True) - group.add_argument( - "--src", action="store_true", help="Dump the challenge source" - ) - group.add_argument( - "--bin", action="store_true", help="Dump the challenge binary" - ) - group.add_argument( - "--verify", action="store_true", help="Verify that the challenge works" - ) - group.add_argument( - "--yml", action="store_true", help="Dump the challenge yaml" - ) - group.add_argument( - "--dojo-insert", action="store_true", help="Insert challenge yaml into dojo spec." - ) + # where to write + for c in [ command_render, command_build ]: + c.add_argument( + "-O", + "--out", + type=argparse.FileType('w'), + default='-', + help="change the output destination" + ) - parser.add_argument( - "-w", - "--walkthrough", - action="store_true", - help="enable challenge walkthrough mode", - ) + # common arguments + for subparser in commands.choices.values(): + subparser.add_argument( + "-s", + "--seed", + required=False, + default=random.randrange(0, 999999), + help="the seed from which to generate the challenge (default: random)", + ) - parser.add_argument( - "--lpath", - help="Location to store needed library files", - ) + subparser.add_argument( + "-d", + "--debug", + action="store_true", + help="print out debugging information", + ) + + subparser.add_argument( + "-w", + "--walkthrough", + action="store_true", + help="enable challenge walkthrough mode", + ) + + subparser.add_argument("challenge", help="the challenge, as ChallengeClassName or ModuleName:level_number") + + 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() if getattr(args, "import", None): - imports = getattr(args, "import") + imports = glob.glob(getattr(args, "import")) for i in imports: sys.path.append(os.path.realpath(os.path.dirname(i))) try: @@ -135,110 +109,42 @@ def main(): finally: sys.path.pop() - if args.debug: - pwn.context.log_level = "DEBUG" - - if args.flag: - with open("/flag", "wb") as f: - f.write(args.flag.encode()) - - if (args.yml or args.dojo_insert) and args.module and not args.level: - assert args.module in pwnshop.MODULE_LEVELS, "Uknown module specified!" - module = getattr(pwnshop.challenges, args.module) - num_test = module.NUM_TESTING - metadata = [ ] - for i,C in enumerate(pwnshop.MODULE_LEVELS[args.module], start=1): - if hasattr(module, "CHOOSE_LEVELS") and i not in module.CHOOSE_LEVELS: - continue - - cm = C(seed=args.seed, walkthrough=1).metadata() - cm["name"] = f"level{i}" if num_test == 0 else f"level{i}.0" - cm["category"] = args.module - del cm["class_name"] - metadata.append(cm) - - for j in range(num_test): - cm = C(seed=args.seed, walkthrough=0).metadata() - cm["name"] = f"level{i}.{j+1}" - cm["category"] = args.module - del cm["class_name"] - metadata.append(cm) - - if args.yml: - args.out.write(yaml.dump(metadata)) - else: - # using ruamel to preserve structure, comments, etc - from ruamel.yaml import YAML - y = YAML() - y.indent(offset=2, sequence=4, mapping=2) - y.preserve_quotes = True - y.width = 4096 - - # thanks to https://stackoverflow.com/questions/8640959/how-can-i-control-what-scalar-form-pyyaml-uses-for-my-data - #def str_presenter(dumper, data): - # if len(data.splitlines()) > 1: # check for multiline string - # return dumper.represent_scalar('tag:yaml.org,2002:str', data, style='|') - # return dumper.represent_scalar('tag:yaml.org,2002:str', data) - #yaml.add_representer(str, str_presenter) - - dojo = y.load(open(args.dojo, "r").read()) - try: - dojo_module = next(m for m in dojo["modules"] if m["id"] == module.DOJO_MODULE) - except StopIteration: - print("Can't find module with specified challenge category in dojo spec.") - sys.exit(1) - dojo_module["challenges"] = metadata - y.dump(dojo, open(args.dojo, "w")) - elif args.challenge or (args.module and args.level): - challenge = make_challenge( - challenge=args.challenge, module=args.module, level=args.level, seed=args.seed, walkthrough=args.walkthrough - ) - - if args.src: - src = challenge.generate_source() - if not args.lineno: - args.out.write(src+"\n") - else: - for i, line in enumerate(src.splitlines()): - args.out.write(f"{i + 1}\t{line}\n") - if os.path.isfile(args.out.name): - os.chmod(args.out.name, 0o644) - - if args.bin: - binary, libs = challenge.build_binary() - args.out.buffer.write(binary) - if os.path.isfile(args.out.name): - os.chmod(args.out.name, 0o755) - - if args.lpath and libs: - os.makedirs(args.lpath, exist_ok=True) - for lib_name, lib_bytes in libs: - lib_path = args.lpath + f'/{lib_name}' - with open(lib_path, 'wb+') as f: - f.write(lib_bytes) - os.chmod(lib_path, 0o755) - - - if args.verify: - # TODO: this was an optimization which may be critical for verifying against 1000s of program runs - # binary = challenge.build_binary() - # challenge.verify(binary=binary, strace=args.strace) - challenge.verify() - - if args.yml: - args.out.write(yaml.dump(challenge.metadata())) - - elif args.src or args.bin or args.verify: - print("Improper challenge specification (need challenge or module&level).") - sys.exit(1) - elif args.yml: - print("Improper challenge specification (need challenge or module&level or module).") - sys.exit(1) - else: - print("Unexpected configuration.") - sys.exit(1) - + challenge = make_challenge(challenge=args.challenge, seed=args.seed, walkthrough=args.walkthrough) + if args.ACTION == "render": + src = challenge.generate_source() + if not args.lineno: + args.out.write(src+"\n") + else: + for i, line in enumerate(src.splitlines()): + args.out.write(f"{i + 1}\t{line}\n") + if os.path.isfile(args.out.name): + os.chmod(args.out.name, 0o644) + + if args.ACTION == "build": + binary, libs = challenge.build_binary() + args.out.buffer.write(binary) + if os.path.isfile(args.out.name): + os.chmod(args.out.name, 0o755) + + if args.lpath and libs: + os.makedirs(args.lpath, exist_ok=True) + for lib_name, lib_bytes in libs: + lib_path = args.lpath + f'/{lib_name}' + with open(lib_path, 'wb+') as f: + f.write(lib_bytes) + os.chmod(lib_path, 0o755) + + + if args.ACTION == "verify": + if args.debug: + pwn.context.log_level = "DEBUG" + + if args.flag: + with open("/flag", "wb") as f: + f.write(args.flag.encode()) + + challenge.verify(strace=args.strace) if __name__ == "__main__": main()