#!/bin/bash

set -e  # exit on error

TARGET_FOLDER="./rpi_image"
IMG_FILE_BASENAME="2021-10-30-raspios-bullseye-armhf-lite"
DOWNLOAD_URL="https://downloads.raspberrypi.org/raspios_lite_armhf/images/raspios_lite_armhf-2021-11-08/${IMG_FILE_BASENAME}.zip"

echo "This script downloads raspian lite, and modifies the image to  enable SSH and set hostname"

if [[ $EUID -ne 0 ]]; then
   echo "This script must be run as root"
   exit 1
fi

echo -n Hostname of new pi
read -p "Hostname of new pi [newrpi]: " RPI_HOSTNAME
read -p "WIFI password (leave empty for non-wifi setup) [] " WIFI_PASSWORD
RPI_HOSTNAME=${RPI_HOSTNAME:-newrpi}

SCRIPT_DIR=`pwd`
mkdir -p $TARGET_FOLDER
cd $TARGET_FOLDER


echo "Downloading image"
wget ${DOWNLOAD_URL}
wget ${DOWNLOAD_URL}.sha256

echo "Checksum verification"
sha256sum -c ${IMG_FILE_BASENAME}.zip.sha256

echo "Unpack image"
unzip ${IMG_FILE_BASENAME}.zip
rm ${IMG_FILE_BASENAME}.zip


echo "Mounting image"
mkdir mounted_image
mkdir mounted_image/boot
mkdir mounted_image/system
losetup -P /dev/loop42 ${IMG_FILE_BASENAME}.img
mount /dev/loop42p1 mounted_image/boot
mount /dev/loop42p2 mounted_image/system


echo "Enabling SSH and writing hostname"
echo $RPI_HOSTNAME > mounted_image/system/etc/hostname
touch mounted_image/boot/ssh  # startup ssh
sed -i "/^#PermitRootLogin/ cPermitRootLogin prohibit-password" mounted_image/system/etc/ssh/sshd_config
mkdir -p mounted_image/system/root/.ssh
cat $SCRIPT_DIR/*.pub > mounted_image/system/root/.ssh/authorized_keys
chmod 700 mounted_image/system/root/.ssh
chmod 600 mounted_image/system/root/.ssh/authorized_keys


if [ -z "$WIFI_PASSWORD" ]
then
   echo "Skipping WIFI setup"
else
   echo "Setting up initial wifi config"
   /bin/cat <<EOM >mounted_image/boot/wpa_supplicant.conf
country=DE
ctrl_interface=DIR=/var/run/wpa_supplicant GROUP=netdev
update_config=1

network={
   ssid="WLAN"
   psk="${WIFI_PASSWORD}"
}
EOM

fi


echo "Unmounting image and other cleanup"
sleep 3 # to prevent "device is busy"
umount /dev/loop42p1
umount /dev/loop42p2
losetup -d /dev/loop42
rmdir mounted_image/boot
rmdir mounted_image/system
rmdir mounted_image
rm ${IMG_FILE_BASENAME}.zip.sha256
chmod a+rw ${IMG_FILE_BASENAME}.img .

echo ""
echo ""
echo "The image is ready in folder ${TARGET_FOLDER}"
echo "copy to SD card with"
echo "  dd bs=4M status=progress if=the_image of=/dev/your/sdcard"
echo "  use e.g. /dev/sdb not /dev/sdb1"