From b27df71969d60ad3c38a3af127e89e4a7d9206e4 Mon Sep 17 00:00:00 2001
From: 0xInfection <0xinfection@gmail.com>
Date: Thu, 29 Feb 2024 15:38:42 +0530
Subject: [PATCH] add rate limit

---
 arjun/__main__.py       | 1 +
 arjun/core/requester.py | 3 +++
 setup.py                | 3 ++-
 3 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/arjun/__main__.py b/arjun/__main__.py
index 1dd14fd..23081da 100644
--- a/arjun/__main__.py
+++ b/arjun/__main__.py
@@ -32,6 +32,7 @@ parser.add_argument('-i', help='Import target URLs from file.', dest='import_fil
 parser.add_argument('-T', help='HTTP request timeout in seconds. (default: 15)', dest='timeout', type=float, default=15)
 parser.add_argument('-c', help='Chunk size. The number of parameters to be sent at once', type=int, dest='chunks', default=250)
 parser.add_argument('-q', help='Quiet mode. No output.', dest='quiet', action='store_true')
+parser.add_argument('--rate-limit', help='Max number of requests to be sent out per second (default: 9999)', dest='rate_limit', type=int, default=9999)
 parser.add_argument('--headers', help='Add headers. Separate multiple headers with a new line.', dest='headers', nargs='?', const=True)
 parser.add_argument('--passive', help='Collect parameter names from passive sources like wayback, commoncrawl and otx.', dest='passive', nargs='?', const='-')
 parser.add_argument('--stable', help='Prefer stability over speed.', dest='stable', action='store_true')
diff --git a/arjun/core/requester.py b/arjun/core/requester.py
index fdbf0b3..b3742d3 100644
--- a/arjun/core/requester.py
+++ b/arjun/core/requester.py
@@ -6,10 +6,13 @@ import warnings
 
 import arjun.core.config as mem
 
+from ratelimit import limits, sleep_and_retry
 from arjun.core.utils import dict_to_xml
 
 warnings.filterwarnings('ignore') # Disable SSL related warnings
 
+@sleep_and_retry
+@limits(calls=mem.var['rate_limit'], period=1)
 def requester(request, payload={}):
     """
     central function for making http requests
diff --git a/setup.py b/setup.py
index d521aea..cf3bae8 100644
--- a/setup.py
+++ b/setup.py
@@ -24,7 +24,8 @@ setup(
     package_data={'arjun': ['db/*']},
     install_requires=[
         'requests',
-        'dicttoxml'
+        'dicttoxml',
+        'ratelimit'
     ],
     classifiers=[
         'Development Status :: 5 - Production/Stable',