From 63072a7af1e70a7bf1612ba1082638284c3e22b9 Mon Sep 17 00:00:00 2001
From: Jussi-Pekka Erkkila <jp.erkkila@gmail.com>
Date: Sun, 15 Dec 2024 12:11:16 +0200
Subject: [PATCH] cmd output tuning for multi target scan

---
 secheaders/cmd_utils.py  |  3 ++-
 secheaders/secheaders.py | 25 ++++++++++++++++++-------
 2 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/secheaders/cmd_utils.py b/secheaders/cmd_utils.py
index 6ae7ec6..81276e0 100644
--- a/secheaders/cmd_utils.py
+++ b/secheaders/cmd_utils.py
@@ -22,7 +22,7 @@ def get_eval_output(warn, no_color):
 
 def output_text(target_url, headers, https, no_color=False, verbose=False) -> str:
     terminal_width = shutil.get_terminal_size().columns
-    output_str = f"Scanning target {target_url} ...\n\n"
+    output_str = f"Scanning target {target_url} ...\n"
 
     # If the stdout is not going into terminal, disable colors
     no_color = no_color or not sys.stdout.isatty()
@@ -66,4 +66,5 @@ def output_text(target_url, headers, https, no_color=False, verbose=False) -> st
 
         output_str += output
 
+    output_str += '\n'
     return output_str
diff --git a/secheaders/secheaders.py b/secheaders/secheaders.py
index d836c20..1e5730e 100644
--- a/secheaders/secheaders.py
+++ b/secheaders/secheaders.py
@@ -59,9 +59,8 @@ def main():
 
 def async_scan_done(scan):
     try:
-        res = scan.result()
-        print(cmd_utils.output_text(res['target'], res['headers'], res['https']))
-        print("========================\n")
+        res, args = scan.result()
+        print(cmd_utils.output_text(res['target'], res['headers'], res['https'], args.no_color, args.verbose))
     except SecurityHeadersException as e:
         print(e, file=sys.stderr)
 
@@ -81,22 +80,34 @@ def scan_target(url, args):
     return {'target': header_check.get_full_url(), 'headers': headers, 'https': https}
 
 
+def scan_target_wrapper(url, args):
+    # A bit of a dirty hack to pass args to the done callback
+    return scan_target(url, args), args
+
+
 async def scan_multiple_targets(args):
     with open(args.target_list, encoding='utf-8') as file:
         targets = [line.rstrip() for line in file]
 
     loop = asyncio.get_event_loop()
     tasks = []
-    for t in targets:
-        task = loop.run_in_executor(None, scan_target, t, args)
-        task.add_done_callback(async_scan_done)
+    for target in targets:
+        if args.json:
+            task = loop.run_in_executor(None, scan_target, target, args)
+        else:
+            task = loop.run_in_executor(None, scan_target_wrapper, target, args)
+            task.add_done_callback(async_scan_done)
         tasks.append(task)
 
+    res = []
     for task in tasks:
         await task
 
-    print("ALL COMPLETED!")
+    if args.json:
+        for t in tasks:
+            res.append(t.result())
 
+        print(str(res))
 
 if __name__ == "__main__":
     main()