Skip to content

Commit bcf631c

Browse files
shaymargolisshaymargolis
committed
tests: Convert test to use compiler arch
1 parent 4338aa6 commit bcf631c

4 files changed

Lines changed: 49 additions & 41 deletions

File tree

‎tests/test_goto.py‎

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
import pytest
22

3-
from unicorn import Uc, UC_ARCH_MIPS, UC_MODE_32, UC_MODE_BIG_ENDIAN
43
from unicorn.mips_const import UC_MIPS_REG_PC
54

65
from shellblocks.shellcode_step import ShellcodeStep
@@ -16,7 +15,7 @@
1615
(0xbc000000, 0xbcf00010),
1716
(0x91000000, 0x91000118),
1817
])
19-
def test_goto_sanity(temp_dir_path, goto_page_and_address):
18+
def test_goto_sanity(get_mu, temp_dir_path, compiler_arch, goto_page_and_address):
2019
# Generate shellcode
2120
# ------------------
2221
shellcode_address = 0xbfc00000
@@ -31,13 +30,13 @@ def test_goto_sanity(temp_dir_path, goto_page_and_address):
3130
0x1000
3231
)
3332

34-
out_file = step.generate(temp_dir_path / step.nickname)
33+
out_file = step.generate(temp_dir_path / step.nickname, compiler_arch)
3534
shellcode = out_file.read_bytes()
3635

3736
# Try to run shellcode
3837
# --------------------
3938

40-
mu = Uc(UC_ARCH_MIPS, UC_MODE_32 | UC_MODE_BIG_ENDIAN)
39+
mu = get_mu()
4140
mu.mem_map(shellcode_address, 0x2000)
4241

4342
# write machine code to be emulated to memory
@@ -54,7 +53,7 @@ def test_goto_sanity(temp_dir_path, goto_page_and_address):
5453
(0xbcf00010),
5554
(0x91000118),
5655
])
57-
def test_goto_is_pic(temp_dir_path, shellcode_run_addr):
56+
def test_goto_is_pic(get_mu, temp_dir_path, compiler_arch, shellcode_run_addr):
5857
# Generate shellcode
5958
# ------------------
6059
shellcode_address = 0xbfc00000
@@ -71,13 +70,13 @@ def test_goto_is_pic(temp_dir_path, shellcode_run_addr):
7170
0x1000
7271
)
7372

74-
out_file = step.generate(temp_dir_path / step.nickname)
73+
out_file = step.generate(temp_dir_path / step.nickname, compiler_arch)
7574
shellcode = out_file.read_bytes()
7675

7776
# Try to run shellcode
7877
# --------------------
7978

80-
mu = Uc(UC_ARCH_MIPS, UC_MODE_32 | UC_MODE_BIG_ENDIAN)
79+
mu = get_mu()
8180
mu.mem_map(shellcode_run_sector, 0x2000)
8281

8382
# write machine code to be emulated to memory

‎tests/test_jump_hook.py‎

Lines changed: 12 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,5 @@
11
import pytest
22

3-
from unicorn import Uc, UC_ARCH_MIPS, UC_MODE_32, UC_MODE_BIG_ENDIAN
4-
53
from shellblocks.shellcode_step import ShellcodeStep
64
from shellblocks.primitives.jump_hook import ShellcodePrimitiveJumpHook
75

@@ -27,7 +25,15 @@
2725
0xbcf00070,
2826
0x910f0218,
2927
])
30-
def test_jump_hook_sanity(temp_dir_path, shellcode_run_addr, jump_hook_location, jump_hook_goto):
28+
def test_jump_hook_sanity(
29+
get_mu,
30+
arch_helper,
31+
temp_dir_path,
32+
compiler_arch,
33+
shellcode_run_addr,
34+
jump_hook_location,
35+
jump_hook_goto
36+
):
3137
# Generate shellcode
3238
# ------------------
3339
shellcode_address = 0xbfc00000
@@ -47,23 +53,15 @@ def test_jump_hook_sanity(temp_dir_path, shellcode_run_addr, jump_hook_location,
4753
0x1000
4854
)
4955

50-
out_file = step.generate(temp_dir_path / step.nickname)
56+
out_file = step.generate(temp_dir_path / step.nickname, compiler_arch)
5157
shellcode = out_file.read_bytes()
5258

53-
EXPECTED_HOOK = b"".join(map(
54-
lambda x: x.to_bytes(4, 'big'),
55-
[
56-
0x3c020000 + (jump_hook_goto >> 16),
57-
0x24420000 + (jump_hook_goto & 0xffff),
58-
0x00400008,
59-
0x00000000,
60-
]
61-
))
59+
EXPECTED_HOOK = arch_helper.get_jump_hook_bytes(jump_hook_goto)
6260

6361
# Try to run shellcode
6462
# --------------------
6563

66-
mu = Uc(UC_ARCH_MIPS, UC_MODE_32 | UC_MODE_BIG_ENDIAN)
64+
mu = get_mu()
6765
mu.mem_map(shellcode_run_sector, 0x2000)
6866
mu.mem_map(jump_hook_sector, 0x2000)
6967

‎tests/test_memcpy.py‎

Lines changed: 14 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,5 @@
11
import pytest
22

3-
from unicorn import Uc, UC_ARCH_MIPS, UC_MODE_32, UC_MODE_BIG_ENDIAN
4-
53
from shellblocks.shellcode_step import ShellcodeStep
64
from shellblocks.primitives.memcpy import ShellcodePrimitiveMemcpy
75

@@ -11,6 +9,7 @@
119

1210
class UcMemcpyHelper:
1311
def __init__(self,
12+
get_mu,
1413
shellcode_address,
1514
first_copy_addr,
1615
second_copy_addr):
@@ -22,7 +21,7 @@ def __init__(self,
2221
self.second_copy_addr = second_copy_addr
2322
self.second_sector = int(self.second_copy_addr/SECTOR_SIZE) * SECTOR_SIZE
2423

25-
self.mu = Uc(UC_ARCH_MIPS, UC_MODE_32 | UC_MODE_BIG_ENDIAN)
24+
self.mu = get_mu()
2625
self.mu.mem_map(self.shellcode_address, SECTOR_SIZE)
2726
self.mu.mem_map(self.first_sector, SECTOR_SIZE)
2827
self.mu.mem_map(self.second_sector, SECTOR_SIZE)
@@ -32,15 +31,16 @@ def write_shellcode(self, shellcode):
3231

3332

3433
@pytest.fixture(scope='function')
35-
def default_memcpy_helper():
34+
def default_memcpy_helper(get_mu):
3635
return UcMemcpyHelper(
36+
get_mu,
3737
0xbfc00000,
3838
0x81000010,
3939
0x82000010
4040
)
4141

4242

43-
def memcpy_get_shellcode(temp_dir_path, memcpy_helper, copy_len):
43+
def memcpy_get_shellcode(temp_dir_path, compiler_arch, memcpy_helper, copy_len):
4444
helper = memcpy_helper
4545

4646
step = ShellcodeStep(
@@ -57,7 +57,7 @@ def memcpy_get_shellcode(temp_dir_path, memcpy_helper, copy_len):
5757
0x1000
5858
)
5959

60-
out_file = step.generate(temp_dir_path / step.nickname)
60+
out_file = step.generate(temp_dir_path / step.nickname, compiler_arch)
6161
shellcode = out_file.read_bytes()
6262

6363
return shellcode
@@ -71,9 +71,9 @@ def memcpy_get_shellcode(temp_dir_path, memcpy_helper, copy_len):
7171
3,
7272
100
7373
])
74-
def test_memcpy_sanity(temp_dir_path, default_memcpy_helper, copy_len):
74+
def test_memcpy_sanity(temp_dir_path, compiler_arch, default_memcpy_helper, copy_len):
7575
helper = default_memcpy_helper
76-
shellcode = memcpy_get_shellcode(temp_dir_path, helper, copy_len)
76+
shellcode = memcpy_get_shellcode(temp_dir_path, compiler_arch, helper, copy_len)
7777

7878
# Try to run shellcode
7979
# --------------------
@@ -96,10 +96,10 @@ def test_memcpy_sanity(temp_dir_path, default_memcpy_helper, copy_len):
9696
(0xbcf00010),
9797
(0x91000118),
9898
])
99-
def test_memcpy_is_pic(temp_dir_path, shellcode_run_addr, default_memcpy_helper):
99+
def test_memcpy_is_pic(temp_dir_path, compiler_arch, shellcode_run_addr, default_memcpy_helper):
100100
copy_len = 0x1000
101101
helper = default_memcpy_helper
102-
shellcode = memcpy_get_shellcode(temp_dir_path, helper, copy_len)
102+
shellcode = memcpy_get_shellcode(temp_dir_path, compiler_arch, helper, copy_len)
103103

104104
shellcode_run_sector = int(shellcode_run_addr/SECTOR_SIZE) * SECTOR_SIZE
105105
helper.mu.mem_map(shellcode_run_sector, SECTOR_SIZE)
@@ -127,10 +127,10 @@ def test_memcpy_is_pic(temp_dir_path, shellcode_run_addr, default_memcpy_helper)
127127
3,
128128
100
129129
])
130-
def test_memcpy_short(temp_dir_path, default_memcpy_helper, copy_len):
130+
def test_memcpy_short(temp_dir_path, compiler_arch, default_memcpy_helper, copy_len):
131131

132132
helper = default_memcpy_helper
133-
shellcode = memcpy_get_shellcode(temp_dir_path, helper, copy_len)
133+
shellcode = memcpy_get_shellcode(temp_dir_path, compiler_arch, helper, copy_len)
134134

135135
# Try to run shellcode
136136
# --------------------
@@ -155,11 +155,11 @@ def test_memcpy_short(temp_dir_path, default_memcpy_helper, copy_len):
155155
4,
156156
100
157157
])
158-
def test_memcpy_two_halves(temp_dir_path, default_memcpy_helper, copy_len):
158+
def test_memcpy_two_halves(temp_dir_path, compiler_arch, default_memcpy_helper, copy_len):
159159
half_copy_len = int(copy_len/2)
160160

161161
helper = default_memcpy_helper
162-
shellcode = memcpy_get_shellcode(temp_dir_path, helper, copy_len)
162+
shellcode = memcpy_get_shellcode(temp_dir_path, compiler_arch, helper, copy_len)
163163

164164
# Try to run shellcode
165165
# --------------------

‎tests/test_print.py‎

Lines changed: 17 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
import pytest
22

3-
from unicorn import Uc, UC_ARCH_MIPS, UC_MODE_32, UC_MODE_BIG_ENDIAN
43
from unicorn.mips_const import UC_MIPS_REG_PC, UC_MIPS_REG_29, UC_MIPS_REG_4
54

65
from shellblocks.shellcode_step import ShellcodeStep
@@ -40,7 +39,7 @@ def string_to_print(request):
4039

4140

4241
@pytest.fixture()
43-
def print_shellcode(temp_dir_path, print_function_addr, string_to_print):
42+
def print_shellcode(compiler_arch, temp_dir_path, print_function_addr, string_to_print):
4443
# Generate shellcode
4544
# ------------------
4645
shellcode_address = 0xbfc00000
@@ -54,13 +53,15 @@ def print_shellcode(temp_dir_path, print_function_addr, string_to_print):
5453
0x1000
5554
)
5655

57-
out_file = step.generate(temp_dir_path / step.nickname)
56+
out_file = step.generate(temp_dir_path / step.nickname, compiler_arch)
5857
shellcode = out_file.read_bytes()
5958

6059
return shellcode, shellcode_address
6160

6261

6362
def get_print_mu(
63+
get_mu,
64+
arch_helper,
6465
print_shellcode,
6566
shellcode_run_addr,
6667
print_function_addr,
@@ -74,7 +75,7 @@ def get_print_mu(
7475
# Try to run shellcode
7576
# --------------------
7677

77-
mu = Uc(UC_ARCH_MIPS, UC_MODE_32 | UC_MODE_BIG_ENDIAN)
78+
mu = get_mu()
7879

7980
# Print function uses the stack pointer
8081
mu.reg_write(UC_MIPS_REG_29, stack_address + 0x2000)
@@ -85,16 +86,18 @@ def get_print_mu(
8586

8687
# write machine code to be emulated to memory
8788
mu.mem_write(shellcode_run_addr, shellcode)
88-
mu.mem_write(print_function_addr, (0x03e00008).to_bytes(4, 'big')) # "jr $ra" in MIPS
89+
mu.mem_write(print_function_addr, arch_helper.get_ret_bytes())
8990

9091
return mu
9192

9293

9394
def test_print_reaches_print_function(
94-
print_shellcode, string_to_print, print_function_addr, stack_address
95+
get_mu, arch_helper, print_shellcode, string_to_print, print_function_addr, stack_address
9596
):
9697
shellcode, shellcode_address = print_shellcode
9798
print_mu = get_print_mu(
99+
get_mu,
100+
arch_helper,
98101
print_shellcode,
99102
shellcode_address,
100103
print_function_addr,
@@ -115,12 +118,16 @@ def test_print_reaches_print_function(
115118

116119

117120
def test_print_reaches_end(
121+
get_mu,
122+
arch_helper,
118123
print_shellcode,
119124
print_function_addr,
120125
string_to_print,
121126
stack_address):
122127
shellcode, shellcode_address = print_shellcode
123128
print_mu = get_print_mu(
129+
get_mu,
130+
arch_helper,
124131
print_shellcode,
125132
shellcode_address,
126133
print_function_addr,
@@ -140,6 +147,8 @@ def test_print_reaches_end(
140147
(0x92000118),
141148
])
142149
def test_print_is_pic(
150+
get_mu,
151+
arch_helper,
143152
shellcode_run_addr,
144153
print_shellcode,
145154
print_function_addr,
@@ -148,6 +157,8 @@ def test_print_is_pic(
148157
):
149158
shellcode, shellcode_address = print_shellcode
150159
print_mu = get_print_mu(
160+
get_mu,
161+
arch_helper,
151162
print_shellcode,
152163
shellcode_run_addr,
153164
print_function_addr,

0 commit comments

Comments
 (0)