utils/raspberrypi_setup

#!/bin/bash
# Copyright 2021 The ChromiumOS Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

# A tool to help set up a new Raspberry Pi (with Raspberry Pi OS installed).
# This script is meant to be executed OUTSIDE the chroot in a regular Linux
# system environment.
# Usage: raspberrypi_setup PI_IP_ADDR ROOT_PASSWORD IF_PI, where:
#   PI_IP_ADDR: IP of the Raspberry Pi this script is setting up.
#   ROOT_PASSWORD: The root password will set to the Raspberry Pi.
#   IF_PI: This tells the script if it is initiated from the host, leave it
#          blank while execute from the host (autotest server).
# e.g., raspberrypi_setup 100.100.100.100 test0000

SSH_KEY_TARBALL="https://chromium.googlesource.com/chromiumos/chromite/+archive/main/ssh_keys.tar.gz"

PI_IP="$1"
IF_PI="$3"

prepare_host() {
    if ping -W 10 -c 1 "${PI_IP}"; then
        echo "Raspberry Pi Found - $(date)"
        # Buffer time to allow SSH up after pinable.
        sleep 5
    else
        echo "Failed to Ping Raspberry Pi"
        exit 1
    fi

    echo "Install required packets"
    sudo apt-get -y install sshpass > /dev/null 2>&1

    ssh-keygen -R pi@"${PI_IP}" > /dev/null 2>&1
}

prepare_pi() {
    echo "Copy setup files to the target Raspberry Pi"
    if ! sshpass -p 'raspberry' scp -o StrictHostKeyChecking=no \
        ~/.ssh/ssh_keys.tar.gz "$0" pi@"${PI_IP}":/tmp; then
        echo "Failed copy setup files to the target Raspberry Pi"
        exit 1
    fi
}

setup_host() {
    echo "Setup autotest server (local) SSH keys"

    if ! wget -O ~/.ssh/ssh_keys.tar.gz "${SSH_KEY_TARBALL}" > /dev/null 2>&1;
    then
        echo "Failed to download SSH keys"
        exit 1
    fi
    tar -xf ~/.ssh/ssh_keys.tar.gz -C ~/.ssh/
    chmod 0600 ~/.ssh/testing_rsa

    echo "Add Raspberry Pi to SSH config"
    cat << EOF >> ~/.ssh/config
HOST ${PI_IP}
    User root
    IdentityFile ~/.ssh/testing_rsa
EOF

    echo "Apply new SSH configuration"
    ssh-agent > /dev/null 2>&1
    ssh-add > /dev/null 2>&1
    ssh-add -l > /dev/null 2>&1
}

setup_pi() {
    PI_ROOT_PASSWORD="$1"

    echo "Updating SSH settings"
    systemctl enable ssh
    sed -i 's/.*#PermitRootLogin .*/PermitRootLogin yes/' /etc/ssh/sshd_config
    sed -i 's/.*#PubkeyAuthentication .*/PubkeyAuthentication yes/' \
        /etc/ssh/sshd_config
    sed -i 's/.*#PasswordAuthentication .*/PasswordAuthentication yes/' \
        /etc/ssh/sshd_config
    sed -i \
        's~.*AuthorizedKeysFile.*~AuthorizedKeysFile .ssh/authorized_keys~' \
        /etc/ssh/sshd_config

    echo "Updating root password"
    printf "%s\n%s" "${PI_ROOT_PASSWORD}" "${PI_ROOT_PASSWORD}" | passwd

    echo "Setup Raspberry Pi (remote) SSH keys"
    mkdir -p ~/.ssh

    tar -xf /tmp/ssh_keys.tar.gz -C ~/.ssh/
    chmod 0600 /root/.ssh/testing_rsa
    mv ~/.ssh/testing_rsa.pub ~/.ssh/authorized_keys
    rm /tmp/ssh_keys.tar.gz

    echo "Restart Raspberry Pi SSH service"
    service ssh restart
}

if [ -z "${IF_PI}" ]; then
    prepare_host
    setup_host
    prepare_pi
    sshpass -p "raspberry" ssh pi@"${PI_IP}" \
        "sudo sh /tmp/raspberrypi_setup ${1} ${2} yes"
    echo "Raspberry Pi Setup Finished"
else
    setup_pi "$2"
fi