Merge pull request #54 from kwoodson/ossh

Adding ssh dynamic hosts capabilities

5 files changed, 354 insertions, 0 deletions

diff --git a/bin/README_SHELL_COMPLETION b/bin/README_SHELL_COMPLETION
new file mode 100644
index 000000000..e17b4b205
--- /dev/null
+++ b/bin/README_SHELL_COMPLETION
@@ -0,0 +1,33 @@
+# ossh is an ssh replacement.
+
+
+ossh uses a dynamic inventory cache in order to lookup hostnames and translate them
+to something meaningful such as an IP address or dns name.
+
+This allows us to treat our servers as cattle and not as pets.
+
+If you have not run the ossh command and it has not laid down
+a cache file the completions will not be available.
+
+You can populate the cache by running `ossh --list`.  This
+will populate the cache file and the completions should
+become available.
+
+This script will look at the cached version of your
+multi_ec2 results in ~/.ansible/tmp/multi_ec2_inventory.cache.
+It will then parse a few {host}.{env} out of the json
+and return them to be completable.
+
+# BASH
+In order to setup bash completion, source the following script:
+/path/to/repository/openshift-online-ansible/bin/ossh_bash_completion
+
+# ZSH
+In order to setup zsh completion, you will need to verify
+that the _ossh_zsh_completion script is somewhere in the path
+of $fpath.
+
+Once $fpath includes the _ossh_zsh_completion script then you should
+run `exec zsh`.  This will then allow you to call `ossh host[TAB]`
+for a list of completions.
+
diff --git a/bin/ansibleutil.py b/bin/ansibleutil.py
new file mode 100644
index 000000000..b12b7b447
--- /dev/null
+++ b/bin/ansibleutil.py
@@ -0,0 +1,69 @@
+# vim: expandtab:tabstop=4:shiftwidth=4
+
+import subprocess
+import sys
+import os
+import json
+import re
+
+class AnsibleUtil(object):
+    def __init__(self):
+        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'))
+
+    def get_inventory(self,args=[]):
+        cmd = [self.multi_ec2_path]
+
+        if args:
+            cmd.extend(args)
+
+        env = {}
+        p = subprocess.Popen(cmd, stderr=subprocess.PIPE,
+                         stdout=subprocess.PIPE, env=env)
+
+        out,err = p.communicate()
+
+        if p.returncode != 0:
+            raise RuntimeError(err)
+
+        return json.loads(out.strip())
+
+    def get_environments(self):
+        pattern = re.compile(r'^tag_environment_(.*)')
+
+        envs = []
+        inv = self.get_inventory()
+        for key in inv.keys():
+            m = pattern.match(key)
+            if m:
+                envs.append(m.group(1))
+
+        return envs
+
+    def get_security_groups(self):
+        pattern = re.compile(r'^security_group_(.*)')
+
+        groups = []
+        inv = self.get_inventory()
+        for key in inv.keys():
+            m = pattern.match(key)
+            if m:
+                groups.append(m.group(1))
+
+        return groups
+
+    def build_host_dict(self, args=[]):
+        inv = self.get_inventory(args)
+
+        inst_by_env = {}
+        for dns, host in inv['_meta']['hostvars'].items():
+            if host['ec2_tag_environment'] not in inst_by_env:
+                inst_by_env[host['ec2_tag_environment']] = {}
+            host_id = "%s:%s" % (host['ec2_tag_Name'],host['ec2_id'])
+            inst_by_env[host['ec2_tag_environment']][host_id] = host
+
+
+        return inst_by_env
+
+
+
diff --git a/bin/ossh b/bin/ossh
new file mode 100755
index 000000000..0b48ca46c
--- /dev/null
+++ b/bin/ossh
@@ -0,0 +1,210 @@
+#!/usr/bin/env python
+# vim: expandtab:tabstop=4:shiftwidth=4
+
+import argparse
+import ansibleutil
+import traceback
+import sys
+import os
+import re
+
+class Ossh(object):
+    def __init__(self):
+        self.file_path = os.path.join(os.path.dirname(os.path.realpath(__file__)))
+        self.parse_cli_args()
+
+        self.ansible = ansibleutil.AnsibleUtil()
+
+        # get a dict of host inventory
+        if self.args.list:
+            self.get_hosts()
+        else:
+            self.get_hosts(True)
+
+        # parse host and user
+        self.process_host()
+
+        if not self.args.list and not self.env:
+            print "Please specify an environment."
+            return
+
+        if self.args.host == '' and not self.args.list:
+            self.parser.print_help()
+            return
+
+
+        if self.args.debug:
+            print self.args
+
+        # perform the SSH
+        if self.args.list:
+            self.list_hosts()
+        else:
+            self.ssh()
+
+    def parse_cli_args(self):
+        parser = argparse.ArgumentParser(description='Openshift Online SSH Tool.')
+        parser.add_argument('-e', '--env', action="store",
+                          help="Which environment to search for the host ")
+        parser.add_argument('-d', '--debug', default=False,
+                          action="store_true", help="debug mode")
+        parser.add_argument('-v', '--verbose', default=False,
+                          action="store_true", help="Verbose?")
+        parser.add_argument('--list', default=False,
+                          action="store_true", help="list out hosts")
+        parser.add_argument('-c', '--command', action='store',
+                            help='Command to run on remote host')
+        parser.add_argument('-l', '--login_name', action='store',
+                            help='User in which to ssh as')
+
+        parser.add_argument('-o', '--ssh_opts', action='store',
+                            help='options to pass to SSH.\n \
+                                  "-oForwardX11=yes,TCPKeepAlive=yes"')
+        parser.add_argument('host', nargs='?', default='')
+
+        self.args = parser.parse_args()
+        self.parser = parser
+
+
+    def process_host(self):
+        '''Determine host name and user name for SSH.
+        '''
+        self.env = None
+        self.user = None
+
+        re_env = re.compile("\.(" + "|".join(self.host_inventory.keys()) + ")")
+        search = re_env.search(self.args.host)
+        if self.args.env:
+            self.env = self.args.env
+        elif search:
+            # take the first?
+            self.env = search.groups()[0]
+
+        # remove env from hostname command line arg if found
+        if search:
+            self.args.host = re_env.split(self.args.host)[0]
+
+        # parse username if passed
+        if '@' in self.args.host:
+            self.user, self.host = self.args.host.split('@')
+        else:
+            self.host = self.args.host
+            if self.args.login_name:
+                self.user = self.args.login_name
+
+    def get_hosts(self, cache_only=False):
+        '''Query our host inventory and return a dict where the format
+           equals:
+
+           dict['servername'] = dns_name
+        '''
+        if cache_only:
+            self.host_inventory = self.ansible.build_host_dict(['--cache-only'])
+        else:
+            self.host_inventory = self.ansible.build_host_dict()
+
+    def select_host(self, regex=False):
+        '''select host attempts to match the host specified
+           on the command line with a list of hosts.
+
+           if regex is specified then we will attempt to match
+           all *{host_string}* equivalents.
+        '''
+        re_host = re.compile(self.host)
+
+        results = []
+        for hostname, server_info in self.host_inventory[self.env].items():
+            if hostname.split(':')[0] == self.host:
+                # an exact match, return it!
+                return [(hostname, server_info)]
+            elif re_host.search(hostname):
+                results.append((hostname, server_info))
+
+        if results:
+            return results
+        else:
+            print "Could not find specified host: %s in %s" % (self.host, self.env)
+
+        # default - no results found.
+        return None
+
+    def list_hosts(self, limit=None):
+        '''Function to print out the host inventory.
+
+           Takes a single parameter to limit the number of hosts printed.
+        '''
+
+        if self.env:
+            results = self.select_host(True)
+            if len(results) == 1:
+                hostname, server_info = results[0]
+                sorted_keys = server_info.keys()
+                sorted_keys.sort()
+                for key in sorted_keys:
+                    print '{0:<35} {1}'.format(key, server_info[key])
+            else:
+                for host_id, server_info in results[:limit]:
+                    name = server_info['ec2_tag_Name']
+                    ec2_id = server_info['ec2_id']
+                    ip = server_info['ec2_ip_address']
+                    print '{ec2_tag_Name:<35} {ec2_tag_environment:<8} {ec2_id:<15} {ec2_ip_address}'.format(**server_info)
+
+                if limit:
+                    print
+                    print 'Showing only the first %d results...' % limit
+                    print
+
+        else:
+            for env, host_ids in self.host_inventory.items():
+                for host_id, server_info in host_ids.items():
+                    name = server_info['ec2_tag_Name']
+                    ec2_id = server_info['ec2_id']
+                    ip = server_info['ec2_ip_address']
+                    print '{ec2_tag_Name:<35} {ec2_tag_environment:<5} {ec2_id:<15} {ec2_ip_address}'.format(**server_info)
+
+    def ssh(self):
+        '''SSH to a specified host
+        '''
+        try:
+            # shell args start with the program name in position 1
+            ssh_args = ['/usr/bin/ssh']
+
+            if self.user:
+                ssh_args.append('-l%s' % self.user)
+
+            if self.args.verbose:
+                ssh_args.append('-vvv')
+
+            if self.args.ssh_opts:
+                for arg in self.args.ssh_opts.split(","):
+                    ssh_args.append("-o%s" % arg)
+
+            result = self.select_host()
+            if not result:
+                return # early exit, no results
+
+            if len(result) > 1:
+                self.list_hosts(10)
+                return # early exit, too many results
+
+            # Assume we have one and only one.
+            hostname, server_info = result[0]
+            dns = server_info['ec2_public_dns_name']
+
+            ssh_args.append(dns)
+
+            #last argument
+            if self.args.command:
+                ssh_args.append("%s" % self.args.command)
+
+            print "Running: %s\n" % ' '.join(ssh_args)
+
+            os.execve('/usr/bin/ssh', ssh_args, os.environ)
+        except:
+            print traceback.print_exc()
+            print sys.exc_info()
+
+
+if __name__ == '__main__':
+    ossh = Ossh()
+
diff --git a/bin/ossh_bash_completion b/bin/ossh_bash_completion
new file mode 100755
index 000000000..0d0bdb0e6
--- /dev/null
+++ b/bin/ossh_bash_completion
@@ -0,0 +1,18 @@
+__ossh_known_hosts(){
+    if [[ -f ~/.ansible/tmp/multi_ec2_inventory.cache ]]; then
+      /usr/bin/python -c 'import json,os; z = json.loads(open("%s"%os.path.expanduser("~/.ansible/tmp/multi_ec2_inventory.cache")).read()); print "\n".join(["%s.%s" % (host["ec2_tag_Name"],host["ec2_tag_environment"]) for dns, host in z["_meta"]["hostvars"].items()])'
+    fi
+}
+
+_ossh()
+{
+    local cur prev known_hosts
+    COMPREPLY=()
+    cur="${COMP_WORDS[COMP_CWORD]}"
+    prev="${COMP_WORDS[COMP_CWORD-1]}"
+    known_hosts="$(__ossh_known_hosts)"
+    COMPREPLY=( $(compgen -W "${known_hosts}" -- ${cur}))
+
+    return 0
+}
+complete -F _ossh ossh
diff --git a/bin/ossh_zsh_completion b/bin/ossh_zsh_completion
new file mode 100644
index 000000000..f057ca8ce
--- /dev/null
+++ b/bin/ossh_zsh_completion
@@ -0,0 +1,24 @@
+#compdef ossh
+
+_ossh_known_hosts(){
+  if [[ -f ~/.ansible/tmp/multi_ec2_inventory.cache ]]; then
+    print $(/usr/bin/python -c 'import json,os; z = json.loads(open("%s"%os.path.expanduser("~/.ansible/tmp/multi_ec2_inventory.cache")).read()); print "\n".join(["%s.%s" % (host["ec2_tag_Name"],host["ec2_tag_environment"]) for dns, host in z["_meta"]["hostvars"].items()])')
+  fi
+
+}
+_ossh(){
+  local curcontext="$curcontext" state line
+  typeset -A opt_args
+
+#_arguments "*:Hosts:_ossh_known_hosts"
+    _arguments -s : \
+        "*:hosts:->hosts"
+
+    case "$state" in
+        hosts)
+            _values 'hosts' $(_ossh_known_hosts)
+            ;;
+    esac
+
+}
+_ossh "$@"