#!/usr/bin/env python
"""Tests for grr.lib.throttle."""

from absl import app

from grr_response_core.lib import rdfvalue
from grr_response_core.lib.rdfvalues import file_finder as rdf_file_finder
from grr_response_server import throttle
from grr_response_server.flows.general import file_finder
from grr.test_lib import flow_test_lib
from grr.test_lib import test_lib


class ThrottleTest(test_lib.GRRBaseTest):
  BASE_TIME = 1439501002

  def setUp(self):
    super().setUp()
    self.client_id = self.SetupClient(0)

  def testCheckFlowRequestLimit(self):
    # Create a flow
    with test_lib.FakeTime(self.BASE_TIME):
      flow_test_lib.StartFlow(
          flow_test_lib.DummyLogFlow,
          client_id=self.client_id,
          creator=self.test_username,
      )

    # One day + 1s later
    with test_lib.FakeTime(self.BASE_TIME + 86400 + 1):
      flow_test_lib.StartFlow(
          flow_cls=flow_test_lib.DummyLogFlow,
          client_id=self.client_id,
          creator=self.test_username,
      )

      # Disable the dup interval checking by setting it to 0.
      throttler = throttle.FlowThrottler(
          daily_req_limit=2,
          dup_interval=rdfvalue.Duration.From(0, rdfvalue.SECONDS),
      )

      # Should succeed, only one flow present in the 1 day window.
      throttler.EnforceLimits(
          self.client_id,
          self.test_username,
          flow_test_lib.DummyLogFlow.__name__,
          None,
      )

      # Start some more flows with a different user
      flow_test_lib.StartFlow(
          flow_cls=flow_test_lib.DummyLogFlow,
          client_id=self.client_id,
          creator="test2",
      )
      flow_test_lib.StartFlow(
          flow_test_lib.DummyLogFlow, client_id=self.client_id, creator="test2"
      )

      # Should still succeed, since we count per-user
      throttler.EnforceLimits(
          self.client_id,
          self.test_username,
          flow_test_lib.DummyLogFlow.__name__,
          None,
      )

      # Add another flow at current time
      flow_test_lib.StartFlow(
          flow_test_lib.DummyLogFlow,
          client_id=self.client_id,
          creator=self.test_username,
      )

      with self.assertRaises(throttle.DailyFlowRequestLimitExceededError):
        throttler.EnforceLimits(
            self.client_id,
            self.test_username,
            flow_test_lib.DummyLogFlow.__name__,
            None,
        )

  def testFlowDuplicateLimit(self):
    # Disable the request limit checking by setting it to 0.
    throttler = throttle.FlowThrottler(
        daily_req_limit=0,
        dup_interval=rdfvalue.Duration.From(1200, rdfvalue.SECONDS),
    )

    # Running the same flow immediately should fail
    with test_lib.FakeTime(self.BASE_TIME):
      throttler.EnforceLimits(
          self.client_id,
          self.test_username,
          flow_test_lib.DummyLogFlow.__name__,
          None,
      )

      flow_test_lib.StartFlow(
          flow_cls=flow_test_lib.DummyLogFlow,
          client_id=self.client_id,
          creator=self.test_username,
      )

      with self.assertRaises(throttle.DuplicateFlowError):
        throttler.EnforceLimits(
            self.client_id,
            self.test_username,
            flow_test_lib.DummyLogFlow.__name__,
            None,
        )

    # Doing the same outside the window should work
    with test_lib.FakeTime(self.BASE_TIME + 1200 + 1):
      throttler.EnforceLimits(
          self.client_id,
          self.test_username,
          flow_test_lib.DummyLogFlow.__name__,
          None,
      )

      flow_test_lib.StartFlow(
          flow_test_lib.DummyLogFlow,
          client_id=self.client_id,
          creator=self.test_username,
      )

      with self.assertRaises(throttle.DuplicateFlowError):
        throttler.EnforceLimits(
            self.client_id,
            self.test_username,
            flow_test_lib.DummyLogFlow.__name__,
            None,
        )

    # Now try a flow with more complicated args
    args = rdf_file_finder.FileFinderArgs(
        paths=["/tmp/1", "/tmp/2"],
        action=rdf_file_finder.FileFinderAction(action_type="STAT"),
    )

    with test_lib.FakeTime(self.BASE_TIME):
      throttler.EnforceLimits(
          self.client_id,
          self.test_username,
          file_finder.FileFinder.__name__,
          args,
      )

      new_args = rdf_file_finder.FileFinderArgs(
          paths=["/tmp/1", "/tmp/2"],
          action=rdf_file_finder.FileFinderAction(action_type="STAT"),
      )

      flow_test_lib.StartFlow(
          flow_cls=file_finder.FileFinder,
          client_id=self.client_id,
          creator=self.test_username,
          flow_args=new_args,
      )

      with self.assertRaises(throttle.DuplicateFlowError):
        throttler.EnforceLimits(
            self.client_id,
            self.test_username,
            file_finder.FileFinder.__name__,
            args,
        )

      # Different args should succeed.
      args = rdf_file_finder.FileFinderArgs(
          paths=["/tmp/1", "/tmp/3"],
          action=rdf_file_finder.FileFinderAction(action_type="STAT"),
      )

      throttler.EnforceLimits(
          self.client_id,
          self.test_username,
          file_finder.FileFinder.__name__,
          args,
      )


def main(argv):
  # Run the full test suite
  test_lib.main(argv)


if __name__ == "__main__":
  app.run(main)