Graph EC2 security groups

graph-security-groups.py

import json
import argparse
import subprocess
import tempfile

def main():
    parser = argparse.ArgumentParser(description='Visualize AWS Security Groups')
    parser.add_argument('--file', dest='file', action='store', help='JSON output of `aws ec2 describe-security-groups`')
    parser.add_argument('--output', dest='output', action='store', help='SVG output filename')
    args = parser.parse_args()

    out = args.output or "output.svg"

    if args.file:
        print("Importing JSON file...")
        jsdata = open(args.file, "r").read()
        c = conv(jsdata.decode())
    else:
        print("Running `aws ec2 describe-security-groups`...")
        sp = subprocess.Popen(["aws", "ec2", "describe-security-groups"], stdin=subprocess.PIPE, stdout=subprocess.PIPE)
        output, err = sp.communicate()
        out = output.decode()
        if len(out) == 0 or out[0] != "{":
            print("Error running `aws ec2 describe-security-groups`...")
            print(out)
            return
        c = conv(out)

    tosvg(c, out)

def conv(jsdata):
    js = json.loads(jsdata)
    d = js["SecurityGroups"]
    for i in d:
        name = None
        for j in i["Tags"]:
            if j["Key"] == "Name":
                i["GroupName"] = j["Value"]

    tf = tempfile.NamedTemporaryFile(delete=False)
    tf.write(json.dumps(js).encode())
    return tf.name

def tosvg(fname, out):
    print("Running `aws_security_viz`...")
    subprocess.Popen(["aws_security_viz", "-o", fname, "-f", out]).communicate()
    print("Saved as", out)

if __name__ == '__main__':
    main()