<s> import re <EOL> import random <EOL> import urllib <EOL> import operator <EOL> from util import scanTokens , scanNgrams <EOL> from Levenshtein import distance <EOL> from random import choice <EOL> LEXER_TOKENS = "<STR_LIT>" <EOL> LEXER_NGRAMS = "<STR_LIT>" <EOL> OFS_TRANSITION = <NUM_LIT:3> <EOL> def pickIndex ( counts ) : <EOL> total = sum ( counts ) <EOL> dice = random . randint ( <NUM_LIT:1> , total ) <EOL> countsSoFar = <NUM_LIT:0> <EOL> for ( ind , count ) in enumerate ( counts ) : <EOL> countsSoFar += count <EOL> if countsSoFar >= dice : <EOL> ret = ind <EOL> break <EOL> return ret <EOL> class MarkovModel : <EOL> def __init__ ( self , model_file ) : <EOL> self . load_markov_model ( model_file ) <EOL> self . init ( ) <EOL> self . simplifiedModel = { } <EOL> for ( state , suc ) in self . markov_model . iteritems ( ) : <EOL> self . simplifiedModel [ state ] = { } . fromkeys ( [ s [ <NUM_LIT:0> ] for s in suc ] , <EOL> True ) <EOL> self . allStates = self . simplifiedModel . copy ( ) <EOL> for trans in self . simplifiedModel . values ( ) : <EOL> self . allStates . update ( trans ) <EOL> self . allStates = self . allStates . keys ( ) <EOL> self . fuzzer = "<STR_LIT>" <EOL> def init ( self ) : <EOL> """<STR_LIT>""" <EOL> self . state = "<STR_LIT:|>" . join ( [ "<STR_LIT>" ] * self . order ) <EOL> def set_fuzzer ( self , fuzzer ) : <EOL> self . fuzzer = fuzzer <EOL> def getSimplifiedModel ( self , stateFilter = None ) : <EOL> """<STR_LIT>""" <EOL> if stateFilter is None : <EOL> ret = self . simplifiedModel <EOL> else : <EOL> ret = { } <EOL> for k in self . simplifiedModel . keys ( ) : <EOL> if k in stateFilter : <EOL> ret [ k ] = dict ( [ kv for kv in self . simplifiedModel [ k ] . items ( ) <EOL> if kv [ <NUM_LIT:0> ] in stateFilter ] ) <EOL> return ret <EOL> def getAllStates ( self ) : <EOL> return self . allStates <EOL> def getMarkovMatrix ( self ) : <EOL> allStates = self . getAllStates ( ) <EOL> nStates = len ( allStates ) <EOL> ret = [ [ <NUM_LIT:0> ] * nStates for _ in xrange ( nStates ) ] <EOL> for ( ind , state ) in enumerate ( allStates ) : <EOL> counts = self . markov_model . get ( state , [ ] ) <EOL> for ( s , c ) in counts : <EOL> ret [ ind ] [ allStates . index ( s ) ] = c <EOL> return ret <EOL> def exportMarkovMartrix ( self , fileName ) : <EOL> f = file ( fileName , "<STR_LIT:w>" ) <EOL> f . write ( "<STR_LIT>" % "<STR_LIT:\t>" . join ( self . allStates ) ) <EOL> for ( s , l ) in zip ( self . allStates , self . getMarkovMatrix ( ) ) : <EOL> f . write ( "<STR_LIT>" % ( s , "<STR_LIT:\t>" . join ( [ str ( v ) for v in l ] ) ) ) <EOL> f . close ( ) <EOL> def update_state ( self , state ) : <EOL> """<STR_LIT>""" <EOL> self . state = state <EOL> def getState ( self ) : <EOL> return self . state <EOL> def getOrder ( self ) : <EOL> return self . order <EOL> def load_markov_model ( self , model_file ) : <EOL> """<STR_LIT>""" <EOL> fd = open ( model_file , "<STR_LIT:r>" ) <EOL> text_markov_model = fd . read ( ) <EOL> fd . close ( ) <EOL> re_markov_model = re . findall ( "<STR_LIT>" , <EOL> text_markov_model ) <EOL> markov_model = { } <EOL> for edge in re_markov_model : <EOL> if edge [ <NUM_LIT:0> ] in markov_model : <EOL> ( markov_model [ edge [ <NUM_LIT:0> ] ] ) . extend ( [ ( edge [ <NUM_LIT:1> ] , int ( edge [ <NUM_LIT:2> ] ) ) ] ) <EOL> else : <EOL> markov_model [ edge [ <NUM_LIT:0> ] ] = [ ( edge [ <NUM_LIT:1> ] , int ( edge [ <NUM_LIT:2> ] ) ) ] <EOL> self . order = edge [ <NUM_LIT:0> ] . count ( "<STR_LIT:|>" ) + <NUM_LIT:1> <EOL> self . markov_model = markov_model <EOL> def get_next_states ( self , role ) : <EOL> """<STR_LIT>""" <EOL> allStates = self . markov_model [ self . state ] <EOL> if role is not None : <EOL> ret = [ s for s in allStates if s [ <NUM_LIT:0> ] . endswith ( role ) or <EOL> s [ <NUM_LIT:0> ] . endswith ( "<STR_LIT>" ) ] <EOL> else : <EOL> ret = allStates <EOL> return ret <EOL> def choose_most_probable_next_state ( self , role ) : <EOL> """<STR_LIT>""" <EOL> next_states = self . get_next_states ( role ) <EOL> print "<STR_LIT>" . format ( next_states ) <EOL> sorted_states = [ ( s , c ) for s , c in sorted ( next_states , <EOL> key = lambda ( k , v ) : ( v , k ) ) ] <EOL> return sorted_states [ - <NUM_LIT:1> ] <EOL> def choose_least_probable_next_state ( self , role ) : <EOL> """<STR_LIT>""" <EOL> next_states = self . get_next_states ( role ) <EOL> print "<STR_LIT>" . format ( next_states ) <EOL> sorted_states = [ ( s , c ) for s , c in sorted ( next_states , <EOL> key = lambda ( k , v ) : ( v , k ) ) ] <EOL> state = sorted_states [ <NUM_LIT:0> ] <EOL> print "<STR_LIT>" . format ( state ) <EOL> return state <EOL> def choose_random_next_state ( self , role ) : <EOL> """<STR_LIT>""" <EOL> next_states = self . get_next_states ( role ) <EOL> print "<STR_LIT>" . format ( next_states ) <EOL> state = choice ( next_states ) <EOL> print "<STR_LIT>" . format ( state ) <EOL> return state <EOL> def choose_next_state_OFS ( self , role , lens , transition = OFS_TRANSITION ) : <EOL> """<STR_LIT>""" <EOL> states = { s : <NUM_LIT:0> for s in self . get_next_states ( role ) } <EOL> for s in states : <EOL> states [ s ] += self . get_subtree_weight_OFS ( lens , s , transition ) <EOL> print "<STR_LIT>" <EOL> state = choice ( [ k for k , v in states . items ( ) if v is max ( states . values ( ) ) ] ) <EOL> print "<STR_LIT>" . format ( state ) <EOL> return state <EOL> def get_subtree_weight_OFS ( self , lens , state , transition = OFS_TRANSITION ) : <EOL> """<STR_LIT>""" <EOL> w = self . get_state_weight_OFS ( state [ <NUM_LIT:0> ] , lens ) <EOL> if transition > <NUM_LIT:1> : <EOL> try : <EOL> states = self . markov_model [ state [ <NUM_LIT:0> ] ] <EOL> except KeyError : <EOL> states = [ ] <EOL> for s in states : <EOL> w += self . get_subtree_weight_OFS ( lens , s , transition - <NUM_LIT:1> ) <EOL> if transition == <NUM_LIT:1> : <EOL> return w <EOL> return w <EOL> def get_state_weight_OFS ( self , state , lens ) : <EOL> """<STR_LIT>""" <EOL> w = <NUM_LIT:0> <EOL> try : <EOL> ids = lens . tl . state_id [ state ] <EOL> except KeyError : <EOL> return <NUM_LIT:0> <EOL> for i in ids : <EOL> try : <EOL> w += lens . fuzzer . tracker [ i ] <EOL> except KeyError : <EOL> pass <EOL> return w <EOL> class TemplateList : <EOL> def __init__ ( self , templates_file , lexerStyle = LEXER_TOKENS ) : <EOL> self . template_list = self . load_templates ( templates_file ) <EOL> self . template_dic = { } <EOL> self . state_id = { } <EOL> self . lexerStyle = lexerStyle <EOL> for template in self . template_list : <EOL> self . template_dic [ template . ID ] = template <EOL> if template . state in self . state_id : <EOL> self . state_id [ template . state ] . extend ( [ template . ID ] ) <EOL> else : <EOL> self . state_id [ template . state ] = [ template . ID ] <EOL> def load_templates ( self , templates_file ) : <EOL> """<STR_LIT>""" <EOL> fd = open ( templates_file , "<STR_LIT:r>" ) <EOL> text_templates = fd . read ( ) <EOL> fd . close ( ) <EOL> templates_str = text_templates . split ( "<STR_LIT>" ) <EOL> r = re . compile ( "<STR_LIT>" , re . DOTALL ) <EOL> template_list = [ ] <EOL> for i in range ( <NUM_LIT:1> , len ( templates_str ) ) : <EOL> template_data = r . findall ( templates_str [ i ] ) <EOL> ID , state , count , ntokens , fields , content = template_data [ <NUM_LIT:0> ] <EOL> if content == "<STR_LIT:\n>" and ntokens == "<STR_LIT:1>" : <EOL> content = [ "<STR_LIT>" ] <EOL> else : <EOL> content = content . split ( "<STR_LIT:\n>" ) [ <NUM_LIT:0> : - <NUM_LIT:1> ] <EOL> new_template = Template ( ID , state , count , ntokens , fields , content ) <EOL> template_list . append ( new_template ) <EOL> return template_list <EOL> def getPossibleNextTemplates ( self , next_possible_states ) : <EOL> """<STR_LIT>""" <EOL> ret = [ ] <EOL> for state in next_possible_states : <EOL> try : <EOL> id_list = self . state_id [ state [ <NUM_LIT:0> ] ] <EOL> except KeyError : <EOL> continue <EOL> for id in id_list : <EOL> template = self . template_dic [ id ] <EOL> ret . append ( template ) <EOL> ret . sort ( key = lambda t : t . count , reverse = True ) <EOL> return [ "<STR_LIT>" % ( t . display ( ) , t . count ) for t in ret ] <EOL> def find_template ( self , next_possible_states , message ) : <EOL> """<STR_LIT>""" <EOL> for state in next_possible_states : <EOL> try : <EOL> id_list = self . state_id [ state [ <NUM_LIT:0> ] ] <EOL> except KeyError : <EOL> continue <EOL> for id in id_list : <EOL> template = self . template_dic [ id ] <EOL> matcher = TemplateMessageMatcher ( template , message , <EOL> self . lexerStyle ) <EOL> fields = matcher . match ( ) <EOL> if fields is not None : <EOL> return ( template , fields ) <EOL> return ( None , None ) <EOL> def find_similar_template ( self , next_possible_states , message ) : <EOL> """<STR_LIT>""" <EOL> for state in next_possible_states : <EOL> try : <EOL> id_list = self . state_id [ state [ <NUM_LIT:0> ] ] <EOL> except KeyError : <EOL> continue <EOL> templates = [ ] <EOL> for id in id_list : <EOL> template = self . template_dic [ id ] <EOL> matcher = SimilarTemplateMessageMatcher ( template , message , <EOL> self . lexerStyle ) <EOL> fields , dist = matcher . match ( ) <EOL> templates += [ ( template , fields , distance ) ] <EOL> sorted_templates = [ ( t , f , d ) for t , f , d in sorted ( templates , <EOL> key = lambda ( t , f , d ) : ( d , t , f ) ) ] <EOL> template , fields = sorted_templates [ <NUM_LIT:0> ] [ : <NUM_LIT:2> ] <EOL> return ( template , fields ) <EOL> def find_next_template ( self , state , rl , previousTemplates , interactive ) : <EOL> """<STR_LIT>""" <EOL> try : <EOL> id_list = self . state_id [ state [ <NUM_LIT:0> ] ] <EOL> except KeyError : <EOL> return None <EOL> templates_with_rules = [ ] <EOL> for template_id in id_list : <EOL> new_template = self . template_dic [ template_id ] <EOL> if ( not new_template . has_fields ( ) <EOL> or rl . transition_exists ( previousTemplates , new_template ) ) : <EOL> templates_with_rules . append ( new_template ) <EOL> if len ( templates_with_rules ) > <NUM_LIT:0> : <EOL> fuzz_templates = [ t for t in templates_with_rules if t . has_fields ( ) ] <EOL> if fuzz_templates : <EOL> templates = fuzz_templates <EOL> else : <EOL> templates = templates_with_rules <EOL> templates_ids = [ t . ID for t in templates ] <EOL> if interactive : <EOL> choice = "<STR_LIT>" <EOL> while choice not in templates_ids : <EOL> print "<STR_LIT>" . format ( templates_ids ) <EOL> if len ( templates_ids ) == <NUM_LIT:1> : <EOL> choice = templates_ids [ <NUM_LIT:0> ] <EOL> else : <EOL> choice = raw_input ( '<STR_LIT>' ) <EOL> ind = templates_ids . index ( choice ) <EOL> else : <EOL> print "<STR_LIT>" . format ( templates_ids ) <EOL> ind = pickIndex ( [ t . count for t in templates ] ) <EOL> return templates [ ind ] <EOL> else : <EOL> return None <EOL> def find_next_template_no_rules ( self , state , templates_no_fields , interactive ) : <EOL> try : <EOL> id_list = self . state_id [ state [ <NUM_LIT:0> ] ] <EOL> except KeyError : <EOL> return None <EOL> templates = [ ] <EOL> fuzz_templates = [ ] <EOL> for template_id in id_list : <EOL> templates += [ self . template_dic [ template_id ] ] <EOL> for t in templates : <EOL> if t . has_fields ( ) : <EOL> fuzz_templates += [ t ] <EOL> if not fuzz_templates or templates_no_fields : <EOL> fuzz_templates = templates <EOL> fuzz_templates_ids = [ t . ID for t in fuzz_templates ] <EOL> if interactive : <EOL> choice = "<STR_LIT>" <EOL> while choice not in fuzz_templates_ids : <EOL> print "<STR_LIT>" . format ( fuzz_templates_ids ) <EOL> if len ( fuzz_templates_ids ) == <NUM_LIT:1> : <EOL> choice = fuzz_templates_ids [ <NUM_LIT:0> ] <EOL> else : <EOL> choice = raw_input ( '<STR_LIT>' ) <EOL> ind = fuzz_templates_ids . index ( choice ) <EOL> else : <EOL> print "<STR_LIT>" . format ( fuzz_templates_ids ) <EOL> ind = pickIndex ( [ t . count for t in fuzz_templates ] ) <EOL> return fuzz_templates [ ind ] <EOL> class TemplateMessageMatcher : <EOL> def __init__ ( self , template , msg , lexerType = LEXER_TOKENS , <EOL> whitespace = "<STR_LIT>" ) : <EOL> self . template = template <EOL> self . msg = msg <EOL> self . lexerType = lexerType <EOL> self . ws = whitespace <EOL> def match ( self ) : <EOL> if self . lexerType == LEXER_TOKENS : <EOL> tokens = scanTokens ( self . msg , self . ws ) <EOL> else : <EOL> tokens = scanNgrams ( self . msg ) <EOL> if len ( tokens ) != len ( self . template . content ) : <EOL> return None <EOL> fields = [ ] <EOL> match = True <EOL> for ( tok , msg ) in zip ( self . template . content , tokens ) : <EOL> curVal = msg [ <NUM_LIT:1> ] <EOL> if tok == "<STR_LIT>" : <EOL> fields . append ( curVal ) <EOL> else : <EOL> if urllib . quote ( curVal ) != tok : <EOL> match = False <EOL> break <EOL> if match : <EOL> ret = fields <EOL> else : <EOL> ret = None <EOL> return ret <EOL> class SimilarTemplateMessageMatcher : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , template , msg , lexerType = LEXER_TOKENS , <EOL> whitespace = "<STR_LIT>" ) : <EOL> self . template = template <EOL> self . msg = msg <EOL> self . lexerType = lexerType <EOL> self . ws = whitespace <EOL> def match ( self ) : <EOL> if self . lexerType == LEXER_TOKENS : <EOL> tokens = scanTokens ( self . msg , self . ws ) <EOL> else : <EOL> tokens = scanNgrams ( self . msg ) <EOL> fields = [ ] <EOL> for ( tok , msg ) in zip ( self . template . content , tokens ) : <EOL> curVal = msg [ <NUM_LIT:1> ] <EOL> if tok == "<STR_LIT>" : <EOL> fields . append ( curVal ) <EOL> template_str = '<STR_LIT>' . join ( [ t for t in self . template . content ] ) <EOL> d = distance ( self . msg , template_str ) <EOL> return fields , d <EOL> class Template : <EOL> def __init__ ( self , ID = "<STR_LIT>" , state = "<STR_LIT>" , count = "<STR_LIT:0>" , ntokens = "<STR_LIT:0>" , <EOL> fields = "<STR_LIT>" , content = [ ] ) : <EOL> """<STR_LIT>""" <EOL> self . ID = ID <EOL> self . state = state <EOL> self . count = int ( count ) <EOL> self . ntokens = int ( ntokens ) <EOL> if fields == "<STR_LIT>" : <EOL> self . fields = [ ] <EOL> else : <EOL> self . fields = [ int ( f ) for f in fields . split ( "<STR_LIT:U+002C>" ) ] <EOL> self . content = content <EOL> def display ( self ) : <EOL> return str ( self ) . strip ( ) <EOL> def __str__ ( self ) : <EOL> tokens = [ urllib . unquote ( t ) if t != "<STR_LIT>" else "<STR_LIT>" for t in self . content ] <EOL> return "<STR_LIT>" . join ( tokens ) <EOL> def fillFields ( self , fields ) : <EOL> """<STR_LIT>""" <EOL> fuzz = None <EOL> if self . ID in self . fuzzer . tracker : <EOL> print self . fuzzer . tracker <EOL> if self . fuzzer . interactive : <EOL> while fuzz is None : <EOL> inp = raw_input ( '<STR_LIT>' . format ( self . ID ) ) <EOL> if inp == '<STR_LIT:y>' : <EOL> fuzz = True <EOL> elif inp == '<STR_LIT:n>' : <EOL> fuzz = False <EOL> else : <EOL> fuzz = True <EOL> if fuzz : <EOL> self . new_fields = fields <EOL> ret = self . fuzzer . fuzz_template ( self ) <EOL> if not fuzz : <EOL> ret = self . content [ : ] <EOL> for ( k , v ) in fields . iteritems ( ) : <EOL> try : <EOL> ret [ self . fields [ k ] ] = v <EOL> except IndexError : <EOL> print "<STR_LIT>" <EOL> pass <EOL> return ret <EOL> def fillFields ( self , fields ) : <EOL> """<STR_LIT>""" <EOL> ret = self . content [ : ] <EOL> for ( k , v ) in fields . iteritems ( ) : <EOL> try : <EOL> ret [ self . fields [ k ] ] = v <EOL> except IndexError : <EOL> print "<STR_LIT>" <EOL> pass <EOL> return ret <EOL> def has_fields ( self ) : <EOL> return len ( self . fields ) is not <NUM_LIT:0> <EOL> class RuleList : <EOL> def __init__ ( self , rules_file ) : <EOL> self . load_rules ( rules_file ) <EOL> def replaceData ( self , transition , dst_field , newData ) : <EOL> for r in self . rules [ transition ] : <EOL> r . replaceData ( transition , dst_field , newData ) <EOL> def load_rules ( self , rules_file ) : <EOL> """<STR_LIT>""" <EOL> fd = open ( rules_file , "<STR_LIT:r>" ) <EOL> text_rules = fd . read ( ) <EOL> fd . close ( ) <EOL> rules_str = text_rules . split ( "<STR_LIT>" ) <EOL> r = re . compile ( "<STR_LIT>" , re . DOTALL ) <EOL> self . rules = { } <EOL> for i in range ( <NUM_LIT:1> , len ( rules_str ) ) : <EOL> rule_data = r . findall ( rules_str [ i ] ) <EOL> transition , src_id , src_field , dst_field , typ , data = rule_data [ <NUM_LIT:0> ] <EOL> new_rule = Rule ( transition , src_id , src_field , <EOL> dst_field , typ , data ) <EOL> self . rules . setdefault ( transition , [ ] ) . append ( new_rule ) <EOL> def create_message ( self , previousTemplates , previousFields , next_template ) : <EOL> """<STR_LIT>""" <EOL> transition = self . generateTransitionLabel ( previousTemplates , <EOL> next_template ) <EOL> fields = dict ( zip ( [ str ( ind ) for ind in range ( - len ( previousFields ) - <NUM_LIT:1> , <EOL> - <NUM_LIT:1> ) ] , previousFields ) ) <EOL> if transition not in self . rules : <EOL> if not next_template . has_fields ( ) : <EOL> current_rules = [ ] <EOL> else : <EOL> return ( None , None , None ) <EOL> else : <EOL> current_rules = self . rules [ transition ] <EOL> dst_fields = { } <EOL> fields [ "<STR_LIT>" ] = dst_fields <EOL> for rule in current_rules : <EOL> src_id = rule . src_id <EOL> dst = rule . dst_field <EOL> src = rule . src_field <EOL> if rule . typ == "<STR_LIT>" : <EOL> dst_fields [ dst ] = fields [ src_id ] [ src ] <EOL> continue <EOL> if rule . typ == "<STR_LIT>" : <EOL> dst_fields [ dst ] = random . choice ( rule . data [ <NUM_LIT:5> : ] . split ( "<STR_LIT:U+002C>" ) ) <EOL> continue <EOL> if rule . typ == "<STR_LIT>" : <EOL> if rule . data [ <NUM_LIT:6> : <NUM_LIT:20> ] == "<STR_LIT>" : <EOL> dst_fields [ dst ] = fields [ src_id ] [ src ] + random . choice ( rule . data [ <NUM_LIT> : ] . split ( "<STR_LIT:U+002C>" ) ) <EOL> else : <EOL> dst_fields [ dst ] = random . choice ( rule . data [ <NUM_LIT> : ] . split ( "<STR_LIT:U+002C>" ) ) + fields [ src_id ] [ src ] <EOL> continue <EOL> if rule . typ == "<STR_LIT>" : <EOL> split_data = fields [ src_id ] [ src ] . split ( rule . data [ <NUM_LIT> : ] , <NUM_LIT:1> ) <EOL> if rule . data [ <NUM_LIT:6> : <NUM_LIT:20> ] == "<STR_LIT>" : <EOL> if len ( split_data ) == <NUM_LIT:2> : <EOL> dst_fields [ dst ] = split_data [ <NUM_LIT:0> ] <EOL> else : <EOL> dst_fields [ dst ] = "<STR_LIT>" <EOL> else : <EOL> if len ( split_data ) == <NUM_LIT:2> : <EOL> dst_fields [ dst ] = split_data [ <NUM_LIT:1> ] <EOL> else : <EOL> dst_fields [ dst ] = "<STR_LIT>" <EOL> continue <EOL> if rule . typ == "<STR_LIT>" : <EOL> try : <EOL> dst_fields [ dst ] = int ( fields [ src_id ] [ src ] ) + int ( rule . data [ <NUM_LIT:5> : ] ) <EOL> except : <EOL> dst_fields [ dst ] = <NUM_LIT:0> <EOL> continue <EOL> message = next_template . fillFields ( dst_fields ) <EOL> fields = dst_fields . items ( ) <EOL> fields . sort ( key = operator . itemgetter ( <NUM_LIT:0> ) ) <EOL> msg = urllib . unquote ( '<STR_LIT>' . join ( str ( n ) for n in message ) ) <EOL> return ( msg , [ f [ <NUM_LIT:1> ] for f in fields ] , transition ) <EOL> def create_fuzzed_message ( self , previousTemplates , previousFields , <EOL> next_template , fuzzer ) : <EOL> """<STR_LIT>""" <EOL> try : <EOL> transition = self . generateTransitionLabel ( previousTemplates , <EOL> next_template ) <EOL> except : <EOL> transition = None <EOL> fields = dict ( zip ( [ str ( ind ) for ind in range ( - len ( previousFields ) - <NUM_LIT:1> , <EOL> - <NUM_LIT:1> ) ] , previousFields ) ) <EOL> template_rules = [ r for t , r in self . rules . items ( ) if t . endswith ( next_template . ID ) ] <EOL> rules = [ item for sublist in template_rules for item in sublist ] <EOL> dst_fields = { } <EOL> fields [ "<STR_LIT>" ] = dst_fields <EOL> fuzz_fields = fuzzer . get_fuzz_fields ( next_template . ID ) <EOL> non_fuzzed_values = [ ] <EOL> for rule in rules : <EOL> src_id = rule . src_id <EOL> dst = rule . dst_field <EOL> src = rule . src_field <EOL> if rule . typ == "<STR_LIT>" : <EOL> try : <EOL> data = fields [ src_id ] [ src ] <EOL> except : <EOL> data = "<STR_LIT>" <EOL> dst_fields [ dst ] = fuzzer . fuzz ( rule , data , fuzz_fields ) <EOL> continue <EOL> if rule . typ == "<STR_LIT>" : <EOL> try : <EOL> data = random . choice ( rule . data [ <NUM_LIT:5> : ] . split ( "<STR_LIT:U+002C>" ) ) <EOL> except : <EOL> data = "<STR_LIT>" <EOL> dst_fields [ dst ] = fuzzer . fuzz ( rule , data , fuzz_fields ) <EOL> continue <EOL> if rule . typ == "<STR_LIT>" : <EOL> try : <EOL> if rule . data [ <NUM_LIT:6> : <NUM_LIT:20> ] == "<STR_LIT>" : <EOL> data = fields [ src_id ] [ src ] + random . choice ( rule . data [ <NUM_LIT> : ] . split ( "<STR_LIT:U+002C>" ) ) <EOL> else : <EOL> data = random . choice ( rule . data [ <NUM_LIT> : ] . split ( "<STR_LIT:U+002C>" ) ) + fields [ src_id ] [ src ] <EOL> except : <EOL> data = "<STR_LIT>" <EOL> dst_fields [ dst ] = fuzzer . fuzz ( rule , data , fuzz_fields ) <EOL> continue <EOL> if rule . typ == "<STR_LIT>" : <EOL> try : <EOL> split_data = fields [ src_id ] [ src ] . split ( rule . data [ <NUM_LIT> : ] , <NUM_LIT:1> ) <EOL> if rule . data [ <NUM_LIT:6> : <NUM_LIT:20> ] == "<STR_LIT>" : <EOL> if len ( split_data ) == <NUM_LIT:2> : <EOL> data = split_data [ <NUM_LIT:0> ] <EOL> else : <EOL> data = "<STR_LIT>" <EOL> else : <EOL> if len ( split_data ) == <NUM_LIT:2> : <EOL> data = split_data [ <NUM_LIT:1> ] <EOL> else : <EOL> data = "<STR_LIT>" <EOL> except : <EOL> data = "<STR_LIT>" <EOL> dst_fields [ dst ] = fuzzer . fuzz ( rule , data , fuzz_fields ) <EOL> continue <EOL> if rule . typ == "<STR_LIT>" : <EOL> try : <EOL> data = int ( fields [ src_id ] [ src ] ) + int ( rule . data [ <NUM_LIT:5> : ] ) <EOL> except : <EOL> data = <NUM_LIT:0> <EOL> dst_fields [ dst ] = fuzzer . fuzz ( rule , data , fuzz_fields ) <EOL> continue <EOL> if dst_fields [ dst ] == data : <EOL> non_fuzzed_values += [ data ] <EOL> message = next_template . fillFields ( dst_fields ) <EOL> fields = dst_fields . items ( ) <EOL> fields . sort ( key = operator . itemgetter ( <NUM_LIT:0> ) ) <EOL> fuzzer . update_tracker ( next_template . ID ) <EOL> fuzzer . trace = list ( fuzz_fields ) , non_fuzzed_values <EOL> print "<STR_LIT>" , fuzz_fields <EOL> print "<STR_LIT>" <EOL> return ( urllib . unquote ( '<STR_LIT>' . join ( str ( n ) for n in message ) ) , <EOL> [ f [ <NUM_LIT:1> ] for f in fields ] , transition ) <EOL> def generateTransitionLabel ( self , previousTemplates , nextTemplate ) : <EOL> if len ( previousTemplates ) == <NUM_LIT:0> : <EOL> transition = nextTemplate . ID <EOL> else : <EOL> transition = "<STR_LIT>" % ( "<STR_LIT:;>" . join ( [ p . ID for p in previousTemplates ] ) , nextTemplate . ID ) <EOL> return transition <EOL> def transition_exists ( self , previousTemplates , nextTemplate ) : <EOL> """<STR_LIT>""" <EOL> transition = self . generateTransitionLabel ( previousTemplates , <EOL> nextTemplate ) <EOL> return transition in self . rules <EOL> class Rule : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , transition , src_id , src_field , dst_field , typ , data ) : <EOL> self . transition = transition <EOL> self . src_id = src_id <EOL> self . src_field = int ( src_field ) <EOL> self . dst_field = int ( dst_field ) <EOL> self . typ = typ <EOL> self . data = data . rstrip ( ) <EOL> def replaceData ( self , transition , dst_field , newData ) : <EOL> if ( self . transition == transition <EOL> and self . dst_field == dst_field <EOL> and self . src_id == "<STR_LIT:0>" <EOL> and self . src_field == - <NUM_LIT:1> <EOL> and self . typ == "<STR_LIT>" ) : <EOL> self . data = "<STR_LIT>" % newData <EOL> class RingBuffer : <EOL> def __init__ ( self , size_max ) : <EOL> self . max = size_max <EOL> self . data = [ ] <EOL> def getLastEntries ( self , n ) : <EOL> if n == <NUM_LIT:0> : <EOL> return [ ] <EOL> else : <EOL> return self . data [ - n : ] <EOL> def append ( self , x ) : <EOL> """<STR_LIT>""" <EOL> self . data . append ( x ) <EOL> if len ( self . data ) == self . max : <EOL> self . cur = <NUM_LIT:0> <EOL> self . __class__ = RingBuffer . __Full <EOL> def get ( self ) : <EOL> """<STR_LIT>""" <EOL> return self . data <EOL> class __Full : <EOL> def __init__ ( self , n ) : <EOL> raise "<STR_LIT>" <EOL> def append ( self , x ) : <EOL> self . data [ self . cur ] = x <EOL> self . cur = ( self . cur + <NUM_LIT:1> ) % self . max <EOL> def get ( self ) : <EOL> return self . data [ self . cur : ] + self . data [ : self . cur ] <EOL> def getLastEntries ( self , n ) : <EOL> if n == <NUM_LIT:0> : <EOL> return [ ] <EOL> else : <EOL> if self . cur >= n : <EOL> return self . data [ ( self . cur - n ) : self . cur ] <EOL> else : <EOL> return self . get ( ) [ - n : ] <EOL> STATUS_OK = "<STR_LIT:OK>" <EOL> STATUS_NO_TRANSITION = "<STR_LIT>" <EOL> STATUS_END = "<STR_LIT>" <EOL> MODE_MOST_PROBABLE_TRANSITION = <NUM_LIT:0> <EOL> MODE_LEAST_PROBABLE_TRANSITION = <NUM_LIT:1> <EOL> MODE_RANDOM_TRANSITION = <NUM_LIT:2> <EOL> MODE_FUZZING_TRANSITION = <NUM_LIT:3> <EOL> class Lens : <EOL> def __init__ ( self , modelPath , role , sim_search = False , mode = MODE_RANDOM_TRANSITION , <EOL> lexerStyle = LEXER_TOKENS , templates_no_fields = False ) : <EOL> markov_model , rule_list , template_list = ( "<STR_LIT>" % modelPath , <EOL> "<STR_LIT>" % modelPath , <EOL> "<STR_LIT>" % modelPath ) <EOL> self . modelPath = modelPath <EOL> self . mm = MarkovModel ( markov_model ) <EOL> self . rl = RuleList ( rule_list ) <EOL> self . tl = TemplateList ( template_list , lexerStyle ) <EOL> self . mode = int ( mode ) <EOL> self . fuzzer = "<STR_LIT>" <EOL> self . role = role <EOL> if role == "<STR_LIT>" : <EOL> self . otherside = "<STR_LIT>" <EOL> else : <EOL> self . otherside = "<STR_LIT>" <EOL> self . templates = RingBuffer ( <NUM_LIT:20> ) <EOL> self . fields = RingBuffer ( <NUM_LIT:20> ) <EOL> self . reset_model ( ) <EOL> self . similarity_search = sim_search <EOL> self . sim_matched_template = False <EOL> self . templates_no_fields = templates_no_fields <EOL> def set_fuzzer ( self , fuzzer ) : <EOL> """<STR_LIT>""" <EOL> self . fuzzer = fuzzer <EOL> self . mm . set_fuzzer ( fuzzer ) <EOL> def getLastFields ( self ) : <EOL> ret = self . fields . getLastEntries ( <NUM_LIT:1> ) <EOL> if len ( ret ) == <NUM_LIT:1> : <EOL> return ret [ <NUM_LIT:0> ] <EOL> else : <EOL> return [ ] <EOL> def replaceData ( self , transition , dst_field , newData ) : <EOL> self . rl . replaceData ( transition , dst_field , newData ) <EOL> def reset_model ( self ) : <EOL> self . mm . init ( ) <EOL> for o in range ( self . mm . getOrder ( ) ) : <EOL> self . templates . append ( Template ( ) ) <EOL> self . fields . append ( [ ] ) <EOL> def consumeOtherSide ( self , message ) : <EOL> """<STR_LIT>""" <EOL> return self . consume ( message , self . otherside ) <EOL> def transitionSelf ( self ) : <EOL> """<STR_LIT>""" <EOL> return self . transition ( self . role , True ) <EOL> def consume ( self , message , role ) : <EOL> """<STR_LIT>""" <EOL> status = STATUS_NO_TRANSITION <EOL> if not ( self . mm . getState ( ) . endswith ( "<STR_LIT>" ) ) : <EOL> next_states = self . mm . get_next_states ( role ) <EOL> if next_states : <EOL> ( template , fields ) = self . tl . find_template ( next_states , message ) <EOL> else : <EOL> if role == self . role : <EOL> other_role = self . otherside <EOL> else : <EOL> other_role = self . role <EOL> next_states = self . mm . get_next_states ( other_role ) <EOL> template = None <EOL> if self . similarity_search : <EOL> if template is not None : <EOL> self . sim_matched_template = False <EOL> print "<STR_LIT>" . format ( template . ID ) <EOL> elif template is None : <EOL> try : <EOL> next_OFS_states = [ ] <EOL> for state in self . mm . get_next_states ( role ) : <EOL> if self . mm . get_subtree_weight_OFS ( self , <EOL> state ) > <NUM_LIT:0> : <EOL> next_OFS_states += [ state ] <EOL> if not next_OFS_states : <EOL> next_OFS_states = self . mm . get_next_states ( role ) <EOL> ( template , fields ) = self . tl . find_similar_template ( next_OFS_states , <EOL> message ) <EOL> except TypeError : <EOL> ( template , fields ) = ( None , None ) <EOL> status = STATUS_END <EOL> if template is not None : <EOL> self . sim_matched_template = True <EOL> print "<STR_LIT>" . format ( template . ID ) <EOL> if template is not None : <EOL> self . templates . append ( template ) <EOL> self . fields . append ( fields ) <EOL> self . mm . update_state ( template . state ) <EOL> print "<STR_LIT>" . format ( self . mm . state ) <EOL> if template . state . endswith ( "<STR_LIT>" ) : <EOL> status = STATUS_END <EOL> else : <EOL> status = STATUS_OK <EOL> else : <EOL> if status == STATUS_END : <EOL> print "<STR_LIT>" <EOL> else : <EOL> print "<STR_LIT>" <EOL> else : <EOL> status = STATUS_END <EOL> return status <EOL> def transition ( self , role , chooseNextStateIndependentFromRole = False ) : <EOL> """<STR_LIT>""" <EOL> status = STATUS_NO_TRANSITION <EOL> if chooseNextStateIndependentFromRole : <EOL> roleForNextState = None <EOL> else : <EOL> roleForNextState = role <EOL> if self . mode == MODE_MOST_PROBABLE_TRANSITION : <EOL> next_state = self . mm . choose_most_probable_next_state ( roleForNextState ) <EOL> elif self . mode == MODE_LEAST_PROBABLE_TRANSITION : <EOL> next_state = self . mm . choose_least_probable_next_state ( roleForNextState ) <EOL> elif self . mode == MODE_RANDOM_TRANSITION : <EOL> next_state = self . mm . choose_random_next_state ( roleForNextState ) <EOL> elif self . mode == MODE_FUZZING_TRANSITION : <EOL> next_state = self . mm . choose_next_state_OFS ( roleForNextState , self ) <EOL> else : <EOL> raise "<STR_LIT>" % self . mode <EOL> if not next_state [ <NUM_LIT:0> ] . endswith ( "<STR_LIT>" ) : <EOL> if chooseNextStateIndependentFromRole and not ( next_state [ <NUM_LIT:0> ] . endswith ( role ) ) : <EOL> next_state = None <EOL> ( state , template , msg , fields , transition ) = self . _transition_to_state ( next_state ) <EOL> if state is not None and state . endswith ( "<STR_LIT>" ) : <EOL> status = STATUS_END <EOL> if msg is not None : <EOL> self . templates . append ( template ) <EOL> print "<STR_LIT>" . format ( template . ID ) <EOL> self . fields . append ( fields ) <EOL> self . mm . update_state ( state ) <EOL> print "<STR_LIT>" . format ( self . mm . state ) <EOL> status = STATUS_OK <EOL> return ( status , msg , transition ) <EOL> def _transition_to_state ( self , next_state ) : <EOL> if next_state is not None : <EOL> for theOrder in range ( self . mm . getOrder ( ) , - <NUM_LIT:1> , - <NUM_LIT:1> ) : <EOL> previousTemplates = self . templates . getLastEntries ( theOrder ) <EOL> previousFields = self . fields . getLastEntries ( theOrder ) <EOL> if not self . sim_matched_template : <EOL> next_template = self . tl . find_next_template ( next_state , <EOL> self . rl , <EOL> previousTemplates , <EOL> self . fuzzer . interactive ) <EOL> else : <EOL> print "<STR_LIT>" <EOL> next_template = self . tl . find_next_template_no_rules ( next_state , <EOL> self . templates_no_fields , <EOL> self . fuzzer . interactive ) <EOL> if next_template is not None : <EOL> break <EOL> if next_template is not None : <EOL> if self . sim_matched_template or self . do_fuzz ( next_template ) : <EOL> ( message , fields , <EOL> transition ) = self . rl . create_fuzzed_message ( previousTemplates , <EOL> previousFields , <EOL> next_template , <EOL> self . fuzzer ) <EOL> else : <EOL> ( message , fields , <EOL> transition ) = self . rl . create_message ( previousTemplates , <EOL> previousFields , <EOL> next_template ) <EOL> else : <EOL> message = None <EOL> fields = None <EOL> transition = None <EOL> ret = ( next_state [ <NUM_LIT:0> ] , next_template , message , fields , transition ) <EOL> else : <EOL> ret = ( None , None , None , None , None ) <EOL> return ret <EOL> def waiting_response ( self ) : <EOL> return not bool ( self . mm . get_next_states ( self . role ) ) <EOL> def getPossibleNextTemplates ( self ) : <EOL> return self . tl . getPossibleNextTemplates ( self . mm . get_next_states ( self . otherside ) ) <EOL> def do_fuzz ( self , template ) : <EOL> fuzz = None <EOL> if template . ID in self . fuzzer . tracker : <EOL> if self . fuzzer . interactive : <EOL> while fuzz is None : <EOL> inp = raw_input ( '<STR_LIT>' . format ( template . ID ) ) <EOL> if inp == '<STR_LIT:y>' : <EOL> fuzz = True <EOL> elif inp == '<STR_LIT:n>' : <EOL> fuzz = False <EOL> return fuzz <EOL> else : <EOL> return True <EOL> else : <EOL> return False </s>
