#!/bin/sh

# CPE reverse SSH agent
# Supports:
#   - RDK-B using OpenSSH client
#   - Homeware/OpenWRT using Dropbear dbclient v2019.78
#
# Server-side expected command:
#   /usr/local/bin/register-cpe.sh '<SERIAL>' '<PORT>' '<DEVICE>' '<TYPE>'

RELAY_HOST="cpe.aquaplace.co.uk"
RELAY_PORT="44322"
RELAY_USER="cpetunnel"

OPENSSH_KEY="/opt/SSHServer_Private"
DROPBEAR_KEY_PRIMARY="/root/SSHServer_Private.dropbear"
DROPBEAR_KEY_ALT="/opt/SSHServer_Private.dropbear"
DROPBEAR_KEY_FALLBACK="/root/SSHServer_Private"


LOCAL_SSH_PORT="22"
RECONNECT_SLEEP="10"
KEEPALIVE_SECONDS="5"
KEEPALIVE_COUNT="2"

# Set to 1 when testing interactively. For background/daemon use, keep 0.
LOG_TO_CONSOLE="1"

log()
{
    MSG="$*"

    if [ "$LOG_TO_CONSOLE" = "1" ]; then
        echo "$MSG"
    fi

    logger -t cpe-agent "$MSG" 2>/dev/null
}

safe_value()
{
    echo "$1" | tr ' ' '_' 
}

get_free_port()
{
    awk 'BEGIN {
        srand()
        print int(30000 + rand() * 20000)
    }'
}

get_device_info()
{
    # RDK-B
    if command -v ba-cli >/dev/null 2>&1; then
        SERIAL="$(ba-cli -l Device.DeviceInfo.SerialNumber? 2>/dev/null |grep -E '.+' )"

        DEVICE="$(ba-cli -l Device.DeviceInfo.Description? 2>/dev/null |grep -E '.+' )"
        DEVICE="$(safe_value "$DEVICE")"
        TYPE="RDK"

        [ -z "$SERIAL" ] && SERIAL="UNKNOWN"
        [ -z "$DEVICE" ] && DEVICE="UNKNOWN"

        echo "${SERIAL}|${DEVICE}|${TYPE}"
        return 0
    fi

    # Homeware / OpenWRT
    if command -v uci >/dev/null 2>&1; then
        SERIAL="$(uci get env.var.serial 2>/dev/null)"
        DEVICE="$(uci get env.var.prod_friendly_name 2>/dev/null)"

        SERIAL="$(safe_value "$SERIAL")"
        DEVICE="$(safe_value "$DEVICE")"
        TYPE="OpenWRT"

        [ -z "$SERIAL" ] && SERIAL="UNKNOWN"
        [ -z "$DEVICE" ] && DEVICE="UNKNOWN"

        echo "${SERIAL}|${DEVICE}|${TYPE}"
        return 0
    fi

    echo "UNKNOWN|UNKNOWN|UNKNOWN"
}

get_ssh_client()
{
    if command -v dbclient >/dev/null 2>&1; then
        if dbclient -V 2>&1 | grep -qi Dropbear; then
            echo "dropbear"
            return 0
        fi
    fi

    if command -v ssh >/dev/null 2>&1; then
        if ssh -V 2>&1 | grep -qi Dropbear; then
            echo "dropbear-ssh"
            return 0
        fi

        echo "openssh"
        return 0
    fi

    echo "unknown"
}

get_dropbear_key()
{
    if [ -f "$DROPBEAR_KEY_PRIMARY" ]; then
        echo "$DROPBEAR_KEY_PRIMARY"
        return 0
    fi

    if [ -f "$DROPBEAR_KEY_ALT" ]; then
        echo "$DROPBEAR_KEY_ALT"
        return 0
    fi

    if [ -f "$DROPBEAR_KEY_FALLBACK" ]; then
        echo "$DROPBEAR_KEY_FALLBACK"
        return 0
    fi

    echo "$DROPBEAR_KEY_PRIMARY"
}

run_openssh_tunnel()
{
    log "Using OpenSSH client"
    log "Key: $OPENSSH_KEY"

    ssh \
        -n \
        -q \
        -T \
        -p "$RELAY_PORT" \
        -i "$OPENSSH_KEY" \
        -o SetEnv="PORT=$REMOTE_PORT SERIAL=$SERIAL DEVICE=$DEVICE TYPE=$TYPE" \
        -o ServerAliveInterval=$KEEPALIVE_SECONDS \
        -o ServerAliveCountMax=$KEEPALIVE_COUNT \
        -o ExitOnForwardFailure=yes \
        -o StrictHostKeyChecking=no \
        -o UserKnownHostsFile=/dev/null \
        -R 127.0.0.1:${REMOTE_PORT}:127.0.0.1:${LOCAL_SSH_PORT} \
        ${RELAY_USER}@${RELAY_HOST} \
        "/usr/local/bin/register-cpe.sh $SERIAL $REMOTE_PORT $DEVICE $TYPE"
}

run_dropbear_tunnel()
{
    DBCLIENT="dbclient"

    if [ "$SSH_CLIENT_TYPE" = "dropbear-ssh" ]; then
        DBCLIENT="ssh"
    fi

    DROPBEAR_KEY="$(get_dropbear_key)"
	DROPBEAR_LAN_SECTION="$(uci show dropbear 2>/dev/null | awk -F'.' '/=dropbear/ {print $2}' | awk -F= '{print $1}')"
    log "Using Dropbear client: $DBCLIENT"
    log "Key: $DROPBEAR_KEY"

    # Do not use -f here. The agent loop owns reconnect handling.
    # Dropbear -f would return immediately and cause repeated tunnel creation.
    $DBCLIENT \
        -T \
        -y \
        -K "$KEEPALIVE_SECONDS" \
        -i "$DROPBEAR_KEY" \
        -p "$RELAY_PORT" \
        -R ${REMOTE_PORT}:$(ip -4 addr show br-lan | awk '/inet / {print $2}' | cut -d/ -f1):$(uci get dropbear.${DROPBEAR_LAN_SECTION}.Port) \
        ${RELAY_USER}@${RELAY_HOST} \
        "/usr/local/bin/register-cpe.sh $SERIAL $REMOTE_PORT $DEVICE $TYPE"
}

REMOTE_PORT="$(get_free_port)"

check_openwrt_prerequisites()
{
    command -v uci >/dev/null 2>&1 || return 0

    CURRENT="$(uci get dropbear.lan.AllowLocalForwarding 2>/dev/null)"
	DROPBEAR_LAN_SECTION="$(uci show dropbear 2>/dev/null | awk -F'.' '/=dropbear/ {print $2}' | awk -F= '{print $1}')"
    [ "$CURRENT" = "1" ] && return 0

    echo
    echo "OpenWRT prerequisite check"
    echo "=========================="
    echo
    echo "dropbear.${DROPBEAR_LAN_SECTION}.AllowLocalForwarding is currently:"
    echo "    $CURRENT"
    echo

    #
    # Interactive shell?
    #
    if [ -t 0 ]; then

        printf "Enable forwarding and restart dropbear? (Initial run can no be ni the background)[Y/N]: "
        read ANSWER

        case "$ANSWER" in
            y|Y|yes|YES)

                uci set dropbear.${DROPBEAR_LAN_SECTION}.AllowLocalForwarding='1'

                if uci get dropbear.wan.enable >/dev/null 2>&1; then
                    uci set dropbear.wan.AllowLocalForwarding='1'
                fi

                uci commit dropbear
                /etc/init.d/dropbear restart

                echo
                echo "Dropbear updated."
                ;;
        esac

    else

        echo "Non-interactive mode detected."
        echo "Automatically enabling forwarding."

        uci set dropbear.${DROPBEAR_LAN_SECTION}.AllowLocalForwarding='1'

        if uci get dropbear.wan.enable >/dev/null 2>&1; then
            uci set dropbear.wan.AllowLocalForwarding='1'
        fi

        uci commit dropbear
        /etc/init.d/dropbear restart

    fi
}

IFS='|' read -r SERIAL DEVICE TYPE <<EOF
$(get_device_info)
EOF

if [ "$TYPE" = "OpenWRT" ]; then
    check_openwrt_prerequisites
fi

SERIAL="$(safe_value "$SERIAL")"
DEVICE="$(safe_value "$DEVICE")"
TYPE="$(safe_value "$TYPE")"

[ -z "$SERIAL" ] && SERIAL="UNKNOWN"
[ -z "$DEVICE" ] && DEVICE="UNKNOWN"
[ -z "$TYPE" ] && TYPE="UNKNOWN"

log "Detected Serial: $SERIAL"
log "Detected Device: $DEVICE"
log "Detected Type: $TYPE"
log "Remote Port: $REMOTE_PORT"

if [ "$SERIAL" = "UNKNOWN" ]; then
    log "Unable to detect serial number, exiting"
    exit 1
fi

SSH_CLIENT_TYPE="$(get_ssh_client)"

log "Detected SSH client type: $SSH_CLIENT_TYPE"
log "Starting SSH Agent"

while true
do
    log "Connecting Serial=$SERIAL Port=$REMOTE_PORT Device=$DEVICE Type=$TYPE"
    log "Tunnel target=${RELAY_USER}@${RELAY_HOST}:${RELAY_PORT}"

    case "$SSH_CLIENT_TYPE" in
        openssh)
            run_openssh_tunnel
            RC=$?
            ;;

        dropbear|dropbear-ssh)
            run_dropbear_tunnel
            RC=$?
            ;;

        *)
            log "No supported SSH client found"
            exit 1
            ;;
    esac

    log "Tunnel disconnected Serial=$SERIAL Port=$REMOTE_PORT rc=$RC"
    sleep "$RECONNECT_SLEEP"
done