#!/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
#
# Blog post about this script: https://kalgan.cc/blog/posts/On_demand_ssh-add/

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)