okd/bin/ossh

#!/usr/bin/env python2
# vim: expandtab:tabstop=4:shiftwidth=4

import argparse
import traceback
import sys
import os
import re
import ConfigParser

from openshift_ansible import awsutil

CONFIG_MAIN_SECTION = 'main'

class Ossh(object):
    def __init__(self):
        self.user = None
        self.host = None
        self.file_path = os.path.join(os.path.dirname(os.path.realpath(__file__)))

        # Default the config path to /etc
        self.config_path = os.path.join(os.path.sep, 'etc',  \
                                        'openshift_ansible', \
                                        'openshift_ansible.conf')

        self.parse_cli_args()
        self.parse_config_file()

        self.aws = awsutil.AwsUtil()

        if self.args.refresh_cache:
            self.get_hosts(True)
        else:
            self.get_hosts()

        # parse host and user
        self.process_host()

        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_config_file(self):
        if os.path.isfile(self.config_path):
            config = ConfigParser.ConfigParser()
            config.read(self.config_path)

    def parse_cli_args(self):
        parser = argparse.ArgumentParser(description='OpenShift Online SSH Tool.')
        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('--refresh-cache', default=False,
                            action="store_true", help="Force a refresh on the host cache.")
        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('-A', default=False, action="store_true",
                            help='Forward authentication agent')
        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.
        '''

        parts = self.args.host.split('@')

        # parse username if passed
        if len(parts) > 1:
            self.user = parts[0]
            self.host = parts[1]
        else:
            self.host = parts[0]

            if self.args.login_name:
                self.user = self.args.login_name


    def get_hosts(self, refresh_cache=False):
        '''Query our host inventory and return a dict where the format '''
        if refresh_cache:
            self.host_inventory = self.aws.get_inventory(['--refresh-cache'])['_meta']['hostvars']
        else:
            self.host_inventory = self.aws.get_inventory()['_meta']['hostvars']

    def select_host(self):
        '''select host attempts to match the host specified
           on the command line with a list of hosts.
        '''
        results = None
        if self.host_inventory.has_key(self.host):
           results = (self.host, self.host_inventory[self.host])
        else:
            print "Could not find specified host: %s." % self.host

        # default - no results found.
        return results

    def list_hosts(self, limit=None):
        '''Function to print out the host inventory.

           Takes a single parameter to limit the number of hosts printed.
        '''
        for host_id, server_info in self.host_inventory.items():
            print '{oo_name:<35} {oo_clusterid:<10} {oo_environment:<8} ' \
                '{oo_id:<15} {oo_public_ip:<18} {oo_private_ip:<18}'.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.A:
                ssh_args.append('-A')

            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)

            results = self.select_host()
            if not results:
                return # early exit, no results

            # Assume we have one and only one.
            server_info = results[1]

            ssh_args.append(server_info['oo_public_ip'])

            #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()