<s> """<STR_LIT>""" <EOL> from grr . lib import aff4 <EOL> from grr . lib import config_lib <EOL> from grr . lib import flow <EOL> from grr . lib import rdfvalue <EOL> class Error ( Exception ) : <EOL> """<STR_LIT>""" <EOL> class ErrorDailyFlowRequestLimitExceeded ( Error ) : <EOL> """<STR_LIT>""" <EOL> class ErrorFlowDuplicate ( Error ) : <EOL> """<STR_LIT>""" <EOL> class FlowThrottler ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , daily_req_limit = None , dup_interval = None ) : <EOL> """<STR_LIT>""" <EOL> if dup_interval is None : <EOL> self . daily_req_limit = config_lib . CONFIG . Get ( "<STR_LIT>" ) <EOL> else : <EOL> self . daily_req_limit = daily_req_limit <EOL> if dup_interval is None : <EOL> self . dup_interval = config_lib . CONFIG . Get ( "<STR_LIT>" ) <EOL> else : <EOL> self . dup_interval = dup_interval <EOL> def EnforceLimits ( self , client_id , user , flow_name , flow_args , token = None ) : <EOL> """<STR_LIT>""" <EOL> if self . dup_interval is None and self . daily_req_limit == <NUM_LIT:0> : <EOL> return <EOL> flows_dir = aff4 . FACTORY . Open ( client_id . Add ( "<STR_LIT>" ) , token = token ) <EOL> now = rdfvalue . RDFDatetime ( ) . Now ( ) <EOL> earlier = now - rdfvalue . Duration ( "<STR_LIT>" ) <EOL> dup_boundary = now - self . dup_interval <EOL> flow_count = <NUM_LIT:0> <EOL> flow_list = flows_dir . ListChildren ( age = ( <EOL> earlier . AsMicroSecondsFromEpoch ( ) , now . AsMicroSecondsFromEpoch ( ) ) ) <EOL> for flow_obj in aff4 . FACTORY . MultiOpen ( flow_list , token = token ) : <EOL> flow_state = flow_obj . Get ( flow_obj . Schema . FLOW_STATE ) <EOL> flow_context = flow_state . context <EOL> if self . dup_interval and flow_context . create_time > dup_boundary : <EOL> if flow_context . args . flow_name == flow_name : <EOL> if flow_state . args == flow_args or ( <EOL> isinstance ( flow_state . args , flow . EmptyFlowArgs ) and <EOL> flow_args is None ) : <EOL> raise ErrorFlowDuplicate ( "<STR_LIT>" % ( <EOL> flow_name , client_id , flow_context . create_time ) ) <EOL> if flow_context . creator == user and flow_context . create_time > earlier : <EOL> flow_count += <NUM_LIT:1> <EOL> if self . daily_req_limit and flow_count >= self . daily_req_limit : <EOL> raise ErrorDailyFlowRequestLimitExceeded ( <EOL> "<STR_LIT>" % ( <EOL> flow_count , earlier , self . daily_req_limit ) ) </s>
