From 305b4829ef3994b526458175b656b4d3e5376b44 Mon Sep 17 00:00:00 2001
From: Anna <anna.madej2@tutanota.com>
Date: Thu, 26 Oct 2023 21:33:33 +0200
Subject: [PATCH] Improved request configuration (#88) (#205)

---
 artemis/config.py            |  4 ++--
 artemis/modules/nuclei.py    |  2 +-
 artemis/utils.py             | 23 ++++++++++++++++-------
 docker-compose.test-e2e.yaml |  2 +-
 docker-compose.test.yaml     |  2 +-
 5 files changed, 21 insertions(+), 12 deletions(-)

diff --git a/artemis/config.py b/artemis/config.py
index b788fccc6..1c2cdd205 100644
--- a/artemis/config.py
+++ b/artemis/config.py
@@ -129,12 +129,12 @@ class Limits:
             "E.g. when set to 100, Artemis will send no more than 100 port scanning packets per seconds per port scanner instance.",
         ] = get_config("SCANNING_PACKETS_PER_SECOND", default=100, cast=int)
 
-        SECONDS_PER_REQUEST: Annotated[
+        REQUEST_PER_SECOND: Annotated[
             int,
             """
             E.g. when set to 2, Artemis will make sure no HTTP/MySQL connect/... request takes less than 2 seconds, sleeping if needed.
             """,
-        ] = get_config("SECONDS_PER_REQUEST", default=0, cast=int)
+        ] = get_config("REQUEST_PER_SECOND", default=1, cast=float)
 
     class Miscellaneous:
         BLOCKLIST_FILE: Annotated[
diff --git a/artemis/modules/nuclei.py b/artemis/modules/nuclei.py
index 91fce5f4b..394daad09 100644
--- a/artemis/modules/nuclei.py
+++ b/artemis/modules/nuclei.py
@@ -104,7 +104,7 @@ def run_multiple(self, tasks: List[Task]) -> None:
             "-headless-bulk-size",
             str(len(tasks_filtered)),
             "-milliseconds-per-request",
-            str(int(Config.Limits.SECONDS_PER_REQUEST * 1000.0 / len(tasks_filtered))),
+            str(int((1 / Config.Limits.REQUEST_PER_SECOND) * 1000.0 / len(tasks_filtered))),
         ] + additional_configuration
 
         targets = []
diff --git a/artemis/utils.py b/artemis/utils.py
index 663f76738..2f2e7cca2 100644
--- a/artemis/utils.py
+++ b/artemis/utils.py
@@ -1,4 +1,5 @@
 import logging
+import math
 import subprocess
 import time
 import urllib.parse
@@ -43,13 +44,21 @@ def is_directory_index(content: str) -> bool:
 
 
 def throttle_request(f: Callable[[], Any]) -> Any:
-    time_start = time.time()
-    try:
-        return f()
-    finally:
-        time_elapsed = time.time() - time_start
-        if time_elapsed < Config.Limits.SECONDS_PER_REQUEST:
-            time.sleep(Config.Limits.SECONDS_PER_REQUEST - time_elapsed)
+    request_per_second = Config.Limits.REQUEST_PER_SECOND
+    if request_per_second >= 1:
+        average_time_per_request = 1 / request_per_second
+        f_start = time.time()
+        f()
+        func_time = time.time() - f_start
+        if func_time < average_time_per_request:
+            time.sleep(average_time_per_request - func_time)
+    elif request_per_second < 1:
+        seconds_for_req = math.floor(1 / request_per_second)
+        f_start = time.time()
+        f()
+        func_time = time.time() - f_start
+        if func_time < seconds_for_req:
+            time.sleep(func_time - f_start)
 
 
 def get_host_from_url(url: str) -> str:
diff --git a/docker-compose.test-e2e.yaml b/docker-compose.test-e2e.yaml
index c79dcac59..3b067720b 100644
--- a/docker-compose.test-e2e.yaml
+++ b/docker-compose.test-e2e.yaml
@@ -6,7 +6,7 @@ services:
 
   karton-bruter:
     environment:
-      SECONDS_PER_REQUEST: 0
+      REQUEST_PER_SECOND: 10
   redis:
     volumes:
       - data-test-redis:/data
diff --git a/docker-compose.test.yaml b/docker-compose.test.yaml
index 2fad0743c..52cb4f51d 100644
--- a/docker-compose.test.yaml
+++ b/docker-compose.test.yaml
@@ -33,7 +33,7 @@ services:
       POSTMAN_MAIL_FROM: artemis@localhost.com
       POSTMAN_MAIL_TO: artemis@localhost.com
 
-      SECONDS_PER_REQUEST: 0
+      REQUEST_PER_SECOND: 10
       SCANNING_PACKETS_PER_SECOND: 5
       CUSTOM_PORT_SCANNER_PORTS: 21,80,6379
       NUCLEI_CHECK_TEMPLATE_LIST: False