#!/bin/sh
# nanokvm Rev3.1

wait_for_exit() {
    process="$1"
    timeout="$2"
    elapsed=0

    while pidof "$process" >/dev/null 2>&1; do
        if [ "$elapsed" -ge "$timeout" ]; then
            return 1
        fi
        sleep 1
        elapsed=$((elapsed + 1))
    done
    return 0
}

stop_process() {
    process="$1"

    if ! pidof "$process" >/dev/null 2>&1; then
        return
    fi

    killall -INT "$process" 2>/dev/null || true
    if wait_for_exit "$process" 20; then
        return
    fi

    killall -TERM "$process" 2>/dev/null || true
    if wait_for_exit "$process" 5; then
        return
    fi

    killall -KILL "$process" 2>/dev/null || true
    wait_for_exit "$process" 2 || true
}

stop_services() {
    # The server owns VI/VPSS. Let it release MMF before stopping kvm_system.
    stop_process NanoKVM-Server
    stop_process kvm_system
    rm -rf /tmp/kvm_system /tmp/server
}

start_services() {
    cp -r /kvmapp/kvm_system /tmp/
    if [ -f /kvmapp/kvm_new_app ]; then
        touch /tmp/.nanokvm_migrating
        /tmp/kvm_system/kvm_system &

        # The legacy migration invokes this init script recursively and starts
        # the server itself. Wait for that single migration path to finish.
        migration_wait=0
        while [ "$migration_wait" -lt 30 ]; do
            if [ ! -f /kvmapp/kvm_new_app ] && pidof NanoKVM-Server >/dev/null 2>&1; then
                rm -f /tmp/.nanokvm_migrating
                return
            fi
            sleep 1
            migration_wait=$((migration_wait + 1))
        done
        rm -f /tmp/.nanokvm_migrating
    else
        /tmp/kvm_system/kvm_system &
    fi

    if ! pidof NanoKVM-Server >/dev/null 2>&1; then
        rm -rf /tmp/server
        cp -r /kvmapp/server /tmp/
        /tmp/server/NanoKVM-Server &
    fi
}

case "$1" in
  start)
    echo -n kvm > /boot/hostname.prefix
    [ -e /root/.profile ] || printf '%s\n' 'export PS1="\u@\h:\w# "' > /root/.profile

    # Copy sensor config if the file exists
    [ -f /mnt/data/sensor_cfg.ini.LT ] && cp /mnt/data/sensor_cfg.ini.LT /mnt/data/sensor_cfg.ini

    # Generate unique device key
    if [ -f /sys/class/cvi-base/base_uid ]; then
        str_value=$(awk '{print $2}' /sys/class/cvi-base/base_uid)
        first_uint=$(echo "$str_value" | cut -d'_' -f1)
        second_uint=$(echo "$str_value" | cut -d'_' -f2)
        echo "$first_uint$second_uint" > /device_key
    fi

    # Set iptables rules (skip if already present)
    iptables -C INPUT -i eth0 -p tcp --dport 80 -m state --state NEW,ESTABLISHED -j ACCEPT 2>/dev/null || \
    iptables -A INPUT -i eth0 -p tcp --dport 80 -m state --state NEW,ESTABLISHED -j ACCEPT

    iptables -C OUTPUT -o eth0 -p tcp --sport 80 -m state --state ESTABLISHED -j ACCEPT 2>/dev/null || \
    iptables -A OUTPUT -o eth0 -p tcp --sport 80 -m state --state ESTABLISHED -j ACCEPT

    iptables -C INPUT -i eth0 -p tcp --sport 22 -m state --state NEW,ESTABLISHED -j ACCEPT 2>/dev/null || \
    iptables -A INPUT -i eth0 -p tcp --sport 22 -m state --state NEW,ESTABLISHED -j ACCEPT

    iptables -C OUTPUT -o eth0 -p tcp --dport 22 -m state --state ESTABLISHED -j ACCEPT 2>/dev/null || \
    iptables -A OUTPUT -o eth0 -p tcp --dport 22 -m state --state ESTABLISHED -j ACCEPT

    iptables -C OUTPUT -o eth0 -p tcp --sport 8000 -m state --state ESTABLISHED -j DROP 2>/dev/null || \
    iptables -A OUTPUT -o eth0 -p tcp --sport 8000 -m state --state ESTABLISHED -j DROP

    # Logo selection
    if [ -f /boot/logo.ico ]; then
        if [ ! -f /kvmapp/server/web/original.ico ]; then
            mv /kvmapp/server/web/sipeed.ico /kvmapp/server/web/original.ico
            cp /boot/logo.ico /kvmapp/server/web/sipeed.ico
        fi
    else
        if [ -f /kvmapp/server/web/original.ico ]; then
            mv /kvmapp/server/web/original.ico /kvmapp/server/web/sipeed.ico
        fi
    fi

    # Start services
    start_services
   ;;

  stop)
    stop_services
    echo "OK"
   ;;

  restart)
    if [ -f /tmp/.nanokvm_migrating ]; then
        exit 0
    fi
    stop_services
    start_services

    sync
    echo "OK"
   ;;
    
  *)
    echo "Usage: $0 {start|stop|restart}"
    exit 1
  ;;
esac
