Kaarsemaker.net

  1 #!/usr/bin/python
  2 #
  3 # Simple wrapper around SSH to achieve the following:
  4 # - Add keys to a running ssh-agent if no keys have been added yet
  5 # - Set up a controlmaster connection the first time you connect to a host
  6 # - my_ssh -Q kills all controlmaster connections
  7 # Works best with 'alias ssh=my_ssh' in your ~/.bashrc (or the equivalent for
  8 # your shell)
  9 #
 10 # (c)2007 Dennis Kaarsemaker <dennis@kaarsemaker.net>
 11 #
 12 # Log:
 13 # 2007/06/20
 14 #  - Initial version
 15 # 2007/06/21
 16 #  - Use paramiko if possible
 17 # 2007/09/18
 18 #  - Don't use a ControlMaster if an explicit forwarding is requested
 19 #  - Guard against KeyboardInterrupt whilst running subprocesses
 20 #  - Move argument error handling up
 21 # 2007/09/22
 22 #  - Really don't use controlmaster when -L / -R is used
 23 
 24 import os, sys, subprocess, getopt
 25 
 26 # Check whether the keys are added to the agent. Prefer paramiko, but fall back
 27 # to opening a pipe to ssh-add -l
 28 try:
 29     import paramiko.agent
 30     a = paramiko.agent.Agent()
 31     if not a.keys:
 32         try:
 33             ret = subprocess.Popen(c_args).wait()
 34             if ret:
 35                 sys.exit(ret)
 36         except KeyboardInterrupt:
 37             sys.exit(130)
 38     a.close()
 39 except ImportError:
 40     if 'identities' in subprocess.Popen(['/usr/bin/ssh-add','-l'], stdout=subprocess.PIPE).communicate()[0]:
 41         try:
 42             ret = subprocess.Popen(c_args).wait()
 43             if ret:
 44                 sys.exit(ret)
 45         except KeyboardInterrupt:
 46             sys.exit(130)
 47 
 48 # Parse arguments
 49 c_args = ['/usr/bin/ssh','-f','-N','-n']
 50 
 51 try:
 52     (opts, args) = getopt.getopt(sys.argv[1:], 'Q1246AaCcfgKkMNnqsTtVvXxYb:c:D::F:i:L:l:m:O:o:p:R:S:w')
 53 except getopt.GetoptError:
 54     # Let SSH do the error reporting
 55     sys.argv[0] = '/usr/bin/ssh'
 56     os.execve('/usr/bin/ssh', sys.argv, os.environ)
 57 
 58 if ('-Q','') in opts:
 59     import glob, struct, socket
 60     # Kill all existing ssh controlmasters
 61     fd = open(os.path.expanduser('~/.ssh/config'))
 62     for line in fd:
 63         line = line.strip()
 64         if line.startswith('ControlPath'):
 65             path = line.split()[1].replace('%h','*').replace('%p','*').replace('%h','*').replace('%r','*')
 66             sockets = glob.glob(os.path.expanduser(path))
 67             msg = struct.pack('>IBII',9,1,3,0) # Read the source luke :) (payload length, type, exit command, flags)
 68             for sock in sockets:
 69                 try:
 70                     s = socket.socket(socket.AF_UNIX)
 71                     s.connect(sock)
 72                     s.send(msg)
 73                     s.recv()
 74                 except:
 75                     pass
 76     sys.exit(0)
 77 
 78 for a in 'ADKLRXY':
 79     if '-' + a in [x[0] for x in opts]:
 80         # Don't use a controlmaster when requesting any forwarding explicitely
 81         sys.argv.insert(1, '-oControlPath=none')
 82         break
 83 else:
 84     if len(args) and ('-O' not in [x[0] for x in opts]):
 85         host = args[0]
 86         c_args += sys.argv[1:sys.argv.index(host)+1]
 87         # Start the control connection if needed
 88         fd = open('/dev/null','w')
 89         if subprocess.Popen(['/usr/bin/ssh','-O','check',host], stdout=fd, stderr=fd).wait():
 90             try:
 91                 ret = subprocess.Popen(c_args).wait()
 92                 if ret:
 93                     sys.exit(ret)
 94             except KeyboardInterrupt:
 95                 sys.exit(130)
 96         fd.close()
 97 
 98 # Don't use subprocess here -- we want to replace ourselves
 99 sys.argv[0] = '/usr/bin/ssh'
100 os.execve('/usr/bin/ssh', sys.argv, os.environ)