From a397e1215fe3c25faa5038dc4def2d0b05776420 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?F=C3=BCckert=2C=20Oliver?= Date: Tue, 28 Jan 2020 22:00:09 +0100 Subject: [PATCH] Cleanup Code --- aci2dot.py | 186 +++++++++++++++++++++++++++++------------------------ 1 file changed, 101 insertions(+), 85 deletions(-) diff --git a/aci2dot.py b/aci2dot.py index 7343de3..7a96f1a 100755 --- a/aci2dot.py +++ b/aci2dot.py @@ -20,14 +20,13 @@ simple = False show_attributes = True - gformat = ''' graph [ size="8.27"; ratio="1"; nodesep="0.15"; - ranksep="0.5"; - #splines="false"; + ranksep="0.5"; + #splines="false"; rankdir=LR; bgcolor="transparent"; ]; @@ -42,103 +41,120 @@ ] ''' -def format_attr(policy,attr): - if not show_attributes: - return '<{0}>'.format(policy) +def format_attr(policy, attr): + + if not show_attributes: + return '<{0}>'.format(policy) - attr_table = '<' - attr_table = attr_table + ''.format(policy) - attr_table = attr_table + '' + attr_table = '<
{0}
' + attr_table = attr_table + \ + ''.format( + policy) + attr_table = attr_table + '' - for k, v in attr.items(): - if k not in ['status'] and v: - attr_table = attr_table + ''.format(k,(v[:20] + '..') if len(v) > 20 else v) + for k, v in attr.items(): + if k not in ['status'] and v: + attr_table = attr_table + \ + ''.format( + k, (v[:20] + '..') if len(v) > 20 else v) - attr_table = attr_table + '
{0}
{0}: {1}
{0}: {1}
>' + attr_table = attr_table + '>' + + return attr_table - return attr_table def iterd(d, i): - #print('# Index: {0}, {1}'.format(i, d)) - for k, v in d.items(): + for k, v in d.items(): + + if isinstance(v, dict): + attr = v.get("attributes") - if isinstance(v, dict): - attr = v.get("attributes") - - # Write Node Properties - if attr: - print('{0} [label={1};]'.format(k + str(i), format_attr(k, attr))) - else: - print('{0} [label="{0}";]'.format(k + str(i))) + if attr: + print('{0} [label={1};]'.format( + k + str(i), format_attr(k, attr))) + else: + print('{0} [label="{0}";]'.format(k + str(i))) - children = v.get("children") + children = v.get("children") + + if children: + childi = 0 + for child in children: + if not simple: + childi = childi + 1 + + for childk, childv in child.items(): + print('{0} -> {1}'.format(k + + str(i), childk + str(childi))) + iterd(child, childi) - if children: - i2=0 - for child in children: - if not simple: i2 = i2 + 1 - for childk, childv in child.items(): - print('{0} -> {1}'.format(k + str(i), childk + str(i2))) - iterd(child, i2) def write_gformat(): - with open(".aci2dot", "wt") as text_file: - text_file.write(gformat) + with open(".aci2dot", "wt") as text_file: + text_file.write(gformat) + def main(): - global simple - global show_attributes - global gformat - parser = argparse.ArgumentParser(description='Create DOT formatted Graph from JSON formatted ACI policy export.') - parser.add_argument('policy_file', type=argparse.FileType('r'), help='JSON ACI Policy Filename') - parser.add_argument('--nr', action='store_true', help='Suppress redundant children') - parser.add_argument('--na', action='store_true', help="Don't show attributes") - parser.add_argument('--write', action='store_true', help="Write config template to .aci2dot and exit") - group = parser.add_mutually_exclusive_group() - group.add_argument('--stdout', action='store_true', help="Write to STDOUT instead of to file") - choices = ["svg", "png", "pdf"] - group.add_argument('--dot', choices=choices, help="Also write SVG/PNG/PDF. 'dot' needs to be installed.") - - args = parser.parse_args() - - #print(args.dot) - #sys.exit() - - if args.write: - write_gformat() - sys.exit('Config template written to .aci2dot') - - simple = args.nr - show_attributes = not args.na if not args.nr else False - base_name = (os.path.splitext(args.policy_file.name)[0]) - - try: - data = json.load(args.policy_file) - except: - sys.exit('JSON File can not be read.') - - try: - with open(".aci2dot", 'r') as config_file: - gformat = config_file.read() - print("Graph config read from .aci2dot", file=sys.stderr) - except IOError: - pass - - if not args.stdout: - sys.stdout = open('{0}.dot'.format(base_name), 'w') - - print('strict digraph Policy {') - print(gformat) - iterd(data, 0) - print('}', flush=True) - - if args.dot: - os.system("dot -T{0} -o{1}.{0} {1}.dot".format(args.dot, base_name)) - #print(("dot -T{0} -o{1}.svg {1}.dot".format(args.dot, base_name))) - print("{0} exported to {1}.{0}".format(args.dot, base_name), file=sys.stderr) + global simple + global show_attributes + global gformat + + parser = argparse.ArgumentParser( + description='Create DOT formatted Graph from JSON formatted ACI policy export.') + parser.add_argument('policy_file', type=argparse.FileType( + 'r'), help='JSON ACI Policy Filename') + parser.add_argument('--nr', action='store_true', + help='Suppress redundant children') + parser.add_argument('--na', action='store_true', + help="Don't show attributes") + parser.add_argument('--write', action='store_true', + help="Write config template to .aci2dot and exit") + group = parser.add_mutually_exclusive_group() + group.add_argument('--stdout', action='store_true', + help="Write to STDOUT instead of to file") + choices = ["svg", "png", "pdf"] + group.add_argument('--dot', choices=choices, + help="Also write SVG/PNG/PDF. 'dot' needs to be installed.") + + args = parser.parse_args() + + if args.write: + write_gformat() + sys.exit('Config template written to .aci2dot') + + simple = args.nr + show_attributes = not args.na if not args.nr else False + + try: + data = json.load(args.policy_file) + except json.decoder.JSONDecodeError: + sys.exit('JSON File can not be read.') + + base_name = (os.path.splitext(args.policy_file.name)[0]) + + try: + with open(".aci2dot", 'r') as config_file: + gformat = config_file.read() + print("Graph config read from .aci2dot", file=sys.stderr) + except IOError: + pass + + if not args.stdout: + sys.stdout = open('{0}.dot'.format(base_name), 'w') + + print('strict digraph Policy {') + print(gformat) + iterd(data, 0) + print('}', flush=True) + + if args.dot: + os.system("dot -T{0} -o{1}.{0} {1}.dot".format(args.dot, base_name)) + print("{0} exported to {1}.{0}".format( + args.dot, base_name), file=sys.stderr) + if __name__ == '__main__': - main() \ No newline at end of file + main()