#!/usr/bin/env python
#
# This script is an inventory-aware ssh wrapper that eliminates the need
# to repeat information from the inventory in .ssh/config.

from ansible.cli import CLI
from ansible.inventory import Inventory
from ansible.parsing.dataloader import DataLoader
from ansible.vars import VariableManager

import ConfigParser
import os
import sys


def main(argv):
    # We require at least one host pattern on the command line (but we
    # may accept other options in future).
    if len(argv) < 2:
        print "You must specify at least one hostname/pattern"
        return -1

    # Basic Ansible initialisation
    loader = DataLoader()
    variable_manager = VariableManager()
    parser = CLI.base_parser(
        usage='%prog <host-pattern>',
        runtask_opts=True,
        inventory_opts=True,
    )
    options, args = parser.parse_args()
    pattern = args[0]

    # Load the inventory and find hosts matching the specified pattern.
    inventory = Inventory(
        loader=loader,
        variable_manager=variable_manager,
        host_list=options.inventory)
    variable_manager.set_inventory(inventory)
    hosts = inventory.list_hosts(pattern)

    # We can't ssh to more than one matching host.
    if len(hosts) != 1:
        print "Your host pattern matched %d hosts" % len(hosts)
        return -1

    # Build ssh arguments for this host
    host = hosts[0]

    # append global vars
    config = ConfigParser.ConfigParser(allow_no_value=True)
    config.read(options.inventory)
    if config.has_section('{0}:vars'.format(pattern)):
        data = config.items('{0}:vars'.format(pattern))
        for key, value in dict(data).items():
            if key in host.vars:
                continue
            host.vars[key] = value

    ssh_args = ['ssh']

    if 'ansible_ssh_user' in host.vars:
        host.vars['ansible_user'] = host.vars['ansible_ssh_user']
        del host.vars['ansible_ssh_user']

    if 'ansible_ssh_port' in host.vars:
        host.vars['ansible_port'] = host.vars['ansible_ssh_port']
        del host.vars['ansible_ssh_port']

    if host.vars.get('ansible_ssh_private_key_file', None) is not None:
        private_key_file = host.vars['ansible_ssh_private_key_file']
        private_key_path = os.path.expanduser(private_key_file)
        ssh_args += ("-o", "IdentityFile=\"{0}\"".format(private_key_path))

    if host.vars.get('ansible_port', None) is not None:
        ssh_args += ("-o", "Port={0}".format(host.vars['ansible_port']))

    if host.vars.get('ansible_user', None) is not None:
        ssh_args += ("-o", "User={0}".format(host.vars['ansible_user']))

    ssh_args.append(host.address)

    ansible_ssh_executable = '/usr/bin/ssh'
    if 'ansible_ssh_executable' in host.vars:
        ansible_ssh_executable = host.vars['ansible_ssh_executable']

    # Launch ssh
    os.execl(ansible_ssh_executable, *ssh_args)


if __name__ == '__main__':
    main(sys.argv)
