<s> """<STR_LIT>""" <EOL> import functools <EOL> import operator <EOL> import time <EOL> from M2Crypto import X509 <EOL> import logging <EOL> from grr . client import actions <EOL> from grr . lib import access_control <EOL> from grr . lib import aff4 <EOL> from grr . lib import communicator <EOL> from grr . lib import config_lib <EOL> from grr . lib import data_store <EOL> from grr . lib import flow_runner <EOL> from grr . lib import queue_manager <EOL> from grr . lib import queues <EOL> from grr . lib import rdfvalue <EOL> from grr . lib import registry <EOL> from grr . lib import stats <EOL> from grr . lib import threadpool <EOL> from grr . lib import type_info <EOL> from grr . lib import utils <EOL> from grr . lib . rdfvalues import client as rdf_client <EOL> from grr . lib . rdfvalues import flows as rdf_flows <EOL> from grr . lib . rdfvalues import protodict as rdf_protodict <EOL> from grr . lib . rdfvalues import structs as rdf_structs <EOL> from grr . proto import jobs_pb2 <EOL> class AuditEvent ( rdf_structs . RDFProtoStruct ) : <EOL> protobuf = jobs_pb2 . AuditEvent <EOL> def __init__ ( self , initializer = None , age = None , ** kwargs ) : <EOL> super ( AuditEvent , self ) . __init__ ( initializer = initializer , age = age , <EOL> ** kwargs ) <EOL> if not self . id : <EOL> self . id = utils . PRNG . GetULong ( ) <EOL> if not self . timestamp : <EOL> self . timestamp = rdfvalue . RDFDatetime ( ) . Now ( ) <EOL> class FlowError ( Exception ) : <EOL> """<STR_LIT>""" <EOL> FlowRunnerArgs = flow_runner . FlowRunnerArgs <EOL> class Responses ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , request = None , responses = None , auth_required = True , <EOL> next_states = None ) : <EOL> self . status = None <EOL> self . success = True <EOL> self . next_states = next_states <EOL> self . request = request <EOL> self . _auth_required = auth_required <EOL> if request : <EOL> self . request_data = rdf_protodict . Dict ( request . data ) <EOL> self . _responses = [ ] <EOL> self . _dropped_responses = [ ] <EOL> if responses : <EOL> responses . sort ( key = operator . attrgetter ( "<STR_LIT>" ) ) <EOL> self . iterator = None <EOL> for msg in responses : <EOL> if msg . auth_state == msg . AuthorizationState . DESYNCHRONIZED or ( <EOL> self . _auth_required and <EOL> msg . auth_state != msg . AuthorizationState . AUTHENTICATED ) : <EOL> logging . warning ( "<STR_LIT>" , <EOL> msg . session_id , msg . auth_state ) <EOL> self . _dropped_responses . append ( msg ) <EOL> continue <EOL> if msg . type == msg . Type . ITERATOR : <EOL> self . iterator = rdf_client . Iterator ( msg . payload ) <EOL> continue <EOL> if msg . type == msg . Type . STATUS : <EOL> self . status = rdf_flows . GrrStatus ( msg . payload ) <EOL> self . success = self . status . status == self . status . ReturnedStatus . OK <EOL> break <EOL> self . _responses . append ( msg ) <EOL> if self . status is None : <EOL> if self . _dropped_responses : <EOL> logging . error ( "<STR_LIT>" + "<STR_LIT:\n>" . join ( <EOL> [ utils . SmartUnicode ( x ) for x in self . _dropped_responses ] ) ) <EOL> if responses : <EOL> self . _LogFlowState ( responses ) <EOL> raise FlowError ( "<STR_LIT>" ) <EOL> self . message = None <EOL> def __iter__ ( self ) : <EOL> """<STR_LIT>""" <EOL> old_response_id = None <EOL> action_registry = actions . ActionPlugin . classes <EOL> expected_response_classes = [ ] <EOL> is_client_request = False <EOL> if self . request . HasField ( "<STR_LIT>" ) : <EOL> is_client_request = True <EOL> client_action_name = self . request . request . name <EOL> if client_action_name not in action_registry : <EOL> raise RuntimeError ( "<STR_LIT>" % <EOL> client_action_name ) <EOL> expected_response_classes = action_registry [ <EOL> client_action_name ] . out_rdfvalues <EOL> for message in self . _responses : <EOL> self . message = rdf_flows . GrrMessage ( message ) <EOL> if self . message . response_id == old_response_id : <EOL> continue <EOL> else : <EOL> old_response_id = self . message . response_id <EOL> if self . message . type == self . message . Type . MESSAGE : <EOL> if is_client_request : <EOL> if not expected_response_classes : <EOL> raise RuntimeError ( <EOL> "<STR_LIT>" % <EOL> client_action_name ) <EOL> else : <EOL> args_rdf_name = self . message . args_rdf_name <EOL> if not args_rdf_name : <EOL> raise RuntimeError ( "<STR_LIT>" <EOL> "<STR_LIT>" ) <EOL> elif args_rdf_name not in [ <EOL> x . __name__ for x in expected_response_classes ] : <EOL> raise RuntimeError ( <EOL> "<STR_LIT>" % <EOL> ( args_rdf_name , expected_response_classes , <EOL> client_action_name ) ) <EOL> yield self . message . payload <EOL> def First ( self ) : <EOL> """<STR_LIT>""" <EOL> for x in self : <EOL> return x <EOL> def __len__ ( self ) : <EOL> return len ( self . _responses ) <EOL> def __nonzero__ ( self ) : <EOL> return bool ( self . _responses ) <EOL> def _LogFlowState ( self , responses ) : <EOL> session_id = responses [ <NUM_LIT:0> ] . session_id <EOL> token = access_control . ACLToken ( username = "<STR_LIT>" , reason = "<STR_LIT>" ) <EOL> token . supervisor = True <EOL> logging . error ( <EOL> "<STR_LIT>" , <EOL> data_store . DB . ResolvePrefix ( <EOL> session_id . Add ( "<STR_LIT:state>" ) , <EOL> "<STR_LIT>" , token = token ) , <EOL> data_store . DB . ResolvePrefix ( <EOL> session_id . Add ( "<STR_LIT>" % responses [ <NUM_LIT:0> ] . request_id ) , <EOL> "<STR_LIT>" , token = token ) , <EOL> data_store . DB . ResolvePrefix ( <EOL> queues . FLOWS , "<STR_LIT>" % session_id , token = token ) ) <EOL> class FakeResponses ( Responses ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , messages , request_data ) : <EOL> super ( FakeResponses , self ) . __init__ ( ) <EOL> self . success = True <EOL> self . _responses = messages or [ ] <EOL> self . request_data = request_data <EOL> self . iterator = None <EOL> def __iter__ ( self ) : <EOL> return iter ( self . _responses ) <EOL> def StateHandler ( next_state = "<STR_LIT>" , auth_required = True ) : <EOL> """<STR_LIT>""" <EOL> if not isinstance ( next_state , ( basestring , list , tuple ) ) : <EOL> raise RuntimeError ( "<STR_LIT>" ) <EOL> def Decorator ( f ) : <EOL> """<STR_LIT>""" <EOL> if isinstance ( next_state , basestring ) : <EOL> next_states = set ( [ next_state ] ) <EOL> else : <EOL> next_states = set ( next_state ) <EOL> @ functools . wraps ( f ) <EOL> def Decorated ( self , responses = None , request = None , direct_response = None ) : <EOL> """<STR_LIT>""" <EOL> if "<STR_LIT:r>" in self . mode : <EOL> pending_termination = self . Get ( self . Schema . PENDING_TERMINATION ) <EOL> if pending_termination : <EOL> self . Error ( pending_termination . reason ) <EOL> return <EOL> runner = self . GetRunner ( ) <EOL> next_states = Decorated . next_states <EOL> old_next_states = runner . GetAllowedFollowUpStates ( ) <EOL> try : <EOL> if direct_response is not None : <EOL> runner . SetAllowedFollowUpStates ( next_states ) <EOL> return f ( self , direct_response ) <EOL> if isinstance ( responses , Responses ) : <EOL> next_states . update ( responses . next_states ) <EOL> else : <EOL> responses = Responses ( request = request , <EOL> next_states = next_states , <EOL> responses = responses , <EOL> auth_required = auth_required ) <EOL> if responses . status : <EOL> runner . SaveResourceUsage ( request , responses ) <EOL> stats . STATS . IncrementCounter ( "<STR_LIT>" ) <EOL> runner . SetAllowedFollowUpStates ( next_states ) <EOL> if f . __name__ == "<STR_LIT>" : <EOL> stats . STATS . IncrementCounter ( "<STR_LIT>" , <EOL> fields = [ self . Name ( ) ] ) <EOL> args = [ self , responses ] <EOL> res = f ( * args [ : f . func_code . co_argcount ] ) <EOL> return res <EOL> finally : <EOL> runner . SetAllowedFollowUpStates ( old_next_states ) <EOL> Decorated . next_states = next_states <EOL> return Decorated <EOL> return Decorator <EOL> class PendingFlowTermination ( rdf_structs . RDFProtoStruct ) : <EOL> """<STR_LIT>""" <EOL> protobuf = jobs_pb2 . PendingFlowTermination <EOL> class EmptyFlowArgs ( rdf_structs . RDFProtoStruct ) : <EOL> """<STR_LIT>""" <EOL> protobuf = jobs_pb2 . EmptyMessage <EOL> class Behaviour ( object ) : <EOL> """<STR_LIT>""" <EOL> LEXICON = { } <EOL> def __init__ ( self , * args ) : <EOL> self . set = set ( ) <EOL> for arg in args : <EOL> if arg not in self . LEXICON : <EOL> raise ValueError ( "<STR_LIT>" % arg ) <EOL> self . set . add ( str ( arg ) ) <EOL> def __add__ ( self , other ) : <EOL> other = str ( other ) <EOL> if other not in self . LEXICON : <EOL> raise ValueError ( "<STR_LIT>" % other ) <EOL> return self . __class__ ( other , * list ( self . set ) ) <EOL> def __sub__ ( self , other ) : <EOL> other = str ( other ) <EOL> result = self . set . copy ( ) <EOL> result . discard ( other ) <EOL> return self . __class__ ( * list ( result ) ) <EOL> def __iter__ ( self ) : <EOL> return iter ( self . set ) <EOL> def IsSupported ( self , other ) : <EOL> """<STR_LIT>""" <EOL> if not isinstance ( other , self . __class__ ) : <EOL> raise TypeError ( "<STR_LIT>" % self . __class__ ) <EOL> return self . set . issubset ( other . set ) <EOL> class FlowBehaviour ( Behaviour ) : <EOL> LEXICON = { <EOL> "<STR_LIT>" : ( "<STR_LIT>" <EOL> "<STR_LIT>" ) , <EOL> "<STR_LIT>" : ( "<STR_LIT>" <EOL> "<STR_LIT>" ) , <EOL> "<STR_LIT>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : "<STR_LIT>" , <EOL> } <EOL> class GRRFlow ( aff4 . AFF4Volume ) : <EOL> """<STR_LIT>""" <EOL> class SchemaCls ( aff4 . AFF4Volume . SchemaCls ) : <EOL> """<STR_LIT>""" <EOL> FLOW_STATE = aff4 . Attribute ( "<STR_LIT>" , <EOL> rdf_flows . FlowState , <EOL> "<STR_LIT>" , <EOL> "<STR_LIT>" , versioned = False , <EOL> creates_new_object_version = False ) <EOL> NOTIFICATION = aff4 . Attribute ( "<STR_LIT>" , rdf_flows . Notification , <EOL> "<STR_LIT>" ) <EOL> CLIENT_CRASH = aff4 . Attribute ( "<STR_LIT>" , rdf_client . ClientCrash , <EOL> "<STR_LIT>" , <EOL> default = None , <EOL> creates_new_object_version = False ) <EOL> PENDING_TERMINATION = aff4 . Attribute ( "<STR_LIT>" , <EOL> PendingFlowTermination , <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" , <EOL> creates_new_object_version = False ) <EOL> category = "<STR_LIT>" <EOL> friendly_name = None <EOL> AUTHORIZED_LABELS = [ ] <EOL> __metaclass__ = registry . MetaclassRegistry <EOL> ACL_ENFORCED = True <EOL> behaviours = FlowBehaviour ( "<STR_LIT>" , "<STR_LIT>" ) <EOL> args_type = EmptyFlowArgs <EOL> runner = None <EOL> state = None <EOL> runner_cls = flow_runner . FlowRunner <EOL> handles_crashes = False <EOL> def Initialize ( self ) : <EOL> """<STR_LIT>""" <EOL> if "<STR_LIT:r>" in self . mode : <EOL> self . state = self . Get ( self . Schema . FLOW_STATE ) <EOL> if self . state : <EOL> self . Load ( ) <EOL> self . args = self . state . get ( "<STR_LIT:args>" ) <EOL> if self . state is None : <EOL> self . state = self . Schema . FLOW_STATE ( ) <EOL> elif self . state . errors : <EOL> logging . warning ( "<STR_LIT>" , <EOL> self . urn ) <EOL> self . mode = "<STR_LIT:r>" <EOL> @ classmethod <EOL> def GetDefaultArgs ( cls , token = None ) : <EOL> """<STR_LIT>""" <EOL> _ = token <EOL> return cls . args_type ( ) <EOL> @ classmethod <EOL> def FilterArgsFromSemanticProtobuf ( cls , protobuf , kwargs ) : <EOL> """<STR_LIT>""" <EOL> for descriptor in protobuf . type_infos : <EOL> value = kwargs . pop ( descriptor . name , None ) <EOL> if value is not None : <EOL> setattr ( protobuf , descriptor . name , value ) <EOL> def UpdateKillNotification ( self ) : <EOL> if ( self . runner . schedule_kill_notifications and <EOL> self . runner . context . kill_timestamp ) : <EOL> with queue_manager . QueueManager ( token = self . token ) as manager : <EOL> manager . DeleteNotification ( <EOL> self . session_id , <EOL> start = self . runner . context . kill_timestamp , <EOL> end = self . runner . context . kill_timestamp + <EOL> rdfvalue . Duration ( "<STR_LIT>" ) ) <EOL> stuck_flows_timeout = rdfvalue . Duration ( <EOL> config_lib . CONFIG [ "<STR_LIT>" ] ) <EOL> self . runner . context . kill_timestamp = ( rdfvalue . RDFDatetime ( ) . Now ( ) + <EOL> stuck_flows_timeout ) <EOL> manager . QueueNotification ( <EOL> session_id = self . session_id , in_progress = True , <EOL> timestamp = self . runner . context . kill_timestamp ) <EOL> def HeartBeat ( self ) : <EOL> if self . locked : <EOL> lease_time = config_lib . CONFIG [ "<STR_LIT>" ] <EOL> if self . CheckLease ( ) < lease_time / <NUM_LIT:2> : <EOL> logging . info ( "<STR_LIT>" , self . session_id ) <EOL> self . UpdateLease ( lease_time ) <EOL> self . UpdateKillNotification ( ) <EOL> else : <EOL> logging . warning ( "<STR_LIT>" , self . urn ) <EOL> def WriteState ( self ) : <EOL> if "<STR_LIT:w>" in self . mode : <EOL> if self . state . Empty ( ) : <EOL> raise IOError ( "<STR_LIT>" % <EOL> self . urn ) <EOL> self . Set ( self . Schema . FLOW_STATE ( self . state ) ) <EOL> def FlushMessages ( self ) : <EOL> """<STR_LIT>""" <EOL> self . GetRunner ( ) . FlushMessages ( ) <EOL> def Flush ( self , sync = True ) : <EOL> """<STR_LIT>""" <EOL> self . CheckLease ( ) <EOL> self . Save ( ) <EOL> self . WriteState ( ) <EOL> self . Load ( ) <EOL> super ( GRRFlow , self ) . Flush ( sync = sync ) <EOL> self . FlushMessages ( ) <EOL> def Close ( self , sync = True ) : <EOL> """<STR_LIT>""" <EOL> self . CheckLease ( ) <EOL> self . Save ( ) <EOL> self . WriteState ( ) <EOL> super ( GRRFlow , self ) . Close ( sync = sync ) <EOL> self . FlushMessages ( ) <EOL> def CreateRunner ( self , parent_runner = None , runner_args = None , ** kw ) : <EOL> """<STR_LIT>""" <EOL> self . runner = self . runner_cls ( <EOL> self , token = self . token , parent_runner = parent_runner , <EOL> runner_args = runner_args , ** kw ) <EOL> return self . runner <EOL> def GetRunner ( self ) : <EOL> if self . runner : <EOL> return self . runner <EOL> self . runner = self . runner_cls ( self ) <EOL> return self . runner <EOL> @ StateHandler ( ) <EOL> def End ( self ) : <EOL> """<STR_LIT>""" <EOL> if self . runner . output is not None : <EOL> self . Notify ( <EOL> "<STR_LIT>" , self . runner . output . urn , <EOL> u"<STR_LIT>" . format ( len ( self . runner . output ) ) ) <EOL> else : <EOL> self . Notify ( "<STR_LIT>" , self . client_id , <EOL> "<STR_LIT>" % self . __class__ . __name__ ) <EOL> @ StateHandler ( ) <EOL> def Start ( self , unused_message = None ) : <EOL> """<STR_LIT>""" <EOL> pass <EOL> def Load ( self ) : <EOL> """<STR_LIT>""" <EOL> def Save ( self ) : <EOL> """<STR_LIT>""" <EOL> def Log ( self , format_str , * args ) : <EOL> """<STR_LIT>""" <EOL> self . GetRunner ( ) . Log ( format_str , * args ) <EOL> def GetLog ( self ) : <EOL> return self . GetRunner ( ) . GetLog ( ) <EOL> def Status ( self , format_str , * args ) : <EOL> """<STR_LIT>""" <EOL> self . GetRunner ( ) . Status ( format_str , * args ) <EOL> def Notify ( self , message_type , subject , msg ) : <EOL> """<STR_LIT>""" <EOL> self . GetRunner ( ) . Notify ( message_type , subject , msg ) <EOL> def Publish ( self , event_name , message = None , session_id = None , <EOL> priority = rdf_flows . GrrMessage . Priority . MEDIUM_PRIORITY ) : <EOL> """<STR_LIT>""" <EOL> result = message <EOL> logging . debug ( "<STR_LIT>" , <EOL> utils . SmartUnicode ( message ) [ : <NUM_LIT:100> ] , event_name ) <EOL> if not isinstance ( message , rdf_flows . GrrMessage ) : <EOL> result = rdf_flows . GrrMessage ( payload = message ) <EOL> result . session_id = session_id or self . session_id <EOL> result . auth_state = rdf_flows . GrrMessage . AuthorizationState . AUTHENTICATED <EOL> result . source = self . session_id <EOL> result . priority = priority <EOL> self . runner . Publish ( event_name , result ) <EOL> def CallClient ( self , action_name , request = None , next_state = None , <EOL> request_data = None , ** kwargs ) : <EOL> return self . runner . CallClient ( <EOL> action_name = action_name , request = request , next_state = next_state , <EOL> request_data = request_data , ** kwargs ) <EOL> def CallStateInline ( self , messages = None , next_state = "<STR_LIT>" , request_data = None , <EOL> responses = None ) : <EOL> if responses is None : <EOL> responses = FakeResponses ( messages , request_data ) <EOL> getattr ( self , next_state ) ( self . runner , direct_response = responses ) <EOL> def CallState ( self , messages = None , next_state = "<STR_LIT>" , request_data = None , <EOL> start_time = None ) : <EOL> return self . runner . CallState ( messages = messages , next_state = next_state , <EOL> request_data = request_data , <EOL> start_time = start_time ) <EOL> def CallFlow ( self , flow_name , next_state = None , request_data = None , ** kwargs ) : <EOL> return self . runner . CallFlow ( flow_name , next_state = next_state , <EOL> request_data = request_data , ** kwargs ) <EOL> def SendReply ( self , response ) : <EOL> return self . runner . SendReply ( response ) <EOL> def Error ( self , backtrace , client_id = None ) : <EOL> return self . runner . Error ( backtrace , client_id = client_id ) <EOL> def Terminate ( self ) : <EOL> return self . runner . Terminate ( ) <EOL> @ property <EOL> def session_id ( self ) : <EOL> return self . state . context . session_id <EOL> @ property <EOL> def client_id ( self ) : <EOL> return self . state . context . args . client_id <EOL> def Name ( self ) : <EOL> return self . __class__ . __name__ <EOL> @ classmethod <EOL> def StartFlow ( cls , args = None , runner_args = None , <EOL> parent_flow = None , _store = None , sync = True , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> if runner_args is None : <EOL> runner_args = flow_runner . FlowRunnerArgs ( ) <EOL> cls . FilterArgsFromSemanticProtobuf ( runner_args , kwargs ) <EOL> if runner_args . start_time : <EOL> sync = False <EOL> if runner_args . flow_name not in GRRFlow . classes : <EOL> stats . STATS . IncrementCounter ( "<STR_LIT>" ) <EOL> raise RuntimeError ( "<STR_LIT>" % runner_args . flow_name ) <EOL> if not runner_args . HasField ( "<STR_LIT>" ) : <EOL> if data_store . default_token is None : <EOL> raise access_control . UnauthorizedAccess ( "<STR_LIT>" ) <EOL> runner_args . token = data_store . default_token . Copy ( ) <EOL> data_store . DB . security_manager . CheckIfCanStartFlow ( <EOL> runner_args . token , runner_args . flow_name , <EOL> with_client_id = runner_args . client_id ) <EOL> flow_cls = GRRFlow . GetPlugin ( runner_args . flow_name ) <EOL> if flow_cls . ACL_ENFORCED and runner_args . client_id : <EOL> data_store . DB . security_manager . CheckClientAccess ( runner_args . token , <EOL> runner_args . client_id ) <EOL> token = runner_args . token . SetUID ( ) <EOL> token . expiry = <NUM_LIT> * rdfvalue . RDFDatetime . converter <EOL> flow_obj = aff4 . FACTORY . Create ( None , runner_args . flow_name , <EOL> token = token ) <EOL> if args is None : <EOL> args = flow_obj . args_type ( ) <EOL> cls . FilterArgsFromSemanticProtobuf ( args , kwargs ) <EOL> args . Validate ( ) <EOL> flow_obj . state . Register ( "<STR_LIT:args>" , args ) <EOL> flow_obj . args = args <EOL> if kwargs : <EOL> raise type_info . UnknownArg ( <EOL> "<STR_LIT>" % kwargs ) <EOL> if parent_flow : <EOL> parent_runner = parent_flow . runner <EOL> else : <EOL> parent_runner = None <EOL> runner = flow_obj . CreateRunner ( parent_runner = parent_runner , <EOL> runner_args = runner_args , <EOL> _store = _store or data_store . DB ) <EOL> logging . info ( u"<STR_LIT>" , flow_obj . urn , <EOL> runner_args . flow_name , runner_args . client_id ) <EOL> if sync : <EOL> flow_obj . Start ( ) <EOL> else : <EOL> runner . CallState ( next_state = "<STR_LIT>" , start_time = runner_args . start_time ) <EOL> if not runner . OutstandingRequests ( ) : <EOL> flow_obj . Terminate ( ) <EOL> flow_obj . Close ( ) <EOL> if parent_flow is None : <EOL> Events . PublishEvent ( "<STR_LIT>" , <EOL> AuditEvent ( user = token . username , <EOL> action = "<STR_LIT>" , <EOL> flow_name = runner_args . flow_name , <EOL> urn = flow_obj . urn , <EOL> client = runner_args . client_id ) , <EOL> token = token ) <EOL> return flow_obj . urn <EOL> @ classmethod <EOL> def MarkForTermination ( cls , flow_urn , reason = None , sync = False , token = None ) : <EOL> """<STR_LIT>""" <EOL> data_store . DB . Set ( flow_urn , <EOL> cls . SchemaCls . PENDING_TERMINATION . predicate , <EOL> PendingFlowTermination ( reason = reason ) , <EOL> replace = False , sync = sync , token = token ) <EOL> @ classmethod <EOL> def TerminateFlow ( cls , flow_id , reason = None , status = None , token = None , <EOL> force = False ) : <EOL> """<STR_LIT>""" <EOL> if not force : <EOL> flow_obj = aff4 . FACTORY . OpenWithLock ( flow_id , aff4_type = "<STR_LIT>" , <EOL> blocking = True , token = token ) <EOL> else : <EOL> flow_obj = aff4 . FACTORY . Open ( flow_id , aff4_type = "<STR_LIT>" , mode = "<STR_LIT>" , <EOL> token = token ) <EOL> if not flow_obj : <EOL> raise FlowError ( "<STR_LIT>" % flow_id ) <EOL> with flow_obj : <EOL> runner = flow_obj . GetRunner ( ) <EOL> if not runner . IsRunning ( ) : <EOL> return <EOL> if token is None : <EOL> token = access_control . ACLToken ( ) <EOL> if reason is None : <EOL> reason = "<STR_LIT>" <EOL> data_store . DB . security_manager . CheckIfCanStartFlow ( <EOL> token , flow_obj . Name ( ) , with_client_id = runner . args . client_id ) <EOL> runner . Error ( reason , status = status ) <EOL> flow_obj . Log ( "<STR_LIT>" . format ( <EOL> token . username , reason ) ) <EOL> super_token = token . SetUID ( ) <EOL> children_to_kill = aff4 . FACTORY . MultiOpen ( <EOL> flow_obj . ListChildren ( ) , token = super_token , aff4_type = "<STR_LIT>" ) <EOL> for child_obj in children_to_kill : <EOL> cls . TerminateFlow ( child_obj . urn , reason = "<STR_LIT>" , <EOL> token = super_token , force = force ) <EOL> @ classmethod <EOL> def PrintArgsHelp ( cls ) : <EOL> print cls . GetArgsHelpAsString ( ) <EOL> @ classmethod <EOL> def _ClsHelpEpilog ( cls ) : <EOL> return cls . GetArgsHelpAsString ( ) <EOL> @ classmethod <EOL> def GetCallingPrototypeAsString ( cls ) : <EOL> """<STR_LIT>""" <EOL> output = [ ] <EOL> output . append ( "<STR_LIT>" ) <EOL> output . append ( "<STR_LIT>" % cls . __name__ ) <EOL> prototypes = [ ] <EOL> if cls . args_type : <EOL> for type_descriptor in cls . args_type . type_infos : <EOL> if not type_descriptor . hidden : <EOL> prototypes . append ( "<STR_LIT>" % ( type_descriptor . name , <EOL> type_descriptor . name ) ) <EOL> output . append ( "<STR_LIT:U+002CU+0020>" . join ( prototypes ) ) <EOL> output . append ( "<STR_LIT:)>" ) <EOL> return "<STR_LIT>" . join ( output ) <EOL> @ classmethod <EOL> def GetArgs ( cls ) : <EOL> """<STR_LIT>""" <EOL> args = { } <EOL> if cls . args_type : <EOL> for type_descriptor in cls . args_type . type_infos : <EOL> if not type_descriptor . hidden : <EOL> args [ type_descriptor . name ] = { <EOL> "<STR_LIT:description>" : type_descriptor . description , <EOL> "<STR_LIT:default>" : type_descriptor . default , <EOL> "<STR_LIT:type>" : "<STR_LIT>" , <EOL> } <EOL> if type_descriptor . type : <EOL> args [ type_descriptor . name ] [ "<STR_LIT:type>" ] = type_descriptor . type . __name__ <EOL> return args <EOL> @ classmethod <EOL> def GetArgsHelpAsString ( cls ) : <EOL> """<STR_LIT>""" <EOL> output = [ "<STR_LIT>" , "<STR_LIT>" % cls . GetCallingPrototypeAsString ( ) , "<STR_LIT>" ] <EOL> arg_list = sorted ( cls . GetArgs ( ) . items ( ) , key = lambda x : x [ <NUM_LIT:0> ] ) <EOL> if not arg_list : <EOL> output . append ( "<STR_LIT>" ) <EOL> else : <EOL> output . append ( "<STR_LIT>" ) <EOL> for arg , val in arg_list : <EOL> output . append ( "<STR_LIT>" % arg ) <EOL> output . append ( "<STR_LIT>" % val [ "<STR_LIT:description>" ] ) <EOL> output . append ( "<STR_LIT>" % val [ "<STR_LIT:type>" ] ) <EOL> output . append ( "<STR_LIT>" % val [ "<STR_LIT:default>" ] ) <EOL> output . append ( "<STR_LIT>" ) <EOL> return "<STR_LIT:\n>" . join ( output ) <EOL> @ staticmethod <EOL> def GetFlowRequests ( flow_urns , token = None ) : <EOL> """<STR_LIT>""" <EOL> flow_requests = { } <EOL> flow_request_urns = [ flow_urn . Add ( "<STR_LIT:state>" ) for flow_urn in flow_urns ] <EOL> for flow_urn , values in data_store . DB . MultiResolvePrefix ( <EOL> flow_request_urns , "<STR_LIT>" , token = token ) : <EOL> for subject , serialized , _ in values : <EOL> try : <EOL> if "<STR_LIT:status>" in subject : <EOL> msg = rdf_flows . GrrMessage ( serialized ) <EOL> else : <EOL> msg = rdf_flows . RequestState ( serialized ) <EOL> except Exception as e : <EOL> logging . warn ( "<STR_LIT>" , e ) <EOL> continue <EOL> flow_requests . setdefault ( flow_urn , [ ] ) . append ( msg ) <EOL> return flow_requests <EOL> class GRRGlobalFlow ( GRRFlow ) : <EOL> """<STR_LIT>""" <EOL> behaviours = GRRFlow . behaviours + "<STR_LIT>" - "<STR_LIT>" <EOL> class WellKnownFlow ( GRRFlow ) : <EOL> """<STR_LIT>""" <EOL> well_known_session_id = None <EOL> category = None <EOL> @ classmethod <EOL> def GetAllWellKnownFlows ( cls , token = None ) : <EOL> """<STR_LIT>""" <EOL> well_known_flows = { } <EOL> for cls in GRRFlow . classes . values ( ) : <EOL> if aff4 . issubclass ( cls , WellKnownFlow ) and cls . well_known_session_id : <EOL> well_known_flow = cls ( cls . well_known_session_id , <EOL> mode = "<STR_LIT>" , token = token ) <EOL> well_known_flows [ cls . well_known_session_id . FlowName ( ) ] = well_known_flow <EOL> return well_known_flows <EOL> def _SafeProcessMessage ( self , * args , ** kwargs ) : <EOL> try : <EOL> self . ProcessMessage ( * args , ** kwargs ) <EOL> except Exception as e : <EOL> logging . exception ( "<STR_LIT>" , e ) <EOL> stats . STATS . IncrementCounter ( "<STR_LIT>" , <EOL> fields = [ str ( self . session_id ) ] ) <EOL> def CallState ( self , messages = None , next_state = None , delay = <NUM_LIT:0> ) : <EOL> """<STR_LIT>""" <EOL> pass <EOL> @ property <EOL> def session_id ( self ) : <EOL> return self . well_known_session_id <EOL> def OutstandingRequests ( self ) : <EOL> return <NUM_LIT:1> <EOL> def FlushMessages ( self ) : <EOL> """<STR_LIT>""" <EOL> pass <EOL> def FetchAndRemoveRequestsAndResponses ( self , session_id ) : <EOL> """<STR_LIT>""" <EOL> messages = [ ] <EOL> with queue_manager . WellKnownQueueManager ( <EOL> token = self . token ) as manager : <EOL> for _ , responses in manager . FetchRequestsAndResponses ( session_id ) : <EOL> messages . extend ( responses ) <EOL> manager . DeleteWellKnownFlowResponses ( session_id , responses ) <EOL> return messages <EOL> def ProcessResponses ( self , responses , thread_pool ) : <EOL> """<STR_LIT>""" <EOL> for response in responses : <EOL> thread_pool . AddTask ( target = self . _SafeProcessMessage , <EOL> args = ( response , ) , name = self . __class__ . __name__ ) <EOL> def ProcessMessages ( self , msgs ) : <EOL> for msg in msgs : <EOL> self . ProcessMessage ( msg ) <EOL> self . HeartBeat ( ) <EOL> def ProcessMessage ( self , msg ) : <EOL> """<STR_LIT>""" <EOL> def CallClient ( self , client_id , action_name , request = None , <EOL> response_session_id = None , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> if client_id is None : <EOL> raise FlowError ( "<STR_LIT>" ) <EOL> client_id = rdf_client . ClientURN ( client_id ) <EOL> try : <EOL> action = actions . ActionPlugin . classes [ action_name ] <EOL> except KeyError : <EOL> raise RuntimeError ( "<STR_LIT>" % action_name ) <EOL> if action . in_rdfvalue is None : <EOL> if request : <EOL> raise RuntimeError ( "<STR_LIT>" % <EOL> action_name ) <EOL> else : <EOL> if request is None : <EOL> request = action . in_rdfvalue ( ** kwargs ) <EOL> else : <EOL> if not isinstance ( request , action . in_rdfvalue ) : <EOL> raise RuntimeError ( "<STR_LIT>" % ( <EOL> action . in_rdfvalue , type ( request ) ) ) <EOL> if response_session_id is None : <EOL> cls = GRRFlow . classes [ "<STR_LIT>" ] <EOL> response_session_id = cls . well_known_session_id <EOL> msg = rdf_flows . GrrMessage ( <EOL> session_id = utils . SmartUnicode ( response_session_id ) , <EOL> name = action_name , request_id = <NUM_LIT:0> , queue = client_id . Queue ( ) , <EOL> payload = request ) <EOL> queue_manager . QueueManager ( token = self . token ) . Schedule ( msg ) <EOL> def WriteState ( self ) : <EOL> if "<STR_LIT:w>" in self . mode : <EOL> self . Set ( self . Schema . FLOW_STATE ( self . state ) ) <EOL> def UpdateKillNotification ( self ) : <EOL> pass <EOL> def EventHandler ( source_restriction = False , auth_required = True , <EOL> allow_client_access = False ) : <EOL> """<STR_LIT>""" <EOL> def Decorator ( f ) : <EOL> """<STR_LIT>""" <EOL> @ functools . wraps ( f ) <EOL> def Decorated ( self , msg ) : <EOL> """<STR_LIT>""" <EOL> if ( auth_required and <EOL> msg . auth_state != msg . AuthorizationState . AUTHENTICATED ) : <EOL> raise RuntimeError ( "<STR_LIT>" % msg . source ) <EOL> if ( not allow_client_access and msg . source and <EOL> rdf_client . ClientURN . Validate ( msg . source ) ) : <EOL> raise RuntimeError ( "<STR_LIT>" ) <EOL> if source_restriction : <EOL> source_check_method = getattr ( self , "<STR_LIT>" ) <EOL> if not source_check_method : <EOL> raise RuntimeError ( "<STR_LIT>" ) <EOL> if not source_check_method ( msg . source ) : <EOL> raise RuntimeError ( "<STR_LIT>" ) <EOL> stats . STATS . IncrementCounter ( "<STR_LIT>" ) <EOL> rdf_msg = rdf_flows . GrrMessage ( msg ) <EOL> res = f ( self , message = rdf_msg , event = rdf_msg . payload ) <EOL> return res <EOL> return Decorated <EOL> return Decorator <EOL> class EventListener ( WellKnownFlow ) : <EOL> """<STR_LIT>""" <EOL> EVENTS = [ ] <EOL> __metaclass__ = registry . EventRegistry <EOL> @ EventHandler ( auth_required = True ) <EOL> def ProcessMessage ( self , message = None , event = None ) : <EOL> """<STR_LIT>""" <EOL> class Events ( object ) : <EOL> """<STR_LIT>""" <EOL> @ classmethod <EOL> def PublishEvent ( cls , event_name , msg , token = None ) : <EOL> """<STR_LIT>""" <EOL> cls . PublishMultipleEvents ( { event_name : [ msg ] } , token = token ) <EOL> @ classmethod <EOL> def PublishMultipleEvents ( cls , events , token = None ) : <EOL> """<STR_LIT>""" <EOL> with queue_manager . WellKnownQueueManager ( token = token ) as manager : <EOL> event_name_map = registry . EventRegistry . EVENT_NAME_MAP <EOL> for event_name , messages in events . iteritems ( ) : <EOL> handler_urns = [ ] <EOL> if isinstance ( event_name , basestring ) : <EOL> for event_cls in event_name_map . get ( event_name , [ ] ) : <EOL> if event_cls . well_known_session_id is None : <EOL> logging . error ( "<STR_LIT>" , <EOL> event_cls . __name__ ) <EOL> else : <EOL> handler_urns . append ( event_cls . well_known_session_id ) <EOL> else : <EOL> handler_urns . append ( event_name ) <EOL> for msg in messages : <EOL> if not isinstance ( msg , rdfvalue . RDFValue ) : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> if not isinstance ( msg , rdf_flows . GrrMessage ) : <EOL> msg = rdf_flows . GrrMessage ( payload = msg ) <EOL> msg . response_id = msg . task_id = msg . GenerateTaskID ( ) <EOL> msg . request_id = <NUM_LIT:0> <EOL> for event_urn in handler_urns : <EOL> manager . QueueResponse ( event_urn , msg ) <EOL> manager . QueueNotification ( <EOL> rdf_flows . GrrNotification ( session_id = event_urn , <EOL> priority = msg . priority ) ) <EOL> @ classmethod <EOL> def PublishEventInline ( cls , event_name , msg , token = None ) : <EOL> """<STR_LIT>""" <EOL> if not isinstance ( msg , rdfvalue . RDFValue ) : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> if not isinstance ( msg , rdf_flows . GrrMessage ) : <EOL> msg = rdf_flows . GrrMessage ( payload = msg ) <EOL> if not isinstance ( event_name , basestring ) : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> event_name_map = registry . EventRegistry . EVENT_NAME_MAP <EOL> for event_cls in event_name_map . get ( event_name , [ ] ) : <EOL> event_obj = event_cls ( event_cls . well_known_session_id , <EOL> mode = "<STR_LIT>" , token = token ) <EOL> event_obj . ProcessMessage ( msg ) <EOL> class ServerPubKeyCache ( communicator . PubKeyCache ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , client_cache , token = None ) : <EOL> self . client_cache = client_cache <EOL> self . token = token <EOL> def GetRSAPublicKey ( self , common_name = "<STR_LIT>" ) : <EOL> """<STR_LIT>""" <EOL> common_name = str ( common_name ) <EOL> try : <EOL> client = self . client_cache . Get ( common_name ) <EOL> cert = client . Get ( client . Schema . CERT ) <EOL> return cert . GetPubKey ( ) <EOL> except ( KeyError , AttributeError ) : <EOL> client = aff4 . FACTORY . Create ( common_name , "<STR_LIT>" , mode = "<STR_LIT>" , <EOL> token = self . token , ignore_cache = True ) <EOL> cert = client . Get ( client . Schema . CERT ) <EOL> if not cert : <EOL> stats . STATS . IncrementCounter ( "<STR_LIT>" ) <EOL> raise communicator . UnknownClientCert ( "<STR_LIT>" ) <EOL> if rdfvalue . RDFURN ( cert . common_name ) != rdfvalue . RDFURN ( common_name ) : <EOL> logging . error ( "<STR_LIT>" , common_name ) <EOL> raise communicator . UnknownClientCert ( "<STR_LIT>" ) <EOL> self . client_cache . Put ( common_name , client ) <EOL> stats . STATS . SetGaugeValue ( "<STR_LIT>" , <EOL> len ( self . client_cache ) ) <EOL> return cert . GetPubKey ( ) <EOL> class ServerCommunicator ( communicator . Communicator ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , certificate , private_key , token = None ) : <EOL> self . client_cache = utils . FastStore ( <NUM_LIT:1000> ) <EOL> self . token = token <EOL> super ( ServerCommunicator , self ) . __init__ ( certificate = certificate , <EOL> private_key = private_key ) <EOL> self . pub_key_cache = ServerPubKeyCache ( self . client_cache , token = token ) <EOL> def _LoadOurCertificate ( self ) : <EOL> """<STR_LIT>""" <EOL> self . cert = X509 . load_cert_string ( str ( self . certificate ) ) <EOL> self . common_name = self . pub_key_cache . GetCNFromCert ( self . cert ) <EOL> def VerifyMessageSignature ( self , response_comms , signed_message_list , <EOL> cipher , api_version ) : <EOL> """<STR_LIT>""" <EOL> result = rdf_flows . GrrMessage . AuthorizationState . UNAUTHENTICATED <EOL> try : <EOL> if cipher . signature_verified : <EOL> result = rdf_flows . GrrMessage . AuthorizationState . AUTHENTICATED <EOL> if result == rdf_flows . GrrMessage . AuthorizationState . AUTHENTICATED : <EOL> client_id = cipher . cipher_metadata . source <EOL> try : <EOL> client = self . client_cache . Get ( client_id ) <EOL> except KeyError : <EOL> client = aff4 . FACTORY . Create ( client_id , <EOL> "<STR_LIT>" , mode = "<STR_LIT>" , <EOL> token = self . token ) <EOL> self . client_cache . Put ( client_id , client ) <EOL> stats . STATS . SetGaugeValue ( "<STR_LIT>" , <EOL> len ( self . client_cache ) ) <EOL> ip = response_comms . orig_request . source_ip <EOL> client . Set ( client . Schema . CLIENT_IP ( ip ) ) <EOL> remote_time = client . Get ( client . Schema . CLOCK ) or <NUM_LIT:0> <EOL> client_time = signed_message_list . timestamp or <NUM_LIT:0> <EOL> if client_time > long ( remote_time ) : <EOL> stats . STATS . IncrementCounter ( "<STR_LIT>" ) <EOL> client . Set ( client . Schema . CLOCK , rdfvalue . RDFDatetime ( client_time ) ) <EOL> client . Set ( client . Schema . PING , rdfvalue . RDFDatetime ( ) . Now ( ) ) <EOL> else : <EOL> logging . debug ( "<STR_LIT>" , <EOL> client_id , long ( remote_time ) , int ( client_time ) ) <EOL> return rdf_flows . GrrMessage . AuthorizationState . DESYNCHRONIZED <EOL> client . Flush ( ) <EOL> except communicator . UnknownClientCert : <EOL> pass <EOL> if result != rdf_flows . GrrMessage . AuthorizationState . AUTHENTICATED : <EOL> stats . STATS . IncrementCounter ( "<STR_LIT>" ) <EOL> return result <EOL> class FrontEndServer ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , certificate , private_key , max_queue_size = <NUM_LIT:50> , <EOL> message_expiry_time = <NUM_LIT> , max_retransmission_time = <NUM_LIT:10> , store = None , <EOL> threadpool_prefix = "<STR_LIT>" ) : <EOL> self . token = access_control . ACLToken ( username = "<STR_LIT>" , <EOL> reason = "<STR_LIT>" ) <EOL> self . token . supervisor = True <EOL> self . throttle_callback = lambda : True <EOL> self . SetThrottleBundlesRatio ( None ) <EOL> self . _communicator = ServerCommunicator ( <EOL> certificate = certificate , private_key = private_key , token = self . token ) <EOL> self . data_store = store or data_store . DB <EOL> self . receive_thread_pool = { } <EOL> self . message_expiry_time = message_expiry_time <EOL> self . max_retransmission_time = max_retransmission_time <EOL> self . max_queue_size = max_queue_size <EOL> self . thread_pool = threadpool . ThreadPool . Factory ( <EOL> threadpool_prefix , min_threads = <NUM_LIT:2> , <EOL> max_threads = config_lib . CONFIG [ "<STR_LIT>" ] ) <EOL> self . thread_pool . Start ( ) <EOL> self . well_known_flows = ( <EOL> WellKnownFlow . GetAllWellKnownFlows ( token = self . token ) ) <EOL> well_known_flow_names = self . well_known_flows . keys ( ) <EOL> for well_known_flow in well_known_flow_names : <EOL> if well_known_flow not in config_lib . CONFIG [ "<STR_LIT>" ] : <EOL> del self . well_known_flows [ well_known_flow ] <EOL> def SetThrottleCallBack ( self , callback ) : <EOL> self . throttle_callback = callback <EOL> def SetThrottleBundlesRatio ( self , throttle_bundles_ratio ) : <EOL> """<STR_LIT>""" <EOL> self . throttle_bundles_ratio = throttle_bundles_ratio <EOL> if throttle_bundles_ratio is None : <EOL> self . handled_bundles = [ ] <EOL> self . last_not_throttled_bundle_time = <NUM_LIT:0> <EOL> stats . STATS . SetGaugeValue ( "<STR_LIT>" , <EOL> str ( throttle_bundles_ratio ) ) <EOL> def UpdateAndCheckIfShouldThrottle ( self , bundle_time ) : <EOL> """<STR_LIT>""" <EOL> if self . throttle_bundles_ratio is None : <EOL> return False <EOL> self . handled_bundles . append ( bundle_time ) <EOL> oldest_limit = bundle_time - config_lib . CONFIG [ <EOL> "<STR_LIT>" ] <EOL> try : <EOL> oldest_index = next ( i for i , v in enumerate ( self . handled_bundles ) <EOL> if v > oldest_limit ) <EOL> self . handled_bundles = self . handled_bundles [ oldest_index : ] <EOL> except StopIteration : <EOL> self . handled_bundles = [ ] <EOL> blen = len ( self . handled_bundles ) <EOL> if blen > <NUM_LIT:1> : <EOL> interval = ( self . handled_bundles [ - <NUM_LIT:1> ] - <EOL> self . handled_bundles [ <NUM_LIT:0> ] ) / float ( blen - <NUM_LIT:1> ) <EOL> else : <EOL> return self . throttle_bundles_ratio == <NUM_LIT:0> <EOL> should_throttle = ( bundle_time - self . last_not_throttled_bundle_time < <EOL> interval / max ( <NUM_LIT> , <EOL> float ( self . throttle_bundles_ratio ) ) ) <EOL> if not should_throttle : <EOL> self . last_not_throttled_bundle_time = bundle_time <EOL> return should_throttle <EOL> @ stats . Counted ( "<STR_LIT>" ) <EOL> @ stats . Timed ( "<STR_LIT>" ) <EOL> def HandleMessageBundles ( self , request_comms , response_comms ) : <EOL> """<STR_LIT>""" <EOL> messages , source , timestamp = self . _communicator . DecodeMessages ( <EOL> request_comms ) <EOL> now = time . time ( ) <EOL> if messages : <EOL> self . ReceiveMessages ( source , messages ) <EOL> required_count = max ( <NUM_LIT:0> , self . max_queue_size - request_comms . queue_size ) <EOL> tasks = [ ] <EOL> message_list = rdf_flows . MessageList ( ) <EOL> if self . UpdateAndCheckIfShouldThrottle ( time . time ( ) ) : <EOL> stats . STATS . IncrementCounter ( "<STR_LIT>" ) <EOL> elif self . throttle_callback ( ) : <EOL> if time . time ( ) - now < <NUM_LIT:10> : <EOL> tasks = self . DrainTaskSchedulerQueueForClient ( source , required_count , <EOL> message_list ) <EOL> else : <EOL> stats . STATS . IncrementCounter ( "<STR_LIT>" ) <EOL> try : <EOL> self . _communicator . EncodeMessages ( <EOL> message_list , response_comms , destination = str ( source ) , <EOL> timestamp = timestamp , api_version = request_comms . api_version ) <EOL> except communicator . UnknownClientCert : <EOL> queue_manager . QueueManager ( token = self . token ) . Schedule ( tasks ) <EOL> raise <EOL> return source , len ( messages ) <EOL> def DrainTaskSchedulerQueueForClient ( self , client , max_count , <EOL> response_message ) : <EOL> """<STR_LIT>""" <EOL> if max_count <= <NUM_LIT:0> : <EOL> return [ ] <EOL> client = rdf_client . ClientURN ( client ) <EOL> start_time = time . time ( ) <EOL> new_tasks = queue_manager . QueueManager ( token = self . token ) . QueryAndOwn ( <EOL> queue = client . Queue ( ) , limit = max_count , <EOL> lease_seconds = self . message_expiry_time ) <EOL> initial_ttl = rdf_flows . GrrMessage ( ) . task_ttl <EOL> check_before_sending = [ ] <EOL> result = [ ] <EOL> for task in new_tasks : <EOL> if task . task_ttl < initial_ttl - <NUM_LIT:1> : <EOL> check_before_sending . append ( task ) <EOL> else : <EOL> response_message . job . Append ( task ) <EOL> result . append ( task ) <EOL> if check_before_sending : <EOL> with queue_manager . QueueManager ( token = self . token ) as manager : <EOL> status_found = manager . MultiCheckStatus ( check_before_sending ) <EOL> for task in check_before_sending : <EOL> if task not in status_found : <EOL> result . append ( task ) <EOL> else : <EOL> manager . DeQueueClientRequest ( client , task . task_id ) <EOL> stats . STATS . IncrementCounter ( "<STR_LIT>" , len ( result ) ) <EOL> logging . debug ( "<STR_LIT>" , <EOL> len ( result ) , client , time . time ( ) - start_time ) <EOL> return result <EOL> def ReceiveMessages ( self , client_id , messages ) : <EOL> """<STR_LIT>""" <EOL> now = time . time ( ) <EOL> with queue_manager . QueueManager ( <EOL> token = self . token , store = self . data_store ) as manager : <EOL> sessions_handled = [ ] <EOL> for session_id , msgs in utils . GroupBy ( <EOL> messages , operator . attrgetter ( "<STR_LIT>" ) ) . iteritems ( ) : <EOL> unprocessed_msgs = self . HandleWellKnownFlows ( msgs ) <EOL> if not unprocessed_msgs : continue <EOL> sessions_handled . append ( session_id ) <EOL> for msg in unprocessed_msgs : <EOL> manager . QueueResponse ( session_id , msg ) <EOL> for msg in unprocessed_msgs : <EOL> if msg . request_id == <NUM_LIT:0> : <EOL> manager . QueueNotification ( session_id = msg . session_id , <EOL> priority = msg . priority ) <EOL> break <EOL> elif msg . type == rdf_flows . GrrMessage . Type . STATUS : <EOL> manager . DeQueueClientRequest ( client_id , msg . task_id ) <EOL> manager . QueueNotification ( session_id = msg . session_id , <EOL> priority = msg . priority , <EOL> last_status = msg . request_id ) <EOL> stat = rdf_flows . GrrStatus ( msg . payload ) <EOL> if stat . status == rdf_flows . GrrStatus . ReturnedStatus . CLIENT_KILLED : <EOL> Events . PublishEvent ( "<STR_LIT>" , rdf_flows . GrrMessage ( msg ) , <EOL> token = self . token ) <EOL> logging . debug ( "<STR_LIT>" , len ( messages ) , <EOL> time . time ( ) - now ) <EOL> def HandleWellKnownFlows ( self , messages ) : <EOL> """<STR_LIT>""" <EOL> msgs_by_wkf = { } <EOL> result = [ ] <EOL> for msg in messages : <EOL> if msg . response_id != <NUM_LIT:0> : <EOL> result . append ( msg ) <EOL> continue <EOL> flow_name = msg . session_id . FlowName ( ) <EOL> if flow_name in self . well_known_flows : <EOL> msgs_by_wkf . setdefault ( flow_name , [ ] ) . append ( msg ) <EOL> stats . STATS . IncrementCounter ( "<STR_LIT>" ) <EOL> stats . STATS . IncrementCounter ( "<STR_LIT>" , <EOL> fields = [ str ( msg . session_id ) ] ) <EOL> else : <EOL> msg . response_id = utils . PRNG . GetULong ( ) <EOL> result . append ( msg ) <EOL> for flow_name , msg_list in msgs_by_wkf . iteritems ( ) : <EOL> flow = self . well_known_flows [ flow_name ] <EOL> flow . ProcessMessages ( msg_list ) <EOL> return result <EOL> def ProcessCompletedRequests ( flow_obj , thread_pool , reqs ) : <EOL> flow_obj . ProcessCompletedRequests ( thread_pool , reqs ) <EOL> class FlowInit ( registry . InitHook ) : <EOL> """<STR_LIT>""" <EOL> pre = [ "<STR_LIT>" ] <EOL> def RunOnce ( self ) : <EOL> stats . STATS . RegisterGaugeMetric ( <EOL> "<STR_LIT>" , int , fields = [ ( "<STR_LIT:source>" , str ) ] ) <EOL> stats . STATS . RegisterGaugeMetric ( <EOL> "<STR_LIT>" , int ) <EOL> stats . STATS . RegisterCounterMetric ( <EOL> "<STR_LIT>" , fields = [ ( "<STR_LIT:source>" , str ) ] ) <EOL> stats . STATS . RegisterCounterMetric ( <EOL> "<STR_LIT>" , fields = [ ( "<STR_LIT:source>" , str ) ] ) <EOL> stats . STATS . RegisterCounterMetric ( <EOL> "<STR_LIT>" , fields = [ ( "<STR_LIT:source>" , str ) ] ) <EOL> stats . STATS . RegisterCounterMetric ( <EOL> "<STR_LIT>" , fields = [ ( "<STR_LIT:source>" , str ) ] ) <EOL> stats . STATS . RegisterEventMetric ( <EOL> "<STR_LIT>" , fields = [ ( "<STR_LIT:source>" , str ) ] ) <EOL> stats . STATS . RegisterCounterMetric ( "<STR_LIT>" ) <EOL> stats . STATS . RegisterCounterMetric ( "<STR_LIT>" ) <EOL> stats . STATS . RegisterCounterMetric ( "<STR_LIT>" ) <EOL> stats . STATS . RegisterCounterMetric ( "<STR_LIT>" ) <EOL> stats . STATS . RegisterEventMetric ( "<STR_LIT>" ) <EOL> stats . STATS . RegisterCounterMetric ( "<STR_LIT>" ) <EOL> stats . STATS . RegisterCounterMetric ( "<STR_LIT>" ) <EOL> stats . STATS . RegisterCounterMetric ( "<STR_LIT>" ) <EOL> stats . STATS . RegisterCounterMetric ( "<STR_LIT>" ) <EOL> stats . STATS . RegisterCounterMetric ( "<STR_LIT>" ) <EOL> stats . STATS . RegisterCounterMetric ( "<STR_LIT>" ) <EOL> stats . STATS . RegisterCounterMetric ( "<STR_LIT>" ) <EOL> stats . STATS . RegisterCounterMetric ( "<STR_LIT>" ) <EOL> stats . STATS . RegisterCounterMetric ( "<STR_LIT>" ) <EOL> stats . STATS . RegisterCounterMetric ( "<STR_LIT>" ) <EOL> stats . STATS . RegisterGaugeMetric ( "<STR_LIT>" , str ) <EOL> stats . STATS . RegisterGaugeMetric ( "<STR_LIT>" , int ) <EOL> stats . STATS . RegisterCounterMetric ( "<STR_LIT>" , <EOL> fields = [ ( "<STR_LIT>" , str ) ] ) <EOL> stats . STATS . RegisterCounterMetric ( "<STR_LIT>" , <EOL> fields = [ ( "<STR_LIT>" , str ) ] ) <EOL> stats . STATS . RegisterCounterMetric ( "<STR_LIT>" , <EOL> fields = [ ( "<STR_LIT>" , str ) ] ) <EOL> stats . STATS . RegisterCounterMetric ( "<STR_LIT>" , <EOL> fields = [ ( "<STR_LIT>" , str ) ] ) <EOL> stats . STATS . RegisterCounterMetric ( "<STR_LIT>" , <EOL> fields = [ ( "<STR_LIT>" , str ) ] ) </s>
