1import binascii
2
3import pyvex
4import archinfo
5import angr
6
7
8def test_strict_block_ends_cbz():
9    # ldr     r3, [sp, #4]
10    # cbz     r3, #0x8149
11    # mov.w   r2, #0x10000000
12    # ldr     r3, [pc, #0x38]
13    # str     r2, [r3]
14    # add     sp, #8
15    # pop     {r4, r5, r6, pc}
16
17    p = angr.load_shellcode(b'\x01\x9b\x1b\xb1O\xf0\x80R\x0eK\x1a`\x02\xb0p\xbd', 'arm')
18    assert p.factory.block(1, strict_block_end=False).instructions == 7
19    assert p.factory.block(1, strict_block_end=True).instructions == 2
20    p.factory.default_engine.default_strict_block_end = False
21    assert p.factory.block(1).instructions == 7
22    p.factory.default_engine.default_strict_block_end = True
23    assert p.factory.block(1).instructions == 2
24
25
26def test_strict_block_ends_with_size_x86():
27    # cmovnz  cx, dx
28    # pop     eax
29    # jecxz   short loc_4010D7  ; the first block should end here
30    # sub     edi, 2
31    # loop    loc_4010c9
32    # nop
33    # nop
34
35
36    b = b"\x66\x0f\x45\xca\x58\xe3\x07\x83\xef\x02\xe2\xf4\x90\x90"
37    p = angr.load_shellcode(b, 'x86', load_address=0x4010c9)
38
39    # jecxz
40    assert p.factory.block(0x4010c9, size=len(b), strict_block_end=False).instructions == 7
41    assert p.factory.block(0x4010c9, strict_block_end=True).instructions == 3
42    assert p.factory.block(0x4010c9, size=len(b), strict_block_end=True).instructions == 3
43
44    # loop
45    assert p.factory.block(0x4010d0, strict_block_end=False).instructions == 4
46    assert p.factory.block(0x4010d0, strict_block_end=True).instructions == 2
47    assert p.factory.block(0x4010d0, size=7, strict_block_end=True).instructions == 2
48
49
50def test_strict_block_ends_with_size_amd64():
51    # cmovnz  cx, dx
52    # pop     rax
53    # jrcxz   short loc_4010D7  ; the first block should end here
54    # sub     edi, 2
55    # loop    loc_4010c9
56    # nop
57    # nop
58
59
60    b = b"\x66\x0f\x45\xca\x58\xe3\x07\x83\xef\x02\xe2\xf4\x90\x90"
61    p = angr.load_shellcode(b, 'amd64', load_address=0x4010c9)
62
63    # jrcxz
64    assert p.factory.block(0x4010c9, size=len(b), strict_block_end=False).instructions == 7
65    assert p.factory.block(0x4010c9, strict_block_end=True).instructions == 3
66    assert p.factory.block(0x4010c9, size=len(b), strict_block_end=True).instructions == 3
67
68    # loop
69    assert p.factory.block(0x4010d0, strict_block_end=False).instructions == 4
70    assert p.factory.block(0x4010d0, strict_block_end=True).instructions == 2
71    assert p.factory.block(0x4010d0, size=7, strict_block_end=True).instructions == 2
72
73
74def test_no_cross_insn_boundary_opt_amd64():
75
76    # 0x4020f8:       sub     rsp, 8
77    # 0x4020fc:       mov     rax, qword ptr [rip + 0x221ef5]
78    # 0x402103:       test    rax, rax
79    # 0x402106:       je      0x40210d
80
81    b = binascii.unhexlify("4883ec08488b05f51e22004885c07405")
82    p = angr.load_shellcode(b, 'amd64', load_address=0x4020f8)
83
84    # No optimization
85    block = p.factory.block(0x4020f8, size=len(b), opt_level=0)
86    assert len(block.vex.statements) == 32
87    # Full level-1 optimization
88    block = p.factory.block(0x4020f8, size=len(b), opt_level=1, cross_insn_opt=True)
89    assert len(block.vex.statements) == 20
90    # Level-1 optimization within each instruction
91    block = p.factory.block(0x4020f8, size=len(b), opt_level=1, cross_insn_opt=False)
92    stmts = block.vex.statements
93    assert len(stmts) == 25
94    # 12 | ------ IMark(0x402103, 3, 0) ------
95    assert isinstance(stmts[12], pyvex.IRStmt.IMark)
96    assert stmts[12].addr == 0x402103
97    # 13 | t6 = GET:I64(rax)
98    assert isinstance(stmts[13], pyvex.IRStmt.WrTmp)
99    assert isinstance(stmts[13].data, pyvex.IRExpr.Get)
100    assert stmts[13].data.offset == archinfo.arch_from_id('amd64').registers['rax'][0]
101    # 14 | PUT(cc_op) = 0x0000000000000014
102    assert isinstance(stmts[14], pyvex.IRStmt.Put)
103    assert stmts[14].offset == archinfo.arch_from_id('amd64').registers['cc_op'][0]
104    assert isinstance(stmts[14].data, pyvex.IRExpr.Const)
105    assert stmts[14].data.con.value == 0x14
106    # 15 | PUT(cc_dep1) = t6
107    assert isinstance(stmts[15], pyvex.IRStmt.Put)
108    assert stmts[15].offset == archinfo.arch_from_id('amd64').registers['cc_dep1'][0]
109    # 16 | PUT(cc_dep2) = 0x0000000000000000
110    assert isinstance(stmts[16], pyvex.IRStmt.Put)
111    assert stmts[16].offset == archinfo.arch_from_id('amd64').registers['cc_dep2'][0]
112    assert isinstance(stmts[16].data, pyvex.IRExpr.Const)
113    assert stmts[16].data.con.value == 0
114    # 17 | PUT(rip) = 0x0000000000402106
115    assert isinstance(stmts[17], pyvex.IRStmt.Put)
116    assert stmts[17].offset == archinfo.arch_from_id('amd64').registers['rip'][0]
117    assert isinstance(stmts[17].data, pyvex.IRExpr.Const)
118    assert stmts[17].data.con.value == 0x402106
119    # 18 | ------ IMark(0x402106, 2, 0) ------
120    assert isinstance(stmts[18], pyvex.IRStmt.IMark)
121    assert stmts[18].addr == 0x402106
122
123
124if __name__ == '__main__':
125    test_strict_block_ends_cbz()
126    test_strict_block_ends_with_size_x86()
127    test_strict_block_ends_with_size_amd64()
128    test_no_cross_insn_boundary_opt_amd64()
129