#! /bin/bash
#
# Copyright (C) 2023 The Android Open Source Project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

set -e

ART_TEST_ON_VM=true . "$(dirname $0)/buildbot-utils.sh"
SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )

known_actions="create|boot|geniso|install-keys|setup-ssh|connect|quit"

if [[ -z $ANDROID_BUILD_TOP ]]; then
    msgfatal "ANDROID_BUILD_TOP is not set"
elif [[ ( $# -ne 1 ) || ! ( "$1" =~ ^($known_actions)$ ) ]]; then
    msgfatal "usage: $0 <$known_actions>"
fi

action="$1"

get_stable_binary() {
    mkdir tmp && cd tmp
    wget "http://security.ubuntu.com/ubuntu/pool/main/$1"
    7z x "$(basename $1)" && zstd -d data.tar.zst && tar -xf data.tar
    mv "$2" ..
    cd .. && rm -rf tmp
}

if [[ $action = create ]]; then
(
    rm -rf "$ART_TEST_VM_DIR"
    mkdir -p "$ART_TEST_VM_DIR"
    cd "$ART_TEST_VM_DIR"

    # sudo apt install qemu-system-<arch> qemu-efi cloud-image-utils

    # Get the cloud image for Ubunty 23.10 (Mantic Minotaur)
    wget "http://cloud-images.ubuntu.com/releases/23.10/release/$ART_TEST_VM_IMG"

    if [[ "$TARGET_ARCH" = "riscv64" ]]; then
        # Get U-Boot for Ubuntu 22.04 (Jammy)
        get_stable_binary \
            u/u-boot/u-boot-qemu_2024.01+dfsg-5ubuntu2_all.deb \
            usr/lib/u-boot/qemu-riscv64_smode/uboot.elf

        # Get OpenSBI for Ubuntu 22.04 (Jammy)
        get_stable_binary \
            o/opensbi/opensbi_1.3-1ubuntu0.23.04.2_all.deb \
            usr/lib/riscv64-linux-gnu/opensbi/generic/fw_jump.elf

    elif [[ "$TARGET_ARCH" = "arm64" ]]; then
        # Get EFI (ARM64)
        get_stable_binary \
            e/edk2/qemu-efi-aarch64_2024.05-2ubuntu0.1_all.deb \
            usr/share/qemu-efi-aarch64/QEMU_EFI.fd

        dd if=/dev/zero of=flash0.img bs=1M count=64
        dd if=QEMU_EFI.fd of=flash0.img conv=notrunc
        dd if=/dev/zero of=flash1.img bs=1M count=64
    fi

    qemu-img resize "$ART_TEST_VM_IMG" +128G
)
elif [[ $action = geniso ]]; then
(
    #https://help.ubuntu.com/community/CloudInit
    cat >user-data <<EOF
#cloud-config
ssh_pwauth: true
chpasswd:
  expire: false
  users:
    - name: $ART_TEST_SSH_USER
      password: ubuntu
      type: text
users:
  - default
  - name: $ART_TEST_SSH_USER
    ssh-authorized-keys:
      - ssh-rsa [payload removed: 544 characters] ubuntu_user@example.com
    sudo: ALL=(ALL) NOPASSWD:ALL
    groups: users, admin
EOF
    # meta-data is necessary, even if empty.
    cat >meta-data <<EOF
EOF
    genisoimage -output user-data.img -volid cidata -joliet -rock user-data meta-data
    mv user-data.img "$(dirname $0)/user-data.img"
    rm user-data meta-data
)
elif [[ $action = boot ]]; then
(
    cp "$(dirname $0)/user-data.img" "$ART_TEST_VM_DIR/user-data.img"
    cd "$ART_TEST_VM_DIR"
    if [[ "$TARGET_ARCH" = "riscv64" ]]; then
        ($ANDROID_BUILD_TOP/device/google/cuttlefish_vmm/qemu/x86_64-linux-gnu/bin/qemu-system-riscv64 \
            -M virt \
            -nographic \
            -m 16G \
            -smp 8 \
            -cpu rv64,v=true,elen=64,vlen=128,zba=true,zbb=true,zbs=true \
            -bios fw_jump.elf \
            -kernel uboot.elf \
            -drive file="$ART_TEST_VM_IMG",if=virtio \
            -drive file=user-data.img,format=raw,if=virtio \
            -device virtio-net-device,netdev=usernet \
            -netdev user,id=usernet,hostfwd=tcp::$ART_TEST_SSH_PORT-:22 > $SCRIPT_DIR/boot.out &)
        echo "Now listening for successful boot"
        finish_str='.*finished at.*'
        while IFS= read -d $'\0' -n 1 a ; do
            line+="${a}"
            if [[ "$line" =~ $finish_str ]] ; then
                echo $line
                echo "VM Successfully booted!"
                exit 0
            elif [[ $a = $'\n' ]]
            then
                echo $line
                unset line
            fi
        done < <(tail -f $SCRIPT_DIR/boot.out)

    elif [[ "$TARGET_ARCH" = "arm64" ]]; then
        (qemu-system-aarch64 \
            -m 16G \
            -smp 8 \
            -cpu cortex-a710,sve=on \
            -M virt \
            -nographic \
            -drive if=none,file="$ART_TEST_VM_IMG",id=hd0 \
            -pflash flash0.img \
            -pflash flash1.img \
            -drive file=user-data.img,format=raw,id=cloud \
            -device virtio-blk-device,drive=hd0 \
            -device virtio-net-device,netdev=usernet \
            -netdev user,id=usernet,hostfwd=tcp::$ART_TEST_SSH_PORT-:22 > $SCRIPT_DIR/boot.out &)
        echo "Now listening for successful boot"
        finish_str='.*finished at.*'
        while IFS= read -d $'\0' -n 1 a ; do
            line+="${a}"
            if [[ "$line" =~ $finish_str ]] ; then
                echo $line
                echo "VM Successfully booted!"
                exit 0
            elif [[ $a = $'\n' ]]
            then
                echo $line
                unset line
            fi
        done < <(tail -f $SCRIPT_DIR/boot.out)
    fi

)
elif [[ $action = setup-ssh ]]; then
    # Clean up mentions of this VM from known_hosts
    sed -i -E "/\[$ART_TEST_SSH_HOST.*\]:$ART_TEST_SSH_PORT .*/d" $HOME/.ssh/known_hosts
    ssh-copy-id -p "$ART_TEST_SSH_PORT" -o IdentityAgent=none -o StrictHostKeyChecking=no "$ART_TEST_SSH_USER@$ART_TEST_SSH_HOST"

elif [[ $action = install-keys ]]; then
    if [ -f "$HOME/.ssh/known_hosts" ]; then
        sed -i -E "/\[$ART_TEST_SSH_HOST.*\]:$ART_TEST_SSH_PORT .*/d" $HOME/.ssh/known_hosts
    fi
    # This key is only used to authorize access to a local test VM and does
    # not pose any security risk.
    echo "ssh-rsa [payload removed: 544 characters] ubuntu_user@example.com" > ~/.ssh/ubuntu.pub
    echo "-----BEGIN OPENSSH PRIVATE KEY-----
[payload removed: 2516 characters in 36 lines]
-----END OPENSSH PRIVATE KEY-----" > ~/.ssh/ubuntu
    chmod 600 ~/.ssh/ubuntu

elif [[ $action = connect ]]; then
    $ART_SSH_CMD

elif [[ $action = quit ]]; then
    $ART_SSH_CMD "sudo poweroff"

fi
