#!/usr/bin/env bash
# This script is used to setup various things to step a machine in l3. This
# script needs five parameters
# 1. Cobbler user's ldap password to be used to fetch ssh keys for the robot name
# 2. Robot name to be used for setting up the robot
# 3. Machine name to be used to determine if the machine is vision or main-ai
# 4. Total number of machines in the env
# 5. Compute major version: Major version of the compute
# 6. Compute_minor_version: Minor version of the compute
# 7. Operator Version: version of the operator
# 8. Robot has operator: 1 to indicate that operator exists and can be imaged as part of robot otherwise 0

# Note: Explicitly not setting a set -eou flag as this script should not return
# error

exec > /var/log/zooxpxeinstall.log 2>&1

vault_cobbler_password="$1"
robot_name="$2"
machine_name="$3"
total_machine_count="$4"
compute_major_version="$5"
compute_minor_version="$6"
operator_version="$7"
robot_has_operator="$8"
imaging_server="{{ inventory_hostname }}"

devuser=zoox
auto_login_robots=" {{ cobbler_auto_login_robots|join(' ') }} "

chown ${devuser}:${devuser} /var/log/zooxpxeinstall.log

http_server={{ cobbler_dhcp_router_ip }}
sshd_config=/etc/ssh/sshd_config
ssh_config_out=$(mktemp)
host=$(hostname)
vault_version=1.3.0
vault_publick_key_id=51852D87348FFC4C
vault_url="https://releases.hashicorp.com/vault/${vault_version}/vault_${vault_version}_linux_amd64.zip"
vault_shasum_url="https://releases.hashicorp.com/vault/${vault_version}/vault_${vault_version}_SHA256SUMS"
vault_signature_url="https://releases.hashicorp.com/vault/${vault_version}/vault_${vault_version}_SHA256SUMS.sig"
vault_sha256="d89b8a317831b06f2a32c56cb86071d058b09d9317b416bb509ce3d01e912eb3"
vault_finger_print="91A6 E7F8 5D05 C656 30BE F189 5185 2D87 348F FC4C"

ki_branch_info_url="{{ cobbler_ki_branch_info_url }}branch_info/${robot_name}"

ki_status_url="{{ cobbler_ki_status_url }}imaging_status/"

cobbler_vault_ssh_namespace="secret/robot/ssh"

auto_login_user_string=""

if [[ "${machine_name}" == "operator" && "${robot_has_operator}" == "0" ]];
then
    cobbler_vault_ssh_namespace="secret/operator/ssh"
fi

#################### install base software ####################

apt-get update
apt-get --assume-yes install linux-firmware {{ cobbler_linux_header }} ubuntu-desktop git \
    sshpass openssh-server vim-nox emacs24-nox lm-sensors htop iftop iotop curl cpufrequtils jq dpkg

if [[ "${compute_major_version}" == "3" && "${machine_name}" != "operator" ]];
then
    if [[ "${compute_minor_version}" == "2" ]];
    then

        mkdir --parents /etc/X11/
        cp /var/tmp/xorg.conf /etc/X11/xorg.conf
        apt-get --assume-yes install xserver-xorg-core-lts-trusty xserver-xorg-input-all \
            xserver-xorg xserver-xorg-lts-trusty
    fi
fi

# If it is main-ai and robot is in the list of auto login robots list then update the auto_login_user_string
if [[ "${auto_login_robots}" =~ " ${robot_name} " && "$machine_name" != "operator" && "$machine_name" != vision* ]];
then
    auto_login_user_string="autologin-user=zoox"
    # Set it as the default session
mkdir --parents /etc/lightdm/lightdm.conf.d
echo "[SeatDefaults]
${auto_login_user_string}
" > /etc/lightdm/lightdm.conf.d/50-myconfig.conf
fi


gpg --recv ${vault_publick_key_id}

curl -Os "$vault_url"
curl -Os "$vault_shasum_url"
curl -Os "$vault_signature_url"

gpg --fingerprint "$vault_finger_print" --verify vault_"${vault_version}"_SHA256SUMS.sig vault_"${vault_version}"_SHA256SUMS
grep linux_amd64.zip "vault_${vault_version}_SHA256SUMS" | shasum --algorithm 256 --check -

check_sum_result="$?"

if [[ ${check_sum_result} -ne 0 ]]; then
    echo "Invalid vault checksum found. Stopping PXE! "
    exit 0 # we should not return error otherwise pxe script would not be debuggable.
fi

unzip -o "vault_${vault_version}_linux_amd64.zip" -d /var/tmp/
vault_path="/var/tmp/vault"

#################### update grub robot, user and ssh details ####################
# Update Grub to make networking as eth* instead of em*
sed -i 's|GRUB_CMDLINE_LINUX_DEFAULT=\"\(.*\)\"|GRUB_CMDLINE_LINUX_DEFAULT=\"net.ifnames=0 biosdevname=0 \1\"|g' /etc/default/grub
update-grub

# Update robot name
if [[ "$machine_name" != "operator" || "${robot_has_operator}" == "1" ]];
then
    echo ${robot_name} | sudo tee /etc/robot_name
fi

# User sudo access
sudo adduser $devuser sudo
echo "$devuser ALL=NOPASSWD:ALL" > /etc/sudoers.d/$devuser

# Enable SSH Login via Password
cp $sshd_config "$sshd_config"-bk
sed -e '0,/PasswordAuthentication no/s/PasswordAuthentication no/PasswordAuthentication yes/' <"$sshd_config"> $ssh_config_out
mv $ssh_config_out $sshd_config

# Setup SSH for zoox user
sudo su ${devuser} --command "mkdir /home/$devuser/.ssh"
VAULT_ADDR=https://vault.zooxlabs.com:8200 ${vault_path} login -method=ldap username=cobbler password=${vault_cobbler_password}
VAULT_ADDR=https://vault.zooxlabs.com:8200 ${vault_path} read -field=ssh_public_key "${cobbler_vault_ssh_namespace}/${robot_name}" > /home/${devuser}/.ssh/id_rsa.pub
VAULT_ADDR=https://vault.zooxlabs.com:8200 ${vault_path} read -field=ssh_private_key "${cobbler_vault_ssh_namespace}/${robot_name}" > /home/${devuser}/.ssh/id_rsa
chown -R $devuser /home/$devuser/.ssh
chmod 644 /home/$devuser/.ssh/id_rsa.pub
chmod 600 /home/$devuser/.ssh/id_rsa
cat /home/$devuser/.ssh/id_rsa.pub > /home/$devuser/.ssh/authorized_keys

# Set AD password
mkdir /etc/zoox
VAULT_ADDR=https://vault.zooxlabs.com:8200 ${vault_path} read -field=password secret/ad_readonly_user > /etc/zoox/ad_readonly_password
chown -R $devuser:$devuser /etc/zoox/
chmod 600 /etc/zoox/ad_readonly_password
rm /root/.vault-token

sudo su ${devuser} --command "VAULT_ADDR=https://vault.zooxlabs.com:8200 ${vault_path} login -method=ldap username=cobbler password=${vault_cobbler_password}"

# Pull Setup Scripts from Github
driving_branch=$(curl --silent "${ki_branch_info_url}" | jq .branch | sed "s/\"//g")
echo "Using the driving branch ${driving_branch}"

sudo su ${devuser} --command "ssh-keyscan git.zooxlabs.com >> /home/$devuser/.ssh/known_hosts"

# Checkout the calibrations repo
sudo su ${devuser} --command "git clone git@git.zooxlabs.com:zooxco/calibration.git /etc/zoox/calibration"

# Checkout the driving repo
sudo su ${devuser} --command "git clone git@git.zooxlabs.com:zooxco/driving.git /home/$devuser/driving"
sudo su ${devuser} --command "cd /home/$devuser/driving; git checkout ${driving_branch}"
driving_sha=$(sudo su ${devuser} --command "cd /home/$devuser/driving; git rev-parse HEAD")

##################### run setup ####################

should_reboot="sudo reboot"
if [[ "${machine_name}" =~ ^vision.* ]];
then
    echo "##################### vision machine ###############################"
    mkdir -p /etc/udev/rules.d/
    cp /home/${devuser}/driving/scripts/ansible/roles/vision_network/files/10-ethernet-${compute_major_version}-${compute_minor_version}.rules /etc/udev/rules.d/10-ethernet.rules
    rm /etc/udev/rules.d/70-persistent-net.rules
    # vision should not reboot as main-ai would do the reboot.
    should_reboot=""

elif [[ "${machine_name}" == "operator" && "${robot_has_operator}" == "1" ]];
then

    echo "##################### robot's operator machine ###############################"
    mkdir -p /etc/udev/rules.d/
    cp /home/${devuser}/driving/scripts/ansible/roles/network_operator/files/10-ethernet-vulcan.rules /etc/udev/rules.d/10-ethernet.rules
    rm /etc/udev/rules.d/70-persistent-net.rules
    # operator should not reboot as main-ai would do the reboot.
    should_reboot=""

elif [[ "${compute_major_version}" == "3" && "${machine_name}" != "operator" ]];
then
    echo "##################### AI machine ###############################"
    mkdir -p /etc/udev/rules.d/

    if [[ "${compute_minor_version}" == "2" ]];
    then

        framegrabber_status=$(lspci | grep Xilinx | wc -l)

        if [[ "${framegrabber_status}" == "0" ]];
        then
            echo "##################### Robot does not have framegrabber ###############################"
            udev_rule_file_name="10-ethernet-3-nofg.rules"
        else
            echo "##################### Robot has framegrabber ###############################"
            udev_rule_file_name="10-ethernet-3-fg.rules"
        fi
    elif [[ "${compute_minor_version}" == "5" ]];
    then
        udev_rule_file_name="10-ethernet-3-5.rules"
    fi

    cp /home/${devuser}/driving/scripts/ansible/roles/network_ai/files/${udev_rule_file_name} /etc/udev/rules.d/10-ethernet.rules
    rm /etc/udev/rules.d/70-persistent-net.rules

fi

echo "#!/usr/bin/env bash
# No set -euo pipefail as this script should not fail
/var/tmp/zoox_post_init.sh \"${devuser}\" \"${robot_name}\" ${total_machine_count} \"${driving_sha}\" \"${compute_major_version}\" \"${compute_minor_version}\"  \"${operator_version}\" \"${robot_has_operator}\"
mv /etc/zoox_first_boot.sh /var/tmp/
${should_reboot}
" > /etc/zoox_first_boot.sh
chmod 0744 /etc/zoox_first_boot.sh
sed --in-place 's|^exit 0|[ -f /etc/zoox_first_boot.sh ] \&\& /etc/zoox_first_boot.sh;\nmodprobe nvidia_uvm;\nnvidia-smi;\nexit 0|g' /etc/rc.local

if [[ "$machine_name" != "operator" && "$machine_name" != vision* ]];
then

    # Update the KI status information as initialized
    curl  --verbose --request POST \
    --header "Content-Type: application/json" \
    --header "Accept: application/json" \
    --data "{\"robot\":\"${robot_name}\",\"imagingServer\":\"${imaging_server}\",\"status\":\"INITIALIZED\"}" \
    ${ki_status_url}
fi

echo "####################### Finished init ##################################"
