Merge pull request #29 from juerkkil/cmd_utils

move cmd related utils to a separate file

secheaders/cmd_utils.py

@@ -0,0 +1,69 @@
+import shutil
+import sys
+import textwrap
+
+from .constants import WARN_COLOR, OK_COLOR, END_COLOR, COLUMN_WIDTH_R
+
+
+def get_eval_output(warn, no_color):
+    color_start = OK_COLOR
+    color_end = END_COLOR
+    eval_result = "OK"
+    if warn:
+        color_start = WARN_COLOR
+        eval_result = "WARN"
+
+    if no_color:
+        color_start = ""
+        color_end = ""
+
+    return f"[ {color_start}{eval_result}{color_end} ]"
+
+
+def output_text(target_url, headers, https, args) -> str:
+    terminal_width = shutil.get_terminal_size().columns
+    output_str = f"Scan target: {target_url}\n"
+
+    # If the stdout is not going into terminal, disable colors
+    no_color = args.no_color or not sys.stdout.isatty()
+    for header, value in headers.items():
+        truncated = False
+        if not value['defined']:
+            output = f"Header '{header}' is missing"
+        else:
+            output = f"{header}: {value['contents']}"
+            if len(output) > terminal_width - COLUMN_WIDTH_R:
+                truncated = True
+                output = f"{output[0:(terminal_width - COLUMN_WIDTH_R - 3)]}..."
+
+        eval_value = get_eval_output(value['warn'], no_color)
+
+        if no_color:
+            output_str += f"{output:<{terminal_width - COLUMN_WIDTH_R}}{eval_value:^{COLUMN_WIDTH_R}}\n"
+        else:
+            # This is a dirty hack required to align ANSI-colored str correctly
+            output_str += f"{output:<{terminal_width - COLUMN_WIDTH_R}}{eval_value:^{COLUMN_WIDTH_R + 9}}\n"
+
+        if truncated and args.verbose:
+            output_str += f"Full header contents: {value['contents']}\n"
+        for note in value['notes']:
+            output_str += textwrap.fill(f" * {note}", terminal_width - COLUMN_WIDTH_R, subsequent_indent='   ')
+            output_str += "\n"
+
+    msg_map = {
+        'supported': 'HTTPS supported',
+        'certvalid': 'HTTPS valid certificate',
+        'redirect': 'HTTP -> HTTPS automatic redirect',
+    }
+    for key in https:
+        output = f"{msg_map[key]}"
+        eval_value = get_eval_output(not https[key], no_color)
+        if no_color:
+            output = f"{output:<{terminal_width - COLUMN_WIDTH_R}}{eval_value:^{COLUMN_WIDTH_R}}"
+        else:
+            # This is a dirty hack required to align ANSI-colored str correctly
+            output = f"{output:<{terminal_width - COLUMN_WIDTH_R}}{eval_value:^{COLUMN_WIDTH_R + 9}}"
+
+        output_str += output
+
+    return output_str

secheaders/securityheaders.py

@@ -2,17 +2,15 @@
 import http.client
 import json
 import re
-import shutil
 import socket
 import ssl
 import sys
-import textwrap
 from typing import Union
 from urllib.parse import ParseResult, urlparse
 
-from . import utils
+from . import utils, cmd_utils
 from .constants import DEFAULT_TIMEOUT, DEFAULT_URL_SCHEME, EVAL_WARN, REQUEST_HEADERS, HEADER_STRUCTURED_LIST, \
-        SERVER_VERSION_HEADERS, COLUMN_WIDTH_R
+        SERVER_VERSION_HEADERS
 from .exceptions import SecurityHeadersException, InvalidTargetURL, UnableToConnect
 
 
@@ -196,52 +194,6 @@ def get_full_url(self) -> str:
         return f"{self.protocol_scheme}://{self.hostname}{self.path}"
 
 
-def output_text(headers, https, verbose=False, no_color=False) -> None:
-    terminal_width = shutil.get_terminal_size().columns
-
-    # If the stdout is not going into terminal, disable colors
-    no_color = no_color or not sys.stdout.isatty()
-    for header, value in headers.items():
-        truncated = False
-        header_contents = value['contents']
-        if not value['defined']:
-            output_str = f"Header '{header}' is missing"
-        else:
-            output_str = f"{header}: {header_contents}"
-            if len(output_str) > terminal_width- COLUMN_WIDTH_R:
-                truncated = True
-                output_str = f"{output_str[0:(terminal_width - COLUMN_WIDTH_R - 3)]}..."
-
-        eval_value = utils.get_eval_output(value['warn'], no_color)
-
-        if no_color:
-            print(f"{output_str:<{terminal_width - COLUMN_WIDTH_R}}{eval_value:^{COLUMN_WIDTH_R}}")
-        else:
-            # This is a dirty hack required to align ANSI-colored str correctly
-            print(f"{output_str:<{terminal_width - COLUMN_WIDTH_R}}{eval_value:^{COLUMN_WIDTH_R + 9}}")
-
-        if truncated and verbose:
-            print((f"Full header contents: {header_contents}"))
-        for note in value['notes']:
-            print(textwrap.fill(f" * {note}", terminal_width - COLUMN_WIDTH_R, subsequent_indent='   '))
-
-    msg_map = {
-        'supported': 'HTTPS supported',
-        'certvalid': 'HTTPS valid certificate',
-        'redirect': 'HTTP -> HTTPS automatic redirect',
-    }
-    for key in https:
-        output_str = f"{msg_map[key]}"
-        eval_value = utils.get_eval_output(not https[key], no_color)
-        if no_color:
-            output_str = f"{output_str:<{terminal_width - COLUMN_WIDTH_R}}{eval_value:^{COLUMN_WIDTH_R}}"
-        else:
-            # This is a dirty hack required to align ANSI-colored str correctly
-            output_str = f"{output_str:<{terminal_width - COLUMN_WIDTH_R}}{eval_value:^{COLUMN_WIDTH_R + 9}}"
-
-        print(output_str)
-
-
 def main():
     parser = argparse.ArgumentParser(description='Scan HTTP security headers',
                                      formatter_class=argparse.ArgumentDefaultsHelpFormatter)
@@ -271,7 +223,7 @@ def main():
     if args.json:
         print(json.dumps({'target': header_check.get_full_url(), 'headers': headers, 'https': https}, indent=2))
     else:
-        output_text(headers, https, args.verbose, args.no_color)
+        print(cmd_utils.output_text(header_check.get_full_url(), headers, https, args))
 
 
 if __name__ == "__main__":

secheaders/utils.py

@@ -1,8 +1,7 @@
 import re
 from typing import Tuple
 
-from .constants import EVAL_WARN, EVAL_OK, UNSAFE_CSP_RULES, RESTRICTED_PERM_POLICY_FEATURES, WARN_COLOR, OK_COLOR, \
-    END_COLOR
+from .constants import EVAL_WARN, EVAL_OK, UNSAFE_CSP_RULES, RESTRICTED_PERM_POLICY_FEATURES
 
 
 def eval_x_frame_options(contents: str) -> Tuple[int, list]:
@@ -132,18 +131,3 @@ def permissions_policy_parser(contents: str) -> dict:
             retval[feature] = feature_policy.split()
 
     return retval
-
-
-def get_eval_output(warn, no_color):
-    color_start = OK_COLOR
-    color_end = END_COLOR
-    eval_result = "OK"
-    if warn:
-        color_start = WARN_COLOR
-        eval_result = "WARN"
-
-    if no_color:
-        color_start = ""
-        color_end = ""
-
-    return f"[ {color_start}{eval_result}{color_end} ]"