11from collections import OrderedDict
22from dataclasses import dataclass , field
3+ import time
34from typing import Union , List , Optional
45
56from requests_toolbelt .utils import dump
@@ -254,6 +255,10 @@ def extract_dict(text, line_sep='\n', kv_sep='='):
254255
255256def http_request_generator (request : HttpRequest , dynamic_values : OrderedDict ):
256257 request_count = len (request .path + request .raw )
258+ # Determine the number of requests and modify the req_condition attribute of the HttpRequest object
259+ if request_count > 1 :
260+ request .req_condition = True
261+
257262 for payload_instance in payload_generator (request .payloads , request .attack ):
258263 current_index = 0
259264 dynamic_values .update (payload_instance )
@@ -272,9 +277,14 @@ def http_request_generator(request: HttpRequest, dynamic_values: OrderedDict):
272277 else :
273278 raw = path .strip ()
274279 raws = list (map (lambda x : x .strip (), raw .splitlines ()))
275- method , path , _ = raws [0 ].split (' ' )
276- url = f'{ Marker .ParenthesisOpen } BaseURL{ Marker .ParenthesisClose } { path } '
277-
280+ # Extract timeout value
281+ if raws [0 ].startswith ('@timeout' ):
282+ timeout = Marker .extract_timeout_value (raws [0 ])
283+ del raws [0 ]
284+ method , path , _ = raws [0 ].split (' ' )
285+ kwargs .setdefault ('timeout' , timeout )
286+ else :
287+ method , path , _ = raws [0 ].split (' ' )
278288 if method == "POST" :
279289 index = 0
280290 for i in raws :
@@ -290,6 +300,8 @@ def http_request_generator(request: HttpRequest, dynamic_values: OrderedDict):
290300 else :
291301 headers = extract_dict ('\n ' .join (raws [1 :]), '\n ' , ": " )
292302
303+ url = f'{ Marker .ParenthesisOpen } BaseURL{ Marker .ParenthesisClose } { path } '
304+
293305 kwargs .setdefault ('allow_redirects' , request .redirects )
294306 kwargs .setdefault ('data' , data )
295307 kwargs .setdefault ('headers' , headers )
@@ -324,7 +336,13 @@ def execute_http_request(request: HttpRequest, dynamic_values, interactsh) -> Un
324336 session .max_redirects = request .max_redirects
325337 else :
326338 session .max_redirects = 10
339+
340+ # Calculate response time
341+ start_time = time .time ()
327342 response = session .request (method = method , url = url , ** kwargs )
343+ end_time = time .time ()
344+ resp_time = end_time - start_time
345+
328346 # for debug purpose
329347 try :
330348 logger .debug (dump .dump_all (response ).decode ('utf-8' ))
@@ -337,6 +355,9 @@ def execute_http_request(request: HttpRequest, dynamic_values, interactsh) -> Un
337355 response = None
338356
339357 resp_data = http_response_to_dsl_map (response )
358+ if response is not None :
359+ resp_data ['duration' ] = resp_time
360+
340361 if response :
341362 response .close ()
342363
0 commit comments