From a0fc97186bd08f63e8a5e084aaf4957add768bec Mon Sep 17 00:00:00 2001 From: Jussi-Pekka Erkkila Date: Fri, 29 Nov 2024 01:21:41 +0200 Subject: [PATCH] move cmd related utils to a separate file --- secheaders/cmd_utils.py | 70 +++++++++++++++++++++++++++++++++++ secheaders/securityheaders.py | 54 ++------------------------- secheaders/utils.py | 18 +-------- 3 files changed, 74 insertions(+), 68 deletions(-) create mode 100644 secheaders/cmd_utils.py diff --git a/secheaders/cmd_utils.py b/secheaders/cmd_utils.py new file mode 100644 index 0000000..9b837d5 --- /dev/null +++ b/secheaders/cmd_utils.py @@ -0,0 +1,70 @@ +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(headers, https, verbose=False, no_color=False) -> str: + terminal_width = shutil.get_terminal_size().columns + output_str = "" + + # 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 = f"Header '{header}' is missing" + else: + output = f"{header}: {header_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 verbose: + output_str += f"Full header contents: {header_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 diff --git a/secheaders/securityheaders.py b/secheaders/securityheaders.py index 76208f9..9c331e0 100644 --- a/secheaders/securityheaders.py +++ b/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(headers, https, args.verbose, args.no_color)) if __name__ == "__main__": diff --git a/secheaders/utils.py b/secheaders/utils.py index ff7e2bf..fa6b6a8 100644 --- a/secheaders/utils.py +++ b/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} ]"