#!/usr/bin/python
#
# Simple wrapper around SSH to achieve the following:
# - Add keys to a running ssh-agent if no keys have been added yet
# - Set up a controlmaster connection the first time you connect to a host
# - my_ssh -Q kills all controlmaster connections
# Works best with 'alias ssh=my_ssh' in your ~/.bashrc (or the equivalent for
# your shell)
#
# (c)2007 Dennis Kaarsemaker <dennis@kaarsemaker.net>
#
# Log:
# 2007/06/20
#  - Initial version
# 2007/06/21
#  - Use paramiko if possible
# 2007/09/18
#  - Don't use a ControlMaster if an explicit forwarding is requested
#  - Guard against KeyboardInterrupt whilst running subprocesses
#  - Move argument error handling up
# 2007/09/22
#  - Really don't use controlmaster when -L / -R is used

import os, sys, subprocess, getopt

# Check whether the keys are added to the agent. Prefer paramiko, but fall back
# to opening a pipe to ssh-add -l
try:
    import paramiko.agent
    a = paramiko.agent.Agent()
    if not a.keys:
        try:
            ret = subprocess.Popen(c_args).wait()
            if ret:
                sys.exit(ret)
        except KeyboardInterrupt:
            sys.exit(130)
    a.close()
except ImportError:
    if 'identities' in subprocess.Popen(['/usr/bin/ssh-add','-l'], stdout=subprocess.PIPE).communicate()[0]:
        try:
            ret = subprocess.Popen(c_args).wait()
            if ret:
                sys.exit(ret)
        except KeyboardInterrupt:
            sys.exit(130)

# Parse arguments
c_args = ['/usr/bin/ssh','-f','-N','-n']

try:
    (opts, args) = getopt.getopt(sys.argv[1:], 'Q1246AaCcfgKkMNnqsTtVvXxYb:c:D::F:i:L:l:m:O:o:p:R:S:w')
except getopt.GetoptError:
    # Let SSH do the error reporting
    sys.argv[0] = '/usr/bin/ssh'
    os.execve('/usr/bin/ssh', sys.argv, os.environ)

if ('-Q','') in opts:
    import glob, struct, socket
    # Kill all existing ssh controlmasters
    fd = open(os.path.expanduser('~/.ssh/config'))
    for line in fd:
        line = line.strip()
        if line.startswith('ControlPath'):
            path = line.split()[1].replace('%h','*').replace('%p','*').replace('%h','*').replace('%r','*')
            sockets = glob.glob(os.path.expanduser(path))
            msg = struct.pack('>IBII',9,1,3,0) # Read the source luke :) (payload length, type, exit command, flags)
            for sock in sockets:
                try:
                    s = socket.socket(socket.AF_UNIX)
                    s.connect(sock)
                    s.send(msg)
                    s.recv()
                except:
                    pass
    sys.exit(0)

for a in 'ADKLRXY':
    if '-' + a in [x[0] for x in opts]:
        # Don't use a controlmaster when requesting any forwarding explicitely
        sys.argv.insert(1, '-oControlPath=none')
        break
else:
    if len(args) and ('-O' not in [x[0] for x in opts]):
        host = args[0]
        c_args += sys.argv[1:sys.argv.index(host)+1]
        # Start the control connection if needed
        fd = open('/dev/null','w')
        if subprocess.Popen(['/usr/bin/ssh','-O','check',host], stdout=fd, stderr=fd).wait():
            try:
                ret = subprocess.Popen(c_args).wait()
                if ret:
                    sys.exit(ret)
            except KeyboardInterrupt:
                sys.exit(130)
        fd.close()

# Don't use subprocess here -- we want to replace ourselves
sys.argv[0] = '/usr/bin/ssh'
os.execve('/usr/bin/ssh', sys.argv, os.environ)