Merge pull request #8 from OWASP/fix-bug-sempaphore

change logic for rate limiting

src/offat/http.py

@@ -11,25 +11,30 @@
 
 class AsyncRequests:
     '''
-    AsyncRequests class helps to send HTTP requests.
+    AsyncRequests class helps to send HTTP requests with rate limiting options.
     '''
 
-    def __init__(self, headers: dict = None, proxy:str = None, ssl:bool=True, allow_redirects: bool=True) -> None:
+    def __init__(self, rate_limit:int=None, delay:float=None, headers: dict = None, proxy:str = None, ssl:bool=True, allow_redirects: bool=True) -> None:
         '''AsyncRequests class constructor
         
         Args:
+            rate_limit (int): number of concurrent requests at the same time
+            delay (float): delay between consecutive requests
             headers (dict): overrides default headers while sending HTTP requests
             proxy (str): proxy URL to be used while sending requests
             ssl (bool): ignores few SSL errors if value is False
         
         Returns:
             None
         '''
+        self._rate_limit = rate_limit
+        self._delay = delay
         self._headers = headers
         self._proxy = proxy if proxy else None
         self._ssl = ssl if ssl else None
         self._allow_redirects = allow_redirects
 
+
     async def request(self, url: str, method: str = 'GET', session: ClientSession = None, *args, **kwargs) -> ClientResponse:
         '''Send HTTP requests asynchronously
 
@@ -44,9 +49,7 @@ async def request(self, url: str, method: str = 'GET', session: ClientSession =
             dict: returns request and response data as dict
         '''
         is_new_session = False
-
-
-        connector = TCPConnector(ssl=self._ssl,)
+        connector = TCPConnector(ssl=self._ssl,limit=self._rate_limit,)
 
         if not session:
             session = ClientSession(headers=self._headers, connector=connector)
@@ -86,46 +89,7 @@ async def request(self, url: str, method: str = 'GET', session: ClientSession =
                 await session.close()
                 del session
 
-        return resp_data
-
-
-class AsyncRLRequests(AsyncRequests):
-    '''
-    Send Asynchronous rate limited HTTP requests.
-    '''
-
-    def __init__(self, rate_limit: int = 20, delay: float = 0.05, headers: dict = None, proxy: str = None, ssl:bool = True) -> None:
-        '''AsyncRLRequests constructor
-
-        Args:
-            rate_limit (int): number of concurrent requests at the same time
-            delay (float): delay between consecutive requests
-            headers (dict): overrides default headers while sending HTTP requests
-
-        Returns:
-            None
-        '''
-        assert isinstance(delay, float) or isinstance(delay, int)
-        assert isinstance(rate_limit, float) or isinstance(rate_limit, int)
-
-        self._delay = delay
-        self._semaphore = asyncio.Semaphore(rate_limit)
-        super().__init__(headers=headers, proxy=proxy, ssl=ssl)
-
-
-    async def request(self, url: str, method: str = 'GET', session: ClientSession = None, *args, **kwargs) -> ClientResponse:
-        '''Send HTTP requests asynchronously with rate limit and delay between the requests
-
-        Args:
-            url (str): URL of the webpage/endpoint
-            method (str): HTTP methods (default: GET) supports GET, POST, 
-            PUT, HEAD, OPTIONS, DELETE
-            session (aiohttp.ClientSession): aiohttp Client Session for sending requests
-        
-        Returns:
-            dict: returns request and response data as dict
-        '''
-        async with self._semaphore:
-            response = await super().request(url, method, session, *args, **kwargs)
+        if self._delay:
             await asyncio.sleep(self._delay)
-            return response
+
+        return resp_data

src/offat/tester/test_runner.py

@@ -1,6 +1,6 @@
 from asyncio import ensure_future, gather
 from enum import Enum
-from ..http import AsyncRequests, AsyncRLRequests
+from ..http import AsyncRequests
 from ..logger import create_logger
 
 logger = create_logger(__name__)
@@ -19,11 +19,8 @@ class PayloadFor(Enum):
 
 class TestRunner:
     def __init__(self, rate_limit:int=None, delay:float=None, headers:dict=None, proxy: str = None, ssl: bool = True) -> None:
-        if rate_limit and delay:
-            self._client = AsyncRLRequests(rate_limit=rate_limit, delay=delay, headers=headers, proxy=proxy, ssl=ssl)
-        else:
-            self._client = AsyncRequests(headers=headers, proxy=proxy, ssl=ssl)
-
+        self._client = AsyncRequests(rate_limit=rate_limit, delay=delay, headers=headers, proxy=proxy, ssl=ssl)
+
 
     def _generate_payloads(self, params:list[dict], payload_for:PayloadFor=PayloadFor.BODY):
         '''Generate body payload from passed data for HTTP body and query.