Further simplify minestruct

git-svn-id: https://svn.mcs.anl.gov/repos/bcfg/trunk/bcfg2@4877 ce84e21b-d406-0410-9b95-82705330c041

1 files changed, 23 insertions, 36 deletions

diff --git a/src/lib/Server/Admin/Minestruct.py b/src/lib/Server/Admin/Minestruct.py
index 9463bbc0e..83eaa2c0a 100644
--- a/src/lib/Server/Admin/Minestruct.py
+++ b/src/lib/Server/Admin/Minestruct.py
@@ -1,45 +1,40 @@
 '''Minestruct Admin Mode'''
 import Bcfg2.Server.Admin
-import lxml.etree
+import lxml.etree, sys, getopt
 
 class Minestruct(Bcfg2.Server.Admin.StructureMode):
     '''Pull extra entries out of statistics'''
-    __shorthelp__ = 'bcfg2-admin minestruct <client> [-f file-name] [-g groups]'
+    __shorthelp__ = 'bcfg2-admin minestruct [-f file-name] [-g groups] client'
     __longhelp__ = __shorthelp__ + '\n\tExtract extra entry lists from statistics'
 
     def __call__(self, args):
         Bcfg2.Server.Admin.Mode.__call__(self, args)
         if len(args) == 0:
             self.errExit("No hostname specified (see bcfg2-admin minestruct -h for help)")
-        if "-h" in args:
+        try:
+            (opts, args) = getopt.getopt(args, 'f:g:h')
+        except:
+            self.log.error(self.__shorthelp__)
+            raise SystemExit(1)
+        if "-h" in args or not args:
             print "Usage:"
             print self.__shorthelp__
             raise SystemExit(1)
-        write_to_file = False
-        file_arg = False
-        output_file = None
-        client = None
-        groups_arg = False
+        
+        client = args[0]
+        output = sys.stdout
         groups = []
-        for arg in args:
-            if arg == "-f":
-                file_arg = True
-                groups_arg = False
-                continue
-            elif arg == "-g":
-                groups_arg = True
-                file_arg = False
-                continue                           
-            elif file_arg == True:
-                output_file = arg
-                file_arg = False
-                write_to_file = True
-                continue
-            elif groups_arg == True:
-                groups.append(arg)
-                continue
-            else:
-                client = arg
+
+        for (opt, optarg) in opts:
+            if opt == '-f':
+                try:
+                    output = open(optarg, 'w')
+                except IOError:
+                    self.log.error("Failed to open file: %s" % (optarg))
+                    raise SystemExit(1)
+            elif opt == '-g':
+                groups = optarg.split(':')
+
         extra = self.statistics.GetExtra(client)
         root = lxml.etree.Element("Base")
         self.log.info("Found %d extra entries" % (len(extra)))
@@ -51,13 +46,5 @@ class Minestruct(Bcfg2.Server.Admin.StructureMode):
             lxml.etree.SubElement(add_point, tag, name=name)
 
         tree = lxml.etree.ElementTree(root)
-        if write_to_file == True:
-            try:
-                f = open(output_file, 'w')
-            except IOError:
-                self.log.info("Failed to write to file: %s" % (output_file))
-                raise SystemExit(1)
-            tree.write(f, pretty_print=True)
-        else:
-            print lxml.etree.tostring(tree, pretty_print=True)
+        tree.write(output, pretty_print=True)