1#!/usr/bin/env python
2# coding: utf-8
3"""Test Symbolic Engine."""
4
5import unittest
6import os
7
8from triton import (Instruction, ARCH, CPUSIZE, MemoryAccess, MODE,
9                    TritonContext, REG)
10
11
12def checkAstIntegrity(instruction):
13    """
14    This function check if all ASTs under an Instruction class are still
15    available.
16    """
17    try:
18        for se in instruction.getSymbolicExpressions():
19            str(se.getAst())
20
21        for x, y in instruction.getLoadAccess():
22            str(y)
23
24        for x, y in instruction.getStoreAccess():
25            str(y)
26
27        for x, y in instruction.getReadRegisters():
28            str(y)
29
30        for x, y in instruction.getWrittenRegisters():
31            str(y)
32
33        for x, y in instruction.getReadImmediates():
34            str(y)
35
36        return True
37
38    except:
39        return False
40
41
42class DefCamp2015(object):
43
44    """Test for DefCamp2015 challenge."""
45
46    def emulate(self, pc):
47        """
48        Emulate every opcode from pc.
49
50        * Process instruction until the end and search for constraint
51        resolution on cmp eax, 1 then self.Triton.set the new correct value and keep going.
52        """
53        astCtxt = self.Triton.getAstContext()
54        while pc:
55            # Fetch opcode
56            opcode = self.Triton.getConcreteMemoryAreaValue(pc, 16)
57
58            # Create the Triton instruction
59            instruction = Instruction()
60            instruction.setOpcode(opcode)
61            instruction.setAddress(pc)
62
63            # Process
64            self.Triton.processing(instruction)
65            self.assertTrue(checkAstIntegrity(instruction))
66
67            # 40078B: cmp eax, 1
68            # eax must be equal to 1 at each round.
69            if instruction.getAddress() == 0x40078B:
70                # Slice expressions
71                rax = self.Triton.getSymbolicRegister(self.Triton.registers.rax)
72                eax = astCtxt.extract(31, 0, rax.getAst())
73
74                # Define constraint
75                cstr = astCtxt.land([self.Triton.getPathPredicate(), astCtxt.equal(eax, astCtxt.bv(1, 32))])
76
77                model = self.Triton.getModel(cstr)
78                solution = str()
79                for k, v in list(sorted(model.items())):
80                    value = v.getValue()
81                    solution += chr(value)
82                    self.Triton.setConcreteVariableValue(self.Triton.getSymbolicVariable(k), value)
83
84            # Next
85            pc = self.Triton.getConcreteRegisterValue(self.Triton.registers.rip)
86        return solution
87
88    def load_binary(self, filename):
89        """Load in memory every opcode from an elf program."""
90        import lief
91        #lief.Logger.disable()
92        binary = lief.parse(filename)
93        phdrs  = binary.segments
94        for phdr in phdrs:
95            size   = phdr.physical_size
96            vaddr  = phdr.virtual_address
97            self.Triton.setConcreteMemoryAreaValue(vaddr, phdr.content)
98
99    def test_defcamp_2015(self):
100        """Load binary, self.Triton.setup environment and solve challenge with sym eval."""
101        # Load the binary
102        binary_file = os.path.join(os.path.dirname(__file__), "misc", "defcamp-2015-r100.bin")
103        self.load_binary(binary_file)
104
105        # Define a fake stack
106        self.Triton.setConcreteRegisterValue(self.Triton.registers.rbp, 0x7fffffff)
107        self.Triton.setConcreteRegisterValue(self.Triton.registers.rsp, 0x6fffffff)
108
109        # Define an user input
110        self.Triton.setConcreteRegisterValue(self.Triton.registers.rdi, 0x10000000)
111
112        # Symbolize user inputs (30 bytes)
113        for index in range(30):
114            self.Triton.symbolizeMemory(MemoryAccess(0x10000000+index, CPUSIZE.BYTE))
115
116        # Emulate from the verification function
117        solution = self.emulate(0x4006FD)
118        self.assertEqual(solution, 'Code_Talkers')
119
120
121class SeedCoverage(object):
122
123    """Test for seed coverage."""
124
125    def init_ctxt(self):
126        """Setup memory and register values."""
127        # Define the address of the serial pointer. The address of the serial
128        # pointer must be the same that the one hardcoded into the tarself.Triton.geted
129        # function. However, the serial pointer (here 0x900000) is arbitrary.
130        self.Triton.setConcreteMemoryValue(0x601040, 0x00)
131        self.Triton.setConcreteMemoryValue(0x601041, 0x00)
132        self.Triton.setConcreteMemoryValue(0x601042, 0x90)
133
134        # Define the serial context. We store the serial content located on our
135        # arbitrary # serial pointer (0x900000).
136        self.Triton.setConcreteMemoryValue(0x900000, 0x31)
137        self.Triton.setConcreteMemoryValue(0x900001, 0x3e)
138        self.Triton.setConcreteMemoryValue(0x900002, 0x3d)
139        self.Triton.setConcreteMemoryValue(0x900003, 0x26)
140        self.Triton.setConcreteMemoryValue(0x900004, 0x31)
141
142        # Point RDI on our buffer. The address of our buffer is arbitrary. We
143        # just need to point the RDI register on it as first argument of our
144        # tarself.Triton.geted function.
145        self.Triton.setConcreteRegisterValue(self.Triton.registers.rdi, 0x1000)
146
147        # Setup stack on an abitrary address.
148        self.Triton.setConcreteRegisterValue(self.Triton.registers.rsp, 0x7fffffff)
149        self.Triton.setConcreteRegisterValue(self.Triton.registers.rbp, 0x7fffffff)
150
151    def symbolize_inputs(self, seed):
152        """Add symboles in memory for seed."""
153        self.Triton.concretizeAllRegister()
154        self.Triton.concretizeAllMemory()
155        for address, value in list(seed.items()):
156            self.Triton.setConcreteMemoryValue(address, value)
157            self.Triton.symbolizeMemory(MemoryAccess(address, CPUSIZE.BYTE))
158            self.Triton.symbolizeMemory(MemoryAccess(address+1, CPUSIZE.BYTE))
159
160    def seed_emulate(self, ip):
161        """Emulate one run of the function with already self.Triton.setup memory."""
162        function = {
163            #   <serial> function
164            #   push    rbp
165            0x40056d: b"\x55",
166            #   mov     rbp,rsp
167            0x40056e: b"\x48\x89\xe5",
168            #   mov     QWORD PTR [rbp-0x18],rdi
169            0x400571: b"\x48\x89\x7d\xe8",
170            #   mov     DWORD PTR [rbp-0x4],0x0
171            0x400575: b"\xc7\x45\xfc\x00\x00\x00\x00",
172            #   jmp     4005bd <check+0x50>
173            0x40057c: b"\xeb\x3f",
174            #   mov     eax,DWORD PTR [rbp-0x4]
175            0x40057e: b"\x8b\x45\xfc",
176            #   movsxd  rdx,eax
177            0x400581: b"\x48\x63\xd0",
178            #   mov     rax,QWORD PTR [rbp-0x18]
179            0x400584: b"\x48\x8b\x45\xe8",
180            #   add     rax,rdx
181            0x400588: b"\x48\x01\xd0",
182            #   movzx   eax,BYTE PTR [rax]
183            0x40058b: b"\x0f\xb6\x00",
184            #   movsx   eax,al
185            0x40058e: b"\x0f\xbe\xc0",
186            #   sub     eax,0x1
187            0x400591: b"\x83\xe8\x01",
188            #   xor     eax,0x55
189            0x400594: b"\x83\xf0\x55",
190            #   mov     ecx,eax
191            0x400597: b"\x89\xc1",
192            #   mov     rdx,QWORD PTR [rip+0x200aa0]        # 601040 <serial>
193            0x400599: b"\x48\x8b\x15\xa0\x0a\x20\x00",
194            #   mov     eax,DWORD PTR [rbp-0x4]
195            0x4005a0: b"\x8b\x45\xfc",
196            #   cdqe
197            0x4005a3: b"\x48\x98",
198            #   add     rax,rdx
199            0x4005a5: b"\x48\x01\xd0",
200            #   movzx   eax,BYTE PTR [rax]
201            0x4005a8: b"\x0f\xb6\x00",
202            #   movsx   eax,al
203            0x4005ab: b"\x0f\xbe\xc0",
204            #   cmp     ecx,eax
205            0x4005ae: b"\x39\xc1",
206            #   je      4005b9 <check+0x4c>
207            0x4005b0: b"\x74\x07",
208            #   mov     eax,0x1
209            0x4005b2: b"\xb8\x01\x00\x00\x00",
210            #   jmp     4005c8 <check+0x5b>
211            0x4005b7: b"\xeb\x0f",
212            #   add     DWORD PTR [rbp-0x4],0x1
213            0x4005b9: b"\x83\x45\xfc\x01",
214            #   cmp     DWORD PTR [rbp-0x4],0x4
215            0x4005bd: b"\x83\x7d\xfc\x04",
216            #   jle     40057e <check+0x11>
217            0x4005c1: b"\x7e\xbb",
218            #   mov     eax,0x0
219            0x4005c3: b"\xb8\x00\x00\x00\x00",
220            #   pop     rbp
221            0x4005c8: b"\x5d",
222            #   ret
223            0x4005c9: b"\xc3",
224        }
225        while ip in function:
226            # Build an instruction
227            inst = Instruction()
228
229            # Setup opcode
230            inst.setOpcode(function[ip])
231
232            # Setup Address
233            inst.setAddress(ip)
234
235            # Process everything
236            self.Triton.processing(inst)
237            self.assertTrue(checkAstIntegrity(inst))
238
239            # Next instruction
240            ip = self.Triton.getRegisterAst(self.Triton.registers.rip).evaluate()
241
242    def new_inputs(self):
243        """Look for another branching using current constraints found."""
244        astCtxt = self.Triton.getAstContext()
245
246        # Set of new inputs
247        inputs = list()
248
249        # Get path constraints from the last execution
250        pco = self.Triton.getPathConstraints()
251
252        # We start with any input. T (Top)
253        previousConstraints = astCtxt.equal(astCtxt.bvtrue(), astCtxt.bvtrue())
254
255        # Go through the path constraints
256        for pc in pco:
257            # If there is a condition
258            if pc.isMultipleBranches():
259                # Get all branches
260                branches = pc.getBranchConstraints()
261                for branch in branches:
262                    # Get the constraint of the branch which has been not taken
263                    if branch['isTaken'] == False:
264                        # Ask for a model
265                        models = self.Triton.getModel(astCtxt.land([previousConstraints, branch['constraint']]))
266                        seed = dict()
267                        for k, v in list(models.items()):
268                            # Get the symbolic variable assigned to the model
269                            symVar = self.Triton.getSymbolicVariable(k)
270                            # Save the new input as seed.
271                            seed.update({symVar.getOrigin(): v.getValue()})
272                        if seed:
273                            inputs.append(seed)
274
275            # Update the previous constraints with true branch to keep a good
276            # path.
277            previousConstraints = astCtxt.land([previousConstraints, pc.getTakenPredicate()])
278
279        # Clear the path constraints to be clean at the next execution.
280        self.Triton.clearPathConstraints()
281
282        return inputs
283
284    def test_seed_coverage(self):
285        """Found every seed so that every opcode will be use at least once."""
286        # Define entry point
287        ENTRY = 0x40056d
288
289        # We start the execution with a random value located at 0x1000.
290        lastInput = list()
291        worklist = list([{0x1000: 1}])
292
293        while worklist:
294            # Take the first seed
295            seed = worklist[0]
296
297            # Symbolize inputs
298            self.symbolize_inputs(seed)
299
300            # Init context memory
301            self.init_ctxt()
302
303            # Emulate
304            self.seed_emulate(ENTRY)
305
306            lastInput += [dict(seed)]
307            del worklist[0]
308
309            newInputs = self.new_inputs()
310            for inputs in newInputs:
311                if inputs not in lastInput and inputs not in worklist:
312                    worklist += [dict(inputs)]
313
314        self.assertIn({4096: 101,
315                       4097: 108,
316                       4098: 105,
317                       4099: 116,
318                       4100: 101}, lastInput)
319
320
321class Emu1(object):
322
323    """Test for emu_1.dump."""
324
325    def test_emulate(self, concretize=False):
326        """Run a dumped simulation and check output registers."""
327        # Get dumped data
328        dump = os.path.join(os.path.dirname(__file__), "misc", "emu_1.dump")
329        with open(dump) as f:
330            regs, mems = eval(f.read())
331
332        # Load memory
333        for mem in mems:
334            start = mem['start']
335            if mem['memory'] is not None:
336                self.Triton.setConcreteMemoryAreaValue(start, bytearray(mem['memory']))
337
338        # self.Triton.setup registers
339        for reg_name in ("rax", "rbx", "rcx", "rdx", "rdi", "rsi", "rbp",
340                         "rsp", "rip", "r8", "r9", "r10", "r11", "r12", "r13",
341                         "r14", "eflags", "xmm0", "xmm1", "xmm2", "xmm3",
342                         "xmm4", "xmm5", "xmm6", "xmm7", "xmm8", "xmm9",
343                         "xmm10", "xmm11", "xmm12", "xmm13", "xmm14", "xmm15"):
344            self.Triton.setConcreteRegisterValue(self.Triton.getRegister(getattr(REG.X86_64, reg_name.upper())), regs[reg_name])
345
346        # run the code
347        pc = self.Triton.getConcreteRegisterValue(self.Triton.registers.rip)
348        while pc != 0x409A18:
349            opcode = self.Triton.getConcreteMemoryAreaValue(pc, 20)
350
351            instruction = Instruction()
352            instruction.setOpcode(opcode)
353            instruction.setAddress(pc)
354
355            # Check if triton doesn't supports this instruction
356            self.assertTrue(self.Triton.processing(instruction))
357            self.assertTrue(checkAstIntegrity(instruction))
358
359            pc = self.Triton.getConcreteRegisterValue(self.Triton.registers.rip)
360
361            if concretize:
362                self.Triton.concretizeAllMemory()
363                self.Triton.concretizeAllRegister()
364
365        rax = self.Triton.getConcreteRegisterValue(self.Triton.registers.rax)
366        rbx = self.Triton.getConcreteRegisterValue(self.Triton.registers.rbx)
367        rcx = self.Triton.getConcreteRegisterValue(self.Triton.registers.rcx)
368        rdx = self.Triton.getConcreteRegisterValue(self.Triton.registers.rdx)
369        rsi = self.Triton.getConcreteRegisterValue(self.Triton.registers.rsi)
370
371        self.assertEqual(rax, 0)
372        self.assertEqual(rbx, 0)
373        self.assertEqual(rcx, 0)
374        self.assertEqual(rdx, 0x4d2)
375        self.assertEqual(rsi, 0x3669000000000000)
376
377
378class BaseTestSimulation(DefCamp2015, SeedCoverage, Emu1):
379
380    """BaseClass to test the simulation."""
381
382
383class TestSymbolicEngineNoOptim(BaseTestSimulation, unittest.TestCase):
384
385    """Testing the symbolic emulation engine without optimization."""
386
387    def setUp(self):
388        """Define the arch."""
389        self.Triton = TritonContext()
390        self.Triton.setArchitecture(ARCH.X86_64)
391        super(TestSymbolicEngineNoOptim, self).setUp()
392
393
394class TestSymbolicEngineAligned(BaseTestSimulation, unittest.TestCase):
395
396    """Testing the symbolic emulation engine with ALIGNED_MEMORY."""
397
398    def setUp(self):
399        """Define the arch and modes."""
400        self.Triton = TritonContext()
401        self.Triton.setArchitecture(ARCH.X86_64)
402        self.Triton.setMode(MODE.ALIGNED_MEMORY, True)
403        super(TestSymbolicEngineAligned, self).setUp()
404
405
406class TestSymbolicEngineAlignedAndTaintPtr(BaseTestSimulation, unittest.TestCase):
407
408    """Testing the symbolic emulation engine with ALIGNED_MEMORY."""
409
410    def setUp(self):
411        """Define the arch and modes."""
412        self.Triton = TritonContext()
413        self.Triton.setArchitecture(ARCH.X86_64)
414        self.Triton.setMode(MODE.ALIGNED_MEMORY, True)
415        self.Triton.setMode(MODE.TAINT_THROUGH_POINTERS, True)
416        super(TestSymbolicEngineAlignedAndTaintPtr, self).setUp()
417
418
419class TestSymbolicEngineOnlySymbolized(BaseTestSimulation, unittest.TestCase):
420
421    """Testing the symbolic emulation engine with ONLY_ON_SYMBOLIZED."""
422
423    def setUp(self):
424        """Define the arch and modes."""
425        self.Triton = TritonContext()
426        self.Triton.setArchitecture(ARCH.X86_64)
427        self.Triton.setMode(MODE.ONLY_ON_SYMBOLIZED, True)
428        super(TestSymbolicEngineOnlySymbolized, self).setUp()
429
430
431class TestSymbolicEngineAlignedOnlySymbolized(BaseTestSimulation, unittest.TestCase):
432
433    """Testing the symbolic emulation engine with ALIGNED_MEMORY and ONLY_ON_SYMBOLIZED."""
434
435    def setUp(self):
436        """Define the arch and modes."""
437        self.Triton = TritonContext()
438        self.Triton.setArchitecture(ARCH.X86_64)
439        self.Triton.setMode(MODE.ALIGNED_MEMORY, True)
440        self.Triton.setMode(MODE.ONLY_ON_SYMBOLIZED, True)
441        super(TestSymbolicEngineAlignedOnlySymbolized, self).setUp()
442
443
444class TestSymbolicEngineDisable(BaseTestSimulation, unittest.TestCase):
445
446    """Testing the emulation with the symbolic engine disabled."""
447
448    def setUp(self):
449        """Define the arch and modes."""
450        self.Triton = TritonContext()
451        self.Triton.setArchitecture(ARCH.X86_64)
452        self.Triton.enableSymbolicEngine(False)
453        super(TestSymbolicEngineDisable, self).setUp()
454
455    @unittest.skip("Not possible")
456    def test_seed_coverage(self):
457        pass
458
459    @unittest.skip("Not possible")
460    def test_defcamp_2015(self):
461        pass
462
463
464class TestSymbolicEngineSymOpti(BaseTestSimulation, unittest.TestCase):
465
466    """Testing the symbolic emulation engine without optimization."""
467
468    def setUp(self):
469        """Define the arch."""
470        self.Triton = TritonContext()
471        self.Triton.setArchitecture(ARCH.X86_64)
472        self.Triton.setMode(MODE.AST_OPTIMIZATIONS, True)
473        super(TestSymbolicEngineSymOpti, self).setUp()
474
475
476class TestSymbolicEngineAlignedSymOpti(BaseTestSimulation, unittest.TestCase):
477
478    """Testing the symbolic emulation engine with ALIGNED_MEMORY."""
479
480    def setUp(self):
481        """Define the arch and modes."""
482        self.Triton = TritonContext()
483        self.Triton.setArchitecture(ARCH.X86_64)
484        self.Triton.setMode(MODE.ALIGNED_MEMORY, True)
485        self.Triton.setMode(MODE.AST_OPTIMIZATIONS, True)
486        super(TestSymbolicEngineAlignedSymOpti, self).setUp()
487
488
489class TestSymbolicEngineAlignedOnlySymbolizedSymOpti(BaseTestSimulation, unittest.TestCase):
490
491    """Testing the symbolic emulation engine with ALIGNED_MEMORY, ONLY_ON_SYMBOLIZED and AST_OPTIMIZATIONS."""
492
493    def setUp(self):
494        """Define the arch and modes."""
495        self.Triton = TritonContext()
496        self.Triton.setArchitecture(ARCH.X86_64)
497        self.Triton.setMode(MODE.ALIGNED_MEMORY, True)
498        self.Triton.setMode(MODE.ONLY_ON_SYMBOLIZED, True)
499        self.Triton.setMode(MODE.AST_OPTIMIZATIONS, True)
500        super(TestSymbolicEngineAlignedOnlySymbolizedSymOpti, self).setUp()
501
502
503class TestSymbolicEngineConstantFolding(BaseTestSimulation, unittest.TestCase):
504
505    """Testing the symbolic emulation engine with CONSTANT_FOLDING."""
506
507    def setUp(self):
508        """Define the arch and modes."""
509        self.Triton = TritonContext()
510        self.Triton.setArchitecture(ARCH.X86_64)
511        self.Triton.setMode(MODE.CONSTANT_FOLDING, True)
512        super(TestSymbolicEngineConstantFolding, self).setUp()
513