<s> """<STR_LIT>""" <EOL> import getpass <EOL> import os <EOL> import pdb <EOL> import pickle <EOL> import tempfile <EOL> import time <EOL> from grr . client import actions <EOL> from grr . lib import access_control <EOL> from grr . lib import aff4 <EOL> from grr . lib import flow <EOL> from grr . lib import queue_manager <EOL> from grr . lib import rdfvalue <EOL> from grr . lib import worker <EOL> from grr . lib . rdfvalues import flows as rdf_flows <EOL> from grr . lib . rdfvalues import structs as rdf_structs <EOL> from grr . proto import flows_pb2 <EOL> class ClientActionArgs ( rdf_structs . RDFProtoStruct ) : <EOL> protobuf = flows_pb2 . ClientActionArgs <EOL> def GetActionArgsClass ( self ) : <EOL> if self . action : <EOL> action_cls = actions . ActionPlugin . classes . get ( self . action ) <EOL> if action_cls is None : <EOL> raise ValueError ( "<STR_LIT>" % self . action ) <EOL> return action_cls . in_rdfvalue <EOL> class ClientAction ( flow . GRRFlow ) : <EOL> """<STR_LIT>""" <EOL> args_type = ClientActionArgs <EOL> @ flow . StateHandler ( next_state = "<STR_LIT>" ) <EOL> def Start ( self ) : <EOL> if self . args . save_to : <EOL> if not os . path . isdir ( self . args . save_to ) : <EOL> os . makedirs ( self . args . save_to , <NUM_LIT:0> <NUM_LIT> ) <EOL> self . CallClient ( self . args . action , request = self . args . action_args , <EOL> next_state = "<STR_LIT>" ) <EOL> @ flow . StateHandler ( ) <EOL> def Print ( self , responses ) : <EOL> """<STR_LIT>""" <EOL> if not responses . success : <EOL> self . Log ( "<STR_LIT>" % ( self . args . action , <EOL> responses . status ) ) <EOL> if self . args . break_pdb : <EOL> pdb . set_trace ( ) <EOL> if self . args . save_to : <EOL> self . _SaveResponses ( responses ) <EOL> def _SaveResponses ( self , responses ) : <EOL> """<STR_LIT>""" <EOL> if responses : <EOL> fd = None <EOL> try : <EOL> fdint , fname = tempfile . mkstemp ( prefix = "<STR_LIT>" , <EOL> dir = self . args . save_to ) <EOL> fd = os . fdopen ( fdint , "<STR_LIT:wb>" ) <EOL> pickle . dump ( responses , fd ) <EOL> self . Log ( "<STR_LIT>" , len ( responses ) , fname ) <EOL> finally : <EOL> if fd : <EOL> fd . close ( ) <EOL> class ConsoleDebugFlowArgs ( rdf_structs . RDFProtoStruct ) : <EOL> protobuf = flows_pb2 . ConsoleDebugFlowArgs <EOL> def GetFlowArgsClass ( self ) : <EOL> if self . flow : <EOL> flow_cls = flow . GRRFlow . classes . get ( self . flow ) <EOL> if flow_cls is None : <EOL> raise ValueError ( "<STR_LIT>" % self . flow ) <EOL> return flow_cls . args_type <EOL> class ConsoleDebugFlow ( flow . GRRFlow ) : <EOL> """<STR_LIT>""" <EOL> args_type = ConsoleDebugFlowArgs <EOL> @ flow . StateHandler ( next_state = "<STR_LIT>" ) <EOL> def Start ( self ) : <EOL> if self . args . save_to : <EOL> if not os . path . isdir ( self . args . save_to ) : <EOL> os . makedirs ( self . args . save_to , <NUM_LIT:0> <NUM_LIT> ) <EOL> self . CallFlow ( self . args . flow , next_state = "<STR_LIT>" , <EOL> ** self . args . args . ToDict ( ) ) <EOL> @ flow . StateHandler ( ) <EOL> def Print ( self , responses ) : <EOL> """<STR_LIT>""" <EOL> if not responses . success : <EOL> self . Log ( "<STR_LIT>" % ( self . args . flow , <EOL> responses . status ) ) <EOL> self . Log ( "<STR_LIT>" , len ( responses ) ) <EOL> for response in responses : <EOL> print response <EOL> if self . args . break_pdb : <EOL> pdb . set_trace ( ) <EOL> if self . args . save_to : <EOL> self . _SaveResponses ( responses ) <EOL> def _SaveResponses ( self , responses ) : <EOL> """<STR_LIT>""" <EOL> if responses : <EOL> fd = None <EOL> try : <EOL> fdint , fname = tempfile . mkstemp ( prefix = "<STR_LIT>" , <EOL> dir = self . args . save_to ) <EOL> fd = os . fdopen ( fdint , "<STR_LIT:wb>" ) <EOL> pickle . dump ( responses , fd ) <EOL> self . Log ( "<STR_LIT>" , len ( responses ) , fname ) <EOL> finally : <EOL> if fd : <EOL> fd . close ( ) <EOL> def StartFlowAndWorker ( client_id , flow_name , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> queue = rdfvalue . RDFURN ( "<STR_LIT>" % getpass . getuser ( ) ) <EOL> if "<STR_LIT>" in kwargs : <EOL> token = kwargs . pop ( "<STR_LIT>" ) <EOL> else : <EOL> token = access_control . ACLToken ( username = "<STR_LIT>" ) <EOL> session_id = flow . GRRFlow . StartFlow ( client_id = client_id , <EOL> flow_name = flow_name , queue = queue , <EOL> token = token , ** kwargs ) <EOL> worker_thrd = worker . GRRWorker ( <EOL> queues = [ queue ] , token = token , <EOL> threadpool_size = <NUM_LIT:1> ) <EOL> while True : <EOL> try : <EOL> worker_thrd . RunOnce ( ) <EOL> except KeyboardInterrupt : <EOL> print "<STR_LIT>" <EOL> worker_thrd . thread_pool . Join ( ) <EOL> break <EOL> time . sleep ( <NUM_LIT:2> ) <EOL> with aff4 . FACTORY . Open ( session_id , token = token ) as flow_obj : <EOL> if not flow_obj . GetRunner ( ) . IsRunning ( ) : <EOL> break <EOL> worker_thrd . thread_pool . Join ( ) <EOL> return session_id <EOL> def TestClientActionWithWorker ( client_id , client_action , print_request = False , <EOL> break_pdb = True , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> action_cls = actions . ActionPlugin . classes [ client_action ] <EOL> request = action_cls . in_rdfvalue ( ** kwargs ) <EOL> if print_request : <EOL> print str ( request ) <EOL> StartFlowAndWorker ( client_id , flow_name = "<STR_LIT>" , action = client_action , <EOL> break_pdb = break_pdb , action_args = request ) <EOL> def WakeStuckFlow ( session_id ) : <EOL> """<STR_LIT>""" <EOL> session_id = rdfvalue . SessionID ( session_id ) <EOL> woken = <NUM_LIT:0> <EOL> checked_pending = False <EOL> with queue_manager . QueueManager ( ) as manager : <EOL> for request , responses in manager . FetchRequestsAndResponses ( session_id ) : <EOL> if not checked_pending : <EOL> task = manager . Query ( request . client_id , <EOL> task_id = "<STR_LIT>" % request . request . task_id ) <EOL> if task : <EOL> return <EOL> checked_pending = True <EOL> if ( not responses or responses [ - <NUM_LIT:1> ] . type != <EOL> rdf_flows . GrrMessage . Type . STATUS ) : <EOL> manager . QueueClientMessage ( request . request ) <EOL> woken += <NUM_LIT:1> <EOL> if responses and responses [ - <NUM_LIT:1> ] . type == rdf_flows . GrrMessage . Type . STATUS : <EOL> manager . QueueNotification ( session_id ) <EOL> return woken </s>
