diff options
author | Casey Deccio <casey@deccio.net> | 2016-08-17 17:52:39 +0300 |
---|---|---|
committer | Casey Deccio <casey@deccio.net> | 2016-08-17 17:52:39 +0300 |
commit | 483a6e2ceeb2441c3e0ea42c6e7d7bcc42db29e4 (patch) | |
tree | cf82d116298da7992ff72b50b0168664898a4cb2 | |
parent | 5fabc99631e9ee5f703431b6b7f99455e2a55ae1 (diff) |
Add some color to grok, and make -p the default for TTY
-rw-r--r-- | dnsviz/commands/grok.py | 89 |
1 files changed, 83 insertions, 6 deletions
diff --git a/dnsviz/commands/grok.py b/dnsviz/commands/grok.py index d97f09e..c6e5965 100644 --- a/dnsviz/commands/grok.py +++ b/dnsviz/commands/grok.py @@ -28,6 +28,7 @@ import getopt import io import json import logging +import os import re import sys @@ -51,6 +52,40 @@ except ImportError: logger = logging.getLogger('dnsviz.analysis.offline') +TERM_COLOR_MAP = { + 'BOLD': '\033[1m', + 'RESET': '\033[0m', + 'SECURE': '\033[36m', + 'BOGUS': '\033[31m', + 'INSECURE': '\033[37m', + 'NOERROR': '\033[37m', + 'NXDOMAIN': '\033[37m', + 'INDETERMINATE': '\033[31m', + 'NON_EXISTENT': '\033[37m', + 'VALID': '\033[36m', + 'INDETERMINATE': '\033[37m', + 'INDETERMINATE_NO_DNSKEY': '\033[37m', + 'INDETERMINATE_MATCH_PRE_REVOKE': '\033[37m', + 'INDETERMINATE_UNKNOWN_ALGORITHM': '\033[33m', + 'ALGORITHM_IGNORED': '\033[37m', + 'EXPIRED': '\033[35m', + 'PREMATURE': '\033[35m', + 'INVALID_SIG': '\033[31m', + 'INVALID': '\033[31m', + 'INVALID_DIGEST': '\033[31m', + 'INCOMPLETE': '\033[33m', + 'LAME': '\033[33m', + 'INVALID_TARGET': '\033[31m', + 'ERROR': '\033[31m', + 'WARNING': '\033[33m', +} + +KEY_RE = re.compile(r'^((?P<indent>\s+)")(.+)(": [\[{"])') +ERRORS_RE = re.compile(r'^((?P<indent>\s+)")((?P<level>warning|error)s?)(": \[)$') +ERRORS_CLOSE_RE = re.compile(r'^(?P<indent>\s+)]$') +DESCRIPTION_CODE_RE = re.compile(r'^((?P<indent>\s+)")(?P<name>description|code)(": ")(.+)(",?)$') +STATUS_RE = re.compile(r'^(?P<indent>\s+)("status": ")(?P<status>.+)(",?)') + def usage(err=None): if err is not None: err += '\n\n' @@ -67,6 +102,42 @@ Options: -h - display the usage and exit ''' % (err)) +def color_json(s): + error = None + s1 = '' + + for line in s.split('\n'): + if error is None: + # not in an error object; look for a start + error = ERRORS_RE.search(line) + if error is not None: + # found an error start + line = ERRORS_RE.sub(r'\1%s%s\3%s\5' % (TERM_COLOR_MAP['BOLD'], TERM_COLOR_MAP[error.group('level').upper()], TERM_COLOR_MAP['RESET']), line) + s1 += line + '\n' + continue + + if error is None: + # not in an error object + line = KEY_RE.sub(r'\1%s\3%s\4' % (TERM_COLOR_MAP['BOLD'], TERM_COLOR_MAP['RESET']), line) + m = STATUS_RE.search(line) + if m is not None: + line = STATUS_RE.sub(r'\1\2%s\3%s\4' % (TERM_COLOR_MAP[m.group('status').upper()], TERM_COLOR_MAP['RESET']), line) + s1 += line + '\n' + continue + + # in an error object + m = ERRORS_CLOSE_RE.search(line) + if m is not None and len(m.group('indent')) == len(error.group('indent')): + error = None + s1 += line + '\n' + continue + + line = DESCRIPTION_CODE_RE.sub(r'\1\3\4%s\5%s\6' % (TERM_COLOR_MAP[error.group('level').upper()], TERM_COLOR_MAP['RESET']), line) + line = KEY_RE.sub(r'\1%s\3%s\4' % (TERM_COLOR_MAP['BOLD'], TERM_COLOR_MAP['RESET']), line) + s1 += line + '\n' + + return s1.rstrip() + def test_m2crypto(): try: import M2Crypto @@ -213,11 +284,6 @@ def main(argv): else: names.append(name) - if '-p' in opts: - kwargs = { 'indent': 4, 'separators': (',', ': ') } - else: - kwargs = {} - if '-o' not in opts or opts['-o'] == '-': opts['-o'] = sys.stdout.fileno() try: @@ -226,6 +292,14 @@ def main(argv): logger.error('%s: "%s"' % (e.strerror, opts['-o'])) sys.exit(3) + if fh.isatty(): + opts['-p'] = None + + if '-p' in opts: + kwargs = { 'indent': 4, 'separators': (',', ': ') } + else: + kwargs = {} + # if trusted keys were supplied, check that pygraphviz is installed if trusted_keys: test_pygraphviz() @@ -266,7 +340,10 @@ def main(argv): name_obj.serialize_status(d, loglevel=loglevel) if d: - fh.write(json.dumps(d, ensure_ascii=False, **kwargs).encode('utf-8')) + s = json.dumps(d, ensure_ascii=False, **kwargs) + if fh.isatty() and os.environ.get('TERM', 'dumb') != 'dumb': + s = color_json(s) + fh.write(s.encode('utf-8')) except KeyboardInterrupt: logger.error('Interrupted.') |