zsh-ec2ssh.zsh

function _load_aws_profile() {
    local aws_profile=$1
    if [ -z "${aws_profile}" ]; then
        aws_profile=$AWS_DEFAULT_PROFILE
    fi
    echo $aws_profile
    return
}

function _load_aws_region() {
    local aws_region=$1
    if [ -z "${aws_region}" ]; then
        aws_region=$AWS_DEFAULT_REGION
    fi
    echo $aws_region
    return
}

function _load_user() {
    local user=$1
    if [ -z "${user}" ]; then
        user=$USER
    fi
    echo $user
    return
}

function _load_ssh_private_key_path() {
    local private_key_path=$1
    if [ -z "${private_key_path}" ]; then
        private_key_path="$HOME/.ssh/id_rsa"
    fi
    echo $private_key_path
    return
}

function _load_port() {
    local port=$1
    if [ -z "${port}" ]; then
        port=22
    fi
    echo $port
    return
}

function zsh-ec2ssh() {
    local aws_profile_name=$1
    local aws_region=$2
    local target_user=$3
    local target_private_key_path=$4
    local target_port=$5
    local proxy_host=$6
    local proxy_user=$7
    local proxy_key_path=$8
    local proxy_port=$9

    aws_profile_name=`_load_aws_profile $aws_profile_name`
    aws_region=`_load_aws_region $aws_region`
    target_user=`_load_user $target_user`
    target_private_key_path=`_load_ssh_private_key_path $target_private_key_path`
    target_port=`_load_port $target_port`
    proxy_user=`_load_user $proxy_user`
    proxy_key_path=`_load_ssh_private_key_path $proxy_key_path`
    proxy_port=`_load_port $proxy_port`

    if [ -z "${aws_profile_name}" ]; then
        echo "AWS profile name is required. Please call this function with aws profile name or set AWS_DEFAULT_REGION in evironment variables."
        return
    fi

    if [ -z "${aws_region}" ]; then
        echo "AWS region is required. Please call this function with aws region or set AWS_DEFAULT_REGION in environment variables."
        return
    fi

    if [ -z "${target_user}" ]; then
        echo "User is required. Please call this function with user or set USER in environment variables."
        return
    fi

    echo "Fetching ec2 host..."
    local selected_host=$(myaws ec2 ls --profile=${aws_profile_name} --region=${aws_region} --fields='InstanceId PublicIpAddress PrivateIpAddress LaunchTime Tag:Name' | sort -k5 | peco | cut -f3)
    if [ -n "${selected_host}" ]; then
        if [ -z "${proxy_host}" ]; then
            BUFFER="ssh -i ${target_private_key_path} -p ${target_port} ${target_user}@${selected_host} -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null"
        else
            BUFFER="ssh -i ${proxy_key_path} -p ${proxy_port} -t ${proxy_user}@${proxy_host} ssh ${target_user}@${selected_host} -i ${target_private_key_path} -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null"
        fi
        if zle; then
            zle accept-line
        else
            print -z "$BUFFER"
        fi
    fi
    if zle; then
        zle clear-screen
    fi
}
zle -N zsh-ec2ssh

function zsh-ec2ssh-with-proxy() {
    local aws_profile_name=$1
    local aws_region=$2
    local target_user=$3
    local target_private_key_path=$4
    local target_port=$5
    local proxy_profile=$6
    local proxy_user=$7
    local proxy_key_path=$8
    local proxy_port=$9

    aws_profile_name=`_load_aws_profile $aws_profile_name`
    aws_region=`_load_aws_region $aws_region`
    target_user=`_load_user $target_user`
    target_private_key_path=`_load_ssh_private_key_path $target_private_key_path`
    target_port=`_load_port $target_port`
    proxy_profile=`_load_aws_profile $proxy_profile`
    proxy_user=`_load_user $proxy_user`
    proxy_key_path=`_load_ssh_private_key_path $proxy_key_path`
    proxy_port=`_load_port $proxy_port`

    if [ -z "${aws_profile_name}" -o -z "${proxy_profile}" ]; then
        echo "AWS profile name is required. Please call this function with aws profile name or set AWS_DEFAULT_REGION in evironment variables."
        return
    fi

    if [ -z "${aws_region}" ]; then
        echo "AWS region is required. Please call this function with aws region or set AWS_DEFAULT_REGION in environment variables."
        return
    fi

    if [ -z "${target_user}" -o -z "${proxy_user}" ]; then
        echo "User is required. Please call this function with user or set USER in environment variables."
        return
    fi

    echo "Fetching ec2 host..."
    local selected_proxy=$(myaws ec2 ls --profile=${proxy_profile} --region=${aws_region} --fields='InstanceId PublicIpAddress LaunchTime Tag:Name Tag:attached_asg' | sort -k4 | peco | cut -f2)
    if [ -n "${selected_proxy}" ]; then
        zsh-ec2ssh $aws_profile_name $aws_region $target_user $target_private_key_path $target_port $selected_proxy $proxy_user $proxy_key_path $proxy_port
    fi
    if zle; then
        zle clear-screen
    fi
}
zle -N zsh-ec2ssh-with-proxy