-
Notifications
You must be signed in to change notification settings - Fork 276
/
test_plasma.py
103 lines (80 loc) · 2.51 KB
/
test_plasma.py
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
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
#!/usr/bin/env python3
# to run this script :
# $ nosetest3
# or
# $ python3 test_plasma.py
import os
import sys
from time import time
from contextlib import redirect_stdout
from nose.tools import assert_equal
from pathlib import Path
from io import StringIO
from plasma.lib.api import Api
from plasma.lib import GlobalContext
TESTS = Path('tests')
SYMBOLS = {
TESTS / 'server.bin': ["main", "connection_handler"],
TESTS / 'pendu.bin': ["_main", "___main", "__imp___cexit"],
TESTS / 'shellcode.bin': ["0x0"],
TESTS / 'entryloop1.bin': ["0x4041b0"],
}
OPTIONS = {
TESTS / 'shellcode.bin': ["--raw x86"],
TESTS / 'entryloop1.bin': ["--raw x64", "--rawbase 0x4041b0"],
}
def test_plasma():
for p in TESTS.glob('*.bin'):
for symbol in sorted(SYMBOLS.get(p, [None])):
yield plasma_file, str(p), symbol, OPTIONS.get(p, [])
def plasma_file(filename, symbol, options):
gctx = GlobalContext()
gctx.color = False
gctx.filename = filename
gctx.entry = symbol
gctx.quiet = True
gctx.db_path = None
for o in options:
if o == "--raw x86":
gctx.raw_type = "x86"
elif o == "--raw x64":
gctx.raw_type = "x64"
elif o.startswith("--rawbase"):
gctx.raw_base = int(o.split(" ")[1], 16)
if not gctx.load_file():
die()
gctx.api = Api(gctx, None)
sio = StringIO()
with redirect_stdout(sio):
o = gctx.get_addr_context(gctx.entry).decompile()
if o is not None:
o.print()
postfix = '{0}.rev'.format('' if symbol is None else '_' + symbol)
with open(filename.replace('.bin', postfix)) as f:
assert_equal(sio.getvalue(), f.read())
def color(text, c):
return "\x1b[38;5;" + str(c) + "m" + text + "\x1b[0m"
if __name__ == "__main__":
start = time()
passed = 0
nb = 0
failed = []
for plasma_file, path, symbol, options in test_plasma():
name = os.path.basename(path)
nb += 1
try:
plasma_file(path, symbol, options)
print(".", end="")
passed += 1
except AssertionError:
print(color("F", 1), end="")
failed.append(name)
except Exception as e:
print(color("E", 1), end="")
failed.append(name)
sys.stdout.flush()
elapsed = time()
elapsed = elapsed - start
print("\n%d/%d tests passed successfully in %fs" % (passed, nb, elapsed))
for p in failed:
print("failed:", p)