<s> """<STR_LIT>""" <EOL> from grr . lib import aff4 <EOL> from grr . lib import artifact <EOL> from grr . lib import flow <EOL> from grr . lib import parsers <EOL> from grr . lib . checks import checks <EOL> from grr . lib . flows . general import collectors as _ <EOL> from grr . lib . rdfvalues import anomaly as rdf_anomaly <EOL> from grr . lib . rdfvalues import paths as rdf_paths <EOL> from grr . lib . rdfvalues import structs as rdf_structs <EOL> from grr . proto import flows_pb2 <EOL> class CheckFlowArgs ( rdf_structs . RDFProtoStruct ) : <EOL> protobuf = flows_pb2 . CheckFlowArgs <EOL> class CheckRunner ( flow . GRRFlow ) : <EOL> """<STR_LIT>""" <EOL> friendly_name = "<STR_LIT>" <EOL> category = "<STR_LIT>" <EOL> args_type = CheckFlowArgs <EOL> behaviours = flow . GRRFlow . behaviours + "<STR_LIT>" <EOL> @ flow . StateHandler ( next_state = [ "<STR_LIT>" ] ) <EOL> def Start ( self ) : <EOL> """<STR_LIT>""" <EOL> self . client = aff4 . FACTORY . Open ( self . client_id , token = self . token ) <EOL> self . state . Register ( "<STR_LIT>" , <EOL> self . client . Get ( self . client . Schema . KNOWLEDGE_BASE ) ) <EOL> self . state . Register ( "<STR_LIT>" , self . client . GetLabels ( ) ) <EOL> self . state . Register ( "<STR_LIT>" , { } ) <EOL> self . state . Register ( "<STR_LIT>" , set ( ) ) <EOL> self . state . Register ( "<STR_LIT>" , [ ] ) <EOL> self . state . Register ( "<STR_LIT>" , [ ] ) <EOL> self . state . Register ( "<STR_LIT>" , None ) <EOL> self . state . Register ( "<STR_LIT>" , { } ) <EOL> self . state . Register ( "<STR_LIT>" , rdf_paths . PathSpec . PathType . OS ) <EOL> self . CallState ( next_state = "<STR_LIT>" ) <EOL> @ flow . StateHandler ( next_state = [ "<STR_LIT>" , "<STR_LIT>" ] ) <EOL> def MapArtifactData ( self , responses ) : <EOL> """<STR_LIT>""" <EOL> self . state . artifacts_wanted = checks . CheckRegistry . SelectArtifacts ( <EOL> os_name = self . state . knowledge_base . os , <EOL> restrict_checks = self . args . restrict_checks ) <EOL> for artifact_name in self . state . artifacts_wanted : <EOL> self . CallFlow ( "<STR_LIT>" , artifact_list = [ artifact_name ] , <EOL> apply_parsers = False , <EOL> request_data = { "<STR_LIT>" : artifact_name } , <EOL> next_state = "<STR_LIT>" ) <EOL> self . CallState ( next_state = "<STR_LIT>" ) <EOL> def _ProcessData ( self , processor , responses , artifact_name , source ) : <EOL> """<STR_LIT>""" <EOL> artifact_data = self . state . host_data . get ( artifact_name ) <EOL> result_iterator = artifact . ApplyParserToResponses ( <EOL> processor , responses , source , self . state , self . token ) <EOL> for rdf in result_iterator : <EOL> if isinstance ( rdf , rdf_anomaly . Anomaly ) : <EOL> artifact_data [ "<STR_LIT>" ] . append ( rdf ) <EOL> else : <EOL> artifact_data [ "<STR_LIT>" ] . append ( rdf ) <EOL> def _RunProcessors ( self , artifact_name , responses ) : <EOL> """<STR_LIT>""" <EOL> source = responses . request_data . GetItem ( "<STR_LIT:source>" , None ) <EOL> processors = parsers . Parser . GetClassesByArtifact ( artifact_name ) <EOL> saved_responses = { } <EOL> for response in responses : <EOL> if processors : <EOL> for processor_cls in processors : <EOL> processor = processor_cls ( ) <EOL> if processor . process_together : <EOL> processor_name = processor . __class__ . __name__ <EOL> saved_responses . setdefault ( processor_name , [ ] ) . append ( response ) <EOL> else : <EOL> self . _ProcessData ( processor , response , artifact_name , source ) <EOL> for processor_name , responses_list in saved_responses . items ( ) : <EOL> processor = parsers . Parser . classes [ processor_name ] ( ) <EOL> self . _ProcessData ( processor , responses_list , artifact_name , source ) <EOL> @ flow . StateHandler ( ) <EOL> def AddResponses ( self , responses ) : <EOL> """<STR_LIT>""" <EOL> artifact_name = responses . request_data [ "<STR_LIT>" ] <EOL> artifact_data = self . state . host_data . setdefault ( artifact_name , { } ) <EOL> artifact_data [ "<STR_LIT>" ] = [ ] <EOL> artifact_data [ "<STR_LIT>" ] = [ ] <EOL> artifact_data [ "<STR_LIT>" ] = [ r for r in responses ] <EOL> if responses : <EOL> self . _RunProcessors ( artifact_name , responses ) <EOL> @ flow . StateHandler ( next_state = [ "<STR_LIT>" ] ) <EOL> def RunChecks ( self , responses ) : <EOL> if not responses . success : <EOL> raise RuntimeError ( "<STR_LIT>" ) <EOL> for finding in checks . CheckHost ( self . state . host_data , <EOL> os_name = self . state . knowledge_base . os , <EOL> restrict_checks = self . args . restrict_checks ) : <EOL> self . state . checks_run . append ( finding . check_id ) <EOL> if finding . anomaly : <EOL> self . state . checks_with_findings . append ( finding . check_id ) <EOL> self . SendReply ( finding ) </s>
