#!/bin/sh

export LOG="logger -t Shark [*]"
export LOG_ERR="logger -t Shark -p 3 [!]"

MODE="OFF"
SWITCH_POSITION=$(/usr/bin/SWITCH)

upgrade_leds() {
    /usr/bin/LED OFF
    while true
    do
        echo 1 > /sys/class/leds/shark:red:system/brightness
        sleep 0.2
        echo 0 > /sys/class/leds/shark:red:system/brightness
        echo 1 > /sys/class/leds/shark:blue:system/brightness
        sleep 0.2
        echo 0 > /sys/class/leds/shark:blue:system/brightness
    done
}

# $1: Upgrade file
# $2: MD5 file
# $3: Upgrade file size in bytes
extract_md5() {
    dd if="${1}" of="${2}" skip="${3}" bs=1 count=33
}

# $1: Upgrade file
# $2: MD5 file
verify_md5() {
    expected=$(cat "${2}")
    checksum=$(md5sum "${1}" | awk '{print $1}')
    [ "${expected}" = "${checksum}" ] && {
        return 0
    }
    return 1
}

# $1: Upgrade file
# $2: Upgrade file size in bytes
truncate_upgrade() {
    dd if=/dev/null of="${1}" bs=1 seek="${2}"
}

execute_upgrade() {
    # Check for upgrade file in default location
    $LOG "Checking for firmware upgrade"
    upgrade_file=$(find /root/upgrade-*.bin -type f 2>/dev/null | tail -n1)

    if [ -f "${upgrade_file}" ]; then
        # Upgrade file found
        $LOG "Firmware upgrade found"

        $LOG "Verifying firmware upgrade"
        upgrade_file_size=$(( $(wc -c "${upgrade_file}" | awk '{print $1}') - 33 ))

        # Extract md5sum from upgrade file
        extract_md5 "${upgrade_file}" /tmp/upgrade.md5 "${upgrade_file_size}"

        cp "${upgrade_file}" /tmp/upgrade.bin
        truncate_upgrade /tmp/upgrade.bin "${upgrade_file_size}"

        # Verify upgrade file
        if ! verify_md5 /tmp/upgrade.bin /tmp/upgrade.md5; then
            # Upgrade file not verified; exit
            $LOG "Firmware upgrade not verified. File may be corrupt"
            LED FAIL &
            return 1
        fi
        $LOG "Firmware upgrade verified"
        LED OFF && LED SUCCESS

        # Check battery state first
        $LOG "Checking device power state"
        battery_state=$(/usr/bin/BATTERY)
        if [ "${battery_state}" = "discharging" ] && [ ! -f "/etc/shark/cable" ]; then
            # Device is not plugged in
            $LOG "Device is not powered. Do not attempt firmware upgrade"
            return 1
        fi
        $LOG "Device is powered"

        # Remove upgrade file
        rm -rf "${upgrade_file}"
        sync

        # Upgrade file verified; run upgrade
        $LOG "Executing UPGRADE"
        sleep 2 && upgrade_leds &
        echo "sysupgrade -n /tmp/upgrade.bin" | at now

        exit
    else
        # Upgrade file not found; enter arming mode
        $LOG "Firmware upgrade not found"
        return 1
    fi
}

configure_network() {
    cp "/usr/lib/hak5/shark/config/${SWITCH_POSITION}/network" /etc/config/network
    /etc/init.d/network restart
}

start_http() {
    /etc/init.d/uhttpd start
}

stop_http() {
    /etc/init.d/uhttpd stop
}

start_ssh() {
    /etc/init.d/sshd start
}

stop_ssh() {
    /etc/init.d/sshd stop
}

start_dnsmasq() {
    #remove old dhcp options before setting new
    uci -q delete dhcp.lan.dhcp_option
    #disables gateway, netmode re-enables it
    uci add_list dhcp.lan.dhcp_option="3"
    uci commit
    /etc/init.d/dnsmasq start
}

stop_dnsmasq() {
    /etc/init.d/dnsmasq stop
}

enter_attack_mode() {
    $LOG "Entering ATTACK mode"
    MODE="ATTACK"
    /usr/bin/LED OFF

    stop_dnsmasq
    ip link set eth0 down
    macchanger -r eth0
    configure_network
    #nothing starts this, but ensure it only runs with
    pkill -9 udhcpc
    stop_http
    stop_ssh

    echo "execute_payload" | at now

    enter_idle_mode
}

enter_arming_mode() {
    $LOG "Entering ARMING mode"
    MODE="ARMING"
    /usr/bin/LED OFF

    configure_network
    start_dnsmasq
    start_http
    start_ssh

    enter_idle_mode
}

enter_off_mode() {
    $LOG "Entering OFF mode"
    MODE="OFF"
    /usr/bin/LED OFF

    /etc/init.d/network stop
    stop_dnsmasq
    stop_http
    stop_ssh

    enter_idle_mode
}

enter_idle_mode() {
    $LOG "Entering IDLE mode"
    while true
    do
        if [ "${MODE}" = "ARMING" ] || [ "${MODE}" = "OFF" ]; then
            #ARMING mode is basically always idle which means this overrides user LED
            #status once per second.  Don't override user status with idle status.
            BATTERY_STATE=$(/usr/bin/BATTERY)
            case $BATTERY_STATE in
                "charging")
                    pgrep LED > /dev/null || /usr/bin/LED B SLOW
                    ;;
                "full")
                    pgrep LED > /dev/null || /usr/bin/LED B
                    ;;
                "cable")
                    # Do nothing with battery if we're a sharkjack cable platform
                    ;;
                *)
                    pgrep LED > /dev/null || /usr/bin/LED Y SLOW
                    if [ "${MODE}" = "OFF" ]; then
                        #sleep 2
                        #halt
                        $LOG "battery state other ($BATTERY_STATE) would halt"
                    fi
                    ;;
            esac
        fi

        SWITCH_POSITION=$(/usr/bin/SWITCH)
        case $SWITCH_POSITION in
            "switch3")
                if [ "${MODE}" != "ATTACK" ]; then
                    enter_attack_mode
                fi
                ;;
            "switch2")
                if [ "${MODE}" != "ARMING" ]; then
                    execute_upgrade || enter_arming_mode
                fi
                ;;
            *)
                if [ "${MODE}" != "OFF" ]; then
                    enter_off_mode
                fi
                ;;
        esac
        sleep 1
    done
}


run() {
    case $SWITCH_POSITION in
        "switch3")
            enter_attack_mode
            ;;
        "switch2")
            execute_upgrade || enter_arming_mode
            ;;
        *)
            enter_off_mode
            ;;
    esac
}


# Start framework after a short wait
sleep 2
run > /dev/null 2>&1 &