#!/bin/zsh

# This file is part of xals' zsh utilss.
#
# xals' zsh utils 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.
#
# xals' zsh utils 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 xals' zsh utils.  If not, see
# <http://www.gnu.org/licenses/>.

# _ssh_agent_plugin_debug function
function _ssh_agent_plugin_debug() {
    local _ssh_agent_plugin_verbose
    zstyle -b :xals:plugins:ssh-agent verbose _ssh_agent_plugin_verbose

    if [[ "${_ssh_agent_plugin_verbose}" == "yes" ]]; then
        echo $*
    fi
}

# Function to clean global link
function _ssh_agent_plugin_clean_global_sock() {
    # Clean link to obsolete auth socket. 
    if [ -L "${_ssh_agent_plugin_sock}" ]; then
        _ssh_agent_plugin_debug "Found global link to socket, checking it."

        if [ -S "${_ssh_agent_plugin_sock}" -a -w "${_ssh_agent_plugin_sock}" ]; then
            _ssh_agent_plugin_debug "Global link relating to a valid socket, keeping it."
        else
            _ssh_agent_plugin_debug "Found a broken global link, removing it."

            # Remove obsolete global link
            unlink "${_ssh_agent_plugin_sock}"
        fi
    elif [ -e "${_ssh_agent_plugin_sock}" ]; then
        # Note: here ${_ssh_agent_plugin_sock} may exist, but may be not a symlink, we remove it.
        _ssh_agent_plugin_debug "Found garbage, cleaning up."

        rm -f "${_ssh_agent_plugin_sock}"
    fi
}

function _ssh_agent_plugin_clean_auth_sock() {
    # Cleanup ${SSH_AUTH_SOCK} if not a valid socket.
    if [ -S "${SSH_AUTH_SOCK}" -a -w "${SSH_AUTH_SOCK}" ]; then
        _ssh_agent_plugin_debug "Auth socket is valid."
    else
        if [ -e "${SSH_AUTH_SOCK}" ]; then
            _ssh_agent_plugin_debug "Found garbage, cleaning up."
            
            rm -f "${SSH_AUTH_SOCK}"
        fi

        if [ -n "${SSH_AUTH_SOCK}" ]; then
            _ssh_agent_plugin_debug "Found an invalid \${SSH_AUTH_SOCK}, unsetting it."

            unset SSH_AUTH_SOCK
        fi
    fi
}

function _ssh_agent_plugin_launch_agent() {
    if [ -S "${SSH_AUTH_SOCK}" ]; then
        _ssh_agent_plugin_debug "Found valid \${SSH_AUTH_SOCK}."
    else
        # We don't have an SSH auth socket, we try to create one.
        if [ -x "$(which keychain)" ]; then
            _ssh_agent_plugin_debug "Launching keychain."

            # Launch keychain
            _ssh_agent_plugin_run_keychain
        elif [ -x "$(which ssh-agent)" ]; then
            _ssh_agent_plugin_debug "Lauching classic SSH agent."

            _ssh_agent_plugin_run_classic_agent
        else
            _ssh_agent_plugin_debug "No valid SSH agent manager found."
        fi
    fi
}

function _ssh_agent_plugin_export_global_sock() {
    # Export SSH authentication socket if socket is valid.
    if [ -S "${_ssh_agent_plugin_sock}" -a -w "${_ssh_agent_plugin_sock}" ]; then
        if [ "${_ssh_agent_plugin_sock}" != "${SSH_AUTH_SOCK}" ]; then
            _ssh_agent_plugin_debug "Exporting ${_ssh_agent_plugin_sock} as new \${SSH_AUTH_SOCK}."

            export SSH_AUTH_SOCK="${_ssh_agent_plugin_sock}"
        else
            _ssh_agent_plugin_debug "\${_ssh_agent_plugin_sock} is already the same as \${SSH_AUTH_SOCK}."
        fi
    else
        _ssh_agent_plugin_debug "${_ssh_agent_plugin_sock} is not a valid socket."
    fi
}

function _ssh_agent_plugin_link_global_sock() {
    # Link to existing socket, if we are not already linked to it
    if [ -S "${SSH_AUTH_SOCK}" -a -w "${SSH_AUTH_SOCK}" ]; then
        _ssh_agent_plugin_debug "Found valid \${SSH_AUTH_SOCK}."

        if [ "${SSH_AUTH_SOCK}" != "${_ssh_agent_plugin_sock}" ]; then
            # Update link with new socket if needed
            if [ -S "${SSH_AUTH_SOCK}" -a -w "${SSH_AUTH_SOCK}" ]; then
                _ssh_agent_plugin_debug "Linking to new valid \${SSH_AUTH_SOCK}."

                ln -sf "${SSH_AUTH_SOCK}" "${_ssh_agent_plugin_sock}"
            else
                _ssh_agent_plugin_debug "Don't have a valid \${SSH_AUTH_SOCK} to link to."
            fi
        else
            _ssh_agent_plugin_debug "Already using linked socket."
        fi
    else
        _ssh_agent_plugin_debug "Didn't find valid \${SSH_AUTH_SOCK} to link to."
    fi
}

# launch SSH agent
function _ssh_agent_plugin_init() {
    local -x _ssh_agent_plugin_sock
    zstyle -s :xals:plugins:ssh-agent global_sock _ssh_agent_plugin_sock

    local _ssh_agent_plugin_enable_root
    zstyle -b :xals:plugins:ssh-agent enable_root _ssh_agent_plugin_enable_root

    local -x _ssh_agent_plugin_directory="$HOME/.ssh"

    if [ ! -d "${_ssh_agent_plugin_directory}" ]; then
        _ssh_agent_plugin_debug "Creating ${_ssh_agent_plugin_directory} directory."

        mkdir "${_ssh_agent_plugin_directory}"
        chmod 700 "${_ssh_agent_plugin_directory}"
    fi

    # Reuse user authentication socket if I am root and using sudo.
    if [ -z "${_ssh_agent_plugin_sock}" ]; then
        if [ "${USER}" != "root" ]; then
            _ssh_agent_plugin_sock="/tmp/ssh-agent-$USER"
        elif [ "${_ssh_agent_plugin_enable_root}" == "yes" ]; then
            if [ -n "${SUDO_USER}" ]; then
                _ssh_agent_plugin_sock="/tmp/ssh-agent-$SUDO_USER"
            else
                _ssh_agent_plugin_sock="/tmp/ssh-agent-$USER"
            fi
        fi
    fi
   
    _ssh_agent_plugin_debug "Static socket will be: \"${_ssh_agent_plugin_sock}\"."
 
    # We want to do socket linking only if it is necessary
    if [ -n "${_ssh_agent_plugin_sock}" ]; then
        _ssh_agent_plugin_debug "Global socket is defined"

        # Clean global link.
        _ssh_agent_plugin_clean_global_sock

        # Clean auth sock.
        _ssh_agent_plugin_clean_auth_sock

        # First time export.
        _ssh_agent_plugin_export_global_sock

        # Launch agent if not valid SSH_AUTH_SOCK found.
        _ssh_agent_plugin_launch_agent

        # Link to SSH_AUTH_SOCK.
        _ssh_agent_plugin_link_global_sock 
    
        # Second time export.
        _ssh_agent_plugin_export_global_sock

        # Add identities.
        _ssh_agent_plugin_add_identities
    fi
}

# Use keychain to launch SSH agent.
function _ssh_agent_plugin_run_keychain() {
    if [ -x "$(which keychain)" ]; then
        keychain --quiet --nogui --inherit any --agents ssh

        [ -r $HOME/.keychain/$(hostname)-sh ] &&  source $HOME/.keychain/$(hostname)-sh
    fi
}

# Launch classic SSH agent.
function _ssh_agent_plugin_run_classic_agent() {
    # We launch the agent only if there is not auth socket.
    if [ -S "${SSH_AUTH_SOCK}" -a -w "${SSH_AUTH_SOCK}" ]; then
        # We should not be here, but just in case...
    else
        if [ -e "${SSH_AUTH_SOCK}" ]; then
            # Should not happen... Just in case...
            rm -f "${SSH_AUTH_SOCK}"
        fi

        if [ -x "$(which ssh-agent)" ]; then
            local _ssh_agent_plugin_pidfile="${_ssh_agent_plugin_directory}/agent-pid"

            # Cleanup old PID file
            if [ -r "${_ssh_agent_plugin_pidfile}" ]; then
                local _ssh_agent_plugin_pid=$(< "${_ssh_agent_plugin_pidfile}")

                _ssh_agent_plugin_debug "Found a PID file with value ${_ssh_agent_plugin_pid}."
                if [ ! -e "${_ssh_agent_plugin_pid}" ] ; then
                    _ssh_agent_plugin_debug "Removing remaining PID file."

                    rm -f "${_ssh_agent_plugin_pidfile}"
                fi
            fi

            if [ -r "${_ssh_agent_plugin_pidfile}" ]; then
                _ssh_agent_plugin_debug "Killing existing agent."

                export SSH_AGENT_PID=$(< "${_ssh_agent_plugin_pidfile}")
                eval $(ssh-agent -s -k) > /dev/null
                rm -f "${_ssh_agent_plugin_pidfile}"
            fi

            _ssh_agent_plugin_debug "Launching SSH agent."

            local _ssh_agent_plugin_lifetime
            zstyle -s :xals:plugins:ssh-agent lifetime _ssh_agent_plugin_lifetime

            eval $(ssh-agent -s ${_ssh_agent_plugin_lifetime:+-t} ${_ssh_agent_plugin_lifetime}) > /dev/null
            echo $SSH_AGENT_PID > "${_ssh_agent_plugin_pidfile}"
        fi
    fi
}

# Add SSH key if not already added.
function _ssh_agent_plugin_add_identities() {
    if [ -x "$(which ssh-add)" ]; then
        # "ssh-add -l" return codes:
        # 0: one or more keys were found.
        # 1: no key was found.
        # 2: agent could not be contacted.

        local _ssh_agent_plugin_identities
        zstyle -a :xals:plugins:ssh-agent identities _ssh_agent_plugin_identities

        if [ "${#_ssh_agent_plugin_identities}" -eq 0 ]; then
            _ssh_agent_plugin_identities=("id_rsa")
        fi

        ssh-add -l > /dev/null

        if [ 2 -eq $? ] ; then
            # Agent is not launched
            _ssh_agent_plugin_debug "Agent is not launched. Probably no key available."
        else
            # Agent has identities added, check if the ones we want are present, add them if not.
            local _ssh_agent_plugin_identity
            local _ssh_agent_plugin_identity_path

            for _ssh_agent_plugin_identity in ${^_ssh_agent_plugin_identities}; do
                _ssh_agent_plugin_debug "Handling identity ${_ssh_agent_plugin_identity}."

                _ssh_agent_plugin_identity_path="${_ssh_agent_plugin_directory}/${_ssh_agent_plugin_identity}"

                if [ -r "${_ssh_agent_plugin_identity_path}" ]; then
                    _ssh_agent_plugin_debug "Trying to add identity ${_ssh_agent_plugin_identity_path}"

                    ssh-add -l | awk '{print $3}' | grep -q "${_ssh_agent_plugin_identity_path}"

                    if [ $? -ne 0 ]; then
                        _ssh_agent_plugin_debug "Adding ${_ssh_agent_plugin_identity} identity."

                        ssh-add "${_ssh_agent_plugin_identity_path}"
                    else
                        _ssh_agent_plugin_debug "${_ssh_agent_plugin_identity} identity already added."
                    fi
                else
                    _ssh_agent_plugin_debug "${_ssh_agent_plugin_identity} identity does not exist."
                fi
            done
        fi
    fi
}

_ssh_agent_plugin_init

unfunction _ssh_agent_plugin_init

unfunction _ssh_agent_plugin_clean_global_sock
unfunction _ssh_agent_plugin_clean_auth_sock
unfunction _ssh_agent_plugin_export_global_sock
unfunction _ssh_agent_plugin_launch_agent
unfunction _ssh_agent_plugin_link_global_sock 

unfunction _ssh_agent_plugin_run_keychain
unfunction _ssh_agent_plugin_run_classic_agent

unfunction _ssh_agent_plugin_add_identities

unfunction _ssh_agent_plugin_debug