Enhanced bcfg2-admin viz to allow output to be limited to one client.

By default, bcfg2-admin viz creates a diagram that shows the complete
contents of the repository: all Groups, Bundles, and (optionally)
Hosts/Clients. In a complicated configuration, this can be an
overwhelming amount of information.

This change adds an --only-client option that can be used to limit the
elements on the diagram to those that apply to the named host.

2 files changed, 48 insertions, 21 deletions

diff --git a/src/lib/Server/Admin/Viz.py b/src/lib/Server/Admin/Viz.py
index f39e6d7a8..5d7086e26 100644
--- a/src/lib/Server/Admin/Viz.py
+++ b/src/lib/Server/Admin/Viz.py
@@ -1,5 +1,6 @@
 import getopt
 from subprocess import Popen, PIPE
+import sys
 
 import Bcfg2.Server.Admin
 
@@ -8,18 +9,22 @@ class Viz(Bcfg2.Server.Admin.MetadataCore):
     __shorthelp__ = "Produce graphviz diagrams of metadata structures"
     __longhelp__ = (__shorthelp__ + "\n\nbcfg2-admin viz [--includehosts] "
                                     "[--includebundles] [--includekey] "
+                                    "[--only-client clientname] "
                                     "[-o output.png] [--raw]\n")
     __usage__ = ("bcfg2-admin viz [options]\n\n"
-                 "     %-25s%s\n"
-                 "     %-25s%s\n"
-                 "     %-25s%s\n"
-                 "     %-25s%s\n" %
+                 "     %-32s%s\n"
+                 "     %-32s%s\n"
+                 "     %-32s%s\n"
+                 "     %-32s%s\n"
+                 "     %-32s%s\n" %
                 ("-H, --includehosts",
                  "include hosts in the viz output",
                  "-b, --includebundles",
                  "include bundles in the viz output",
                  "-k, --includekey",
                  "show a key for different digraph shapes",
+                 "-c, --only-client <clientname>",
+                 "show only the groups, bundles for the named client",
                  "-o, --outfile <file>",
                  "write viz output to an output file"))
 
@@ -42,18 +47,21 @@ class Viz(Bcfg2.Server.Admin.MetadataCore):
         Bcfg2.Server.Admin.MetadataCore.__call__(self, args)
         # First get options to the 'viz' subcommand
         try:
-            opts, args = getopt.getopt(args, 'Hbko:',
+            opts, args = getopt.getopt(args, 'Hbkc:o:',
                                        ['includehosts', 'includebundles',
-                                        'includekey', 'outfile='])
+                                        'includekey', 'only-client=', 'outfile='])
         except getopt.GetoptError:
             msg = sys.exc_info()[1]
             print(msg)
+            self.log.error(self.__shorthelp__)
+            raise SystemExit(1)
 
         #FIXME: is this for --raw?
         #rset = False
         hset = False
         bset = False
         kset = False
+        only_client = None
         outputfile = False
         for opt, arg in opts:
             if opt in ("-H", "--includehosts"):
@@ -62,16 +70,18 @@ class Viz(Bcfg2.Server.Admin.MetadataCore):
                 bset = True
             elif opt in ("-k", "--includekey"):
                 kset = True
+            elif opt in ("-c", "--only-client"):
+                only_client = arg
             elif opt in ("-o", "--outfile"):
                 outputfile = arg
 
         data = self.Visualize(self.get_repo_path(), hset, bset,
-                              kset, outputfile)
+                              kset, only_client, outputfile)
         print(data)
         raise SystemExit(0)
 
     def Visualize(self, repopath, hosts=False,
-                  bundles=False, key=False, output=False):
+                  bundles=False, key=False, only_client=None, output=False):
         """Build visualization of groups file."""
         if output:
             format = output.split('.')[-1]
@@ -90,7 +100,7 @@ class Viz(Bcfg2.Server.Admin.MetadataCore):
             raise SystemExit(1)
         dotpipe.stdin.write('\trankdir="LR";\n')
         dotpipe.stdin.write(self.metadata.viz(hosts, bundles,
-                                                key, self.colors))
+                                              key, only_client, self.colors))
         if key:
             dotpipe.stdin.write("\tsubgraph cluster_key {\n")
             dotpipe.stdin.write('''\tstyle="filled";\n''')
diff --git a/src/lib/Server/Plugins/Metadata.py b/src/lib/Server/Plugins/Metadata.py
index 6570f2912..7fc34f178 100644
--- a/src/lib/Server/Plugins/Metadata.py
+++ b/src/lib/Server/Plugins/Metadata.py
@@ -782,8 +782,20 @@ class Metadata(Bcfg2.Server.Plugin.Plugin,
             xdict['xquery'][0].set('auth', 'cert')
             self.clients_xml.write_xml(xdict['filename'], xdict['xmltree'])
 
-    def viz(self, hosts, bundles, key, colors):
+    def viz(self, hosts, bundles, key, only_client, colors):
         """Admin mode viz support."""
+        if only_client:
+            clientmeta = self.core.build_metadata(only_client)
+
+        def include_client(client):
+            return not only_client or client != only_client
+
+        def include_bundle(bundle):
+            return not only_client or bundle in clientmeta.bundles
+
+        def include_group(group):
+            return not only_client or group in clientmeta.groups
+        
         groups_tree = lxml.etree.parse(self.data + "/groups.xml")
         try:
             groups_tree.xinclude()
@@ -791,7 +803,6 @@ class Metadata(Bcfg2.Server.Plugin.Plugin,
             self.logger.error("Failed to process XInclude for file %s" % dest)
         groups = groups_tree.getroot()
         categories = {'default': 'grey83'}
-        instances = {}
         viz_str = ""
         egroups = groups.findall("Group") + groups.findall('.//Groups/Group')
         for group in egroups:
@@ -801,8 +812,11 @@ class Metadata(Bcfg2.Server.Plugin.Plugin,
         if None in categories:
             del categories[None]
         if hosts:
+            instances = {}
             clients = self.clients
             for client, profile in list(clients.items()):
+                if include_client(client):
+                    continue
                 if profile in instances:
                     instances[profile].append(client)
                 else:
@@ -817,7 +831,8 @@ class Metadata(Bcfg2.Server.Plugin.Plugin,
             bundles = []
             [bundles.append(bund.get('name')) \
                  for bund in groups.findall('.//Bundle') \
-                 if bund.get('name') not in bundles]
+                 if bund.get('name') not in bundles \
+                     and include_bundle(bund.get('name'))]
             bundles.sort()
             for bundle in bundles:
                 viz_str += '''\t"bundle-%s" [ label="%s", shape="septagon"];\n''' \
@@ -829,20 +844,22 @@ class Metadata(Bcfg2.Server.Plugin.Plugin,
             else:
                 style = "filled"
             gseen.append(group.get('name'))
-            viz_str += '\t"group-%s" [label="%s", style="%s", fillcolor=%s];\n' % \
-                (group.get('name'), group.get('name'), style, group.get('color'))
-            if bundles:
-                for bundle in group.findall('Bundle'):
-                    viz_str += '\t"group-%s" -> "bundle-%s";\n' % \
-                        (group.get('name'), bundle.get('name'))
+            if include_group(group.get('name')):
+                viz_str += '\t"group-%s" [label="%s", style="%s", fillcolor=%s];\n' % \
+                           (group.get('name'), group.get('name'), style, group.get('color'))
+                if bundles:
+                    for bundle in group.findall('Bundle'):
+                        viz_str += '\t"group-%s" -> "bundle-%s";\n' % \
+                                   (group.get('name'), bundle.get('name'))
         gfmt = '\t"group-%s" [label="%s", style="filled", fillcolor="grey83"];\n'
         for group in egroups:
             for parent in group.findall('Group'):
-                if parent.get('name') not in gseen:
+                if parent.get('name') not in gseen and include_group(parent.get('name')):
                     viz_str += gfmt % (parent.get('name'), parent.get('name'))
                     gseen.append(parent.get("name"))
-                viz_str += '\t"group-%s" -> "group-%s" ;\n' % \
-                    (group.get('name'), parent.get('name'))
+                if include_group(group.get('name')):
+                    viz_str += '\t"group-%s" -> "group-%s" ;\n' % \
+                               (group.get('name'), parent.get('name'))
         if key:
             for category in categories:
                 viz_str += '''\t"''' + category + '''" [label="''' + category + \