added the ability to run opssh and ohi on all hosts in an environment, as well as all hosts of the same host-type regardless of environment

1 files changed, 25 insertions, 11 deletions

diff --git a/bin/ohi b/bin/ohi
index 06a375cdb..408961ee4 100755
--- a/bin/ohi
+++ b/bin/ohi
@@ -12,13 +12,12 @@ import subprocess
 import ConfigParser
 
 from openshift_ansible import awsutil
+from openshift_ansible.awsutil import ArgumentError
 
 CONFIG_MAIN_SECTION = 'main'
 CONFIG_HOST_TYPE_ALIAS_SECTION = 'host_type_aliases'
 CONFIG_INVENTORY_OPTION = 'inventory'
 
-class ArgumentMismatchError(ValueError): pass
-
 class Ohi(object):
     def __init__(self):
         self.inventory = None
@@ -40,15 +39,30 @@ class Ohi(object):
             self.aws.print_host_types()
             return 0
 
-        if self.args.env and \
-           self.args.host_type:
-            hosts = self.aws.get_host_list(self.args.host_type, self.args.env)
-            for host in hosts:
-                print host
-            return 0
+        hosts = None
+        if self.args.host_type is not None and \
+           self.args.env is not None:
+            # Both env and host-type specified
+            hosts = self.aws.get_host_list(host_type=self.args.host_type, \
+                                           env=self.args.env)
+
+        if self.args.host_type is None and \
+           self.args.env is not None:
+            # Only env specified
+            hosts = self.aws.get_host_list(env=self.args.env)
+
+        if self.args.host_type is not None and \
+           self.args.env is None:
+            # Only host-type specified
+            hosts = self.aws.get_host_list(host_type=self.args.host_type)
+
+        if hosts is None:
+            # We weren't able to determine what they wanted to do
+            raise ArgumentError("Invalid combination of arguments")
 
-        # If it makes it here, we weren't able to determine what they wanted to do
-        raise ArgumentMismatchError("Invalid combination of arguments")
+        for host in hosts:
+            print host
+        return 0
 
     def parse_config_file(self):
         if os.path.isfile(self.config_path):
@@ -92,5 +106,5 @@ if __name__ == '__main__':
         ohi = Ohi()
         exitcode = ohi.run()
         sys.exit(exitcode)
-    except ArgumentMismatchError as e:
+    except ArgumentError as e:
         print "\nError: %s\n" % e.message