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