<s> """<STR_LIT>""" <EOL> import logging <EOL> from grr . lib import aff4 <EOL> from grr . lib import flow <EOL> from grr . lib import output_plugin <EOL> from grr . lib import rdfvalue <EOL> from grr . lib import stats <EOL> from grr . lib import utils <EOL> from grr . lib . aff4_objects import cronjobs <EOL> from grr . lib . hunts import results as hunts_results <EOL> from grr . lib . rdfvalues import structs as rdf_structs <EOL> from grr . proto import flows_pb2 <EOL> class ProcessHuntResultCollectionsCronFlowArgs ( rdf_structs . RDFProtoStruct ) : <EOL> protobuf = flows_pb2 . ProcessHuntResultsCronFlowArgs <EOL> class ResultsProcessingError ( Exception ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self ) : <EOL> self . exceptions_by_hunt = { } <EOL> super ( ResultsProcessingError , self ) . __init__ ( ) <EOL> def RegisterSubException ( self , hunt_urn , plugin_name , exception ) : <EOL> self . exceptions_by_hunt . setdefault ( hunt_urn , { } ) . setdefault ( <EOL> plugin_name , [ ] ) . append ( exception ) <EOL> def __repr__ ( self ) : <EOL> messages = [ ] <EOL> for hunt_urn , exceptions_by_plugin in self . exceptions_by_hunt . items ( ) : <EOL> for plugin_name , exception in exceptions_by_plugin . items ( ) : <EOL> messages . append ( "<STR_LIT>" % <EOL> ( hunt_urn , plugin_name , exception ) ) <EOL> return "<STR_LIT:\n>" . join ( messages ) <EOL> class ProcessHuntResultCollectionsCronFlow ( cronjobs . SystemCronFlow ) : <EOL> """<STR_LIT>""" <EOL> frequency = rdfvalue . Duration ( "<STR_LIT>" ) <EOL> lifetime = rdfvalue . Duration ( "<STR_LIT>" ) <EOL> allow_overruns = True <EOL> args_type = ProcessHuntResultCollectionsCronFlowArgs <EOL> DEFAULT_BATCH_SIZE = <NUM_LIT> <EOL> def CheckIfRunningTooLong ( self ) : <EOL> if self . state . args . max_running_time : <EOL> elapsed = ( rdfvalue . RDFDatetime ( ) . Now ( ) . AsSecondsFromEpoch ( ) - <EOL> self . start_time . AsSecondsFromEpoch ( ) ) <EOL> if elapsed > self . state . args . max_running_time : <EOL> return True <EOL> return False <EOL> def LoadPlugins ( self , metadata_obj ) : <EOL> output_plugins = metadata_obj . Get ( metadata_obj . Schema . OUTPUT_PLUGINS ) <EOL> used_plugins = [ ] <EOL> unused_plugins = [ ] <EOL> for _ , ( plugin_def , state ) in output_plugins . data . iteritems ( ) : <EOL> if not hasattr ( plugin_def , "<STR_LIT>" ) : <EOL> logging . error ( "<STR_LIT>" , plugin_def ) <EOL> continue <EOL> used_plugins . append ( ( plugin_def , plugin_def . GetPluginForState ( state ) ) ) <EOL> return output_plugins , used_plugins <EOL> def RunPlugins ( self , hunt_urn , plugins , results , exceptions_by_plugin ) : <EOL> for plugin_def , plugin in plugins : <EOL> try : <EOL> plugin . ProcessResponses ( results ) <EOL> plugin_status = output_plugin . OutputPluginBatchProcessingStatus ( <EOL> plugin_descriptor = plugin_def , <EOL> status = "<STR_LIT>" , <EOL> batch_size = len ( results ) ) <EOL> stats . STATS . IncrementCounter ( "<STR_LIT>" , <EOL> delta = len ( results ) , <EOL> fields = [ plugin_def . plugin_name ] ) <EOL> except Exception as e : <EOL> logging . exception ( e ) <EOL> plugin_status = output_plugin . OutputPluginBatchProcessingStatus ( <EOL> plugin_descriptor = plugin_def , <EOL> status = "<STR_LIT>" , <EOL> summary = utils . SmartStr ( e ) , <EOL> batch_size = len ( results ) ) <EOL> exceptions_by_plugin . setdefault ( plugin_def , [ ] ) . append ( e ) <EOL> aff4 . FACTORY . Open ( <EOL> hunt_urn . Add ( "<STR_LIT>" ) , <EOL> "<STR_LIT>" , <EOL> mode = "<STR_LIT:w>" , <EOL> token = self . token ) . Add ( plugin_status ) <EOL> if plugin_status . status == plugin_status . Status . ERROR : <EOL> aff4 . FACTORY . Open ( <EOL> hunt_urn . Add ( "<STR_LIT>" ) , <EOL> "<STR_LIT>" , <EOL> mode = "<STR_LIT:w>" , <EOL> token = self . token ) . Add ( plugin_status ) <EOL> def FlushPlugins ( self , hunt_urn , plugins , exceptions_by_plugin ) : <EOL> for plugin_def , plugin in plugins : <EOL> try : <EOL> plugin . Flush ( ) <EOL> except Exception as e : <EOL> logging . exception ( "<STR_LIT>" <EOL> "<STR_LIT>" , hunt_urn , str ( plugin ) ) <EOL> self . Log ( "<STR_LIT>" <EOL> "<STR_LIT>" % ( hunt_urn , str ( plugin ) , e ) ) <EOL> exceptions_by_plugin . setdefault ( plugin_def , [ ] ) . append ( e ) <EOL> def ProcessOneHunt ( self , exceptions_by_hunt ) : <EOL> """<STR_LIT>""" <EOL> hunt_results_urn , results = ( <EOL> hunts_results . HuntResultQueue . ClaimNotificationsForCollection ( <EOL> start_time = self . args . start_processing_time , <EOL> token = self . token , <EOL> lease_time = self . lifetime ) ) <EOL> logging . debug ( "<STR_LIT>" , len ( results ) , <EOL> hunt_results_urn ) <EOL> if not results : <EOL> return <NUM_LIT:0> <EOL> hunt_urn = rdfvalue . RDFURN ( hunt_results_urn . Dirname ( ) ) <EOL> batch_size = self . state . args . batch_size or self . DEFAULT_BATCH_SIZE <EOL> metadata_urn = hunt_urn . Add ( "<STR_LIT>" ) <EOL> exceptions_by_plugin = { } <EOL> num_processed_for_hunt = <NUM_LIT:0> <EOL> with aff4 . FACTORY . OpenWithLock ( hunt_results_urn , <EOL> aff4_type = "<STR_LIT>" , <EOL> lease_time = <NUM_LIT> , <EOL> token = self . token ) as collection_obj : <EOL> with aff4 . FACTORY . OpenWithLock ( metadata_urn , <EOL> lease_time = <NUM_LIT> , <EOL> token = self . token ) as metadata_obj : <EOL> all_plugins , used_plugins = self . LoadPlugins ( metadata_obj ) <EOL> num_processed = int ( metadata_obj . Get ( <EOL> metadata_obj . Schema . NUM_PROCESSED_RESULTS ) ) <EOL> for batch in utils . Grouper ( results , batch_size ) : <EOL> results = list ( collection_obj . MultiResolve ( [ ( <EOL> ts , suffix ) for ( _ , ts , suffix ) in batch ] ) ) <EOL> self . RunPlugins ( hunt_urn , used_plugins , results , exceptions_by_plugin ) <EOL> hunts_results . HuntResultQueue . DeleteNotifications ( <EOL> [ record_id for ( record_id , _ , _ ) in batch ] , <EOL> token = self . token ) <EOL> num_processed += len ( batch ) <EOL> num_processed_for_hunt += len ( batch ) <EOL> self . HeartBeat ( ) <EOL> collection_obj . UpdateLease ( <NUM_LIT> ) <EOL> metadata_obj . Set ( metadata_obj . Schema . NUM_PROCESSED_RESULTS ( <EOL> num_processed ) ) <EOL> metadata_obj . UpdateLease ( <NUM_LIT> ) <EOL> if self . CheckIfRunningTooLong ( ) : <EOL> logging . warning ( "<STR_LIT>" ) <EOL> break <EOL> self . FlushPlugins ( hunt_urn , used_plugins , exceptions_by_plugin ) <EOL> metadata_obj . Set ( metadata_obj . Schema . OUTPUT_PLUGINS ( all_plugins ) ) <EOL> metadata_obj . Set ( metadata_obj . Schema . NUM_PROCESSED_RESULTS ( <EOL> num_processed ) ) <EOL> if exceptions_by_plugin : <EOL> for plugin , exceptions in exceptions_by_plugin . items ( ) : <EOL> exceptions_by_hunt . setdefault ( hunt_urn , { } ) . setdefault ( <EOL> plugin , [ ] ) . extend ( exceptions ) <EOL> logging . debug ( "<STR_LIT>" , num_processed_for_hunt ) <EOL> return len ( results ) <EOL> @ flow . StateHandler ( ) <EOL> def Start ( self ) : <EOL> self . start_time = rdfvalue . RDFDatetime ( ) . Now ( ) <EOL> exceptions_by_hunt = { } <EOL> if not self . state . args . max_running_time : <EOL> self . state . args . max_running_time = rdfvalue . Duration ( "<STR_LIT>" % int ( <EOL> ProcessHuntResultCollectionsCronFlow . lifetime . seconds * <NUM_LIT> ) ) <EOL> while not self . CheckIfRunningTooLong ( ) : <EOL> count = self . ProcessOneHunt ( exceptions_by_hunt ) <EOL> if not count : <EOL> break <EOL> if exceptions_by_hunt : <EOL> e = ResultsProcessingError ( ) <EOL> for hunt_urn , exceptions_by_plugin in exceptions_by_hunt . items ( ) : <EOL> for plugin , exceptions in exceptions_by_plugin . items ( ) : <EOL> for exception in exceptions : <EOL> e . RegisterSubException ( hunt_urn , plugin , exception ) <EOL> raise e </s>
