Update docker-py version, use MVC pattern, automatically select file type

.gitignore

@@ -1,4 +1,5 @@
 *.gv
 *.pdf
 *.png
+*.svg
 .idea

Pipfile

@@ -4,10 +4,10 @@ verify_ssl = true
 name = "pypi"
 
 [packages]
-docker-py = "*"
+docker = "*"
 graphviz = "*"
 
 [dev-packages]
 
 [requires]
-python_version = "3"
+python_version = "3.7"

Pipfile.lock

@@ -1,11 +1,11 @@
 {
     "_meta": {
         "hash": {
-            "sha256": "b980cd984b5794dc93711b11dcaadee33a6ee25a811c7aa8c8a6c7db32a78204"
+            "sha256": "9ea88e9a4fb06636a8077b1f1bc6e322fa52659b9159fffa1260d5ad8225ddd4"
         },
         "pipfile-spec": 6,
         "requires": {
-            "python_version": "3"
+            "python_version": "3.7"
         },
         "sources": [
             {
@@ -18,10 +18,10 @@
     "default": {
         "certifi": {
             "hashes": [
-                "sha256:376690d6f16d32f9d1fe8932551d80b23e9d393a8578c5633a2ed39a64861638",
+                "sha256:339dc09518b07e2fa7eda5450740925974815557727d6bd35d319c1524a04a4c",
-                "sha256:456048c7e371c089d0a77a5212fb37a2c2dce1e24146e3b7e0261736aaeaa22a"
+                "sha256:6d58c986d22b038c8c0df30d639f23a3e6d172a05c3583e766f4c0b785c0986a"
             ],
-            "version": "==2018.8.24"
+            "version": "==2018.10.15"
         },
         "chardet": {
             "hashes": [
@@ -30,13 +30,13 @@
             ],
             "version": "==3.0.4"
         },
-        "docker-py": {
+        "docker": {
             "hashes": [
-                "sha256:35b506e95861914fa5ad57a6707e3217b4082843b883be246190f57013948aba",
+                "sha256:31421f16c01ffbd1ea7353c7e7cd7540bf2e5906d6173eb51c8fea4e0ea38b19",
-                "sha256:4c2a75875764d38d67f87bc7d03f7443a3895704efc57962bdf6500b8d4bc415"
+                "sha256:fbe82af9b94ccced752527c8de07fa20267f9634b48674ba478a0bb4000a0b1e"
             ],
             "index": "pypi",
-            "version": "==1.10.6"
+            "version": "==3.5.1"
         },
         "docker-pycreds": {
             "hashes": [
@@ -47,11 +47,11 @@
         },
         "graphviz": {
             "hashes": [
-                "sha256:310bacfb969f0ac7c872610500e017c3e82b24a8abd33d289e99af162de30cb8",
+                "sha256:0e1744a45b0d707bc44f99c7b8e5f25dc22cf96b6aaf2432ac308ed9822a9cb6",
-                "sha256:865afa6ab9775cf29db03abd8e571a164042c726c35a1b3c1e2b8c4c645e2993"
+                "sha256:d311be4fddfe832a56986ac5e1d6e8715d7fcb0208560da79d1bb0f72abef41f"
             ],
             "index": "pypi",
-            "version": "==0.9"
+            "version": "==0.10.1"
         },
         "idna": {
             "hashes": [
@@ -62,10 +62,10 @@
         },
         "requests": {
             "hashes": [
-                "sha256:63b52e3c866428a224f97cab011de738c36aec0185aa91cfacd418b5d58911d1",
+                "sha256:65b3a120e4329e33c9889db89c80976c5272f56ea92d3e74da8a463992e3ff54",
-                "sha256:ec22d826a36ed72a7358ff3fe56cbd4ba69dd7a6718ffd450ff0e9df7a47ce6a"
+                "sha256:ea881206e59f41dbd0bd445437d792e43906703fff75ca8ff43ccdb11f33f263"
             ],
-            "version": "==2.19.1"
+            "version": "==2.20.1"
         },
         "six": {
             "hashes": [
@@ -76,18 +76,17 @@
         },
         "urllib3": {
             "hashes": [
-                "sha256:a68ac5e15e76e7e5dd2b8f94007233e01effe3e50e8daddf69acfd81cb686baf",
+                "sha256:61bf29cada3fc2fbefad4fdf059ea4bd1b4a86d2b6d15e1c7c0b582b9752fe39",
-                "sha256:b5725a0bd4ba422ab0e66e89e030c806576753ea3ee08554382c14e685d117b5"
+                "sha256:de9529817c93f27c8ccbfead6985011db27bd0ddfcdb2d86f3f663385c6a9c22"
             ],
-            "markers": "python_version < '4' and python_version != '3.1.*' and python_version >= '2.6' and python_version != '3.2.*' and python_version != '3.3.*' and python_version != '3.0.*'",
+            "version": "==1.24.1"
-            "version": "==1.23"
         },
         "websocket-client": {
             "hashes": [
-                "sha256:030bbfbf29ac9e315ffb207ed5ed42b6981b5038ea00d1e13b02b872cc95e8f6",
+                "sha256:8c8bf2d4f800c3ed952df206b18c28f7070d9e3dcbd6ca6291127574f57ee786",
-                "sha256:a35bac3d9647c62c1ba3e8a7340385d92981f5486b033557d592138fd4b21b90"
+                "sha256:e51562c91ddb8148e791f0155fdb01325d99bb52c4cdbb291aee7a3563fd0849"
             ],
-            "version": "==0.51.0"
+            "version": "==0.54.0"
         }
     },
     "develop": {}

docker-net-graph.py

@@ -1,10 +1,13 @@
 #!/usr/bin/python3
+
 import os
-import json
 import argparse
 import random
-from docker import Client
+import docker
+import typing
+from dataclasses import dataclass
 from graphviz import Graph
+from graphviz.backend import FORMATS
 
 # colorlover.scales["12"]["qual"]["Paired"] converted to hex strings
 COLORS = ["#1f78b4", "#33a02c", "#e31a1c", "#ff7f00", "#6a3d9a", "#b15928", "#a6cee3", "#b2df8a", "#fdbf6f",
@@ -12,7 +15,36 @@ COLORS = ["#1f78b4", "#33a02c", "#e31a1c", "#ff7f00", "#6a3d9a", "#b15928", "#a6
 i = 0
 
 
-def get_unique_color():
+@dataclass
+class Network:
+    name: str
+    gateway: str
+    internal: bool
+    isolated: bool
+    color: str
+
+
+@dataclass
+class Interface:
+    endpoint_id: str
+    address: str
+
+
+@dataclass
+class Container:
+    container_id: str
+    name: str
+    interfaces: typing.List[Interface]
+
+
+@dataclass
+class Link:
+    container_id: str
+    endpoint_id: str
+    network_name: str
+
+
+def get_unique_color() -> str:
     global i
 
     if i < len(COLORS):
@@ -25,98 +57,135 @@ def get_unique_color():
     return c
 
 
-def generate_graph(verbose: bool, file: str):
+def get_networks(client: docker.DockerClient, verbose: bool) -> typing.Dict[str, Network]:
-    g = Graph(comment="Docker Network Graph", engine="sfdp", format="png",
+    networks: typing.Dict[str, Network] = {}
-              graph_attr=dict(splines="true"))
-
-    docker_client = Client(os.environ.get("DOCKER_HOST", "unix:///var/run/docker.sock"))
-
-    def dump_json(obj):
-        print(json.dumps(obj, indent=4))
-
-    for c in docker_client.containers():
-        name = c["Names"][0][1:]
-        container_id = c["Id"]
-
-        node_id = f"container_{container_id}"
-
-        iface_labels = []
-
-        for net_name, net_info in c["NetworkSettings"]["Networks"].items():
-            label_iface = f"<{net_info['EndpointID']}> {net_info['IPAddress']}"
-
-            iface_labels.append(label_iface)
-
-        labels = "|".join(iface_labels)
-        if verbose:
-            print(labels)
-
-        g.node(node_id,
-               shape="record",
-               label=f"{{ {name} | { {labels} } }}",
-               fillcolor="#ff9999",
-               style="filled"
-               )
-
-    for net in sorted(docker_client.networks(), key=lambda k: k["Name"]):
-        net_name = net["Name"]
-        color = get_unique_color()
 
+    for net in sorted(client.networks.list(), key=lambda k: k.name):
         try:
-            gateway = net["IPAM"]["Config"][0]["Gateway"]
+            gateway = net.attrs["IPAM"]["Config"][0]["Gateway"]
         except (KeyError, IndexError):
             # This network doesn't seem to be used, skip it
             continue
 
-        internal = ""
+        internal = False
         try:
-            if net["Internal"]:
+            if net.attrs["Internal"]:
-                internal = "| Internal"
+                internal = True
         except KeyError:
             pass
 
-        isolated = ""
+        isolated = False
         try:
-            if net["Options"]["com.docker.network.bridge.enable_icc"] == "false":
+            if net.attrs["Options"]["com.docker.network.bridge.enable_icc"] == "false":
-                isolated = "| Containers isolated"
+                isolated = True
         except KeyError:
             pass
 
         if verbose:
-            print(f"Network: {net_name} {internal} gw:{gateway}")
+            print(f"Network: {net.name} {'internal' if internal else ''} {'isolated' if isolated else ''} gw:{gateway}")
 
-        net_node_id = f"net_{net_name}"
+        color = get_unique_color()
+        networks[net.name] = Network(net.name, gateway, internal, isolated, color)
 
-        label = f"{{<gw_iface> {gateway} | {net_name} {internal} {isolated}}}"
+    return networks
 
-        g.node(net_node_id,
-               shape="record",
-               label=label,
-               fillcolor=color,
-               style="filled"
-               )
 
-        if net["Containers"]:
+def get_containers(client: docker.DockerClient, verbose: bool) -> (typing.List[Container], typing.List[Link]):
-            for container_id, container in sorted(net["Containers"].items()):
+    containers: typing.List[Container] = []
-                if verbose:
+    links: typing.List[Link] = []
-                    dump_json(container)
-                    print(" * ", container["Name"], container["IPv4Address"], container["IPv6Address"])
 
-                container_node_id = f"container_{container_id}"
+    for container in client.containers.list():
+        interfaces: typing.List[Interface] = []
 
-                container_iface_ref = f"{container_node_id}:{container['EndpointID']}"
+        # Iterate over container interfaces
+        for net_name, net_info in container.attrs["NetworkSettings"]["Networks"].items():
+            endpoint_id = net_info["EndpointID"]
 
-                g.edge(container_iface_ref, f"{net_node_id}:gw_iface", color=color)
+            interfaces.append(Interface(endpoint_id, net_info['IPAddress']))
+            links.append(Link(container.id, endpoint_id, net_name))
+
+        if verbose:
+            print(f"Container: {container.name} {''.join([iface.address for iface in interfaces])}")
+
+        containers.append(Container(container.id, container.name, interfaces))
+
+    return containers, links
+
+
+def draw_network(g: Graph, net: Network):
+    label = f"{{<gw_iface> {net.gateway} | {net.name}"
+    if net.internal:
+        label += " | Internal"
+    if net.isolated:
+        label += " | Containers isolated"
+    label += "}"
+
+    g.node(f"network_{net.name}",
+           shape="record",
+           label=label,
+           fillcolor=net.color,
+           style="filled"
+           )
+
+
+def draw_container(g: Graph, c: Container):
+    iface_labels = [f"<{iface.endpoint_id}> {iface.address}" for iface in c.interfaces]
+
+    label = f"{{ {c.name} | {{ {'|'.join(iface_labels)} }} }}"
+
+    g.node(f"container_{c.container_id}",
+           shape="record",
+           label=label,
+           fillcolor="#ff9999",
+           style="filled"
+           )
+
+
+def draw_link(g: Graph, networks: typing.Dict[str, Network], link: Link):
+    g.edge(f"container_{link.container_id}:{link.endpoint_id}",
+           f"network_{link.network_name}",
+           color=networks[link.network_name].color
+           )
+
+
+def generate_graph(verbose: bool, file: str):
+    docker_client = docker.from_env()
+
+    networks = get_networks(docker_client, verbose)
+    containers, links = get_containers(docker_client, verbose)
 
-    print(g.source)
     if file:
-        g.render(file)
+        base, ext = os.path.splitext(file)
+        g = Graph(comment="Docker Network Graph", engine="sfdp", format=ext[1:], graph_attr=dict(splines="true"))
+    else:
+        g = Graph(comment="Docker Network Graph", engine="sfdp", graph_attr=dict(splines="true"))
+
+    for _, network in networks.items():
+        draw_network(g, network)
+
+    for container in containers:
+        draw_container(g, container)
+
+    for link in links:
+        draw_link(g, networks, link)
+
+    if file:
+        g.render(base)
+    else:
+        print(g.source)
+
+
+def graphviz_output_file(filename: str):
+    ext = os.path.splitext(filename)[1][1:]
+    if ext.lower() not in FORMATS:
+        raise argparse.ArgumentTypeError("Must be valid graphviz output format")
+    return filename
 
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Generate docker network graph.")
     parser.add_argument("-v", "--verbose", help="Verbose output", action="store_true")
-    parser.add_argument("-o", "--out", help="Write output to file", type=str)
+    parser.add_argument("-o", "--out", help="Write output to file", type=graphviz_output_file)
     args = parser.parse_args()
 
     generate_graph(args.verbose, args.out)