Repository navigation
Expand file tree
/
Copy path__init__.py
More file actions
81 lines (67 loc) 路 2.85 KB
/
Copy path__init__.py
File metadata and controls
81 lines (67 loc) 路 2.85 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
import argparse
import pickle
import sys
import base64
from .compiler import Compiler
from .helper import PickoraError
import ast
def main():
description = "A toy compiler that can convert Python scripts into pickle bytecode."
epilog = "Basic usage: `pickora samples/hello.py` or `pickora --code 'print(\"Hello, world!\")' --extended`"
parser = argparse.ArgumentParser(description=description, epilog=epilog)
parser.add_argument("source", nargs="?", help="source code file")
parser.add_argument("-c", "--code", help="source code string")
parser.add_argument("-p", "--protocol", type=int,
default=pickle.DEFAULT_PROTOCOL, help="pickle protocol")
parser.add_argument("-e", "--extended", action="store_true",
help="enable extended syntax (trigger find_class)")
parser.add_argument("-O", "--optimize", action="store_true",
help="optimize pickle bytecode (with pickletools.optimize)")
parser.add_argument("-o", "--output", help="output file")
parser.add_argument("-d", "--disassemble",
action="store_true", help="disassemble pickle bytecode")
parser.add_argument("-r", "--run", action="store_true",
help="run (load) pickle bytecode immediately")
parser.add_argument("-f", "--format",
choices=["repr", "raw", "hex", "base64", "none"], default="repr", help="output format, none means no output")
args = parser.parse_args()
if args.source and args.code:
parser.error("You can only specify one of source code file or string.")
if args.source:
with open(args.source, "r") as f:
source = f.read()
elif args.code:
source = args.code
else:
parser.error("You must specify source code file or string.")
compiler = Compiler(protocol=args.protocol,
optimize=args.optimize, extended=args.extended)
try:
code = compiler.compile(source, args.source)
except PickoraError as e:
print(e, file=sys.stderr)
sys.exit(1)
if args.disassemble:
import pickletools
try:
pickletools.dis(code)
except Exception as e:
print("[x] Disassemble error:", e, file=sys.stderr)
if args.output:
with open(args.output, "wb") as f:
f.write(code)
else:
if args.format == "repr":
print(repr(code))
elif args.format == "raw":
print(code.decode('latin1'), end="")
elif args.format == "hex":
print(code.hex())
elif args.format == "base64":
print(base64.b64encode(code).decode())
elif args.format == "none":
pass
if args.run:
print("[*] Running pickle bytecode...")
ret = pickle.loads(code)
print("[*] Return value:", repr(ret))