<s> import os <EOL> import shutil <EOL> import ConfigParser <EOL> import urllib <EOL> import data <EOL> import rule <EOL> import DFA <EOL> from ast import literal_eval <EOL> from pulsar . core . harry import Harry <EOL> from pulsar . utils . preprocessing import derrick <EOL> class ModelGenerator ( ) : <EOL> def __init__ ( self , models_dir , path_conf , <EOL> pcaps = "<STR_LIT>" , binaries = "<STR_LIT>" , <EOL> nmf_components = <NUM_LIT:0> , process = False ) : <EOL> self . models_dir = models_dir <EOL> self . path_conf = path_conf <EOL> self . pcaps = pcaps <EOL> self . binaries = binaries <EOL> self . process = process <EOL> self . nmf_components = nmf_components <EOL> config = ConfigParser . RawConfigParser ( ) <EOL> model_conf = os . path . join ( path_conf , "<STR_LIT>" ) <EOL> config . readfp ( open ( model_conf ) ) <EOL> self . prisma_dir = config . get ( '<STR_LIT>' , '<STR_LIT:path>' ) <EOL> self . storage_dir = config . get ( '<STR_LIT:data>' , '<STR_LIT>' ) <EOL> self . binaries_dir = os . path . abspath ( os . path . join ( self . storage_dir , <EOL> '<STR_LIT>' ) ) <EOL> self . analyses_dir = os . path . abspath ( os . path . join ( self . storage_dir , <EOL> '<STR_LIT>' ) ) <EOL> self . analyses_ival = literal_eval ( config . get ( '<STR_LIT:data>' , <EOL> '<STR_LIT>' ) ) <EOL> if self . analyses_ival : <EOL> self . analyses_ival = range ( self . analyses_ival [ <NUM_LIT:0> ] , <EOL> self . analyses_ival [ - <NUM_LIT:1> ] + <NUM_LIT:1> ) <EOL> self . whitespace = config . get ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> self . order = literal_eval ( config . get ( '<STR_LIT>' , '<STR_LIT>' ) ) <EOL> self . ngram = literal_eval ( config . get ( '<STR_LIT>' , '<STR_LIT>' ) ) <EOL> self . dograph = literal_eval ( config . get ( '<STR_LIT>' , '<STR_LIT>' ) ) <EOL> self . minimize = literal_eval ( config . get ( '<STR_LIT>' , '<STR_LIT>' ) ) <EOL> self . subrules = literal_eval ( config . get ( '<STR_LIT>' , '<STR_LIT>' ) ) <EOL> self . colorFieldStates = literal_eval ( config . get ( '<STR_LIT>' , <EOL> '<STR_LIT>' ) ) <EOL> def generate_model ( self ) : <EOL> if self . pcaps : <EOL> self . _generate_model_pcaps ( ) <EOL> elif isinstance ( self . binaries , list ) : <EOL> self . _generate_model_binaries ( ) <EOL> def _generate_model_pcaps ( self ) : <EOL> for pcap_origin in self . pcaps : <EOL> pcap = os . path . basename ( pcap_origin ) <EOL> pcap_noext = pcap . split ( '<STR_LIT:.>' ) [ <NUM_LIT:0> ] <EOL> pcap_named_dir = os . path . join ( self . models_dir , pcap_noext ) <EOL> if not os . path . exists ( pcap_named_dir ) : <EOL> print "<STR_LIT>" . format ( pcap_named_dir ) <EOL> os . makedirs ( pcap_named_dir ) <EOL> pcap_path_model = os . path . join ( pcap_named_dir , pcap ) <EOL> shutil . copyfile ( pcap_origin , pcap_path_model ) <EOL> self . _extract_derrick ( pcap_noext ) <EOL> self . _generate_prisma_input ( pcap_noext ) <EOL> self . _generate_clusters ( pcap_noext ) <EOL> self . _build_model_files ( pcap_noext ) <EOL> def _generate_model_binaries ( self ) : <EOL> if len ( self . binaries ) == <NUM_LIT:0> : <EOL> self . binaries = os . listdir ( self . binaries_dir ) <EOL> pcaps = { b : [ ] for b in self . binaries } <EOL> if not self . analyses_ival : <EOL> for path , dirs , files in os . walk ( os . path . abspath ( self . analyses_dir ) ) : <EOL> for file in files : <EOL> if file . endswith ( "<STR_LIT>" ) : <EOL> binary_link = os . path . join ( path , '<STR_LIT>' ) <EOL> binary_name = os . path . realpath ( binary_link ) . split ( '<STR_LIT:/>' ) [ - <NUM_LIT:1> ] <EOL> if binary_name in self . binaries : <EOL> pcap_path = os . path . join ( path , file ) <EOL> pcaps [ binary_name ] += [ pcap_path ] <EOL> else : <EOL> print "<STR_LIT>" . format ( self . analyses_ival ) <EOL> for d in self . analyses_ival : <EOL> path = os . path . join ( self . analyses_dir , str ( d ) ) <EOL> files = os . listdir ( os . path . abspath ( path ) ) <EOL> for f in files : <EOL> if f . endswith ( "<STR_LIT>" ) : <EOL> binary_link = os . path . join ( path , '<STR_LIT>' ) <EOL> binary_name = os . path . realpath ( binary_link ) . split ( '<STR_LIT:/>' ) [ - <NUM_LIT:1> ] <EOL> if binary_name in self . binaries : <EOL> pcap_path = os . path . join ( path , f ) <EOL> pcaps [ binary_name ] += [ pcap_path ] <EOL> for bin_name , pcap_list in pcaps . items ( ) : <EOL> self . _merge_pcaps ( bin_name , pcap_list ) <EOL> self . _extract_derrick ( bin_name ) <EOL> self . _generate_prisma_input ( bin_name ) <EOL> self . _generate_clusters ( bin_name ) <EOL> self . _build_model_files ( bin_name ) <EOL> def _merge_pcaps ( self , bin_name , pcaps ) : <EOL> """<STR_LIT>""" <EOL> bin_named_dir = os . path . join ( self . models_dir , bin_name ) <EOL> if not os . path . exists ( bin_named_dir ) : <EOL> os . makedirs ( bin_named_dir ) <EOL> bin_path_storage = os . path . join ( self . binaries_dir , bin_name ) <EOL> bin_path_model = os . path . join ( bin_named_dir , bin_name ) <EOL> shutil . copyfile ( bin_path_storage , bin_path_model ) <EOL> merged_pcap_file = os . path . join ( bin_named_dir , <EOL> "<STR_LIT>" . format ( bin_name ) ) <EOL> print "<STR_LIT>" . format ( bin_name ) <EOL> cmd = "<STR_LIT>" . format ( '<STR_LIT:U+0020>' . join ( pcaps ) , merged_pcap_file ) <EOL> os . system ( cmd ) <EOL> def _extract_derrick ( self , bin_name ) : <EOL> """<STR_LIT>""" <EOL> pcap_file = os . path . join ( self . models_dir , bin_name , <EOL> "<STR_LIT>" . format ( bin_name ) ) <EOL> drk_file = os . path . join ( self . models_dir , bin_name , <EOL> "<STR_LIT>" . format ( bin_name ) ) <EOL> print "<STR_LIT>" . format ( bin_name ) <EOL> cmd = "<STR_LIT>" . format ( pcap_file , drk_file ) <EOL> os . system ( cmd ) <EOL> if self . process : <EOL> derrick . process ( drk_file ) <EOL> def _generate_prisma_input ( self , bin_name ) : <EOL> """<STR_LIT>""" <EOL> drk_file = os . path . join ( self . models_dir , bin_name , <EOL> "<STR_LIT>" . format ( bin_name ) ) <EOL> h = Harry ( drk_file , self . path_conf ) <EOL> print "<STR_LIT>" . format ( bin_name ) <EOL> h . generate_prisma_input ( drk_file ) <EOL> def _generate_clusters ( self , bin_name ) : <EOL> """<STR_LIT>""" <EOL> clusters_file = os . path . join ( self . models_dir , bin_name , <EOL> "<STR_LIT>" . format ( bin_name ) ) <EOL> data_dir = os . path . join ( self . models_dir , bin_name , bin_name ) <EOL> core_dir = os . path . dirname ( os . path . abspath ( __file__ ) ) <EOL> cluster_generator_path = os . path . join ( core_dir , <EOL> "<STR_LIT>" ) <EOL> print "<STR_LIT>" <EOL> cmd = "<STR_LIT>" . format ( self . prisma_dir , <EOL> data_dir , <EOL> clusters_file , <EOL> self . nmf_components , <EOL> cluster_generator_path ) <EOL> os . system ( cmd ) <EOL> def _build_model_files ( self , bin_name ) : <EOL> """<STR_LIT>""" <EOL> data_dir = os . path . join ( self . models_dir , bin_name , bin_name ) <EOL> m = ModelFilesGenerator ( data_dir , self . order , self . ngram , <EOL> urllib . unquote ( self . whitespace ) , self . subrules ) <EOL> m . saveModel ( ) <EOL> if self . dograph : <EOL> import pygraphviz <EOL> if self . minimize : <EOL> graph = m . minimizeModel ( ) <EOL> else : <EOL> graph = m . mc <EOL> if self . colorFieldStates : <EOL> m . setColorStates ( ) <EOL> G = pygraphviz . AGraph ( directed = True ) <EOL> cs = set ( m . colorStates ) <EOL> print "<STR_LIT>" . format ( len ( cs ) ) <EOL> for node , neighs in graph . items ( ) : <EOL> if node in cs : <EOL> print node <EOL> if node . endswith ( "<STR_LIT>" ) : <EOL> G . add_node ( node , style = "<STR_LIT>" , fillcolor = "<STR_LIT>" ) <EOL> if node . endswith ( "<STR_LIT>" ) : <EOL> G . add_node ( node , style = "<STR_LIT>" , fillcolor = "<STR_LIT>" ) <EOL> else : <EOL> G . add_node ( node ) <EOL> for node , neighs in graph . items ( ) : <EOL> for neig in neighs : <EOL> G . add_edge ( node , neig ) <EOL> else : <EOL> G = pygraphviz . AGraph ( graph , directed = True ) <EOL> G . draw ( "<STR_LIT>" % data_dir , prog = "<STR_LIT>" ) <EOL> G . write ( "<STR_LIT>" % data_dir ) <EOL> STATE_SEP = "<STR_LIT:|>" <EOL> TEMPLATE_SEP = "<STR_LIT:;>" <EOL> class ModelFilesGenerator : <EOL> def __init__ ( self , datapath , order = <NUM_LIT:2> , ngram = <NUM_LIT:0> , <EOL> whitespace = "<STR_LIT>" , subrules = False ) : <EOL> self . datapath = datapath <EOL> self . ngram = ngram <EOL> self . dh = data . DataHandler ( datapath , ngram , whitespace ) <EOL> self . order = order <EOL> self . _calculateMarkovChain ( ) <EOL> self . _gatherTemplates ( ) <EOL> self . rules = { } <EOL> if subrules : <EOL> subrulesOrder = range ( self . order , - <NUM_LIT:1> , - <NUM_LIT:1> ) <EOL> else : <EOL> subrulesOrder = [ self . order ] <EOL> for theOrder in subrulesOrder : <EOL> templateCombinations = self . _gatherTemplatesPerCommunication ( theOrder ) <EOL> self . _constructRules ( templateCombinations ) <EOL> def minimizeModel ( self ) : <EOL> dfa = DFA . prismaModel2DFA ( self . mc ) <EOL> collapse = dfa . minimize ( ) <EOL> return ( DFA . DFA2prismaModel ( dfa ) ) <EOL> def saveModel ( self ) : <EOL> f = file ( self . datapath + "<STR_LIT>" , "<STR_LIT:w>" ) <EOL> for state , nextStates in self . mc . iteritems ( ) : <EOL> for nextState , count in nextStates . iteritems ( ) : <EOL> f . write ( "<STR_LIT>" % ( state , nextState , count ) ) <EOL> f . close ( ) <EOL> f = file ( self . datapath + "<STR_LIT>" , "<STR_LIT:w>" ) <EOL> for template in self . structureTemplates : <EOL> f . write ( "<STR_LIT>" % <EOL> ( template . templateId , template . state , <EOL> template . numberOfMessages , len ( template . tokens ) , <EOL> "<STR_LIT:U+002C>" . join ( map ( str , template . getFieldIndices ( ) ) ) ) ) <EOL> for tok in template . tokens : <EOL> f . write ( "<STR_LIT>" % urllib . quote ( tok ) ) <EOL> f . close ( ) <EOL> f = file ( self . datapath + "<STR_LIT>" , "<STR_LIT:w>" ) <EOL> for ruleSet in self . rules . itervalues ( ) : <EOL> for r in ruleSet . rules : <EOL> f . write ( "<STR_LIT>" % <EOL> ( ruleSet . fingerprint , r . sourceId , r . sourceField , <EOL> r . targetField , str ( r . __class__ ) ) ) <EOL> f . write ( "<STR_LIT>" % r . getSpecificSaveString ( ) ) <EOL> f . close ( ) <EOL> def setColorStates ( self ) : <EOL> """<STR_LIT>""" <EOL> self . colorStates = [ ] <EOL> for st in self . structureTemplates : <EOL> if st . getFieldIndices ( ) : <EOL> self . colorStates += [ st . state ] <EOL> def setStateLabel ( self , msgIndex , stateLabel ) : <EOL> self . stateLabel4msg [ msgIndex ] = stateLabel <EOL> self . msgs4state . setdefault ( stateLabel , [ ] ) . append ( msgIndex ) <EOL> def getStartLabel ( self ) : <EOL> return STATE_SEP . join ( [ data . START_STATE ] * self . order ) <EOL> def _calculateMarkovChain ( self ) : <EOL> self . mc = { } <EOL> startLabel = self . getStartLabel ( ) <EOL> self . stateLabel4msg = [ "<STR_LIT>" ] * self . dh . getN ( ) <EOL> self . msgs4state = { } <EOL> self . mc [ startLabel ] = { } <EOL> for commId in self . dh . getCommunicationIds ( ) : <EOL> states = self . dh . getStateIdsForComm ( commId , self . order - <NUM_LIT:1> ) <EOL> indices = self . dh . getMsgIndexForComm ( commId ) <EOL> for ind in range ( <NUM_LIT:0> , len ( states ) - self . order ) : <EOL> s1 = STATE_SEP . join ( states [ ind : ( ind + self . order ) ] ) <EOL> s2 = STATE_SEP . join ( states [ ( ind + <NUM_LIT:1> ) : ( ind + self . order + <NUM_LIT:1> ) ] ) <EOL> trans = self . mc . setdefault ( s1 , { } ) <EOL> trans [ s2 ] = trans . get ( s2 , <NUM_LIT:0> ) + <NUM_LIT:1> <EOL> if self . order == <NUM_LIT:1> : <EOL> if ind >= <NUM_LIT:1> : <EOL> self . setStateLabel ( indices [ ind - <NUM_LIT:1> ] , s1 ) <EOL> else : <EOL> if ind < len ( indices ) : <EOL> self . setStateLabel ( indices [ ind ] , s1 ) <EOL> if ind == <NUM_LIT:0> : <EOL> trans = self . mc [ startLabel ] <EOL> trans [ s1 ] = trans . get ( s1 , <NUM_LIT:0> ) + <NUM_LIT:1> <EOL> def _gatherTemplates ( self ) : <EOL> self . templateid4msg = [ - <NUM_LIT:1> ] * self . dh . getN ( ) <EOL> self . msg4templateIds = { } <EOL> self . structureTemplates = [ ] <EOL> templateId = <NUM_LIT:0> <EOL> for ( state , msgs ) in self . msgs4state . iteritems ( ) : <EOL> tokens = [ self . dh . getTokensForMsg ( m ) for m in msgs ] <EOL> len2tokens = { } <EOL> for ( tok , msg ) in zip ( tokens , msgs ) : <EOL> len2tokens . setdefault ( len ( tok ) , [ ] ) . append ( ( tok , msg ) ) <EOL> for ( theLen , sameLengthMsgs ) in len2tokens . iteritems ( ) : <EOL> curTemplate = StructureTemplate ( state , templateId , theLen ) <EOL> for ( tok , msg ) in sameLengthMsgs : <EOL> curTemplate . addTokenizedMessage ( tok ) <EOL> self . templateid4msg [ msg ] = templateId <EOL> curTemplate . inferStructure ( ) <EOL> self . structureTemplates . append ( curTemplate ) <EOL> self . msg4templateIds [ templateId ] = [ m [ <NUM_LIT:1> ] for m in sameLengthMsgs ] <EOL> templateId += <NUM_LIT:1> <EOL> def _gatherTemplatesPerCommunication ( self , theOrder ) : <EOL> templateCombinations = { } <EOL> for comm in self . dh . getCommunicationIds ( ) : <EOL> msgs = self . dh . getMsgIndexForComm ( comm ) <EOL> tids = [ self . templateid4msg [ msg ] for msg in msgs ] <EOL> msgs = ( [ - <NUM_LIT:1> ] * theOrder ) + msgs <EOL> tids = ( [ - <NUM_LIT:1> ] * theOrder ) + tids <EOL> for ind in xrange ( len ( msgs ) - theOrder ) : <EOL> curTids = [ str ( t ) for t in tids [ ind : ( ind + theOrder + <NUM_LIT:1> ) ] ] <EOL> curMsgs = msgs [ ind : ( ind + theOrder + <NUM_LIT:1> ) ] <EOL> templateCombinations . setdefault ( TEMPLATE_SEP . join ( curTids ) , <EOL> [ ] ) . append ( tuple ( curMsgs ) ) <EOL> return templateCombinations <EOL> def _constructRules ( self , templateCombinations ) : <EOL> def getTokens ( mid ) : <EOL> if mid == - <NUM_LIT:1> : <EOL> return [ ] <EOL> else : <EOL> return self . dh . getTokensForMsg ( mid ) <EOL> def getFields ( tid ) : <EOL> if tid == - <NUM_LIT:1> : <EOL> return [ ] <EOL> else : <EOL> return self . structureTemplates [ tid ] . getFieldIndices ( ) <EOL> for ( templateIds , msgs ) in templateCombinations . iteritems ( ) : <EOL> tids = [ int ( tid ) for tid in templateIds . split ( TEMPLATE_SEP ) ] <EOL> fieldsForMessages = [ None ] * len ( tids ) <EOL> for ( ind , tid ) in enumerate ( tids ) : <EOL> tokenizedMessages = [ getTokens ( mids [ ind ] ) for mids in msgs ] <EOL> fieldsForMessages [ ind ] = FieldArray ( tokenizedMessages , <EOL> getFields ( tid ) ) <EOL> ruleSet = rule . RuleSet ( templateIds , tids ) <EOL> ruleSet . gatherRules ( fieldsForMessages , self . ngram > <NUM_LIT:0> ) <EOL> self . rules [ templateIds ] = ruleSet <EOL> class FieldArray : <EOL> def __init__ ( self , tokenizedMessages , fieldIndices ) : <EOL> self . array = [ [ tokm [ ind ] [ <NUM_LIT:1> ] for ind in fieldIndices ] <EOL> for tokm in tokenizedMessages ] <EOL> def getValuesForField ( self , index ) : <EOL> return [ m [ index ] for m in self . array ] <EOL> def getNumberOfFields ( self ) : <EOL> if len ( self . array [ <NUM_LIT:0> ] ) == <NUM_LIT:0> : <EOL> return <NUM_LIT:0> <EOL> else : <EOL> return len ( self . array [ <NUM_LIT:0> ] ) <EOL> class StructureTemplate : <EOL> def __init__ ( self , state , templateId , nTokens ) : <EOL> self . state = state <EOL> self . templateId = templateId <EOL> self . tokens = [ set ( ) for _ in xrange ( nTokens ) ] <EOL> self . numberOfMessages = <NUM_LIT:0> <EOL> def addTokenizedMessage ( self , toks ) : <EOL> assert ( len ( toks ) == len ( self . tokens ) ) <EOL> self . numberOfMessages += <NUM_LIT:1> <EOL> for ( ind , tok ) in enumerate ( toks ) : <EOL> self . tokens [ ind ] . add ( tok [ <NUM_LIT:1> ] ) <EOL> def inferStructure ( self ) : <EOL> self . fields = [ ind for ( ind , f ) in enumerate ( self . tokens ) <EOL> if len ( f ) > <NUM_LIT:1> ] <EOL> for ind in xrange ( len ( self . tokens ) ) : <EOL> cur = self . tokens [ ind ] <EOL> if len ( cur ) == <NUM_LIT:1> : <EOL> self . tokens [ ind ] = cur . pop ( ) <EOL> else : <EOL> self . tokens [ ind ] = "<STR_LIT>" <EOL> def getFieldIndices ( self ) : <EOL> return self . fields </s>
