<s> import logging <EOL> from collections import defaultdict <EOL> from androguard . decompiler . dad . basic_blocks import ( build_node_from_block , <EOL> StatementBlock , CondBlock ) <EOL> from androguard . decompiler . dad . util import get_type <EOL> from androguard . decompiler . dad . instruction import Variable <EOL> logger = logging . getLogger ( '<STR_LIT>' ) <EOL> class Graph ( object ) : <EOL> def __init__ ( self ) : <EOL> self . entry = None <EOL> self . exit = None <EOL> self . nodes = list ( ) <EOL> self . rpo = [ ] <EOL> self . edges = defaultdict ( list ) <EOL> self . catch_edges = defaultdict ( list ) <EOL> self . reverse_edges = defaultdict ( list ) <EOL> self . reverse_catch_edges = defaultdict ( list ) <EOL> self . loc_to_ins = None <EOL> self . loc_to_node = None <EOL> def sucs ( self , node ) : <EOL> return self . edges . get ( node , [ ] ) <EOL> def all_sucs ( self , node ) : <EOL> return self . edges . get ( node , [ ] ) + self . catch_edges . get ( node , [ ] ) <EOL> def preds ( self , node ) : <EOL> return [ n for n in self . reverse_edges . get ( node , [ ] ) if not n . in_catch ] <EOL> def all_preds ( self , node ) : <EOL> return ( self . reverse_edges . get ( node , [ ] ) + self . reverse_catch_edges . get ( <EOL> node , [ ] ) ) <EOL> def add_node ( self , node ) : <EOL> self . nodes . append ( node ) <EOL> def add_edge ( self , e1 , e2 ) : <EOL> lsucs = self . edges [ e1 ] <EOL> if e2 not in lsucs : <EOL> lsucs . append ( e2 ) <EOL> lpreds = self . reverse_edges [ e2 ] <EOL> if e1 not in lpreds : <EOL> lpreds . append ( e1 ) <EOL> def add_catch_edge ( self , e1 , e2 ) : <EOL> lsucs = self . catch_edges [ e1 ] <EOL> if e2 not in lsucs : <EOL> lsucs . append ( e2 ) <EOL> lpreds = self . reverse_catch_edges [ e2 ] <EOL> if e1 not in lpreds : <EOL> lpreds . append ( e1 ) <EOL> def remove_node ( self , node ) : <EOL> preds = self . reverse_edges . get ( node , [ ] ) <EOL> for pred in preds : <EOL> self . edges [ pred ] . remove ( node ) <EOL> succs = self . edges . get ( node , [ ] ) <EOL> for suc in succs : <EOL> self . reverse_edges [ suc ] . remove ( node ) <EOL> exc_preds = self . reverse_catch_edges . pop ( node , [ ] ) <EOL> for pred in exc_preds : <EOL> self . catch_edges [ pred ] . remove ( node ) <EOL> exc_succs = self . catch_edges . pop ( node , [ ] ) <EOL> for suc in exc_succs : <EOL> self . reverse_catch_edges [ suc ] . remove ( node ) <EOL> self . nodes . remove ( node ) <EOL> if node in self . rpo : <EOL> self . rpo . remove ( node ) <EOL> del node <EOL> def number_ins ( self ) : <EOL> self . loc_to_ins = { } <EOL> self . loc_to_node = { } <EOL> num = <NUM_LIT:0> <EOL> for node in self . rpo : <EOL> start_node = num <EOL> num = node . number_ins ( num ) <EOL> end_node = num - <NUM_LIT:1> <EOL> self . loc_to_ins . update ( node . get_loc_with_ins ( ) ) <EOL> self . loc_to_node [ start_node , end_node ] = node <EOL> def get_ins_from_loc ( self , loc ) : <EOL> return self . loc_to_ins . get ( loc ) <EOL> def get_node_from_loc ( self , loc ) : <EOL> for ( start , end ) , node in self . loc_to_node . iteritems ( ) : <EOL> if start <= loc <= end : <EOL> return node <EOL> def remove_ins ( self , loc ) : <EOL> ins = self . get_ins_from_loc ( loc ) <EOL> self . get_node_from_loc ( loc ) . remove_ins ( loc , ins ) <EOL> self . loc_to_ins . pop ( loc ) <EOL> def compute_rpo ( self ) : <EOL> '''<STR_LIT>''' <EOL> nb = len ( self . nodes ) + <NUM_LIT:1> <EOL> for node in self . post_order ( ) : <EOL> node . num = nb - node . po <EOL> self . rpo = sorted ( self . nodes , key = lambda n : n . num ) <EOL> def post_order ( self ) : <EOL> '''<STR_LIT>''' <EOL> def _visit ( n , cnt ) : <EOL> visited . add ( n ) <EOL> for suc in self . all_sucs ( n ) : <EOL> if not suc in visited : <EOL> for cnt , s in _visit ( suc , cnt ) : <EOL> yield cnt , s <EOL> n . po = cnt <EOL> yield cnt + <NUM_LIT:1> , n <EOL> visited = set ( ) <EOL> for _ , node in _visit ( self . entry , <NUM_LIT:1> ) : <EOL> yield node <EOL> def draw ( self , name , dname , draw_branches = True ) : <EOL> from pydot import Dot , Edge <EOL> g = Dot ( ) <EOL> g . set_node_defaults ( color = '<STR_LIT>' , <EOL> style = '<STR_LIT>' , <EOL> shape = '<STR_LIT>' , <EOL> fontname = '<STR_LIT>' , <EOL> fontsize = '<STR_LIT>' ) <EOL> for node in sorted ( self . nodes , key = lambda x : x . num ) : <EOL> if draw_branches and node . type . is_cond : <EOL> g . add_edge ( Edge ( str ( node ) , str ( node . true ) , color = '<STR_LIT>' ) ) <EOL> g . add_edge ( Edge ( str ( node ) , str ( node . false ) , color = '<STR_LIT>' ) ) <EOL> else : <EOL> for suc in self . sucs ( node ) : <EOL> g . add_edge ( Edge ( str ( node ) , str ( suc ) , color = '<STR_LIT>' ) ) <EOL> for except_node in self . catch_edges . get ( node , [ ] ) : <EOL> g . add_edge ( Edge ( str ( node ) , <EOL> str ( except_node ) , <EOL> color = '<STR_LIT>' , <EOL> style = '<STR_LIT>' ) ) <EOL> g . write_png ( '<STR_LIT>' % ( dname , name ) ) <EOL> def immediate_dominators ( self ) : <EOL> return dom_lt ( self ) <EOL> def __len__ ( self ) : <EOL> return len ( self . nodes ) <EOL> def __repr__ ( self ) : <EOL> return str ( self . nodes ) <EOL> def __iter__ ( self ) : <EOL> for node in self . nodes : <EOL> yield node <EOL> def split_if_nodes ( graph ) : <EOL> '''<STR_LIT>''' <EOL> node_map = { n : n for n in graph } <EOL> to_update = set ( ) <EOL> for node in graph . nodes [ : ] : <EOL> if node . type . is_cond : <EOL> if len ( node . get_ins ( ) ) > <NUM_LIT:1> : <EOL> pre_ins = node . get_ins ( ) [ : - <NUM_LIT:1> ] <EOL> last_ins = node . get_ins ( ) [ - <NUM_LIT:1> ] <EOL> pre_node = StatementBlock ( '<STR_LIT>' % node . name , pre_ins ) <EOL> cond_node = CondBlock ( '<STR_LIT>' % node . name , [ last_ins ] ) <EOL> node_map [ node ] = pre_node <EOL> node_map [ pre_node ] = pre_node <EOL> node_map [ cond_node ] = cond_node <EOL> pre_node . copy_from ( node ) <EOL> cond_node . copy_from ( node ) <EOL> for var in node . var_to_declare : <EOL> pre_node . add_variable_declaration ( var ) <EOL> pre_node . type . is_stmt = True <EOL> cond_node . true = node . true <EOL> cond_node . false = node . false <EOL> for pred in graph . all_preds ( node ) : <EOL> pred_node = node_map [ pred ] <EOL> if node not in graph . sucs ( pred ) : <EOL> graph . add_catch_edge ( pred_node , pre_node ) <EOL> continue <EOL> if pred is node : <EOL> pred_node = cond_node <EOL> if pred . type . is_cond : <EOL> if pred . true is node : <EOL> pred_node . true = pre_node <EOL> if pred . false is node : <EOL> pred_node . false = pre_node <EOL> graph . add_edge ( pred_node , pre_node ) <EOL> for suc in graph . sucs ( node ) : <EOL> graph . add_edge ( cond_node , node_map [ suc ] ) <EOL> for suc in graph . catch_edges . get ( node , [ ] ) : <EOL> graph . add_catch_edge ( pre_node , node_map [ suc ] ) <EOL> if node is graph . entry : <EOL> graph . entry = pre_node <EOL> graph . add_node ( pre_node ) <EOL> graph . add_node ( cond_node ) <EOL> graph . add_edge ( pre_node , cond_node ) <EOL> pre_node . update_attribute_with ( node_map ) <EOL> cond_node . update_attribute_with ( node_map ) <EOL> graph . remove_node ( node ) <EOL> else : <EOL> to_update . add ( node ) <EOL> for node in to_update : <EOL> node . update_attribute_with ( node_map ) <EOL> def simplify ( graph ) : <EOL> '''<STR_LIT>''' <EOL> redo = True <EOL> while redo : <EOL> redo = False <EOL> node_map = { } <EOL> to_update = set ( ) <EOL> for node in graph . nodes [ : ] : <EOL> if node . type . is_stmt and node in graph : <EOL> sucs = graph . all_sucs ( node ) <EOL> if len ( sucs ) != <NUM_LIT:1> : <EOL> continue <EOL> suc = sucs [ <NUM_LIT:0> ] <EOL> if len ( node . get_ins ( ) ) == <NUM_LIT:0> : <EOL> if any ( pred . type . is_switch <EOL> for pred in graph . all_preds ( node ) ) : <EOL> continue <EOL> if node is suc : <EOL> continue <EOL> node_map [ node ] = suc <EOL> for pred in graph . all_preds ( node ) : <EOL> pred . update_attribute_with ( node_map ) <EOL> if node not in graph . sucs ( pred ) : <EOL> graph . add_catch_edge ( pred , suc ) <EOL> continue <EOL> graph . add_edge ( pred , suc ) <EOL> redo = True <EOL> if node is graph . entry : <EOL> graph . entry = suc <EOL> graph . remove_node ( node ) <EOL> elif ( suc . type . is_stmt and len ( graph . all_preds ( suc ) ) == <NUM_LIT:1> and <EOL> not ( suc in graph . catch_edges ) and not ( <EOL> ( node is suc ) or ( suc is graph . entry ) ) ) : <EOL> ins_to_merge = suc . get_ins ( ) <EOL> node . add_ins ( ins_to_merge ) <EOL> for var in suc . var_to_declare : <EOL> node . add_variable_declaration ( var ) <EOL> new_suc = graph . sucs ( suc ) [ <NUM_LIT:0> ] <EOL> if new_suc : <EOL> graph . add_edge ( node , new_suc ) <EOL> for exception_suc in graph . catch_edges . get ( suc , [ ] ) : <EOL> graph . add_catch_edge ( node , exception_suc ) <EOL> redo = True <EOL> graph . remove_node ( suc ) <EOL> else : <EOL> to_update . add ( node ) <EOL> for node in to_update : <EOL> node . update_attribute_with ( node_map ) <EOL> def dom_lt ( graph ) : <EOL> '''<STR_LIT>''' <EOL> def _dfs ( v , n ) : <EOL> semi [ v ] = n = n + <NUM_LIT:1> <EOL> vertex [ n ] = label [ v ] = v <EOL> ancestor [ v ] = <NUM_LIT:0> <EOL> for w in graph . all_sucs ( v ) : <EOL> if not semi [ w ] : <EOL> parent [ w ] = v <EOL> n = _dfs ( w , n ) <EOL> pred [ w ] . add ( v ) <EOL> return n <EOL> def _compress ( v ) : <EOL> u = ancestor [ v ] <EOL> if ancestor [ u ] : <EOL> _compress ( u ) <EOL> if semi [ label [ u ] ] < semi [ label [ v ] ] : <EOL> label [ v ] = label [ u ] <EOL> ancestor [ v ] = ancestor [ u ] <EOL> def _eval ( v ) : <EOL> if ancestor [ v ] : <EOL> _compress ( v ) <EOL> return label [ v ] <EOL> return v <EOL> def _link ( v , w ) : <EOL> ancestor [ w ] = v <EOL> parent , ancestor , vertex = { } , { } , { } <EOL> label , dom = { } , { } <EOL> pred , bucket = defaultdict ( set ) , defaultdict ( set ) <EOL> semi = { v : <NUM_LIT:0> for v in graph . nodes } <EOL> n = _dfs ( graph . entry , <NUM_LIT:0> ) <EOL> for i in xrange ( n , <NUM_LIT:1> , - <NUM_LIT:1> ) : <EOL> w = vertex [ i ] <EOL> for v in pred [ w ] : <EOL> u = _eval ( v ) <EOL> y = semi [ w ] = min ( semi [ w ] , semi [ u ] ) <EOL> bucket [ vertex [ y ] ] . add ( w ) <EOL> pw = parent [ w ] <EOL> _link ( pw , w ) <EOL> bpw = bucket [ pw ] <EOL> while bpw : <EOL> v = bpw . pop ( ) <EOL> u = _eval ( v ) <EOL> dom [ v ] = u if semi [ u ] < semi [ v ] else pw <EOL> for i in range ( <NUM_LIT:2> , n + <NUM_LIT:1> ) : <EOL> w = vertex [ i ] <EOL> dw = dom [ w ] <EOL> if dw != vertex [ semi [ w ] ] : <EOL> dom [ w ] = dom [ dw ] <EOL> dom [ graph . entry ] = None <EOL> return dom <EOL> def bfs ( start ) : <EOL> to_visit = [ start ] <EOL> visited = set ( [ start ] ) <EOL> while to_visit : <EOL> node = to_visit . pop ( <NUM_LIT:0> ) <EOL> yield node <EOL> if node . exception_analysis : <EOL> for _ , _ , exception in node . exception_analysis . exceptions : <EOL> if exception not in visited : <EOL> to_visit . append ( exception ) <EOL> visited . add ( exception ) <EOL> for _ , _ , child in node . childs : <EOL> if child not in visited : <EOL> to_visit . append ( child ) <EOL> visited . add ( child ) <EOL> class GenInvokeRetName ( object ) : <EOL> def __init__ ( self ) : <EOL> self . num = <NUM_LIT:0> <EOL> self . ret = None <EOL> def new ( self ) : <EOL> self . num += <NUM_LIT:1> <EOL> self . ret = Variable ( '<STR_LIT>' % self . num ) <EOL> return self . ret <EOL> def set_to ( self , ret ) : <EOL> self . ret = ret <EOL> def last ( self ) : <EOL> return self . ret <EOL> def make_node ( graph , block , block_to_node , vmap , gen_ret ) : <EOL> node = block_to_node . get ( block ) <EOL> if node is None : <EOL> node = build_node_from_block ( block , vmap , gen_ret ) <EOL> block_to_node [ block ] = node <EOL> if block . exception_analysis : <EOL> for _type , _ , exception_target in block . exception_analysis . exceptions : <EOL> exception_node = block_to_node . get ( exception_target ) <EOL> if exception_node is None : <EOL> exception_node = build_node_from_block ( exception_target , vmap , <EOL> gen_ret , _type ) <EOL> exception_node . set_catch_type ( _type ) <EOL> exception_node . in_catch = True <EOL> block_to_node [ exception_target ] = exception_node <EOL> graph . add_catch_edge ( node , exception_node ) <EOL> for _ , _ , child_block in block . childs : <EOL> child_node = block_to_node . get ( child_block ) <EOL> if child_node is None : <EOL> child_node = build_node_from_block ( child_block , vmap , gen_ret ) <EOL> block_to_node [ child_block ] = child_node <EOL> graph . add_edge ( node , child_node ) <EOL> if node . type . is_switch : <EOL> node . add_case ( child_node ) <EOL> if node . type . is_cond : <EOL> if_target = ( ( block . end / <NUM_LIT:2> ) - <EOL> ( block . last_length / <NUM_LIT:2> ) + node . off_last_ins ) <EOL> child_addr = child_block . start / <NUM_LIT:2> <EOL> if if_target == child_addr : <EOL> node . true = child_node <EOL> else : <EOL> node . false = child_node <EOL> if node . type . is_cond and node . false is None : <EOL> node . false = node . true <EOL> return node <EOL> def construct ( start_block , vmap , exceptions ) : <EOL> bfs_blocks = bfs ( start_block ) <EOL> graph = Graph ( ) <EOL> gen_ret = GenInvokeRetName ( ) <EOL> block_to_node = { } <EOL> exceptions_start_block = [ ] <EOL> for exception in exceptions : <EOL> for _ , _ , block in exception . exceptions : <EOL> exceptions_start_block . append ( block ) <EOL> for block in bfs_blocks : <EOL> node = make_node ( graph , block , block_to_node , vmap , gen_ret ) <EOL> graph . add_node ( node ) <EOL> graph . entry = block_to_node [ start_block ] <EOL> del block_to_node , bfs_blocks <EOL> graph . compute_rpo ( ) <EOL> graph . number_ins ( ) <EOL> for node in graph . rpo : <EOL> preds = [ pred for pred in graph . all_preds ( node ) if pred . num < node . num ] <EOL> if preds and all ( pred . in_catch for pred in preds ) : <EOL> node . in_catch = True <EOL> lexit_nodes = [ node for node in graph if node . type . is_return ] <EOL> if len ( lexit_nodes ) > <NUM_LIT:1> : <EOL> logger . error ( '<STR_LIT>' ) <EOL> graph . exit = graph . rpo [ - <NUM_LIT:1> ] <EOL> elif len ( lexit_nodes ) < <NUM_LIT:1> : <EOL> logger . debug ( '<STR_LIT>' ) <EOL> else : <EOL> graph . exit = lexit_nodes [ <NUM_LIT:0> ] <EOL> return graph </s>
