On demand ssh-add snippet

1 files changed, 122 insertions, 0 deletions

diff --git a/odsa b/odsa
new file mode 100755
index 0000000..34899d1
--- /dev/null
+++ b/odsa
@@ -0,0 +1,122 @@
+#!/usr/bin/python
+# -*- coding: utf-8 -*-
+#!/usr/bin/env python
+# Copyright (C) 2013 Raúl Benencia <rul@kalgan.cc>
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+#
+# Little script to load SSH keys on demand. You can use it by writing a
+# function like the following one in your shell rc file:
+#
+#    odsa_ssh() {
+#       odsa "$*"
+#       /usr/bin/ssh "$*"
+#    }
+#
+# And then an alias to ssh
+#    alias ssh=odsa_ssh
+
+import signal
+import sys
+
+from os.path import expanduser
+from paramiko import SSHConfig, Agent
+from subprocess import call
+
+def signal_handler(signal, frame):
+        print 'Caught ^C'
+        sys.exit(4)
+
+class ODSA(): # This cryptic name means On Demand Ssh-Add
+    def __init__(self, cfg_filename=expanduser('~/.ssh/config')):
+        self.cfg_filename = cfg_filename
+
+    def keys_for_host(self, host):
+        """
+        Return a list of paths to the private keys configured by the user for
+        logging in on the received host.
+
+        @return: a list of keys configured up in ~/.ssh/config
+        @rtype: list of paths to private keys
+        """
+        try:
+            cfg_file = file(self.cfg_filename)
+        except IOError:
+            print "Couldn't open SSH config file: " + self.cfg_filename
+            sys.exit(3)
+        else:
+            config = SSHConfig()
+            config.parse(cfg_file)
+            keys = config.lookup(host).get('identityfile', None)
+            cfg_file.close()
+
+            if keys:
+                return [expanduser(k) for k in keys]
+            else:
+                return []
+
+    def get_public_key_from_file(self, key_filename):
+        """
+        Return the public key in base64 format.
+
+        @return: a string containing a public key encoded in base64
+        @rtype: str
+        """
+        try:
+            key_file = file(key_filename + ".pub")
+        except IOError:
+            print "Couldn't open the public key file " + key_filename + ".pub"
+            sys.exit(2)
+        else:
+            contents = key_file.read()
+            key_file.close()
+
+            return contents.split()[1]
+
+    def is_key_loaded(self, key_filename):
+        """
+        Return whether a certain SSH key is loaded in the SSH agent
+
+        @return: True if the key is loaded, False otherwise.
+        @rtype: Boolean
+        """
+        key = self.get_public_key_from_file(key_filename)
+
+        return any(key == k.get_base64() for k in Agent().get_keys())
+
+    def load_keys(self, keys):
+        """
+        Spawn ssh_add command for loading the missing keys.
+
+        @return: Returns values of ssh_add command
+        @rtype: Int
+        """
+        keys.insert(0, "/usr/bin/ssh-add")
+
+        return call(keys)
+
+if __name__ == "__main__":
+    if len(sys.argv) == 1:
+        print 'Usage: ' + sys.argv[0] + ' HOST'
+        sys.exit(1)
+
+    signal.signal(signal.SIGINT, signal_handler)
+
+    odsa = ODSA()
+    enabled_keys = odsa.keys_for_host(sys.argv[1])
+    keys_to_load = filter(lambda k: not odsa.is_key_loaded(k), enabled_keys)
+    if keys_to_load:
+        odsa.load_keys(keys_to_load)
+
+    sys.exit(0)