Merge pull request #111 from twiest/opssh

added opssh

2 files changed, 179 insertions, 1 deletions

diff --git a/bin/ansibleutil.py b/bin/ansibleutil.py
index b9f37726e..b527ab140 100644
--- a/bin/ansibleutil.py
+++ b/bin/ansibleutil.py
@@ -1,13 +1,27 @@
 # vim: expandtab:tabstop=4:shiftwidth=4
 
 import subprocess
-import sys
 import os
 import json
 import re
 
 class AnsibleUtil(object):
     def __init__(self):
+        self.host_type_aliases = {
+                'legacy-openshift-broker': ['broker', 'ex-srv'],
+                         'openshift-node': ['node', 'ex-node'],
+                   'openshift-messagebus': ['msg'],
+            'openshift-customer-database': ['mongo'],
+                'openshift-website-proxy': ['proxy'],
+            'openshift-community-website': ['drupal'],
+                         'package-mirror': ['mirror'],
+        }
+
+        self.alias_lookup = {}
+        for key, values in self.host_type_aliases.iteritems():
+            for value in values:
+                self.alias_lookup[value] = key
+
         self.file_path = os.path.join(os.path.dirname(os.path.realpath(__file__)))
         self.multi_ec2_path = os.path.realpath(os.path.join(self.file_path, '..','inventory','multi_ec2.py'))
 
@@ -39,8 +53,22 @@ class AnsibleUtil(object):
             if m:
                 envs.append(m.group(1))
 
+        envs.sort()
         return envs
 
+    def get_host_types(self):
+        pattern = re.compile(r'^tag_host-type_(.*)')
+
+        host_types = []
+        inv = self.get_inventory()
+        for key in inv.keys():
+            m = pattern.match(key)
+            if m:
+                host_types.append(m.group(1))
+
+        host_types.sort()
+        return host_types
+
     def get_security_groups(self):
         pattern = re.compile(r'^security_group_(.*)')
 
@@ -51,6 +79,7 @@ class AnsibleUtil(object):
             if m:
                 groups.append(m.group(1))
 
+        groups.sort()
         return groups
 
     def build_host_dict_by_env(self, args=[]):
@@ -65,3 +94,41 @@ class AnsibleUtil(object):
 
         return inst_by_env
 
+    # Display host_types
+    def print_host_types(self):
+        host_types = self.get_host_types()
+        ht_format_str = "%35s"
+        alias_format_str = "%-20s"
+        combined_format_str = ht_format_str + "    " + alias_format_str
+
+        print
+        print combined_format_str % ('Host Types', 'Aliases')
+        print combined_format_str % ('----------', '-------')
+
+        for ht in host_types:
+            aliases = []
+            if ht in self.host_type_aliases:
+                aliases = self.host_type_aliases[ht]
+                print combined_format_str % (ht, ", ".join(aliases))
+            else:
+                print  ht_format_str % ht
+        print
+
+    # Convert host-type aliases to real a host-type
+    def resolve_host_type(self, host_type):
+        if self.alias_lookup.has_key(host_type):
+            return self.alias_lookup[host_type]
+        return host_type
+
+    def gen_env_host_type_tag(self, host_type, env):
+        """Generate the environment host type tag
+        """
+        host_type = self.resolve_host_type(host_type)
+        return "tag_env-host-type_%s-%s" % (env, host_type)
+
+    def get_host_list(self, host_type, env):
+        """Get the list of hosts from the inventory using host-type and environment
+        """
+        inv = self.get_inventory()
+        host_type_tag = self.gen_env_host_type_tag(host_type, env)
+        return inv[host_type_tag]
diff --git a/bin/opssh b/bin/opssh
new file mode 100755
index 000000000..0836f0354
--- /dev/null
+++ b/bin/opssh
@@ -0,0 +1,111 @@
+#!/usr/bin/env python
+# vim: expandtab:tabstop=4:shiftwidth=4
+
+import argparse
+import ansibleutil
+import traceback
+import sys
+import os
+import re
+import tempfile
+import time
+import subprocess
+
+DEFAULT_PSSH_PAR=200
+PSSH = '/usr/bin/pssh'
+
+class Opssh(object):
+    def __init__(self):
+        self.file_path = os.path.join(os.path.dirname(os.path.realpath(__file__)))
+        self.ansible = ansibleutil.AnsibleUtil()
+
+        self.parse_cli_args()
+
+        if self.args.list_host_types:
+            self.ansible.print_host_types()
+            return
+
+        if self.args.env and \
+           self.args.host_type and \
+           self.args.command:
+            retval = self.run_pssh()
+            if retval != 0:
+                raise ValueError("ansible run failed")
+
+            return
+
+        # If it makes it here, we weren't able to determine what they wanted to do
+        raise ValueError("Invalid combination of arguments")
+
+    def run_pssh(self):
+        """Actually run the pssh command based off of the supplied options
+        """
+
+        # Default set of options
+        pssh_args = [PSSH, '-i', '-t', '0', '-p', str(self.args.par), '--user', self.args.user]
+
+        if self.args.outdir:
+            pssh_args.append("--outdir='%s'" % self.args.outdir)
+
+        if self.args.errdir:
+            pssh_args.append("--errdir='%s'" % self.args.errdir)
+
+        hosts = self.ansible.get_host_list(self.args.host_type, self.args.env)
+        with tempfile.NamedTemporaryFile(prefix='opssh-', delete=True) as f:
+            for h in hosts:
+                f.write(h + os.linesep)
+            f.flush()
+
+            pssh_args.extend(["-h", "%s" % f.name])
+            pssh_args.append("%s" % self.args.command)
+
+            print
+            print "Running: %s" % ' '.join(pssh_args)
+            print
+            return subprocess.call(pssh_args)
+
+        return None
+
+
+    def parse_cli_args(self):
+        """Setup the command line parser with the options we want
+        """
+
+        parser = argparse.ArgumentParser(description='Openshift Online PSSH Tool.')
+
+        parser.add_argument('--list-host-types', default=False, action='store_true',
+                       help='List all of the host types')
+
+        parser.add_argument('-e', '--env', action="store",
+                       help="Which environment to use")
+
+        parser.add_argument('-t', '--host-type', action="store",
+                       help="Which host type to use")
+
+        parser.add_argument('-c', '--command', action='store',
+                       help='Command to run on remote host(s)')
+
+        parser.add_argument('--user', action='store', default='root',
+                       help='username')
+
+        parser.add_argument('-p', '--par', action='store', default=DEFAULT_PSSH_PAR,
+                       help=('max number of parallel threads (default %s)' % DEFAULT_PSSH_PAR))
+
+        parser.add_argument('--outdir', action='store',
+                       help='output directory for stdout files')
+
+        parser.add_argument('--errdir', action='store',
+                       help='output directory for stderr files')
+
+        self.args = parser.parse_args()
+
+
+if __name__ == '__main__':
+    if len(sys.argv) == 1:
+        print "\nError: No options given. Use --help to see the available options\n"
+        sys.exit(0)
+
+    try:
+        opssh = Opssh()
+    except ValueError as e:
+        print "\nError: %s\n" % e.message