diff --git a/func_smash.py b/func_smash.py index 9bb9538..7932a42 100644 --- a/func_smash.py +++ b/func_smash.py @@ -1,10 +1,10 @@ +from argparse import ArgumentParser from code import interact import imp import opcode import os import random import re -import sys import types import prettify @@ -90,18 +90,44 @@ def print_function(func): print def run(): - try: - path, name = os.path.split(sys.argv[1]) - name = re.sub("\.pyc?$", "", name) - except IndexError: - raise RuntimeError("Needs a filename as a command-line argument") + parser = ArgumentParser(prog="func-smash", + description="Smash functions together!") + parser.add_argument("path", metavar="path", help="path to the corpus file") + parser.add_argument("-n", "--no-run", action="store_true", + help="don't run the function after smashing it") + parser.add_argument("-i", "--interpret", action="store_true", + help="open a interpreter session after smashing") + parser.add_argument("-p", "--prettify", action="store_true", + help="prettify the function after smashing it") + args = parser.parse_args() + + path, name = os.path.split(args.path) + name = re.sub("\.pyc?$", "", name) file_obj, path, desc = imp.find_module(name, [path]) try: module = imp.load_module(name, file_obj, path, desc) finally: file_obj.close() - _demo(module.corpus) + corpus = module.corpus + chain = make_chain(corpus) + func = make_function(chain, "func") + print "Using {0}-function corpus.".format(len(corpus)) + print "Smashed function disassembly:" + print_function(func) + + if args.prettify: + print "\nFunction as python code:" + prettify.prettify_function(func, indent=4) + + if not args.no_run: + arg = 12 + print + print "func({0}) =".format(arg), func(arg) + + if args.interpret: + variables = dict(globals().items() + locals().items()) + interact(banner="", local=variables) def _parse_func(func, chain): codeobj = func.__code__ @@ -188,22 +214,5 @@ def _coerce_arg_into_codes(codes, arg): codes.append(arg % 256) codes.append(arg // 256) -def _demo(corpus, arg=12.0): - chain = make_chain(corpus) - func = make_function(chain, "func") - print "Using {0}-function corpus.".format(len(corpus)) - print "Smashed function disassembly:" - print_function(func) - print "\nFunction as python code:" - prettify.prettify_function(func, indent=4) - - if not len(sys.argv) > 2 or "n" not in "".join(sys.argv[2:]): - print - print "func({0}) =".format(arg), func(arg) - - if len(sys.argv) > 2 and "i" in "".join(sys.argv[2:]): - variables = dict(globals().items() + locals().items()) - interact(banner="", local=variables) - if __name__ == "__main__": run()