#!/usr/bin/env bash
# connect to iot services
THIS_DIR=$( cd "$( dirname "${BASH_SOURCE[0]:-${(%):-%x}}" )" && pwd )
source "$THIS_DIR/../tools/common.sh"
RIOT_TRUST_CLIENT=${RIOT_TRUST_CLIENT:-${DFS_TRUST:-0}}
RIOT_TRUST_SERVER=${RIOT_TRUST_SERVER:-0}
EXTRA_SSH_OPTIONS=()
GPG_FORWARD=0

# config
RIOT_CONFIG_FILES=(
    "$DOTFILES/riot-config.sh"
    "$HOME/.config/riot-config.sh"
    "riot-config.sh"
)
for file in "${RIOT_CONFIG_FILES[@]}"; do
    if [[ -f "$file" ]]; then
        source "$file"
    fi
done

# check if port number valid
check_port() {
    [[ "$1" =~ ^[1-9][0-9]{0,4}$ ]] || return 1
    [[ $1 < 65536 && $1 > 0 ]] || return 1
    return 0
}

# check if username valid
check_username() {
    [[ "$1" =~ ^[a-z][-a-z0-9_]*$ ]] || return 1
    return 0
}

# get single server setting
# may be called more than once
get_server_meta() {
    # returns:
    RET_HOSTNAME=""
    RET_TRUST_SERVER=0
    RET_PORT=""  # optional
    RET_USERNAME=""  # optional
    RET_JUMP_SERVER=""  # optional
    # body
    local remote="$1"
    # extract username from user@...
    if [[ "$remote" == *@* ]]; then
        RET_USERNAME=${remote%%@*}
        remote=${remote#*@}
        check_username $RET_USERNAME || fmt_warning \"$RET_USERNAME\" is not a valid unix username
    fi
    # extract port from ...:port
    if [[
        "$remote" =~ ^[^:]+:[1-9][0-9]*$  # contains only one colon
        || "$remote" =~ ^\[.+\]:[1-9][0-9]*$  # in the form of [host]:port
        || "$remote" =~ :::[1-9][0-9]*$  # in the form of :::port
        || "$remote" =~ ^([0-9a-f]{1,4}:){7}[0-9a-f]{1,4}:[1-9][0-9]*$  # full ipv6 address with port
        || "$remote" =~ ^[0-9a-f:]+%.+:[1-9][0-9]*$  # ipv6 address with scope and port
    ]]; then
        RET_PORT=${remote##*:}
        remote=${remote%:*}
        check_port $RET_PORT || fmt_fatal invalid port number \"$RET_PORT\"
    fi
    # remove square brackets
    if [[ "$remote" =~ ^\[.*\]$ ]]; then
        remote=${remote:1:-1}
    fi
    # presets -- match remote
    local remote_func="$remote.remote"
    if is_function "$remote_func"; then
        "$remote_func"
    fi
    # presets -- match domain
    RET_HOSTNAME=${remote}
    local domain=${remote##*.}
    local host=${remote%.*}
    # if it contains no dot and is not ipv6
    if [[ "$remote" != *.* && "$remote" != *:* ]]; then
        domain="default"
    fi
    local domain_func="$domain.domain"
    if is_function "$domain_func"; then
        "$domain_func"
    elif is_function ".domain"; then
        ".domain"
    fi
}

parse_remote() {
    # remote setting, including jump servers
    # called for every remote
    # provides:
    SERVER=""
    TRUST_SERVER=1
    PORT=""  # optional
    USERNAME=""  # optional
    SSH_OPTIONS=("-o" "ServerAliveInterval=60")
    if [[ -t 1 ]]; then
        SSH_OPTIONS+=("-o" "RequestTTY=yes")
    fi
    if [[ "$RIOT_TRUST_CLIENT" == "1" ]]; then
        SSH_OPTIONS+=("-o" "PermitLocalCommand=yes")
        if [[ "$(get_os_type)" != "msys" ]]; then
            test "$DFS_DRY_RUN" = "1" || mkdir -p ~/.ssh/master-socket
            SSH_OPTIONS+=("-o" "ControlMaster=auto" "-o" "ControlPersist=5s" "-o" "ControlPath=~/.ssh/master-socket/%C")
        fi
    fi
    # handle input
    local remote="$1"
    local jump_servers=""
    # loop for jump servers
    while [[ -n $remote ]]; do
        local server=${remote%%/*}
        remote=${remote#*/}
        get_server_meta "$server"
        if [[ -n "$RET_JUMP_SERVER" ]]; then
            jump_servers="$jump_servers${jump_servers:+,}$RET_JUMP_SERVER"
        fi
        # only if all servers are trusted
        TRUST_SERVER=$((TRUST_SERVER*RET_TRUST_SERVER))
        if [[ "$server" == "$remote" || -z "$remote" ]]; then
            SERVER="$RET_HOSTNAME"
            PORT="$RET_PORT"
            USERNAME="$RET_USERNAME"
            remote=""
        else
            jump_servers="$jump_servers${jump_servers:+,}$RET_USERNAME${RET_USERNAME:+@}$RET_HOSTNAME${RET_PORT:+:}$RET_PORT"
        fi
    done
    # construct cmd
    if [[ "$RIOT_TRUST_SERVER" == "1" || "$TRUST_SERVER" == "1" ]]; then
        SSH_OPTIONS+=("-o" "ForwardX11=yes" "-o" "ForwardAgent=yes")
    fi
    if [[ -n "$jump_servers" ]]; then
        SSH_OPTIONS+=("-o" "ProxyJump=$jump_servers")
    fi
}

check_local_gpg_agent() {
    LOCAL_GPG_EXTRA_SOCKET=""
    GPG_FORWARD_ERROR=""

    if ! command -v gpgconf > /dev/null 2>&1 || ! command -v gpg-connect-agent > /dev/null 2>&1; then
        GPG_FORWARD_ERROR="gpgconf or gpg-connect-agent is not available"
        return 1
    fi

    local agent_info
    agent_info=$(gpg-connect-agent --no-autostart 'GETINFO pid' /bye 2>/dev/null || true)
    if ! grep -qE '^D [1-9][0-9]*$' <<< "$agent_info"; then
        GPG_FORWARD_ERROR="local gpg-agent is not running"
        return 1
    fi

    LOCAL_GPG_EXTRA_SOCKET=$(gpgconf --list-dirs agent-extra-socket 2>/dev/null || true)
    if [[ -z "$LOCAL_GPG_EXTRA_SOCKET" || "$LOCAL_GPG_EXTRA_SOCKET" != /* \
        || "$LOCAL_GPG_EXTRA_SOCKET" == *:* || "$LOCAL_GPG_EXTRA_SOCKET" == *$'\r'* \
        || "$LOCAL_GPG_EXTRA_SOCKET" == *$'\n'* || ! -S "$LOCAL_GPG_EXTRA_SOCKET" ]]; then
        GPG_FORWARD_ERROR="local gpg-agent extra socket is unavailable"
        return 1
    fi

    local extra_info
    extra_info=$(gpg-connect-agent --raw-socket "$LOCAL_GPG_EXTRA_SOCKET" 'GETINFO version' /bye 2>/dev/null || true)
    if ! grep -qE '^D [^[:space:]]+' <<< "$extra_info"; then
        GPG_FORWARD_ERROR="local gpg-agent extra socket exists but is not accepting connections"
        return 1
    fi
}

probe_remote_gpg_socket() {
    REMOTE_GPG_SOCKET=""

    # The probe must be an independent connection: reusing or creating a
    # multiplex master here races with the immediately following login.
    local query_cmd=(ssh "-S" "none" "-o" "ClearAllForwardings=yes" "-o" "RequestTTY=no")
    if [[ -n "$PORT" ]]; then
        query_cmd+=("-p" "$PORT")
    fi
    query_cmd+=(
        "${SSH_OPTIONS[@]}"
        "${EXTRA_SSH_OPTIONS[@]}"
        "-T"
        "$USERNAME${USERNAME:+@}$SERVER"
        'socket=$(gpgconf --list-dirs agent-socket 2>/dev/null) || exit 10
case "$socket" in /*/S.gpg-agent) ;; *) exit 11;; esac
case "$socket" in *:*) exit 11;; esac
case "$socket" in *"
"*) exit 11;; esac
carriage_return=$(printf "\r")
case "$socket" in *"$carriage_return"*) exit 11;; esac
if [ -e "$socket" ] && [ ! -S "$socket" ]; then exit 12; fi
if [ -S "$socket" ] && command -v gpg-connect-agent >/dev/null 2>&1; then
    agent_mode=$(gpg-connect-agent --raw-socket "$socket" "GETINFO restricted" /bye 2>/dev/null || true)
    case "$agent_mode" in *"D 1"*) exit 13;; esac
fi
systemd_socket=0
if command -v systemctl >/dev/null 2>&1 && systemctl --user is-active --quiet gpg-agent.socket >/dev/null 2>&1; then
    systemd_socket=1
fi
gpgconf --kill gpg-agent >/dev/null 2>&1 || exit 14
rm -f "$socket" || exit 15
printf "%s\n%s\n" "$socket" "$systemd_socket"'
    )

    local output status
    if output=$("${query_cmd[@]}"); then
        status=0
    else
        status=$?
    fi
    case "$status" in
    0) ;;
    10) GPG_FORWARD_ERROR="gpgconf is unavailable on the remote host" ;;
    11) GPG_FORWARD_ERROR="remote gpgconf returned an invalid agent socket" ;;
    12) GPG_FORWARD_ERROR="refusing to remove the non-socket remote gpg-agent path" ;;
    13) GPG_FORWARD_ERROR="another forwarded gpg-agent is already using the remote socket; close that session first" ;;
    14) GPG_FORWARD_ERROR="failed to stop the remote gpg-agent (the socket may belong to another forwarding session)" ;;
    15) GPG_FORWARD_ERROR="failed to remove the stale remote gpg-agent socket" ;;
    *) GPG_FORWARD_ERROR="failed to query or clean the remote gpg-agent socket (ssh status $status)" ;;
    esac
    if [[ "$status" != "0" ]]; then
        return 1
    fi

    local remote_systemd_socket
    REMOTE_GPG_SOCKET=${output%%$'\n'*}
    remote_systemd_socket=${output#*$'\n'}
    if [[ -z "$REMOTE_GPG_SOCKET" || "$REMOTE_GPG_SOCKET" == *$'\r'* || "$REMOTE_GPG_SOCKET" == *$'\n'* \
        || "$REMOTE_GPG_SOCKET" != /*/S.gpg-agent || "$REMOTE_GPG_SOCKET" == *:* \
        || "$remote_systemd_socket" != "0" && "$remote_systemd_socket" != "1" ]]; then
        GPG_FORWARD_ERROR="remote gpgconf returned an invalid agent socket: $output"
        return 1
    fi
    if [[ "$remote_systemd_socket" == "1" ]]; then
        fmt_warning "remote gpg-agent.socket is active and may race with GPG forwarding; consider disabling its socket activation"
    fi
}

prepare_gpg_forwarding() {
    if [[ "$DFS_DRY_RUN" == "1" ]]; then
        REMOTE_GPG_SOCKET="<remote-gpg-agent-socket>"
        LOCAL_GPG_EXTRA_SOCKET="<local-gpg-agent-extra-socket>"
    else
        if ! check_local_gpg_agent; then
            fmt_fatal "cannot forward gpg-agent: $GPG_FORWARD_ERROR"
        fi
        if ! probe_remote_gpg_socket; then
            fmt_fatal "cannot forward gpg-agent: $GPG_FORWARD_ERROR"
        fi
    fi

    SSH_OPTIONS+=(
        # The probe already removed the old socket. Do not let the main SSH
        # connection unlink a path recreated during the gap between them.
        "-o" "StreamLocalBindUnlink=no"
        "-o" "ExitOnForwardFailure=yes"
        "-R" "$REMOTE_GPG_SOCKET:$LOCAL_GPG_EXTRA_SOCKET"
    )
}

print_cmd() {
    local output=""
    for s in "${CMD[@]}"; do
        if [[ "$s" =~ [\ \\\'\"] ]]; then  # needs to be escaped
            s="${s@Q}"
        fi
        output+="$s "
    done
    fmt_note "--> ${output% }"
}

eval_or_echo() {
    local DO=""
    tmux_win=${tmux_win:-0}
    if [[ "$DFS_DRY_RUN" == "1" ]]; then
        DO=echo
    fi
    if [[ "$USE_TMUX" == "1" ]]; then
        if [[ -z "$TMUX_SESS" ]]; then
            TMUX_SESS=riot-$(date +%s)
            $DO tmux new-session -d -s $TMUX_SESS bash -l
        else
            tmux_win=$((tmux_win+1))
            $DO tmux new-window -t $TMUX_SESS:$tmux_win -d bash -l
        fi
        $DO tmux send-keys -t $TMUX_SESS:$tmux_win $(printf '%s Space ' "${CMD[@]}") Enter
    else
        $DO "${CMD[@]}"
    fi
}

# ssh series
prepare_ssh_cmd() {
    local ssh_bin="${1:-ssh}"
    if [[ "$ssh_bin" == "scp" || "$ssh_bin" == "sftp" ]]; then
        local port_param='-P'
    else
        local port_param='-p'
    fi
    CMD=(
        "$ssh_bin"
        "${PORT:+$port_param}" "$PORT"
        "${SSH_OPTIONS[@]}"
        "${EXTRA_SSH_OPTIONS[@]}"
        "$SCP_SRC"
        "$USERNAME${USERNAME:+@}$SERVER"
        "$SCP_DST"
        "${@:2}"
    )
    for i in ${!CMD[@]}; do if [[ -z "${CMD[i]}" ]]; then unset CMD[i]; fi; done
}

# ssh
run_ssh()
{
    prepare_ssh_cmd "$@"
    print_cmd
    eval_or_echo
}

# sshl
run_sshl()
{
    local arg left right localsock access
    local res="${1//[^:]}"
    local lorr="-L"
    if [[ "${FUNCNAME[1]}" == "run_sshr" ]]; then
        lorr="-R"
    fi
    if [[ ${#res} -eq 2 ]]; then
        arg="$1"
    elif [[ ${#res} -eq 0 ]]; then
        if [[ "$lorr" == "-R" ]]; then
            arg="$1"
        else
            arg="$(get_free_port):localhost:$1"
        fi
    else
        left=${1%%:*}
        right=${1##*:}
        if check_port "$left"; then
            arg="$1"
        elif check_port "$right"; then
            arg="$(get_free_port):$1"
        else
            arg="$1"
            localsock=1
        fi
    fi
    if [[ "$localsock" == "1" ]]; then
        access="unix://${arg%%:*}"
    else
        access="localhost:${arg%%:*}"
    fi
    SSH_OPTIONS+=("-NC" "$lorr" "$arg")
    prepare_ssh_cmd ssh
    print_cmd
    fmt_note "  > please access $access"
    eval_or_echo
}

# sshr
run_sshr()
{
    run_sshl "$1"
}

# sshd
run_sshd()
{
    local port=${1:-$(get_free_port)}
    SSH_OPTIONS+=("-NC" "-D" "$port")
    prepare_ssh_cmd ssh
    print_cmd
    fmt_note "  > please access localhost:$port"
    eval_or_echo
}

# scp
run_scp() {
    local src="$1"
    local dst="$2"
    local dst_is_remote=1
    # whoever is ./*, it can't be the remote; whoever not exists on local, it's possible the remote.
    # it is suggested to use ./* for local files.
    if [[ "$src" != "./"* && ( "$dst" == "./"* || ( ! -e "$src" && -e "$dst" ) ) ]]; then
        dst_is_remote=0
    fi
    if [[ "$dst_is_remote" == "1" ]]; then
        SCP_SRC="$src"
        SERVER="$SERVER":"$dst"
    else
        SERVER="$SERVER":"$src"
        SCP_DST="$dst"
    fi
    SSH_OPTIONS+=("-r")
    prepare_ssh_cmd scp
    print_cmd
    eval_or_echo
}

# ping
run_ping() {
    CMD=(ping)
    if [[ "$1" == "ping4" ]]; then
        CMD+=(-4)
    elif [[ "$1" == "ping6" ]]; then
        CMD+=(-6)
    fi
    CMD+=(-c 4 "$SERVER")
    print_cmd
    eval_or_echo
}

# remove host keys
remove_hostkey() {
    local key
    if [[ -z "$PORT" || "$PORT" == "22" ]]; then
        key=$SERVER
    else
        key="[$SERVER]:$PORT"
    fi
    ssh-keygen -R "$key"
}

# main
print_help()
{
    local pager=(cat)
    if [[ -t 1 ]] && command -v less > /dev/null 2>&1; then
        pager=(less -R)
    fi

    cat <<EOF | "${pager[@]}"
NAME
    riot - connect to remote hosts using SSH presets

SYNOPSIS
    ${0##*/} [OPTION]... REMOTE [COMMAND] [--] [COMMAND-ARG]...

DESCRIPTION
    Connect to REMOTE using the matching configuration from riot-config.sh.
    COMMAND defaults to ssh. Separate multiple remotes with commas and jump
    hosts with slashes.

OPTIONS
    -4
        Force ssh to use IPv4 addresses only.

    -6
        Force ssh to use IPv6 addresses only.

    -D, --dry-run
        Print commands without executing them.

    -d, --dev
        Enable shell execution tracing.

    -g, --gpg
        Forward the local GPG agent during an interactive SSH login. Use only
        with trusted remote hosts.

    -h, --help
        Display this help and exit.

    -l, --lite
        Enable dotfiles lite mode for loaded configuration.

    -o SSH-OPTION
        Pass an option to ssh. This option may be specified multiple times.

    -p, --password
        Use password authentication instead of public-key authentication.

    -q, --quiet
        Enable dotfiles quiet mode for loaded configuration.

    -t, --trust
        Trust the remote and enable X11 and SSH agent forwarding.

    -v
        Enable verbose ssh output.

    --color
        Force colored output.

    --tmux
        Open SSH sessions in tmux windows.

    --
        End option parsing.

COMMANDS
    ssh [SSH-ARG]...
        Open an SSH session. This is the default command.

    tmux [SSH-ARG]...
        Open SSH sessions in multiple tmux windows.

    git [GIT-ARG]...
        Run git on the remote host.

    sshl [LOCAL-PORT:REMOTE-HOST:]REMOTE-PORT
        Create local port forwarding with ssh -L.

    sshr [REMOTE-HOST:]REMOTE-PORT
        Create remote port forwarding with ssh -R.

    sshd [LOCAL-PORT]
        Create dynamic port forwarding with ssh -D.

    zssh [SSH-ARG]...
        Open a zssh session.

    sftp [SFTP-ARG]...
        Open an SFTP session.

    scp SOURCE DESTINATION
        Copy files to or from the remote host.

    rm
        Remove the remote host key from known_hosts.

    ping, ping4, ping6
        Ping the remote host.
EOF
}

router() {
    local positional=()
    while [[ $# > 0 ]]; do
        case "$1" in
            -h|--help )
                print_help
                exit 0
                ;;
            -t|--trust )
                RIOT_TRUST_SERVER=1
                ;;
            -g|--gpg )
                GPG_FORWARD=1
                ;;
            --tmux )
                USE_TMUX=1
                ;;
            -p|--password )
                EXTRA_SSH_OPTIONS+=("-o" "PasswordAuthentication=yes" "-o" "PubkeyAuthentication=no")
                ;;
            -o )
                EXTRA_SSH_OPTIONS+=("-o" "$2")
                shift
                ;;
            -4 )
                EXTRA_SSH_OPTIONS+=("-4")
                ;;
            -6 )
                EXTRA_SSH_OPTIONS+=("-6")
                ;;
            -v )
                EXTRA_SSH_OPTIONS+=("-v")
                ;;
            -- )
                shift
                positional+=("$@")
                break
                ;;
            -* )
                fmt_fatal "unknown option: '$1'. if this option is for the remote command, add '--' before."
                ;;
            * )
                positional+=("$1")
                ;;
        esac
        shift
    done
    IFS=',' read -ra remotes <<< "${positional[0]}"
    for i in ${!remotes[@]}; do if [[ -z "${remotes[i]}" ]]; then unset remotes[i]; fi; done
    if [[ "${#positional[@]}" == "0" || "${#remotes[@]}" == "0" ]]; then
        print_help
        exit 1
    fi
    if [[ "$GPG_FORWARD" == "1" && ( \
        ( "${positional[1]}" != "" && "${positional[1]}" != "ssh" && "${positional[1]}" != "tmux" ) \
        || "${#positional[@]}" -gt 2 ) ]]; then
        fmt_fatal "gpg-agent forwarding is only supported for interactive SSH login"
    fi
    for i in ${!remotes[@]}; do
        remote="${remotes[i]}"
        local batch_func="${remote}.batch"
        if is_function "$batch_func"; then
            "$batch_func"
            continue
        fi
        parse_remote "$remote"
        case "${positional[1]}" in
            ssh|tmux|"" )
                [[ "${positional[1]}" == tmux ]] && USE_TMUX=1
                if [[ "$GPG_FORWARD" == "1" ]]; then
                    prepare_gpg_forwarding
                fi
                run_ssh ssh "${positional[@]:2}"
                ;;
            git )
                if printf '%s\0' "${positional[@]:2}" | grep -Fxqz -- '-C'; then
                    run_ssh ssh git "${positional[@]:2}"
                else
                    run_ssh ssh git -C "$(pwd)" "${positional[@]:2}"
                fi
                ;;
            ping|ping4|ping6 )
                test "${#positional[@]}" -eq 2 || fmt_fatal "ping requires no arguments"
                run_ping "${positional[1]}"
                ;;
            zssh )
                run_ssh zssh "${positional[@]:2}"
                ;;
            sftp )
                run_ssh sftp "${positional[@]:2}"
                ;;
            sshl )
                test -n "${positional[2]}" || fmt_fatal "no target address provided"
                test "${#positional[@]}" -eq 3 || fmt_fatal "sshl requires exactly one argument"
                run_sshl "${positional[2]}"
                ;;
            sshr )
                test -n "${positional[2]}" || fmt_fatal "no target address provided"
                test "${#positional[@]}" -eq 3 || fmt_fatal "sshr requires exactly one argument"
                run_sshr "${positional[2]}"
                ;;
            sshd )
                test "${#positional[@]}" -le 3 || fmt_fatal "sshd requires one or no arguments"
                if [[ "${#positional[@]}" -eq 3 ]]; then
                    check_port "${positional[2]}" || fmt_fatal "invalid port number: ${positional[2]}"
                    run_sshd "${positional[2]}"
                else
                    run_sshd
                fi
                ;;
            scp )
                test "${#positional[@]}" -eq 4 || fmt_fatal "scp requires exactly two arguments: source and destination"
                test -n "${positional[2]}" || fmt_fatal "no source path specified"
                test -n "${positional[3]}" || fmt_fatal "no destination path specified"
                run_scp "${positional[2]}" "${positional[3]}"
                ;;
            rm )
                test "${#positional[@]}" -eq 2 || fmt_fatal "rm requires no arguments"
                remove_hostkey
                ;;
            * )
                print_help
                fmt_fatal "unknown command: ${positional[1]}"
                ;;
        esac
    done

    if [[ -n "$TMUX_SESS" && "$DFS_DRY_RUN" != "1" ]]; then
        tmux attach-session -t $TMUX_SESS
    fi
}

router "${GOT_OPTS[@]}"
