def EnforceLimits(self, client_id, user, flow_name, flow_args=None):
"""Enforce DailyFlowRequestLimit and FlowDuplicateInterval.
Look at the flows that have run on this client recently and check
we aren't exceeding our limits. Raises if limits will be exceeded by running
the specified flow.
Args:
client_id: client URN
user: username string
flow_name: name of the Flow. Only used for FlowDuplicateInterval.
flow_args: flow args rdfvalue for the flow being launched
Raises:
DailyFlowRequestLimitExceededError: if the user has already run
API.DailyFlowRequestLimit on this client in the previous 24h.
DuplicateFlowError: an identical flow was run on this machine by a user
within the API.FlowDuplicateInterval
"""
if not self.dup_interval and not self.daily_req_limit:
return
now = rdfvalue.RDFDatetime.Now()
yesterday = now - rdfvalue.Duration.From(1, rdfvalue.DAYS)
dup_boundary = now - self.dup_interval
min_create_time = min(yesterday, dup_boundary)
flow_count = 0
flow_objs = self._LoadFlows(client_id, min_create_time)
if flow_args is None:
flow_args = rdf_flows.EmptyFlowArgs()
for flow_obj in flow_objs:
if (flow_obj.create_time > dup_boundary
and flow_obj.flow_class_name == flow_name
and flow_obj.args == flow_args):
raise DuplicateFlowError(
"Identical %s already run on %s at %s" %
(flow_name, client_id, flow_obj.create_time),
flow_id=flow_obj.flow_id)
# Filter for flows started by user within the 1 day window.
if flow_obj.creator == user and flow_obj.create_time > yesterday:
flow_count += 1
# If limit is set, enforce it.
if self.daily_req_limit and flow_count >= self.daily_req_limit:
raise DailyFlowRequestLimitExceededError(
"%s flows run since %s, limit: %s" %
(flow_count, yesterday, self.daily_req_limit))
def setUp(self):
super().setUp()
self.handler = hunt_plugin.ApiGetExportedHuntResultsHandler()
self.context = api_call_context.ApiCallContext("test")
self.hunt_id = self.StartHunt(
flow_runner_args=rdf_flow_runner.FlowRunnerArgs(
flow_name=flow_test_lib.DummyFlowWithSingleReply.__name__),
flow_args=rdf_flows.EmptyFlowArgs(),
client_rate=0)
self.client_ids = self.SetupClients(5)
# Ensure that clients are processed sequentially - this way the test won't
# depend on the order of results in the collection (which is normally
# random).
for cid in self.client_ids:
self.RunHunt(client_ids=[cid], failrate=-1)