<s> import bisect <EOL> import logging <EOL> import networkx <EOL> from Queue import Queue <EOL> from pydot import Dot <EOL> from pydot import Edge <EOL> from pydot import Node <EOL> from barf . arch import ARCH_ARM <EOL> from barf . arch import ARCH_X86 <EOL> from barf . core . reil import DualInstruction <EOL> from barf . core . reil import ReilMnemonic <EOL> from barf . core . reil import ReilImmediateOperand <EOL> from barf . arch . arm . armdisassembler import InvalidDisassemblerData , CapstoneOperandNotSupported <EOL> BARF_DISASM_LINEAR = <NUM_LIT:0> <EOL> BARF_DISASM_RECURSIVE = <NUM_LIT:1> <EOL> logger = logging . getLogger ( __name__ ) <EOL> class BasicBlock ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self ) : <EOL> self . _instrs = [ ] <EOL> self . _address = None <EOL> self . _taken_branch = None <EOL> self . _not_taken_branch = None <EOL> self . _direct_branch = None <EOL> @ property <EOL> def instrs ( self ) : <EOL> """<STR_LIT>""" <EOL> return self . _instrs <EOL> @ property <EOL> def address ( self ) : <EOL> """<STR_LIT>""" <EOL> if self . _instrs == [ ] : <EOL> return None <EOL> return self . _instrs [ <NUM_LIT:0> ] . address <EOL> @ property <EOL> def start_address ( self ) : <EOL> """<STR_LIT>""" <EOL> if self . _instrs is [ ] : <EOL> return None <EOL> return self . _instrs [ <NUM_LIT:0> ] . address <EOL> @ property <EOL> def end_address ( self ) : <EOL> """<STR_LIT>""" <EOL> if self . _instrs is [ ] : <EOL> return None <EOL> return self . _instrs [ - <NUM_LIT:1> ] . address + self . _instrs [ - <NUM_LIT:1> ] . asm_instr . size - <NUM_LIT:1> <EOL> @ property <EOL> def size ( self ) : <EOL> """<STR_LIT>""" <EOL> if self . _instrs is [ ] : <EOL> return None <EOL> return sum ( [ dinstr . asm_instr . size for dinstr in self . _instrs ] ) <EOL> @ property <EOL> def taken_branch ( self ) : <EOL> """<STR_LIT>""" <EOL> return self . _taken_branch <EOL> @ taken_branch . setter <EOL> def taken_branch ( self , value ) : <EOL> """<STR_LIT>""" <EOL> self . _taken_branch = value <EOL> @ property <EOL> def not_taken_branch ( self ) : <EOL> """<STR_LIT>""" <EOL> return self . _not_taken_branch <EOL> @ not_taken_branch . setter <EOL> def not_taken_branch ( self , value ) : <EOL> """<STR_LIT>""" <EOL> self . _not_taken_branch = value <EOL> @ property <EOL> def direct_branch ( self ) : <EOL> """<STR_LIT>""" <EOL> return self . _direct_branch <EOL> @ direct_branch . setter <EOL> def direct_branch ( self , value ) : <EOL> """<STR_LIT>""" <EOL> self . _direct_branch = value <EOL> @ property <EOL> def branches ( self ) : <EOL> """<STR_LIT>""" <EOL> branches = [ ] <EOL> if self . _taken_branch : <EOL> branches += [ ( self . _taken_branch , '<STR_LIT>' ) ] <EOL> if self . _not_taken_branch : <EOL> branches += [ ( self . _not_taken_branch , '<STR_LIT>' ) ] <EOL> if self . _direct_branch : <EOL> branches += [ ( self . _direct_branch , '<STR_LIT>' ) ] <EOL> return branches <EOL> def contains ( self , address ) : <EOL> """<STR_LIT>""" <EOL> return self . address <= address and address <= self . end_address <EOL> def empty ( self ) : <EOL> """<STR_LIT>""" <EOL> return len ( self . _instrs ) == <NUM_LIT:0> <EOL> def __str__ ( self ) : <EOL> lines = [ "<STR_LIT>" . format ( self . address if self . address else <NUM_LIT:0> ) ] <EOL> asm_fmt = "<STR_LIT>" <EOL> reil_fmt = "<STR_LIT>" <EOL> for dinstr in self . _instrs : <EOL> asm_instr = dinstr . asm_instr <EOL> lines += [ asm_fmt . format ( asm_instr . address , asm_instr ) ] <EOL> for reil_instr in dinstr . ir_instrs : <EOL> lines += [ reil_fmt . format ( reil_instr . address >> <NUM_LIT> , reil_instr . address & <NUM_LIT> , reil_instr ) ] <EOL> return "<STR_LIT:\n>" . join ( lines ) <EOL> def __eq__ ( self , other ) : <EOL> return self . address == other . address and self . end_address == other . end_address <EOL> def __ne__ ( self , other ) : <EOL> return not self . __eq__ ( other ) <EOL> def __iter__ ( self ) : <EOL> for instr in self . _instrs : <EOL> yield instr <EOL> def __len__ ( self ) : <EOL> return len ( self . _instrs ) <EOL> def __getstate__ ( self ) : <EOL> state = { } <EOL> state [ '<STR_LIT>' ] = self . _instrs <EOL> state [ '<STR_LIT>' ] = self . _address <EOL> state [ '<STR_LIT>' ] = self . _taken_branch <EOL> state [ '<STR_LIT>' ] = self . _not_taken_branch <EOL> state [ '<STR_LIT>' ] = self . _direct_branch <EOL> return state <EOL> def __setstate__ ( self , state ) : <EOL> self . _instrs = state [ '<STR_LIT>' ] <EOL> self . _address = state [ '<STR_LIT>' ] <EOL> self . _taken_branch = state [ '<STR_LIT>' ] <EOL> self . _not_taken_branch = state [ '<STR_LIT>' ] <EOL> self . _direct_branch = state [ '<STR_LIT>' ] <EOL> class BasicBlockGraph ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , basic_blocks ) : <EOL> self . _basic_blocks = sorted ( basic_blocks , key = lambda bb : bb . address ) <EOL> self . _bb_by_addr = dict ( [ ( bb . address , bb ) for bb in basic_blocks ] ) <EOL> self . _graph = self . _build_graph ( ) <EOL> self . _entry_blocks = [ bb . address for bb in basic_blocks <EOL> if len ( self . _graph . in_edges ( bb . address ) ) == <NUM_LIT:0> ] <EOL> self . _exit_blocks = [ bb . address for bb in basic_blocks <EOL> if len ( self . _graph . out_edges ( bb . address ) ) == <NUM_LIT:0> ] <EOL> def all_simple_bb_paths ( self , start_address , end_address ) : <EOL> """<STR_LIT>""" <EOL> bb_start = self . _find_basic_block ( start_address ) <EOL> bb_end = self . _find_basic_block ( end_address ) <EOL> paths = networkx . all_simple_paths ( self . _graph , source = bb_start . address , target = bb_end . address ) <EOL> return ( map ( lambda addr : self . _bb_by_addr [ addr ] , path ) for path in paths ) <EOL> def save ( self , filename , print_ir = False , format = '<STR_LIT>' ) : <EOL> """<STR_LIT>""" <EOL> node_format = { <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> } <EOL> edge_format = { <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> } <EOL> edge_colors = { <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> } <EOL> html_entities = { <EOL> "<STR_LIT:!>" : "<STR_LIT>" , <EOL> "<STR_LIT:#>" : "<STR_LIT>" , <EOL> "<STR_LIT::>" : "<STR_LIT>" , <EOL> "<STR_LIT:{>" : "<STR_LIT>" , <EOL> "<STR_LIT:}>" : "<STR_LIT>" , <EOL> } <EOL> try : <EOL> dot_graph = Dot ( graph_type = "<STR_LIT>" , rankdir = "<STR_LIT>" ) <EOL> nodes = { } <EOL> for bb_addr in self . _bb_by_addr : <EOL> bb_dump = self . _dump_bb ( self . _bb_by_addr [ bb_addr ] , print_ir ) <EOL> for char , encoding in html_entities . items ( ) : <EOL> bb_dump = bb_dump . replace ( char , encoding ) <EOL> label = "<STR_LIT>" . format ( bb_addr , bb_dump ) <EOL> nodes [ bb_addr ] = Node ( bb_addr , label = label , ** node_format ) <EOL> dot_graph . add_node ( nodes [ bb_addr ] ) <EOL> for bb_src_addr in self . _bb_by_addr : <EOL> for bb_dst_addr , branch_type in self . _bb_by_addr [ bb_src_addr ] . branches : <EOL> if not bb_dst_addr in self . _bb_by_addr : <EOL> continue <EOL> dot_graph . add_edge ( <EOL> Edge ( <EOL> nodes [ bb_src_addr ] , <EOL> nodes [ bb_dst_addr ] , <EOL> color = edge_colors [ branch_type ] , <EOL> ** edge_format ) ) <EOL> dot_graph . write ( "<STR_LIT>" . format ( filename , format ) , format = format ) <EOL> except Exception : <EOL> logger . error ( <EOL> "<STR_LIT>" , <EOL> filename , <EOL> format , <EOL> exc_info = True <EOL> ) <EOL> def _build_graph ( self ) : <EOL> graph = networkx . DiGraph ( ) <EOL> for bb_addr in self . _bb_by_addr . keys ( ) : <EOL> graph . add_node ( bb_addr , address = bb_addr ) <EOL> for bb_src_addr in self . _bb_by_addr . keys ( ) : <EOL> for bb_dst_addr , branch_type in self . _bb_by_addr [ bb_src_addr ] . branches : <EOL> graph . add_edge ( bb_src_addr , bb_dst_addr , branch_type = branch_type ) <EOL> return graph <EOL> def _find_basic_block ( self , address ) : <EOL> bb_rv = None <EOL> for bb in self . _basic_blocks : <EOL> if address >= bb . address and address <= bb . end_address : <EOL> bb_rv = bb <EOL> break <EOL> return bb_rv <EOL> def _dump_bb ( self , basic_block , print_ir = False ) : <EOL> lines = [ ] <EOL> asm_fmt = "<STR_LIT>" <EOL> reil_fmt = "<STR_LIT>" <EOL> for dinstr in basic_block : <EOL> lines += [ asm_fmt . format ( dinstr . address , dinstr . asm_instr . size , dinstr . asm_instr ) ] <EOL> if print_ir : <EOL> for ir_instr in dinstr . ir_instrs : <EOL> lines += [ reil_fmt . format ( ir_instr . address & <NUM_LIT> , ir_instr ) ] <EOL> return "<STR_LIT>" . join ( lines ) <EOL> @ property <EOL> def basic_blocks ( self ) : <EOL> return self . _basic_blocks <EOL> def get_basic_block ( self , address ) : <EOL> return self . _bb_by_addr [ address ] <EOL> @ property <EOL> def exit_basic_blocks ( self ) : <EOL> for bb_addr in self . _exit_blocks : <EOL> yield self . _bb_by_addr [ bb_addr ] <EOL> @ property <EOL> def entry_basic_blocks ( self ) : <EOL> for bb_addr in self . _entry_blocks : <EOL> yield self . _bb_by_addr [ bb_addr ] <EOL> def __getstate__ ( self ) : <EOL> state = { } <EOL> state [ '<STR_LIT>' ] = self . _basic_blocks <EOL> return state <EOL> def __setstate__ ( self , state ) : <EOL> basic_blocks = state [ '<STR_LIT>' ] <EOL> self . _basic_blocks = sorted ( basic_blocks , key = lambda bb : bb . address ) <EOL> self . _bb_by_addr = dict ( [ ( bb . address , bb ) for bb in basic_blocks ] ) <EOL> self . _graph = self . _build_graph ( ) <EOL> self . _entry_blocks = [ bb . address for bb in basic_blocks <EOL> if len ( self . _graph . in_edges ( bb . address ) ) == <NUM_LIT:0> ] <EOL> self . _exit_blocks = [ bb . address for bb in basic_blocks <EOL> if len ( self . _graph . out_edges ( bb . address ) ) == <NUM_LIT:0> ] <EOL> class BasicBlockBuilder ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , disassembler , memory , translator , arch_info ) : <EOL> self . _disasm = disassembler <EOL> self . _ir_trans = translator <EOL> self . _lookahead_max = <NUM_LIT:16> <EOL> self . _mem = memory <EOL> self . _arch_info = arch_info <EOL> def build ( self , start_address , end_address , symbols = None ) : <EOL> """<STR_LIT>""" <EOL> if not symbols : <EOL> symbols = { } <EOL> bbs = self . _find_candidate_bbs ( start_address , end_address , symbols = symbols ) <EOL> bbs = self . _refine_bbs ( bbs , symbols ) <EOL> addrs = [ ] <EOL> for bb in bbs : <EOL> for addr , _ in bb . branches : <EOL> if not ( start_address <= addr and addr <= end_address ) : <EOL> addrs . append ( addr ) <EOL> addrs = sorted ( list ( set ( addrs ) ) ) <EOL> bbs_new = [ ] <EOL> for addr in addrs : <EOL> bbs_new += self . _find_candidate_bbs ( addr , end_address , symbols = symbols ) <EOL> bbs = self . _refine_bbs ( bbs_new + bbs , symbols ) <EOL> return bbs <EOL> def _find_candidate_bbs ( self , start_address , end_address , mode = BARF_DISASM_RECURSIVE , symbols = None ) : <EOL> if not symbols : <EOL> symbols = { } <EOL> bbs = [ ] <EOL> addrs_to_process = Queue ( ) <EOL> addrs_processed = set ( ) <EOL> addrs_to_process . put ( start_address ) <EOL> while not addrs_to_process . empty ( ) : <EOL> addr_curr = addrs_to_process . get ( ) <EOL> if addr_curr in addrs_processed or not ( addr_curr >= start_address and addr_curr <= end_address ) : <EOL> continue <EOL> bb = self . _disassemble_bb ( addr_curr , end_address + <NUM_LIT> , symbols ) <EOL> if bb . empty ( ) : <EOL> continue <EOL> bbs += [ bb ] <EOL> addrs_processed . add ( addr_curr ) <EOL> if mode == BARF_DISASM_LINEAR : <EOL> next_addr = bb . address + bb . size <EOL> if not self . _bb_ends_in_direct_jmp ( bb ) and not self . _bb_ends_in_return ( bb ) and not next_addr in addrs_processed : <EOL> addrs_to_process . put ( next_addr ) <EOL> if mode == BARF_DISASM_RECURSIVE : <EOL> for addr , _ in bb . branches : <EOL> if not addr in addrs_processed : <EOL> addrs_to_process . put ( addr ) <EOL> return bbs <EOL> def _bb_ends_in_direct_jmp ( self , bb ) : <EOL> last_instr = bb . instrs [ - <NUM_LIT:1> ] . ir_instrs [ - <NUM_LIT:1> ] <EOL> return last_instr . mnemonic == ReilMnemonic . JCC and isinstance ( last_instr . operands [ <NUM_LIT:0> ] , ReilImmediateOperand ) and last_instr . operands [ <NUM_LIT:0> ] . immediate == <NUM_LIT> <EOL> def _bb_ends_in_return ( self , bb ) : <EOL> last_instr = bb . instrs [ - <NUM_LIT:1> ] . ir_instrs [ - <NUM_LIT:1> ] <EOL> return last_instr . mnemonic == ReilMnemonic . RET <EOL> def _refine_bbs ( self , bbs , symbols ) : <EOL> bbs . sort ( key = lambda x : x . address ) <EOL> bbs_addrs = [ bb . address for bb in bbs ] <EOL> bbs_new = [ ] <EOL> for bb1 in bbs : <EOL> bb_divided = False <EOL> lower = bisect . bisect_left ( bbs_addrs , bb1 . start_address ) <EOL> upper = bisect . bisect_right ( bbs_addrs , bb1 . end_address ) <EOL> for bb2 in bbs [ lower : upper ] : <EOL> if bb1 . contains ( bb2 . address ) and bb1 != bb2 : <EOL> bb_lower , bb_upper = self . _split_bb ( bb1 , bb2 . address , symbols ) <EOL> if not bb_upper . empty ( ) and bb_upper not in bbs_new : <EOL> bbs_new += [ bb_upper ] <EOL> bb1 = bb_lower <EOL> bb_divided = True <EOL> break <EOL> if not bb_divided : <EOL> if not bb1 . empty ( ) and not bb1 in bbs_new : <EOL> bbs_new += [ bb1 ] <EOL> return bbs_new <EOL> def _split_bb ( self , bb , address , symbols ) : <EOL> bb_upper_half = self . _disassemble_bb ( bb . start_address , address , symbols ) <EOL> bb_lower_half = self . _disassemble_bb ( address , bb . end_address , symbols ) <EOL> bb_upper_half . direct_branch = address <EOL> return bb_lower_half , bb_upper_half <EOL> def _disassemble_bb ( self , start_address , end_address , symbols ) : <EOL> bb = BasicBlock ( ) <EOL> addr = start_address <EOL> taken , not_taken , direct = None , None , None <EOL> while addr < end_address : <EOL> try : <EOL> start , end = addr , min ( addr + self . _lookahead_max , end_address ) <EOL> data_chunk = self . _mem [ start : end ] <EOL> except Exception : <EOL> break <EOL> try : <EOL> asm = self . _disasm . disassemble ( data_chunk , addr ) <EOL> except ( InvalidDisassemblerData , CapstoneOperandNotSupported ) : <EOL> break <EOL> if not asm : <EOL> break <EOL> ir = self . _ir_trans . translate ( asm ) <EOL> bb . instrs . append ( DualInstruction ( addr , asm , ir ) ) <EOL> if ir [ - <NUM_LIT:1> ] . mnemonic == ReilMnemonic . RET : <EOL> break <EOL> if asm . mnemonic == "<STR_LIT>" : <EOL> break <EOL> if ir [ - <NUM_LIT:1> ] . mnemonic == ReilMnemonic . JCC and asm . mnemonic == "<STR_LIT>" and isinstance ( ir [ - <NUM_LIT:1> ] . operands [ <NUM_LIT:2> ] , ReilImmediateOperand ) and ( ir [ - <NUM_LIT:1> ] . operands [ <NUM_LIT:2> ] . immediate >> <NUM_LIT> ) in symbols and not symbols [ ir [ - <NUM_LIT:1> ] . operands [ <NUM_LIT:2> ] . immediate >> <NUM_LIT> ] [ <NUM_LIT:2> ] : <EOL> break <EOL> if ir [ - <NUM_LIT:1> ] . mnemonic == ReilMnemonic . JCC and not asm . mnemonic == "<STR_LIT>" and not asm . mnemonic == "<STR_LIT>" and not asm . mnemonic == "<STR_LIT>" and not asm . prefix in [ "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" ] : <EOL> taken , not_taken , direct = self . _extract_branches ( asm , ir ) <EOL> break <EOL> if self . _arch_info == ARCH_ARM and asm . mnemonic == "<STR_LIT>" and "<STR_LIT>" in str ( asm . operands [ <NUM_LIT:1> ] ) : <EOL> break <EOL> if self . _arch_info == ARCH_ARM and asm . mnemonic == "<STR_LIT>" and "<STR_LIT>" in str ( asm . operands [ <NUM_LIT:0> ] ) : <EOL> break <EOL> addr += asm . size <EOL> bb . taken_branch = taken <EOL> bb . not_taken_branch = not_taken <EOL> bb . direct_branch = direct <EOL> return bb <EOL> def _resolve_branch_address ( self , branch_instr , instrs ) : <EOL> target = branch_instr . operands [ <NUM_LIT:2> ] <EOL> if isinstance ( target , ReilImmediateOperand ) : <EOL> return target . immediate >> <NUM_LIT:8> <EOL> else : <EOL> for instr in instrs [ : : - <NUM_LIT:1> ] : <EOL> if instr . mnemonic == ReilMnemonic . STR and isinstance ( instr . operands [ <NUM_LIT:0> ] , ReilImmediateOperand ) and instr . operands [ <NUM_LIT:2> ] == target : <EOL> return instr . operands [ <NUM_LIT:0> ] . immediate >> <NUM_LIT:8> <EOL> def _extract_branches ( self , asm , ir ) : <EOL> taken_branch , not_taken_branch , direct_branch = None , None , None <EOL> instr_last = ir [ - <NUM_LIT:1> ] <EOL> if instr_last . mnemonic == ReilMnemonic . JCC : <EOL> cond = instr_last . operands [ <NUM_LIT:0> ] <EOL> branch_addr = self . _resolve_branch_address ( instr_last , ir ) <EOL> if isinstance ( cond , ReilImmediateOperand ) : <EOL> if cond . immediate == <NUM_LIT> : <EOL> not_taken_branch = branch_addr <EOL> if cond . immediate == <NUM_LIT> : <EOL> if asm . mnemonic == '<STR_LIT>' : <EOL> direct_branch = asm . address + asm . size <EOL> if asm . mnemonic != '<STR_LIT>' : <EOL> direct_branch = branch_addr <EOL> else : <EOL> taken_branch = branch_addr <EOL> not_taken_branch = asm . address + asm . size <EOL> return taken_branch , not_taken_branch , direct_branch </s>
