<s> """<STR_LIT>""" <EOL> import gc <EOL> import logging <EOL> import os <EOL> import pdb <EOL> import threading <EOL> import time <EOL> import traceback <EOL> import psutil <EOL> from grr . client import client_utils <EOL> from grr . lib import config_lib <EOL> from grr . lib import flags <EOL> from grr . lib import rdfvalue <EOL> from grr . lib import registry <EOL> from grr . lib import utils <EOL> from grr . lib . rdfvalues import flows as rdf_flows <EOL> from grr . lib . rdfvalues import protodict as rdf_protodict <EOL> INITIAL_RESPONSE_ID = <NUM_LIT:1> <EOL> class Error ( Exception ) : <EOL> pass <EOL> class CPUExceededError ( Error ) : <EOL> pass <EOL> class NetworkBytesExceededError ( Error ) : <EOL> """<STR_LIT>""" <EOL> class ThreadNotFoundError ( Error ) : <EOL> """<STR_LIT>""" <EOL> class ActionPlugin ( object ) : <EOL> """<STR_LIT>""" <EOL> in_rdfvalue = None <EOL> out_rdfvalues = [ None ] <EOL> _authentication_required = True <EOL> __metaclass__ = registry . MetaclassRegistry <EOL> priority = rdf_flows . GrrMessage . Priority . MEDIUM_PRIORITY <EOL> require_fastpoll = True <EOL> last_progress_time = <NUM_LIT:0> <EOL> def __init__ ( self , grr_worker = None ) : <EOL> """<STR_LIT>""" <EOL> self . grr_worker = grr_worker <EOL> self . response_id = INITIAL_RESPONSE_ID <EOL> self . cpu_used = None <EOL> self . nanny_controller = None <EOL> self . status = rdf_flows . GrrStatus ( <EOL> status = rdf_flows . GrrStatus . ReturnedStatus . OK ) <EOL> self . _last_gc_run = rdfvalue . RDFDatetime ( ) . Now ( ) <EOL> self . _gc_frequency = config_lib . CONFIG [ "<STR_LIT>" ] <EOL> def Execute ( self , message ) : <EOL> """<STR_LIT>""" <EOL> self . message = message <EOL> if message : <EOL> self . priority = message . priority <EOL> self . require_fastpoll = message . require_fastpoll <EOL> args = None <EOL> try : <EOL> if self . message . args_rdf_name : <EOL> if not self . in_rdfvalue : <EOL> raise RuntimeError ( "<STR_LIT>" % <EOL> self . message . args_rdf_name ) <EOL> if self . in_rdfvalue . __name__ != self . message . args_rdf_name : <EOL> raise RuntimeError ( "<STR_LIT>" % <EOL> ( self . message . args_rdf_name , <EOL> self . in_rdfvalue . __name__ ) ) <EOL> args = self . message . payload <EOL> if ( self . _authentication_required and <EOL> self . message . auth_state != <EOL> rdf_flows . GrrMessage . AuthorizationState . AUTHENTICATED ) : <EOL> raise RuntimeError ( "<STR_LIT>" % <EOL> self . message . name ) <EOL> pid = os . getpid ( ) <EOL> self . proc = psutil . Process ( pid ) <EOL> user_start , system_start = self . proc . cpu_times ( ) <EOL> self . cpu_start = ( user_start , system_start ) <EOL> self . cpu_limit = self . message . cpu_limit <EOL> self . network_bytes_limit = self . message . network_bytes_limit <EOL> if getattr ( flags . FLAGS , "<STR_LIT>" , False ) : <EOL> pdb . set_trace ( ) <EOL> try : <EOL> self . Run ( args ) <EOL> finally : <EOL> user_end , system_end = self . proc . cpu_times ( ) <EOL> self . cpu_used = ( user_end - user_start , system_end - system_start ) <EOL> except NetworkBytesExceededError as e : <EOL> self . SetStatus ( rdf_flows . GrrStatus . ReturnedStatus . NETWORK_LIMIT_EXCEEDED , <EOL> "<STR_LIT>" % ( e , e ) , <EOL> traceback . format_exc ( ) ) <EOL> except Exception as e : <EOL> self . SetStatus ( rdf_flows . GrrStatus . ReturnedStatus . GENERIC_ERROR , <EOL> "<STR_LIT>" % ( e , e ) , <EOL> traceback . format_exc ( ) ) <EOL> if flags . FLAGS . debug : <EOL> self . DisableNanny ( ) <EOL> pdb . post_mortem ( ) <EOL> if self . status . status != rdf_flows . GrrStatus . ReturnedStatus . OK : <EOL> logging . info ( "<STR_LIT>" , self . __class__ . __name__ , <EOL> self . status . error_message ) <EOL> if self . status . backtrace : <EOL> logging . debug ( self . status . backtrace ) <EOL> if self . cpu_used : <EOL> self . status . cpu_time_used . user_cpu_time = self . cpu_used [ <NUM_LIT:0> ] <EOL> self . status . cpu_time_used . system_cpu_time = self . cpu_used [ <NUM_LIT:1> ] <EOL> self . SendReply ( self . status , message_type = rdf_flows . GrrMessage . Type . STATUS ) <EOL> self . _RunGC ( ) <EOL> def _RunGC ( self ) : <EOL> now = rdfvalue . RDFDatetime ( ) . Now ( ) <EOL> if now - self . _last_gc_run > self . _gc_frequency : <EOL> gc . collect ( ) <EOL> self . _last_gc_run = now <EOL> def ForceGC ( self ) : <EOL> self . _last_gc_run = rdfvalue . RDFDatetime ( <NUM_LIT:0> ) <EOL> self . _RunGC ( ) <EOL> def Run ( self , unused_args ) : <EOL> """<STR_LIT>""" <EOL> raise KeyError ( "<STR_LIT>" % <EOL> self . message . name ) <EOL> def SetStatus ( self , status , message = "<STR_LIT>" , backtrace = None ) : <EOL> """<STR_LIT>""" <EOL> self . status . status = status <EOL> self . status . error_message = utils . SmartUnicode ( message ) <EOL> if backtrace : <EOL> self . status . backtrace = utils . SmartUnicode ( backtrace ) <EOL> def SendReply ( self , rdf_value = None , <EOL> message_type = rdf_flows . GrrMessage . Type . MESSAGE , ** kw ) : <EOL> """<STR_LIT>""" <EOL> if rdf_value is None : <EOL> rdf_value = self . out_rdfvalues [ <NUM_LIT:0> ] ( ** kw ) <EOL> self . grr_worker . SendReply ( rdf_value , <EOL> name = self . __class__ . __name__ , <EOL> session_id = self . message . session_id , <EOL> response_id = self . response_id , <EOL> request_id = self . message . request_id , <EOL> message_type = message_type , <EOL> task_id = self . message . task_id , <EOL> priority = self . priority , <EOL> require_fastpoll = self . require_fastpoll ) <EOL> self . response_id += <NUM_LIT:1> <EOL> def Progress ( self ) : <EOL> """<STR_LIT>""" <EOL> now = time . time ( ) <EOL> if now - self . last_progress_time <= <NUM_LIT:2> : <EOL> return <EOL> self . last_progress_time = now <EOL> client_utils . KeepAlive ( ) <EOL> if self . nanny_controller is None : <EOL> self . nanny_controller = client_utils . NannyController ( ) <EOL> self . nanny_controller . Heartbeat ( ) <EOL> try : <EOL> user_start , system_start = self . cpu_start <EOL> user_end , system_end = self . proc . cpu_times ( ) <EOL> used_cpu = user_end - user_start + system_end - system_start <EOL> if used_cpu > self . cpu_limit : <EOL> self . grr_worker . SendClientAlert ( "<STR_LIT>" ) <EOL> raise CPUExceededError ( "<STR_LIT>" ) <EOL> except AttributeError : <EOL> pass <EOL> def SyncTransactionLog ( self ) : <EOL> """<STR_LIT>""" <EOL> if self . nanny_controller is None : <EOL> self . nanny_controller = client_utils . NannyController ( ) <EOL> self . nanny_controller . SyncTransactionLog ( ) <EOL> def ChargeBytesToSession ( self , length ) : <EOL> self . grr_worker . ChargeBytesToSession ( self . message . session_id , length , <EOL> limit = self . network_bytes_limit ) <EOL> def DisableNanny ( self ) : <EOL> try : <EOL> self . nanny_controller . nanny . Stop ( ) <EOL> except AttributeError : <EOL> logging . info ( "<STR_LIT>" ) <EOL> class IteratedAction ( ActionPlugin ) : <EOL> """<STR_LIT>""" <EOL> def Run ( self , request ) : <EOL> """<STR_LIT>""" <EOL> client_state = request . iterator . client_state . ToDict ( ) <EOL> self . Iterate ( request , client_state ) <EOL> request . iterator . client_state = rdf_protodict . Dict ( client_state ) <EOL> self . SendReply ( request . iterator , <EOL> message_type = rdf_flows . GrrMessage . Type . ITERATOR ) <EOL> def Iterate ( self , request , client_state ) : <EOL> """<STR_LIT>""" <EOL> class ClientActionWorker ( threading . Thread ) : <EOL> """<STR_LIT>""" <EOL> daemon = True <EOL> def __init__ ( self , action = None , * args , ** kw ) : <EOL> super ( ClientActionWorker , self ) . __init__ ( * args , ** kw ) <EOL> self . cond = threading . Condition ( lock = threading . RLock ( ) ) <EOL> self . id = None <EOL> self . action_obj = action <EOL> self . exception_status = None <EOL> def Resume ( self ) : <EOL> with self . cond : <EOL> self . cond . notify ( ) <EOL> self . cond . wait ( ) <EOL> def Suspend ( self ) : <EOL> with self . cond : <EOL> self . cond . notify ( ) <EOL> self . cond . wait ( ) <EOL> def run ( self ) : <EOL> self . Suspend ( ) <EOL> try : <EOL> self . action_obj . Iterate ( ) <EOL> except Exception : <EOL> if flags . FLAGS . debug : <EOL> pdb . post_mortem ( ) <EOL> self . exception_status = traceback . format_exc ( ) <EOL> raise <EOL> finally : <EOL> self . action_obj . Done ( ) <EOL> with self . cond : <EOL> self . cond . notify ( ) <EOL> class SuspendableAction ( ActionPlugin ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , * args , ** kw ) : <EOL> self . worker_cls = kw . pop ( "<STR_LIT>" , None ) <EOL> super ( SuspendableAction , self ) . __init__ ( * args , ** kw ) <EOL> self . exceptions = [ ] <EOL> self . worker = None <EOL> def Run ( self , request ) : <EOL> """<STR_LIT>""" <EOL> self . response_id = INITIAL_RESPONSE_ID <EOL> self . request = request <EOL> if not self . worker : <EOL> worker_cls = self . worker_cls or ClientActionWorker <EOL> self . worker = worker_cls ( action = self ) <EOL> self . worker . cond . acquire ( ) <EOL> self . worker . start ( ) <EOL> self . worker . Resume ( ) <EOL> self . grr_worker . suspended_actions [ id ( self ) ] = self <EOL> self . request . iterator . suspended_action = id ( self ) <EOL> self . worker . Resume ( ) <EOL> if self . worker . exception_status : <EOL> raise RuntimeError ( "<STR_LIT>" % ( <EOL> self . worker . exception_status ) ) <EOL> self . SendReply ( self . request . iterator , <EOL> message_type = rdf_flows . GrrMessage . Type . ITERATOR ) <EOL> def Done ( self ) : <EOL> self . request . iterator . state = self . request . iterator . State . FINISHED <EOL> del self . grr_worker . suspended_actions [ id ( self ) ] <EOL> def Suspend ( self ) : <EOL> self . worker . Suspend ( ) <EOL> def Iterate ( self ) : <EOL> """<STR_LIT>""" </s>
