#!/bin/sh

BOOT_WAIT_DEFAULT=25
# set GO_PLAY_KIDS to no to perform all tests.
GO_PLAY_KIDS=yes

YESCOUNT=0
NOCOUNT=0

WHISPERSSOCK=/run/whispers/unix-sockets/whispers.sock
VPNSOCK=/run/whispers/vpn/unix-sockets/vpn.sock

REGISTERSLEEP=10
TESTREGISTERSLEEP=4
CONNECTSLEEP=20
TESTCONNECTSLEEP=10

function testcount()
{
    test $1 = $2 && echo SUCCESS || echo FAILED
    test $1 = $2 && YESCOUNT=$(($YESCOUNT + 1)) || NOCOUNT=$(($NOCOUNT + 1))
}

function dirtyssh()
{
    ssh -q \
        -o UserKnownHostsFile=/dev/null \
        -o StrictHostKeyChecking=no \
        -p $1 \
        root@localhost \
        "$2"
}

function herd_service_status()
{
    echo "herd status of $2 on ${VM[$1]}:"
    dirtyssh ${PORT[$1]} "herd status $2"
    echo
}

function herd_test_service_status()
{
    echo "Test if the $2 service is $3 on ${VM[$1]}."
    RESULT=$(herd_service_status $1 $2 | awk -- 'BEGIN{FS="[. ]"} ; /It is '$3'/ {print $5}')
    testcount ${RESULT} $3
}

function whispers_service_status()
{
    echo "herd status of $3/$2 on ${VM[$1]}:"
    dirtyssh ${PORT[$1]} "herd -s \\
                              /run/whispers/$3/unix-sockets/*.sock status $2"
    echo
}

function whispers_test_service_status()
{
    echo "Test if the $3/$2 service is $4 on ${VM[$1]}."
    RESULT=$(whispers_service_status $1 $2 $3 | awk -- 'BEGIN{FS="[. ]"} ; /It is '$4'/ {print $5}')
    testcount ${RESULT} $4
}

function test_forward()
{
    echo "From ${VM[$1]}, take forward opened by
${VM[$2]} to connect to ${VM[$3]} through
sshd of ${VM[$4]}."
    EXITHOST=reset
    EXITHOST=$(dirtyssh \
                   ${PORT[$1]} \
                   "ssh -q \\
                        -o UserKnownHostsFile=/dev/null \\
                        -o StrictHostKeyChecking=no \\
                        -p $5 \\
                        root@$6 \\
                        hostname")
    echo "Host name at forward exit: ${EXITHOST}, expected: ${VM[$3]}"
    testcount ${EXITHOST} ${VM[$3]}
}

function test_forwards()
{
    #   echo "*** REVERSE PORT FORWARD TESTS"
    #   echo
    #   test_forward $WILD01 \
        #                $CLIENT01 \
        #                $CLIENT01 \
        #                $SOCKS \
        #                ${EXTRAPORT_HOST[$SOCKS]} \
        #                10.0.2.2
    #   echo
    #   echo
    echo "*** PORT FORWARD TESTS"
    echo
    test_forward $CLIENT02 \
                 $CLIENT02 \
                 $VPN_SERVER \
                 $VPN_SERVER \
                 36492 \
                 localhost
    echo
    echo
}

function test_book_dummy_user()
{
    echo "Book a hostname $2 in the network state of ${VM[$1]}, test
that hostnme $2 has booked one and one one voucher."
    dirtyssh ${PORT[$1]} "herd -s $VPNSOCK book-client network-rw $2"
    COUNT=reset
    COUNT=$(dirtyssh \
                ${PORT[$1]} \
                "herd -s $VPNSOCK display-network-state network-rw \\
                    | grep -c \"client-hostname $2\"")
    echo "Number of booked vouchers for $2: ${COUNT}, expected: 1"
    testcount ${COUNT} 1
    echo
}

function tests_book_free_dummy_users()
{
    echo "*** TEST DUMMY VOUCHER BOOKING BY ${VM[$VPN_SERVER]}"
    echo
    test_book_dummy_user $1 $2
    test_book_dummy_user $1 $2
    test_book_dummy_user $1 $3
    test_book_dummy_user $1 $2
    test_book_dummy_user $1 $3
    test_book_dummy_user $1 $2
    test_book_dummy_user $1 $4
    test_book_dummy_user $1 $2
    test_book_dummy_user $1 $4
    test_free_dummy_user $1 $3
    test_free_dummy_user $1 $2
    test_book_dummy_user $1 $3
    test_book_dummy_user $1 $3
    test_free_dummy_user $1 $2
    test_free_dummy_user $1 $3
    test_book_dummy_user $1 $4
    test_free_dummy_user $1 $4
    test_free_dummy_user $1 $2
    test_free_dummy_user $1 $2
    test_book_dummy_user $1 $2
    test_free_dummy_user $1 $2
    test_free_dummy_user $1 $3
    test_free_dummy_user $1 $4
    echo
}

function test_free_dummy_user()
{
    echo "Free a hostname $2 in the network state of ${VM[$1]}, test
that hostnme $2 has no booked voucher."
    dirtyssh ${PORT[$1]} "herd -s $VPNSOCK free-client-booking network-rw $2"
    COUNT=reset
    COUNT=$(dirtyssh \
                ${PORT[$1]} \
                "herd -s $VPNSOCK display-network-state network-rw \\
                    | grep -c \"client-hostname $2\"")
    echo "Number of booked vouchers for $2: ${COUNT}, expected: 0"
    testcount ${COUNT} 0
    echo
}

function tests_register()
{
    echo "Register ${VM[$1]} as a prospective client in its configured
VPN network."
    dirtyssh ${PORT[$1]} "herd -s $WHISPERSSOCK register vpn"
    sleep $REGISTERSLEEP
    echo "Test that hostname ${VM[$1]} has booked one and
only one voucher."
    COUNT=reset
    COUNT=$(dirtyssh \
                ${PORT[$1]} \
                "herd -s $VPNSOCK display-network-state network-rw \\
                    | grep -c \"client-hostname ${VM[$1]}\"")
    echo "Number of booked vouchers for ${VM[$1]}: ${COUNT}, expected: 1"
    testcount ${COUNT} 1
    sleep $TESTREGISTERSLEEP
    whispers_test_service_status $1 registered vpn running
    whispers_test_service_status $1 unregistered vpn stopped
    whispers_test_service_status $1 connecting vpn stopped
    whispers_test_service_status $1 disconnecting vpn stopped
    whispers_test_service_status $1 registering vpn stopped
    whispers_test_service_status $1 unregistering vpn stopped
    echo
}

function tests_unregister()
{
    echo "Unregister ${VM[$1]} from its configured VPN network."
    dirtyssh ${PORT[$1]} "herd -s $WHISPERSSOCK unregister vpn"
    sleep $REGISTERSLEEP
    echo "Test that hostname ${VM[$1]} has no booked voucher."
    COUNT=reset
    COUNT=$(dirtyssh \
                ${PORT[$1]} \
                "herd -s $VPNSOCK display-network-state network-rw \\
                    | grep -c \"client-hostname  ${VM[$1]}\"")
    echo "Number of booked vouchers for ${VM[$1]}: ${COUNT}, expected: 0"
    testcount ${COUNT} 0
    sleep $TESTREGISTERSLEEP
    whispers_test_service_status $1 registered vpn stopped
    whispers_test_service_status $1 unregistered vpn running
    whispers_test_service_status $1 connecting vpn stopped
    whispers_test_service_status $1 disconnecting vpn stopped
    whispers_test_service_status $1 registering vpn stopped
    whispers_test_service_status $1 unregistering vpn stopped
    echo
}

function tests_register_unregister()
{
    echo "*** REGISTRATION AND UNREGISTRATION TESTS"
    echo
    tests_register $1
    tests_register $1
    tests_register $2
    tests_register $1
    tests_unregister $1
    tests_unregister $1
    tests_unregister $2
    tests_unregister $1
    tests_register $1
    tests_register $1
    tests_register $2
    tests_register $1
    tests_unregister $1
    tests_unregister $1
    tests_unregister $2
    tests_unregister $1
    echo
}

function test_connect_propagate()
{
    if dirtyssh \
           ${PORT[$2]} \
           "herd -s $VPNSOCK display-network-state network-rw \\
                    | grep ${VM[$2]} \\
                    | grep \"connected. #t\"" > /dev/null
    then
        AUX=${VM[$2]}
    else
        AUX=${VM[$1]}
    fi
    echo "Test that ${VM[$1]} is connected in ${AUX}'s network state."
    COUNT=reset
    COUNT=$(dirtyssh \
                ${PORT[$1]} \
                "herd -s $VPNSOCK display-network-state network-rw \\
                    | grep ${AUX} \\
                    | grep -c \"connected. #t\"")
    echo "Number of connected vouchers for ${AUX}: ${COUNT}, expected: 1"
    testcount ${COUNT} 1
    sleep $TESTCONNECTSLEEP
}

function server_tun()
{
    dirtyssh \
        ${PORT[$VPN_SERVER]}
        "herd -s $VPNSOCK display-network-state network-rw" \
            | grep ${VM[$1]} \
            | sed 's/.*tun-edvice-number ([0-9]*)/tun\1/'
}

function tests_connect()
{
    echo "Connect ${VM[$1]} as a VPN client of the VPN network."
    dirtyssh ${PORT[$1]} "herd -s $WHISPERSSOCK connect vpn"
    sleep $CONNECTSLEEP
    echo "Test that ${VM[$1]} has a tun interface."
    COUNT=reset
    COUNT=$(dirtyssh \
                ${PORT[$1]} \
                "ip addr \\
                    | grep -c \": tun\"")
    echo "Number of tun interfaces for ${VM[$1]}: ${COUNT}, expected: 1"
    testcount ${COUNT} 1
    sleep $TESTCONNECTSLEEP
    test_connect_propagate $1 $1
    test_connect_propagate $1 $2
    whispers_test_service_status $1 connected vpn running
    whispers_test_service_status $1 disconnected vpn stopped
    whispers_test_service_status $1 registered vpn running
    whispers_test_service_status $1 unregistered vpn stopped
    whispers_test_service_status $1 connecting vpn stopped
    whispers_test_service_status $1 disconnecting vpn stopped
    whispers_test_service_status $1 registering vpn stopped
    whispers_test_service_status $1 unregistering vpn stopped
    echo
}

function tests_manual_register_connect
{
    echo "*** REGISTER THEN CONNECT CLIENTS TESTS"
    echo
    tests_register $1
    tests_connect $1 $2
    tests_register $1
    tests_connect $1 $2
    tests_register $2
    tests_connect $2 $1
    tests_register $2
    tests_connect $2 $1
    tests_register $2
    tests_connect $2 $1
    tests_register $1
    tests_connect $1 $2
    echo
}

function tests_disconnect()
{
    echo "Disconnect ${VM[$1]} as a VPN client of the VPN network."
    dirtyssh ${PORT[$1]} "herd -s $WHISPERSSOCK disconnect vpn"
    sleep $CONNECTSLEEP
    whispers_test_service_status $1 connected vpn stopped
    whispers_test_service_status $1 disconnected vpn running
    whispers_test_service_status $1 registered vpn running
    whispers_test_service_status $1 unregistered vpn stopped
    whispers_test_service_status $1 connecting vpn stopped
    whispers_test_service_status $1 disconnecting vpn stopped
    whispers_test_service_status $1 registering vpn stopped
    whispers_test_service_status $1 unregistering vpn stopped
    echo
}

function tests_disconnects_connects
{
    echo "*** DISCONNECT THEN RECONNECT CLIENTS TESTS"
    echo
    tests_disconnect $1
    tests_disconnect $1
    tests_connect $1 $2
    tests_connect $1 $2
    tests_disconnect $2
    tests_disconnect $2
    tests_disconnect $2
    tests_connect $2 $1
    tests_connect $2 $1
    tests_disconnect $2
    tests_disconnect $2
    echo
}

function tests_direct_connects
{
    echo "*** DIRECT CONNECT CLIENTS TESTS"
    echo
    tests_unregister $1
    tests_unregister $2
    tests_connect $1 $2
    tests_connect $2 $1
    echo
}

function tests_disconnects_unregisters
{
    echo "*** DISCONNECT THEN UNREGISTER CLIENTS TESTS"
    echo
    tests_disconnect $1
    tests_disconnect $1
    tests_unregister $1
    tests_unregister $1
    tests_disconnect $2
    tests_disconnect $2
    tests_disconnect $2
    tests_unregister $2
    tests_unregister $2
    echo
}

echo "*** DEFINING VM INSTANCIATION ARRAYS"
echo

I=0
WILD=$I

VM[$I]=wild
echo "* VM: ${VM[$I]}"
EXTRA_SERVICES[$I]=""
PORT[$I]=$((10022+$I))
EXTRAPORT_HOST[$I]=$((${PORT[$I]}+1000))
EXTRAPORT_VM[$I]=$((${PORT[$I]}+2000))
ULTIMAPORT_HOST[$I]=$((${PORT[$I]}+3000))
ULTIMAPORT_VM[$I]=$((${PORT[$I]}+4000))
EXTRA_PACKAGES[$I]=" iproute
                         iptables"
echo

I=$(($I+1))
VPN_SERVER=$I

VM[$I]=vpn
echo "* VM: ${VM[$I]}"
ALLOW_GATEWAY[$I]="yes"
EXTRA_SERVICES[$I]=""
PORT[$I]=$((10022+$I))
EXTRAPORT_HOST[$I]=$((${PORT[$I]}+1000))
EXTRAPORT_VM[$I]=$((${PORT[$I]}+2000))
ULTIMAPORT_HOST[$I]=$((${PORT[$I]}+3000))
ULTIMAPORT_VM[$I]=$((${PORT[$I]}+4000))
ALLOW_TUNDEV[$I]=yes
EXTRA_PACKAGES[$I]=" iproute
                         iptables"
EXTRA_SERVICES[$I]="
           (service whispers-vpn-service-type)"
echo

I=$(($I+1))
SOCKS=$I

VM[$I]=socks
echo "* VM: ${VM[$I]}"
ALLOW_GATEWAY[$I]="yes"
BOOT_WAIT[$I]=25
EXTRA_SERVICES[$I]=""
PORT[$I]=$((10022+$I))
EXTRAPORT_HOST[$I]=$((${PORT[$I]}+1000))
EXTRAPORT_VM[$I]=$((${PORT[$I]}+2000))
ULTIMAPORT_HOST[$I]=$((${PORT[$I]}+3000))
ULTIMAPORT_VM[$I]=$((${PORT[$I]}+4000))
EXTRA_PACKAGES[$I]=" iproute
                         iptables"
echo

I=$(($I+1))
CLIENT01=$I

VM[$I]=client01
echo "* VM: ${VM[$I]}"
PORT[$I]=$((10022+$I))
EXTRAPORT_HOST[$I]=$((${PORT[$I]}+1000))
EXTRAPORT_VM[$I]=$((${PORT[$I]}+2000))
ULTIMAPORT_HOST[$I]=$((${PORT[$I]}+3000))
ULTIMAPORT_VM[$I]=$((${PORT[$I]}+4000))
EXTRA_PACKAGES[$I]=" iproute
                         iptables"
EXTRA_SERVICES[$I]="
           (service elogind-service-type)
           (service whispers-vpn-service-type
            (whispers-vpn-configuration
             (client? #t)
             (server-sshd-host \"10.0.2.2\")
             (server-sshd-port ${PORT[$VPN_SERVER]})
             (forward-exit-port 22)
             (%auto-register? #f)))"
echo

I=$(($I+1))
CLIENT02=$I

VM[$I]=client02
echo "* VM: ${VM[$I]}"
PORT[$I]=$((10022+$I))
EXTRAPORT_HOST[$I]=$((${PORT[$I]}+1000))
EXTRAPORT_VM[$I]=$((${PORT[$I]}+2000))
ULTIMAPORT_HOST[$I]=$((${PORT[$I]}+3000))
ULTIMAPORT_VM[$I]=$((${PORT[$I]}+4000))
EXTRA_PACKAGES[$I]=" iproute
                         iptables"
EXTRA_SERVICES[$I]="
           (service whispers-vpn-service-type
            (whispers-vpn-configuration
             (client? #t)
             (stealth? #t)
             (server-sshd-host \"10.0.2.2\")
             (server-sshd-port ${PORT[$VPN_SERVER]})
             (forward-exit-port 22)
             (proxy-sshd-host \"10.0.2.2\")
             (proxy-sshd-port ${PORT[$VPN_SERVER]})
             (%auto-register? #f)))"
echo

for I in ${!PORT[@]}
do
    EXTRACONTENT[$I]=""
    if [ -v ALLOW_TUNDEV[$I] ]
    then
        EXTRACONTENT[$I]="
                     (extra-content \"
PermitTunnel=point-to-point\")"
    fi
    GATEWAY[$I]=""
    if [ -v ALLOW_GATEWAY[$I] ]
    then
        GATEWAY[$I]="
                     (gateway-ports? #t)"
    fi
    WAIT[$I]=$BOOT_WAIT_DEFAULT
    if [ -v BOOT_WAIT[$I] ]
    then
        WAIT[$I]=${BOOT_WAIT[$I]}
    fi
done
echo

echo "*** CONCATENATING SYSTEM CONFIGURATIONS"
echo
for I in ${!PORT[@]}
do
    echo "* VM: ${VM[$I]}"
    echo "(use-modules (guix records)
                   (gnu)
                   (whispers services whispers vpn)
                   (whispers services ssh-tunneler)
                   (whispers services whispers)
                   (whispers services ssh)
                   (whispers services whispers ssh))
(use-service-modules networking desktop)
(use-package-modules ssh networking linux scm-runciter)

(operating-system
  (host-name \"${VM[$I]}\")
  (timezone \"Asia/Shanghai\")
  (locale \"en_US.utf8\")

  (bootloader (bootloader-configuration
                (bootloader grub-bootloader)
                (targets '(\"/dev/sda\"))))

  (file-systems (cons (file-system
                        (device (file-system-label \"root\"))
                        (mount-point \"/\")
                        (type \"ext4\"))
                      %base-file-systems))

  (users %base-user-accounts)

  (packages (append %base-packages
                    (list${EXTRA_PACKAGES[$I]})))

  (services
    (append
     (list (service dhcp-client-service-type)
           (service openssh-service-type
                    (openssh-configuration
                     (permit-root-login #t)
                     (allow-empty-passwords? #t)
                     (openssh openssh-sans-x)
                     (port-number 22)${GATEWAY[$I]}${EXTRACONTENT[$I]}))
           (service whispers-service-type)
           (service
            whispers-ssh-service-type
            (whispers-ssh-configuration
             (users-groups-keys-forwards
              (append
               (list
                (ssh-user-group-keys-forwards
                (user-and-group (whispers-user-group
                                (user \"root\")
                                (group \"root\")))
                (keys '(\"root/.ssh/id_rsa\"))))))))${EXTRA_SERVICES[$I]})
     %base-services)))" > /tmp/${VM[$I]}.scm
    echo
done
echo

echo "*** INSTANCIATING VMs"
echo
for I in ${!PORT[@]}
do
    echo "* VM: ${VM[$I]}"
    VMRUN[$I]=$(guix system vm /tmp/${VM[$I]}.scm)
    echo
done
echo

echo "*** SILENTLY BOOTING VMs..."
echo
for I in ${!PORT[@]}
do
    NICSTANCE=-"nic user,model=virtio-net-pci,hostfwd=tcp::${PORT[$I]}-:22,hostfwd=tcp::${EXTRAPORT_HOST[$I]}-:${EXTRAPORT_VM[$I]},hostfwd=tcp::${ULTIMAPORT_HOST[$I]}-:${ULTIMAPORT_VM[$I]}"
    if [ $I = $VPN_SERVER ]
    then
        for CURPORT in $VPNPORTS
        do
            NICSTANCE=${NICSTANCE},hostfwd=tcp::${CURPORT}-:${CURPORT}
        done
    fi
    echo "* VM: ${VM[$I]}"
    ${VMRUN[$I]} \
        $NICSTANCE \
        -display none &
    sleep ${WAIT[$I]}
    echo
done
echo

echo "*** SETTING PASSWORDLESS LOGIN FOR WHISPERS USERS"
echo
echo "* VM: ${VM[$VPN_SERVER]}"
dirtyssh ${PORT[$VPN_SERVER]} 'passwd -d whispers'
echo
echo "* VM: ${VM[$CLIENT01]}"
dirtyssh ${PORT[$CLIENT01]} 'passwd -d whispers'
echo
echo "* VM: ${VM[$CLIENT02]}"
dirtyssh ${PORT[$CLIENT02]} 'passwd -d whispers'
echo
echo

function full_sshagent_reports_tests()
{
    echo "*** GENERATING AND ADDING ROOT SSH PRIVATE KEYS"
    echo
    for I in ${!PORT[@]}
    do
        echo "* VM: ${VM[$I]}"
        dirtyssh ${PORT[$I]} 'ssh-keygen -t rsa -N "" -f /root/.ssh/id_rsa'
        dirtyssh ${PORT[$I]} \
                 'herd -s '$WHISPERSSOCK' restart ssh'
        echo
    done
    echo

    echo "*** REPORTING SHEPHERD STATUS"
    echo
    for I in ${!PORT[@]}
    do
        echo "* VM: ${VM[$I]}"
        dirtyssh ${PORT[$I]} 'herd status'
        echo
    done
    echo

    echo "*** REPORTING NETWORK INTERFACES"
    echo
    for I in ${!PORT[@]}
    do
        echo "* VM: ${VM[$I]}"
        dirtyssh ${PORT[$I]} 'ip addr show'
        echo
    done
    echo

    echo "*** REPORTING ROUTES"
    echo
    for I in ${!PORT[@]}
    do
        echo "* VM: ${VM[$I]}"
        dirtyssh ${PORT[$I]} 'ip route'
        echo
    done
    echo

    tests_book_free_dummy_users $VPN_SERVER dummy dump-dumb dumber
    tests_register_unregister $CLIENT02 $CLIENT01
    tests_manual_register_connect $CLIENT01 $CLIENT02
    tests_disconnects_connects $CLIENT01 $CLIENT02
    tests_disconnects_unregisters $CLIENT01 $CLIENT02
    tests_direct_connects $CLIENT01 $CLIENT02
}

if [[ foo$GO_PLAY_KIDS == foono ]]
then
    full_sshagent_reports_tests
    echo
fi

echo "*** CONNECTING CLIENTS"
echo
    dirtyssh ${PORT[$CLIENT01]} "herd -s $WHISPERSSOCK connect vpn"
    sleep $CONNECTSLEEP
echo
    dirtyssh ${PORT[$CLIENT02]} "herd -s $WHISPERSSOCK connect vpn"
    sleep $CONNECTSLEEP
echo
echo

echo "*** TESTS SUMMARRY"
echo
echo Successes: $YESCOUNT
echo Failures: $NOCOUNT
echo
echo

echo "You can ssh into the VMs from another terminal of this host"
for I in ${!PORT[@]}
do
    echo "* VM: ${VM[$I]}, port: ${PORT[$I]}"
done
echo
read -n 1 -r -s -p "When done playing, press any key to halt the VMs..."
echo
echo

echo "*** HALTING VMs"
echo
for I in ${!PORT[@]}
do
    echo "* VM: ${VM[$I]}"
    dirtyssh ${PORT[$I]} halt
    echo
done