#!/bin/bash

FILELOGLEVEL=DEBUG
. $(dirname $0)/../logging
# loggerfactory
LANG=C

usage(){
    cat << EOF

    Usage: $(basename $0) [[-c]|[--create-only]]|[[-t]|[--token-only]]|[[-k]|[--key-only]]|[[-r]|[-f]|[--readd]|[--force]] [<ssh-identity>]

    If started only with <ssh-identity>, the script looks up in configured identity-path \$SSH_IDENTITIES_DIR (${SSH_IDENTITIES_DIR}) if it can find a directory named after <ssh-identity>.
    If no <ssh_identity> is given, the identity is set to \$SSH_DEFAULT_IDENTITY ($SSH_DEFAULT_IDENTITY) configured via Environment.
    IF \$SSH_DEFAULT_IDENTITY is also not set, default is the SSH_DEFAULT_IDENTITY

    The output is the name of the file, where ssh-agent infomations are hold to load it to current shell for further actions.
    Use "$ eval \$(<outputfilenam>)", if you want to load the SSH_AUTH_SOCK and SSH_AGENT_PID in current shell or shorter "$ loadagent [<ssh_identity>]"

    -c|--create-only                Create or restart only the agent. Do not load any 
                                    key or token in it.
                                    The Output is used for loading the agent in the current
                                    shell. (loadagent <identity>)
    -t|--token-only                 To add or renew only configured pkcs11-hardware-token
                                    configured in ${SSH_IDENTITIES_DIR}/<ssh-identity>,
                                    just use this. 
    -k|--key-only                   To add or renew only configured keys configured in
                                    ${SSH_IDENTITIES_DIR}/<ssh-identity>, just use this.
    -r|-f|--readd-token|--force    remove all in ${SSH_IDENTITIES_DIR}/<ssh-identity> 
                                    configured keys and tokens and readd them again.
                                    Depends on -t an -k Option to select wheter only 
                                    keys or tokens only. If no -t and -k is given, all 
                                    keys and token are removed and readded again.
                                    Just to be asked for password again, if you plugged off
                                    hardware-token and plugged it in again.
    --rm|--remove                   remove keys and token instead of adding them.
    -h|--info                       Show this info


EOF

}

check_token(){
    loginfo "a P11M:          ${P11M}"
    loginfo "a PKCS11_MODULE: ${PKCS11_MODULE}"

    [ -n "${P11M:+x}" ] && export PKCS11_MODULE=${P11M}
    loginfo "b P11M:          ${P11M}"
    loginfo "b PKCS11_MODULE: ${PKCS11_MODULE}"
    # it's the same as "token" in functions.sh
    # defined here also, to work also in environments, where functions.sh couldn't be sourced
    [ -z "${PKCS11_MODULE:+x}" ] && { PKCS11_MODULE=$P11M; export PKCS11_MODULE; }

    loginfo "c P11M:          ${P11M}"
    loginfo "c PKCS11_MODULE: ${PKCS11_MODULE}"
    # If DISPLAY is set, ssh-add calls ssh-askpass, and if its in remote-terminal, it wont work
    # So remember and unset DISPLAY, and set it at the end again, if it was set before
    [ $DISPLAY ] && local DISPLAY_ORIG=$DISPLAY
    [ $DISPLAY ] && logtrace "unset DISPLAY: $DISPLAY"
    [ $DISPLAY ] && unset DISPLAY
    
    # Write public keys of all in agent stored keys to a temporary file
    local tmppubkey="$(mktemp -p ${XDG_RUNTIME_DIR} pubkey.XXXXXX.pub)"
    logtrace "tmppubkey: $tmppubkey"
    logdebug "C"
    tmpIFS="${IFS}"
    IFS=$'\n'
    for tmppk in $(ssh-add -L)
    do
        printf "%s" "$tmppk" |tee "${tmppubkey}" || return $?
        #echo "${tmppk}" > $tmppubkey || return $?
        #ssh-add -L > $tmppubkey || return $?
        # Check if public-keys in tmppubkey are working. They are not working, if you removed and add back hardware-token. 
        loginfo "$(ssh-add -T ${tmppubkey}|| { ssh-add -e $PKCS11_MODULE; ssh-add -s $PKCS11_MODULE; } )"
        logdebug "$(ssh-add -l )"
    done
    logdebug "$(rm "${tmppubkey}")"
    IFS=$tmpIFS
    unset tmpIFS
    [ $DISPLAY_ORIG ] && logtrace "reset DISPLAY=$DISPLAY_ORIG"
    [ $DISPLAY_ORIG ] && export DISPLAY=$DISPLAY_ORIG

}

createonly=false
tokenonly=false
readd=false
keyonly=false
remove=false
while :; do
    case $1 in
        -c|--create-only)
            createonly=true
            shift
            ;;
        -t|--token-only)
            tokenonly=true
            shift
            ;;
        -k|--key-only)
            keyonly=true
            shift
            ;;
        -r|-f|--readd-token|--force)
            readd=true
            shift
            ;;
        --rm|--remove)
            remove=true
            shift
            ;;
        -h|--info)
            usage
            exit 0
            ;;
        -*)
            echo "Unknown urgument: »$1«" >&2
            exit 1
            ;;
        *)
            #ssh_identity=${1-${SSH_DEFAULT_IDENTITY}}
            ssh_identity=${1}
            break
            ;;
    esac
done

agent_start_or_restart () {
    
    logtrace "agent_start_or_restart"
    # If DISPLAY is set, ssh-add calls ssh-askpass, and if its in remote-terminal, it wont work
    # So remember and unset DISPLAY, and set it at the end again, if it was set before
    [ $DISPLAY ] && local DISPLAY_ORIG=$DISPLAY
    [ $DISPLAY ] && logtrace "unset DISPLAY: $DISPLAY"
    [ $DISPLAY ] && unset DISPLAY

    local ssh_identity
    local agentfile
    local agentsocket
    local ret

    if [ -n "${1+x}" ]; then
        ssh_identity="${1}"
        identitydir=${SSH_IDENTITIES_DIR}/${ssh_identity}
        loginfo "ssh-identität: ${ssh_identity}" >&2
        if [ -d ${identitydir} ]; then
            [ -e "${identitydir}/config" ] && . "${identitydir}/config"
            agentfile="${SSH_AGENTS_DIR}/agent-${ssh_identity}-$(hostname)"
            agentsocket="${SSH_AGENT_SOCKETS_DIR}/socket-${ssh_identity}-$(hostname)"
            logtrace "agentfile for ${ssh_identity}: $agentfile"
            logtrace "agentsocket for ${ssh_identity}: $agentsocket"
            if (! $keyonly && ! $tokenonly ) && $remove ; then
                logdebug "delete keys and tokens in this ssh-agent"
                logdebug "$(ssh_runinagent $agentfile ssh-add -D 2>&1)"
                createonly=true

            else
                if [ -e $agentfile ]; then

                    local msg
                    # TODO make in runinagent
                    msg="$(/bin/sh -c "unset SSH_AUTH_SOCK SSH_AGENT_PID; . $agentfile >/dev/null 2>&1; ssh-add -l 2>&1")"
                    local ret=$?
                    logtrace "Output from check for running: $msg"
                    case $ret in
                        0)
                            logdebug "agent is running" 
                            ;;
                        1)
                            logdebug "agent is running, but:" 
                            logwarning "$msg"
                            ;;
                        2)
                            logdebug "former agent is not running -> start it" 
                            logdebug "SSH_AGENT_OPTIONS: $SSH_AGENT_OPTIONS"
                            [ -e $agentsocket ] && { logdebug -n "remove socketfile: $( rm -v "$agentsocket" )"; }
                            logtrace "$(ssh-agent -a $agentsocket ${SSH_AGENT_OPTIONS} > $agentfile )"
                            logdebug "agent started" 
                            ;;
                    esac
                else
                    logdebug "agent did not exist -> create it" 
                    logtrace "ssh-agent -a $agentsocket \> $agentfile"
                    logtrace "$(ssh-agent -a $agentsocket $SSH_AGENT_OPTIONS > $agentfile )"
                    logdebug "agent started" 
                fi
            fi

            #logdebug "ssh-agent for identity »$ssh_identity«: $agentfile"
            $createonly && logtrace "current loaded keys after action:
$(ssh_runinagent $agentfile ssh-add -l)"
            echo $agentfile
            ret=0
        else
            logwarning "ssh-identity »$ssh_identity« is not configured. Please create $identitydir and add keys"
            printf "%s" "$agentfile"
            ret=2
        fi
    else
        if which gnome-keyring-daemon >/dev/null 2>&1; then
            logdebug "no identity given -> gnome-keyrings ssh-agent"
            agentfile="${SSH_AGENTS_DIR}/agent-gnome_session-$(hostname)"
            agentsocket="${SSH_AGENT_SOCKETS_DIR}/socket-gnome-session-$(hostname)"
            gnome-keyring-daemon -s > $agentfile
            logdebug "$(cat $agentfile)"
            logdebug "ssh-agent for identity »$ssh_identity«: $agentfile"
#            logdebug "currently loaded keys after action: 
#$(ssh_runinagent $agentfile ssh-add -l)"
            printf "%s" "$agentfile"
            ret=0
        else
            logwarning "no identity given -> exit"
            ret=1
        fi
    fi

    [ $DISPLAY_ORIG ] && logtrace "reset DISPLAY=$DISPLAY_ORIG"
    [ $DISPLAY_ORIG ] && export DISPLAY=$DISPLAY_ORIG
    return $ret

}

agent_load_identity_keys () {

    logtrace "agent_load_identity_keys"
    # If DISPLAY is set, ssh-add calls ssh-askpass, and if its in remote-terminal, it wont work
    # So remember and unset DISPLAY, and set it at the end again, if it was set before
    [ $DISPLAY ] && local DISPLAY_ORIG=$DISPLAY
    [ $DISPLAY ] && logtrace "unset DISPLAY: $DISPLAY"
    [ $DISPLAY ] && unset DISPLAY

    local ssh_identity
    local agentfile
    local agentsocket
    local fingerprints
    declare -a fingerprints
    local pubkeysonly
    declare -a pubkeysonly
    local fingerprint
    local tokenfingerprint

    if [ -n "${1+x}" ]; then
        ssh_identity="${1}"
        identitydir=${SSH_IDENTITIES_DIR}/${ssh_identity}
        if [ -d ${identitydir} ]; then
            [ -e "${identitydir}/config" ] && . "${identitydir}/config"
            agentfile="${SSH_AGENTS_DIR}/agent-${ssh_identity}-$(hostname)"
            agentsocket="${SSH_AGENT_SOCKETS_DIR}/socket-${ssh_identity}-$(hostname)"
            logdebug "SSH_ADD_OPTIONS: $SSH_ADD_OPTIONS"
            logtrace "agentfile: $agentfile" 
            logtrace "agentsocket: $agentsocket" 
            logtrace "identitydir: $identitydir"

            fingerprints=( $(ssh_runinagent $agentfile "ssh-add -l|awk '{print \$2}'") )
            logdebug "fingerprints from loaded keys before action:"
            for f in ${fingerprints[@]};do
                logdebug "$f"
            done
            if ! $tokenonly ; then
                # load keys
                for key in $(ls ${SSH_IDENTITIES_DIR}/${ssh_identity}/id_*|grep -v "pub$\|so$\|config$\|public$"); do
                    fingerprint=$(ssh-keygen -l -f $key|awk '{print $2}')
                    if [[ ${fingerprints[*]} =~ "$fingerprint" ]]; then
                        logdebug "key: $(basename $key) (with fp $fingerprint) is loaded" 
                        if  $readd || $remove ; then
                            $readd && logdebug "re-add key $key" 
                            $remove && logdebug "remove key $key"
                            logdebug "$(ssh_runinagent $agentfile ssh-add ${SSH_ADD_OPTIONS} -d ${key} 2>&1)"
                            $remove || logdebug "$(ssh_runinagent $agentfile ssh-add ${SSH_ADD_OPTIONS} ${key} 2>&1)"
                        fi
                    else
                        logdebug "key: $key is not loaded -> load it" 
                        #logdebug "$(ssh_runinagent $agentfile ssh-add ${SSH_ADD_OPTIONS} -d ${key} 2>&1)"
                        $remove || logdebug "$(ssh_runinagent $agentfile ssh-add ${SSH_ADD_OPTIONS} ${key} 2>&1)"
                    fi
                done
            fi
            if ! $keyonly ; then
                # load tokens
                for pubkey in $(ls ${SSH_IDENTITIES_DIR}/${ssh_identity}/id_*|grep "pub$\|public$"|grep -v "cert.pub"); do
                    if $( ls ${pubkey%.pub} 1>/dev/null 2>&1);then
                        logdebug "pubkey with privkey: $pubkey"
                    else
                        logdebug "pubkey without privkey: $pubkey"
                        pubkeysonly+=($pubkey)
                    fi
                    if [ -e "${pubkey%.pub}-cert.pub" ]; then
                        #logwarning "${pubkey%.pub}-cert.pub: $(date +%s -d $(ssh-keygen -L -f "${pubkey%.pub}-cert.pub"|awk '$1 == "Valid:"{print $5}'))"
                        #logwarning "now: $(date +%s -d NOW)"
                        cert_exp_date=$(ssh-keygen -L -f "${pubkey%.pub}-cert.pub"|awk '$1 == "Valid:"{print $5}')
                        #[ $(date +%s -d $(ssh-keygen -L -f "${pubkey%.pub}-cert.pub"|awk '$1 == "Valid:"{print $5}')) -gt $(date +%s -d NOW) ] \
                        [ $(date +%s -d $cert_exp_date) -gt $(date +%s -d NOW) ] \
                        || logwarning "CERTIFICATE IS NOT VALID ANYMORE: ${pubkey%.pub}-cert.pub"
                        #[ $(date +%s -d $(ssh-keygen -L -f "${pubkey%.pub}-cert.pub"|awk '$1 == "Valid:"{print $5}')) -lt $(date +%s -d "$SSH_CERT_VALIDITY_WARN_SEC") ] \
                        [ $(date +%s -d $cert_exp_date) -lt $(date +%s -d "$SSH_CERT_VALIDITY_WARN_SEC") ] \
                            && logwarning "CERTIFICATE expires in $(echo "scale=0; ( `date -d $cert_exp_date +%s` - `date -d now +%s`) / (24*3600)" | bc -l) days: ${pubkey%.pub}-cert.pub"
                    fi
                done
                loginfo "pubkeysonly: ${pubkeysonly[@]} (count: ${#pubkeysonly[*]})"


                for key in $(ls ${SSH_IDENTITIES_DIR}/${ssh_identity}/*|grep  "\.so$"); do
                    logdebug "PKCS11_MODULE: ${PKCS11_MODULE:-not set} - key: $key"
                    #[ -e "${P11M-x}" ] || [ "$(readlink - f $key)" == "$P11M" ] || key="$P11M"
                    #[ "$(readlink - f $key)" == "$P11M" ] || key="$P11M"
                    [ "$(readlink -f $key)" != "$PKCS11_MODULE" -a "${PKCS11_MODULE:-x}" != "x" ] && key="$PKCS11_MODULE"
                    export PKCS11_MODULE=$(readlink -f $key)
                    export P11M=$PKCS11_MODULE
                    logdebug "PKCS11_MODULE: ${PKCS11_MODULE:-not set} - key: $key"
                    grep -q "PKCS11_MODULE" "${agentfile}" >/dev/null 2>&1 && sed -i -e '/PKCS11_MODULE/d' "${agentfile}"
                    echo "PKCS11_MODULE=$(readlink -f $key)" | tee -a "${agentfile}" >&2

                    tokenfingerprints=($(ssh-keygen -l -D $key|tr -s ' '|awk '{print $2}'))

                    logtrace "fingerprints: ${fingerprints[*]}"
                    logtrace "tokenfingerprints count: ${#tokenfingerprints[@]}"
                    if [ ${#tokenfingerprints[@]} -gt 1 ]; then
                        logerr "Found ${#tokenfingerprints[@]} in $key! Check if only one token is plugged into your computer"
                        
                    else
                        tokenfingerprint=${tokenfingerprints[0]}
                        logtrace "tokenfingerprint: ${tokenfingerprint}"
                        if [[ "${fingerprints[@]}" =~ "$tokenfingerprint" ]]; then
                            logdebug "token: $key ($tokenfingerprint) is loaded"
                            check_pubkeysonly
                            if  $readd || $remove ; then
                                $readd && logdebug "re-add token $key"
                                $remove && logdebug "remove token $key"
                                logtrace "agentfile1 $agentfile \$SSH_ADD_OPTIONS $key "
                                logdebug "$(ssh_runinagent $agentfile ssh-add ${SSH_ADD_OPTIONS} -e "${key}" 2>&1)"
                                $remove || logdebug "$(ssh_runinagent $agentfile ssh-add ${SSH_ADD_OPTIONS} -s "${key}" 2>&1)"
                            fi
                        else
                            logdebug "token: $key is not loaded -> load it" 
                            logdebug "$(ssh_runinagent $agentfile ssh-add -v ${SSH_ADD_OPTIONS} -e ${key} 2>&1)"
                            $remove || logdebug "$(ssh_runinagent $agentfile ssh-add ${SSH_ADD_OPTIONS} -s ${key} 2>&1)"
                        fi
                    fi
                done
            fi
            loginfo "currently loaded keys after action: $(ssh_runinagent $agentfile ssh-add -l|wc -l)"
            logdebug "$(ssh_runinagent $agentfile ssh-add -l)"
        else
            logwarning "ssh-identity $ssh_identity is not configured. Please create $identitydir and add keys"
        fi
    fi

    [ $DISPLAY_ORIG ] && logtrace "reset DISPLAY=$DISPLAY_ORIG"
    [ $DISPLAY_ORIG ] && export DISPLAY=$DISPLAY_ORIG

}

function check_pubkeysonly () { 
    if [ ${#pubkeysonly[*]} -gt 0 ] ; then
        for p in ${pubkeysonly[@]}; do
            logtrace "pubkeyonly: $p"
            logtrace "$(ssh_runinagent $agentfile ssh-add -T ${p} 2>&1)"
            ssh_runinagent $agentfile ssh-add -T ${p} 2>&1 || { $remove || readd=true; break; }
        done
    else
        logwarning "obviously there is no pubkey for the token in ${SSH_IDENTITIES_DIR}/${ssh_identity}/"
        logwarning "you can add the pubkey with"
        logwarning "   ssh-add -L > ${SSH_IDENTITIES_DIR}/${ssh_identity}/id_etoken.pub"
        logwarning "make sure, only the token is loaded into ssh-agent with"
        logwarning "   ssh-add -l"
        logwarning "only one line should be the output"
    fi
    logdebug "readd: $readd"
}

ssh_runinagent () {

    local SSH_AUTH_SOCK
    local SSH_AGENT_PID
    local agentfile
    local command
    local agentfile=${1}
    shift
    local sshcommand=${@}

    logtrace "run command »$sshcommand« in agent $agentfile" 
    if [ -e "$agentfile" ]; then 
        /bin/sh -c "unset SSH_AUTH_SOCK SSH_AGENT_PID; . $agentfile >/dev/null 2>/dev/null; $sshcommand"
        ret=$?
    else
        logwarning "agentfile not existent" 
        ret=99
    fi

    loginfo "SSH_AUTH_SOCK: ${SSH_AUTH_SOCK}"
    return $ret

}

SCRIPTENTRY
[ -z "${SSH_IDENTITIES_DIR+x}" ]    && { SSH_IDENTITIES_DIR="${SSH_IDENTITIES_DEFAULT_DIR-${HOME}/.ssh/identities}"; export SSH_IDENTITIES_DIR; }
[ -z "${SSH_AGENTS_DIR+x}" ]        && { SSH_AGENTS_DIR=${SSH_AGENTS_DEFAULT_DIR-~/.ssh/agents}; export SSH_AGENTS_DIR; }
[ -z "${SSH_AGENT_SOCKETS_DIR+x}" ] && { SSH_AGENT_SOCKETS_DIR=${SSH_AGENT_SOCKETS_DEFAULT_DIR-~/.ssh/agents}; export SSH_AGENT_SOCKETS_DIR; }
[ -z "${SSH_AGENT_OPTIONS+x}" ]     && { SSH_AGENT_OPTIONS=${SSH_AGENT_DEFAULT_OPTIONS--t 7200 }; export SSH_AGENT_OPTIONS; }
logtrace "       SSH_AGENTS_DIR: $SSH_AGENTS_DIR"
logtrace "SSH_AGENT_SOCKETS_DIR: $SSH_AGENT_SOCKETS_DIR"
logtrace "   SSH_IDENTITIES_DIR: $SSH_IDENTITIES_DIR"
[ -z "${SSH_AGENTS_DIR-x}" ]        || mkdir -vp "$SSH_AGENTS_DIR"
[ -z "${SSH_AGENT_SOCKETS_DIR-x}" ] || mkdir -vp "$SSH_AGENT_SOCKETS_DIR"
[ -z "${SSH_IDENTITIES_DIR-x}" ]    || mkdir -vp "$SSH_IDENTITIES_DIR"


if [[ $SSH_TTY || $X2GO_SESSION ]] ; then
    logdebug "Shell running with forwarded ssh-agent. Please add local token manually"
    check_token
    res=1
else
    logdebug "run with local ssh-agent"
    agent_start_or_restart $ssh_identity
    ! $createonly && agent_load_identity_keys $ssh_identity
    res=0
fi

SCRIPTEXIT
exit $res