PythonPickle

pickle.py

import argparse
import pickle
import sys


if sys.version_info < (3, 0, 0):
    _exec = None
    buff = sys.stdout  # Buffer to write binary strings to
else:
    import builtins
    _exec = getattr(builtins, "exec")  # Prevent Py2 SyntaxError exceptions
    buff = sys.stdout.buffer  # Buffer to write bytes to
    raw_input = input


class Exec(object):
    class Compile(object):
        def __init__(self, code):
            self.code = code

        def __reduce__(self):
            return (compile, (self.code,  "<string>", "exec"))

    def __reduce__(self):
        if _exec is None:
            return (eval, (Exec.Compile(self.code),))  # Py2 compile eval
        return (_exec, (self.code,))  # Py3 exec

    def __init__(self, code):
        self.code = code

    @classmethod
    def test(cls):
        test_code = """
            global retval
            print("[+] Running...")
            retval = True
        """

        print("[+] Pickling...")
        pickle.loads(pickle.dumps(cls("if 1:" + test_code)))

        # Py2 evals in the context of the pickle module, Py3 execs in the method scope
        assert getattr(pickle, "retval", globals().get("retval", False)) == True
        print("[+] Success!")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Python code pickler")
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("-u", "--userinput", action="store_true", default=False)
    group.add_argument("-i", "--infile", type=str, default=None)
    group.add_argument("-l", "--loadfile", type=str, default=None)
    group.add_argument("-t", "--test", action="store_true", default=False)
    parser.add_argument("-o", "--outfile", type=str, default=None)
    parser.add_argument("-r", "--run", action="store_true", default=False)
    args = parser.parse_args()

    if args.test:
        Exec.test()

    pickled_code = None
    if args.userinput:
        lines = []
        while True:
            lines.append(raw_input())
            if not lines[-1]:
                break
        pickled_code = pickle.dumps(Exec("\n".join(lines)))
    elif args.infile is not None:
        with open(sys.argv[2], "r") as f:
            pickled_code = pickle.dumps(Exec(f.read()))
    elif args.loadfile is not None:
        with open(args.loadfile, "rb") as f:
            pickled_code = f.read()

    if args.run:
        pickle.loads(pickled_code)
    if args.outfile is not None:
        if args.outfile == "-":
            buff.write(pickled_code)
        else:
            with open(args.outfile, "wb") as f:
                f.write(pickled_code)