fix(nuclei,katana): add -sr flag and write http responses and screens…

…hot to correct folder (#395)
freelabz · May 7, 2024 · 1a51790 · 1a51790
1 parent 4f3ff59
commit 1a51790
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 20 deletions.
diff --git a/secator/tasks/httpx.py b/secator/tasks/httpx.py
@@ -31,6 +31,7 @@ class httpx(Http):
 		'cdn': {'is_flag': True, 'default': False, 'help': 'CDN detection'},
 		'debug_resp': {'is_flag': True, 'default': False, 'help': 'Debug response'},
 		'vhost': {'is_flag': True, 'default': False, 'help': 'Probe and display server supporting VHOST'},
+		'store_responses': {'is_flag': True, 'short': 'sr', 'default': CONFIG.http.store_responses, 'help': 'Save HTTP responses'},  # noqa: E501
 		'screenshot': {'is_flag': True, 'short': 'ss', 'default': False, 'help': 'Screenshot response'},
 		'system_chrome': {'is_flag': True, 'default': False, 'help': 'Use local installed Chrome for screenshot'},
 		'headless_options': {'is_flag': False, 'short': 'ho', 'default': None, 'help': 'Headless Chrome additional options'},
@@ -55,6 +56,7 @@ class httpx(Http):
 		THREADS: 'threads',
 		TIMEOUT: 'timeout',
 		USER_AGENT: OPT_NOT_SUPPORTED,
+		'store_responses': 'sr',
 	}
 	opt_value_map = {
 		DELAY: lambda x: str(x) + 's' if x else None,
@@ -71,15 +73,10 @@ def on_init(self):
 		debug_resp = self.get_opt_value('debug_resp')
 		if debug_resp:
 			self.cmd = self.cmd.replace('-silent', '')
-		if CONFIG.http.store_responses:
-			self.output_response_path = f'{self.reports_folder}/response'
-			self.output_screenshot_path = f'{self.reports_folder}/screenshot'
-			os.makedirs(self.output_response_path, exist_ok=True)
-			os.makedirs(self.output_screenshot_path, exist_ok=True)
-			self.cmd += f' -sr -srd {self.reports_folder}'
-
-		# Remove screenshot bytes and body bytes when screenshot
 		screenshot = self.get_opt_value('screenshot')
+		store_responses = self.get_opt_value('store_responses')
+		if store_responses or screenshot:
+			self.cmd += f' -srd {self.reports_folder}/.outputs'
 		if screenshot:
 			self.cmd += ' -esb -ehb'
 
@@ -98,8 +95,15 @@ def on_item_pre_convert(self, item):
 
 	@staticmethod
 	def on_end(self):
-		if CONFIG.http.store_responses:
-			if os.path.exists(self.output_response_path + '/index.txt'):
-				os.remove(self.output_response_path + '/index.txt')
-			if os.path.exists(self.output_screenshot_path + '/index.txt'):
-				os.remove(self.output_screenshot_path + '/index_screenshot.txt')
+		store_responses = self.get_opt_value('store_responses')
+		response_dir = f'{self.reports_folder}/.outputs'
+		if store_responses:
+			index_rpath = f'{response_dir}/response/index.txt'
+			index_spath = f'{response_dir}/screenshot/index_screenshot.txt'
+			index_spath2 = f'{response_dir}/screenshot/screenshot.html'
+			if os.path.exists(index_rpath):
+				os.remove(index_rpath)
+			if os.path.exists(index_spath):
+				os.remove(index_spath)
+			if os.path.exists(index_spath2):
+				os.remove(index_spath2)
diff --git a/secator/tasks/katana.py b/secator/tasks/katana.py
@@ -29,7 +29,8 @@ class katana(HttpCrawler):
 	opts = {
 		'headless': {'is_flag': True, 'short': 'hl', 'help': 'Headless mode'},
 		'system_chrome': {'is_flag': True, 'short': 'sc', 'help': 'Use local installed chrome browser'},
-		'form_extraction': {'is_flag': True, 'short': 'fx', 'help': 'Detect forms'}
+		'form_extraction': {'is_flag': True, 'short': 'fx', 'help': 'Detect forms'},
+		'store_responses': {'is_flag': True, 'short': 'sr', 'default': CONFIG.http.store_responses, 'help': 'Store responses'}
 	}
 	opt_key_map = {
 		HEADER: 'headers',
@@ -50,7 +51,8 @@ class katana(HttpCrawler):
 		RETRIES: 'retry',
 		THREADS: 'concurrency',
 		TIMEOUT: 'timeout',
-		USER_AGENT: OPT_NOT_SUPPORTED
+		USER_AGENT: OPT_NOT_SUPPORTED,
+		'store_responses': 'sr'
 	}
 	opt_value_map = {
 		DELAY: lambda x: int(x) if isinstance(x, float) else x
@@ -107,14 +109,16 @@ def on_init(self):
 		debug_resp = self.get_opt_value('debug_resp')
 		if debug_resp:
 			self.cmd = self.cmd.replace('-silent', '')
-		if CONFIG.http.store_responses:
-			self.cmd += f' -sr -srd {self.reports_folder}'
+		store_responses = self.get_opt_value('store_responses')
+		if store_responses:
+			self.cmd += f' -srd {self.reports_folder}/.outputs'
 
 	@staticmethod
 	def on_item(self, item):
 		if not isinstance(item, Url):
 			return item
-		if CONFIG.http.store_responses and os.path.exists(item.stored_response_path):
+		store_responses = self.get_opt_value('store_responses')
+		if store_responses and os.path.exists(item.stored_response_path):
 			with open(item.stored_response_path, 'r', encoding='latin-1') as fin:
 				data = fin.read().splitlines(True)
 				first_line = data[0]
@@ -126,5 +130,7 @@ def on_item(self, item):
 
 	@staticmethod
 	def on_end(self):
-		if CONFIG.http.store_responses and os.path.exists(self.reports_folder + '/index.txt'):
-			os.remove(self.reports_folder + '/index.txt')
+		store_responses = self.get_opt_value('store_responses')
+		index_rpath = f'{self.reports_folder}/.outputs/index.txt'
+		if store_responses and os.path.exists(index_rpath):
+			os.remove(index_rpath)