follow_trace.py
#!/usr/bin/env python
'''
A simple trace following execution driver script. Only supports passing symbolic arguments via argv.
'''
import sys
import time
import argparse
import itertools
from manticore import Manticore, issymbolic
from manticore.core.plugin import ExtendedTracer, Plugin
def _partition(pred, iterable):
t1, t2 = itertools.tee(iterable)
return list(itertools.filterfalse(pred, t1)), list(filter(pred, t2))
class TraceReceiver(Plugin):
def __init__(self, tracer):
self._trace = None
self._tracer = tracer
super().__init__()
@property
def trace(self):
return self._trace
def will_generate_testcase_callback(self, state, test_id, msg):
self._trace = state.context[self._tracer.context_key]
instructions, writes = _partition(lambda x: x['type'] == 'regs', self._trace)
total = len(self._trace)
print('Recorded concrete trace: {}/{} instructions, {}/{} writes'.format(
len(instructions), total, len(writes), total))
class Follower(Plugin):
def __init__(self, trace):
self.index = 0
self.trace = trace
self.last_instruction = None
self.symbolic_ranges = []
self.active = True
super().__init__()
def add_symbolic_range(self, pc_start, pc_end):
self.symbolic_ranges.append((pc_start,pc_end))
def get_next(self, type):
event = self.trace[self.index]
assert event['type'] == type
self.index += 1
return event
def did_write_memory_callback(self, state, where, value, size):
if not self.active:
return
write = self.get_next('mem_write')
if not issymbolic(value):
return
assert write['where'] == where and write['size'] == size
state.constrain(value == write['value'])
def did_execute_instruction_callback(self, state, last_pc, pc, insn):
if not self.active:
return
event = self.get_next('regs')
self.last_instruction = event['values']
if issymbolic(pc):
state.constrain(state.cpu.RIP == self.last_instruction['RIP'])
else:
for start, stop in self.symbolic_ranges:
if start <= pc <= stop:
self.active = False
def main():
parser = argparse.ArgumentParser(description='Follow a concrete trace')
parser.add_argument('-f', '--explore_from', help='Value of PC from which to explore symbolically', type=str)
parser.add_argument('-t', '--explore_to', type=str, default=sys.maxsize,
help="Value of PC until which to explore symbolically. (Probably don't want this set)")
parser.add_argument('--verbose', '-v', action='count', default=0, help='Increase verbosity')
parser.add_argument('cmd', type=str, nargs='+',
help='Program and arguments. Use "--" to separate script arguments from target arguments')
args = parser.parse_args(sys.argv[1:])
range = None
if args.explore_from:
range = (args.explore_from, args.explore_to)
# Create a concrete Manticore and record it
m1 = Manticore.linux(args.cmd[0], args.cmd[1:])
t = ExtendedTracer()
r = TraceReceiver(t)
m1.verbosity(args.verbose)
m1.register_plugin(t)
m1.register_plugin(r)
m1.run(procs=1)
time.sleep(3)
# Create a symbolic Manticore and follow last trace
symbolic_args = ['+'*len(arg) for arg in args.cmd[1:]]
m2 = Manticore.linux(args.cmd[0], symbolic_args)
f = Follower(r.trace)
if range:
f.add_symbolic_range(*range)
m2.verbosity(args.verbose)
m2.register_plugin(f)
m2.run()
if __name__=='__main__':
main()