<s> """<STR_LIT>""" <EOL> import ConfigParser <EOL> import glob <EOL> import os <EOL> import os . path <EOL> import random <EOL> import re <EOL> import subprocess <EOL> import sys <EOL> import time <EOL> import addr_util <EOL> import selectors <EOL> import util <EOL> if sys . platform [ : <NUM_LIT:3> ] == '<STR_LIT>' : <EOL> DEFAULT_TIMER = time . clock <EOL> else : <EOL> DEFAULT_TIMER = time . time <EOL> GLOBAL_DATA_CACHE = { } <EOL> DEFAULT_CONFIG_PATH = '<STR_LIT>' <EOL> MAX_NON_UNIQUE_RECORD_COUNT = <NUM_LIT> <EOL> MAX_FILE_MTIME_AGE_DAYS = <NUM_LIT> <EOL> MIN_FILE_SIZE = <NUM_LIT> <EOL> MIN_RECOMMENDED_RECORD_COUNT = <NUM_LIT:200> <EOL> MAX_FQDN_SYNTHESIZE_PERCENT = <NUM_LIT:4> <EOL> class DataSources ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , config_path = DEFAULT_CONFIG_PATH , status_callback = None ) : <EOL> self . source_cache = GLOBAL_DATA_CACHE <EOL> self . source_config = { } <EOL> self . status_callback = status_callback <EOL> self . _LoadConfigFromPath ( config_path ) <EOL> def msg ( self , msg , ** kwargs ) : <EOL> if self . status_callback : <EOL> self . status_callback ( msg , ** kwargs ) <EOL> else : <EOL> print '<STR_LIT>' % msg <EOL> def _LoadConfigFromPath ( self , path ) : <EOL> """<STR_LIT>""" <EOL> conf_file = util . FindDataFile ( path ) <EOL> config = ConfigParser . ConfigParser ( ) <EOL> config . read ( conf_file ) <EOL> for section in config . sections ( ) : <EOL> if section not in self . source_config : <EOL> self . source_config [ section ] = { <EOL> '<STR_LIT:name>' : None , <EOL> '<STR_LIT>' : set ( ) , <EOL> '<STR_LIT>' : True , <EOL> '<STR_LIT>' : False , <EOL> '<STR_LIT>' : False , <EOL> '<STR_LIT>' : MAX_FILE_MTIME_AGE_DAYS <EOL> } <EOL> for ( key , value ) in config . items ( section ) : <EOL> if key == '<STR_LIT:name>' : <EOL> self . source_config [ section ] [ '<STR_LIT:name>' ] = value <EOL> elif key == '<STR_LIT>' and int ( value ) == <NUM_LIT:0> : <EOL> self . source_config [ section ] [ '<STR_LIT>' ] = False <EOL> elif key == '<STR_LIT>' : <EOL> self . source_config [ section ] [ '<STR_LIT>' ] = int ( value ) <EOL> elif key == '<STR_LIT>' : <EOL> self . source_config [ section ] [ '<STR_LIT>' ] = bool ( value ) <EOL> elif key == '<STR_LIT>' : <EOL> self . source_config [ section ] [ '<STR_LIT>' ] = bool ( value ) <EOL> else : <EOL> self . source_config [ section ] [ '<STR_LIT>' ] . add ( value ) <EOL> def ListSourceTypes ( self ) : <EOL> """<STR_LIT>""" <EOL> return sorted ( self . source_config . keys ( ) ) <EOL> def ListSourcesWithDetails ( self ) : <EOL> """<STR_LIT>""" <EOL> for source in self . ListSourceTypes ( ) : <EOL> max_mtime = self . source_config [ source ] [ '<STR_LIT>' ] <EOL> self . _GetHostsFromSource ( source , min_file_size = MIN_FILE_SIZE , <EOL> max_mtime_age_days = max_mtime ) <EOL> details = [ ] <EOL> for source in self . source_cache : <EOL> details . append ( ( source , <EOL> self . source_config [ source ] [ '<STR_LIT:name>' ] , <EOL> self . source_config [ source ] [ '<STR_LIT>' ] , <EOL> len ( self . source_cache [ source ] ) ) ) <EOL> return sorted ( details , key = lambda x : ( x [ <NUM_LIT:2> ] , x [ <NUM_LIT:3> ] * - <NUM_LIT:1> ) ) <EOL> def ListSourceTitles ( self ) : <EOL> """<STR_LIT>""" <EOL> titles = [ ] <EOL> seen_synthetic = False <EOL> seen_organic = False <EOL> for ( unused_type , name , is_synthetic , count ) in self . ListSourcesWithDetails ( ) : <EOL> if not is_synthetic : <EOL> seen_organic = True <EOL> if is_synthetic and seen_organic and not seen_synthetic : <EOL> titles . append ( '<STR_LIT:->' * <NUM_LIT> ) <EOL> seen_synthetic = True <EOL> titles . append ( '<STR_LIT>' % ( name , count ) ) <EOL> return titles <EOL> def ConvertSourceTitleToType ( self , detail ) : <EOL> """<STR_LIT>""" <EOL> for source_type in self . source_config : <EOL> if detail . startswith ( self . source_config [ source_type ] [ '<STR_LIT:name>' ] ) : <EOL> return source_type <EOL> def GetBestSourceDetails ( self ) : <EOL> return self . ListSourcesWithDetails ( ) [ <NUM_LIT:0> ] <EOL> def GetNameForSource ( self , source ) : <EOL> if source in self . source_config : <EOL> return self . source_config [ source ] [ '<STR_LIT:name>' ] <EOL> else : <EOL> return source <EOL> def GetCachedRecordCountForSource ( self , source ) : <EOL> return len ( self . source_cache [ source ] ) <EOL> def _CreateRecordsFromHostEntries ( self , entries , include_duplicates = False ) : <EOL> """<STR_LIT>""" <EOL> last_entry = None <EOL> records = [ ] <EOL> full_host_count = <NUM_LIT:0> <EOL> for entry in entries : <EOL> if entry == last_entry and not include_duplicates : <EOL> continue <EOL> else : <EOL> last_entry = entry <EOL> if '<STR_LIT:U+0020>' in entry : <EOL> ( record_type , host ) = entry . split ( '<STR_LIT:U+0020>' ) <EOL> else : <EOL> record_type = '<STR_LIT:A>' <EOL> host = entry <EOL> if not addr_util . IP_RE . match ( host ) and not addr_util . INTERNAL_RE . search ( host ) : <EOL> if not host . endswith ( '<STR_LIT:.>' ) : <EOL> host += '<STR_LIT:.>' <EOL> records . append ( ( record_type , host ) ) <EOL> if addr_util . FQDN_RE . match ( host ) : <EOL> full_host_count += <NUM_LIT:1> <EOL> if not records : <EOL> raise ValueError ( '<STR_LIT>' % entries ) <EOL> full_host_percent = full_host_count / float ( len ( records ) ) * <NUM_LIT:100> <EOL> if full_host_percent < MAX_FQDN_SYNTHESIZE_PERCENT : <EOL> full_host_names = True <EOL> else : <EOL> full_host_names = False <EOL> return ( records , full_host_names ) <EOL> def GetTestsFromSource ( self , source , count = <NUM_LIT:50> , select_mode = None ) : <EOL> """<STR_LIT>""" <EOL> records = [ ] <EOL> if source in self . source_config : <EOL> include_duplicates = self . source_config [ source ] . get ( '<STR_LIT>' , False ) <EOL> else : <EOL> include_duplicates = False <EOL> records = self . _GetHostsFromSource ( source ) <EOL> if not records : <EOL> raise ValueError ( '<STR_LIT>' % source ) <EOL> self . msg ( '<STR_LIT>' <EOL> % ( self . GetNameForSource ( source ) , len ( records ) , count , select_mode ) ) <EOL> ( records , are_records_fqdn ) = self . _CreateRecordsFromHostEntries ( records , <EOL> include_duplicates = include_duplicates ) <EOL> if select_mode in ( '<STR_LIT>' , '<STR_LIT>' , None ) : <EOL> if include_duplicates : <EOL> if count > len ( records ) : <EOL> select_mode = '<STR_LIT>' <EOL> else : <EOL> select_mode = '<STR_LIT>' <EOL> elif len ( records ) != len ( set ( records ) ) : <EOL> if select_mode == '<STR_LIT>' : <EOL> self . msg ( '<STR_LIT>' % source ) <EOL> select_mode = '<STR_LIT>' <EOL> else : <EOL> select_mode = '<STR_LIT>' <EOL> self . msg ( '<STR_LIT>' % <EOL> ( count , len ( records ) , select_mode ) ) <EOL> if select_mode == '<STR_LIT>' : <EOL> records = selectors . WeightedDistribution ( records , count ) <EOL> elif select_mode == '<STR_LIT>' : <EOL> records = selectors . ChunkSelect ( records , count ) <EOL> elif select_mode == '<STR_LIT>' : <EOL> records = selectors . RandomSelect ( records , count , include_duplicates = include_duplicates ) <EOL> else : <EOL> raise ValueError ( '<STR_LIT>' % select_mode ) <EOL> if source not in self . source_config : <EOL> self . source_config [ source ] = { '<STR_LIT>' : True } <EOL> if are_records_fqdn : <EOL> self . source_config [ source ] [ '<STR_LIT>' ] = False <EOL> self . msg ( '<STR_LIT>' % source ) <EOL> synthesized = [ ] <EOL> for ( req_type , hostname ) in records : <EOL> if not addr_util . FQDN_RE . match ( hostname ) : <EOL> hostname = self . _GenerateRandomHostname ( hostname ) <EOL> synthesized . append ( ( req_type , hostname ) ) <EOL> return synthesized <EOL> else : <EOL> return records <EOL> def _GenerateRandomHostname ( self , domain ) : <EOL> """<STR_LIT>""" <EOL> oracle = random . randint ( <NUM_LIT:0> , <NUM_LIT:100> ) <EOL> if oracle < <NUM_LIT> : <EOL> return '<STR_LIT>' % domain <EOL> elif oracle < <NUM_LIT> : <EOL> return domain <EOL> elif oracle < <NUM_LIT> : <EOL> return '<STR_LIT>' % domain <EOL> else : <EOL> return '<STR_LIT>' % ( random . randint ( <NUM_LIT:0> , <NUM_LIT:10> ) , domain ) <EOL> def _GetHostsFromSource ( self , source , min_file_size = None , max_mtime_age_days = None ) : <EOL> """<STR_LIT>""" <EOL> if source in self . source_cache : <EOL> return self . source_cache [ source ] <EOL> filename = self . _FindBestFileForSource ( source , min_file_size = min_file_size , <EOL> max_mtime_age_days = max_mtime_age_days ) <EOL> if not filename : <EOL> return None <EOL> size_mb = os . path . getsize ( filename ) / <NUM_LIT> / <NUM_LIT> <EOL> if not self . source_config [ source ] [ '<STR_LIT>' ] : <EOL> self . msg ( '<STR_LIT>' % ( self . GetNameForSource ( source ) , filename , size_mb ) ) <EOL> start_clock = DEFAULT_TIMER ( ) <EOL> if filename . endswith ( '<STR_LIT>' ) or filename . endswith ( '<STR_LIT>' ) : <EOL> hosts = self . _ExtractHostsFromPcapFile ( filename ) <EOL> else : <EOL> hosts = self . _ExtractHostsFromHistoryFile ( filename ) <EOL> if not hosts : <EOL> hosts = self . _ReadDataFile ( filename ) <EOL> duration = DEFAULT_TIMER ( ) - start_clock <EOL> if duration > <NUM_LIT:5> : <EOL> self . msg ( '<STR_LIT>' % ( self . GetNameForSource ( source ) , duration ) ) <EOL> self . source_cache [ source ] = hosts <EOL> return hosts <EOL> def _ExtractHostsFromHistoryFile ( self , path ) : <EOL> """<STR_LIT>""" <EOL> parse_re = re . compile ( '<STR_LIT>' ) <EOL> return parse_re . findall ( open ( path , '<STR_LIT:rb>' ) . read ( ) ) <EOL> def _ExtractHostsFromPcapFile ( self , path ) : <EOL> """<STR_LIT>""" <EOL> self . msg ( '<STR_LIT>' ) <EOL> cmd = '<STR_LIT>' % path <EOL> pipe = subprocess . Popen ( cmd , shell = True , stdout = subprocess . PIPE ) . stdout <EOL> parse_re = re . compile ( '<STR_LIT>' ) <EOL> requests = [ ] <EOL> for line in pipe : <EOL> if '<STR_LIT:?>' not in line : <EOL> continue <EOL> match = parse_re . search ( line ) <EOL> if match : <EOL> requests . append ( '<STR_LIT:U+0020>' . join ( match . groups ( ) ) ) <EOL> return requests <EOL> def _ReadDataFile ( self , path ) : <EOL> """<STR_LIT>""" <EOL> records = [ ] <EOL> domains_re = re . compile ( '<STR_LIT>' ) <EOL> requests_re = re . compile ( '<STR_LIT>' ) <EOL> for line in open ( path ) : <EOL> if domains_re . match ( line ) or requests_re . match ( line ) : <EOL> records . append ( line . rstrip ( ) ) <EOL> return records <EOL> def _GetSourceSearchPaths ( self , source ) : <EOL> """<STR_LIT>""" <EOL> if source not in self . source_config : <EOL> return [ source ] <EOL> search_paths = [ ] <EOL> environment_re = re . compile ( '<STR_LIT>' ) <EOL> for path in self . source_config [ source ] [ '<STR_LIT>' ] : <EOL> env_vars = set ( environment_re . findall ( path ) ) <EOL> if env_vars : <EOL> for variable in env_vars : <EOL> env_var = os . getenv ( variable , False ) <EOL> if env_var : <EOL> path = path . replace ( '<STR_LIT>' % variable , env_var ) <EOL> else : <EOL> path = None <EOL> if path : <EOL> path = path . replace ( '<STR_LIT:/>' , os . sep ) <EOL> search_paths . append ( path ) <EOL> if sys . platform [ : <NUM_LIT:3> ] == '<STR_LIT>' : <EOL> keywords = [ '<STR_LIT>' , '<STR_LIT>' ] <EOL> for keyword in keywords : <EOL> if keyword in path : <EOL> replacement = keywords [ keywords . index ( keyword ) - <NUM_LIT:1> ] <EOL> search_paths . append ( path . replace ( '<STR_LIT>' % keyword , '<STR_LIT>' % replacement ) ) <EOL> search_paths . append ( path . replace ( '<STR_LIT>' % keyword , '<STR_LIT>' ) ) <EOL> return search_paths <EOL> def _FindBestFileForSource ( self , source , min_file_size = None , max_mtime_age_days = None ) : <EOL> """<STR_LIT>""" <EOL> found = [ ] <EOL> for path in self . _GetSourceSearchPaths ( source ) : <EOL> if not os . path . isabs ( path ) : <EOL> path = util . FindDataFile ( path ) <EOL> for filename in glob . glob ( path ) : <EOL> if min_file_size and os . path . getsize ( filename ) < min_file_size : <EOL> self . msg ( '<STR_LIT>' % ( filename , os . path . getsize ( filename ) ) ) <EOL> else : <EOL> try : <EOL> fp = open ( filename , '<STR_LIT:rb>' ) <EOL> fp . close ( ) <EOL> found . append ( filename ) <EOL> except IOError : <EOL> self . msg ( '<STR_LIT>' % filename ) <EOL> if found : <EOL> newest = sorted ( found , key = os . path . getmtime ) [ - <NUM_LIT:1> ] <EOL> age_days = ( time . time ( ) - os . path . getmtime ( newest ) ) / <NUM_LIT> <EOL> if max_mtime_age_days and age_days > max_mtime_age_days : <EOL> pass <EOL> else : <EOL> return newest <EOL> else : <EOL> return None <EOL> if __name__ == '<STR_LIT:__main__>' : <EOL> parser = DataSources ( ) <EOL> print parser . ListSourceTypes ( ) <EOL> print parser . ListSourcesWithDetails ( ) <EOL> best = parser . ListSourcesWithDetails ( ) [ <NUM_LIT:0> ] [ <NUM_LIT:0> ] <EOL> print len ( parser . GetRecordsFromSource ( best ) ) </s>
