#!/bin/bash

set -e  # exit on error

SCRIPT_DIR=$(dirname "$(readlink -f "$0")")
TARGET_FOLDER="$SCRIPT_DIR/rpi_image"
#IMG_FILE_BASENAME="2023-12-11-raspios-bookworm-armhf-lite"
#DOWNLOAD_URL="https://downloads.raspberrypi.com/raspios_lite_armhf/images/raspios_lite_armhf-2023-12-11/${IMG_FILE_BASENAME}.img.xz"
IMG_FILE_BASENAME="2026-09-15-raspios-trixie-arm64-lite"
DOWNLOAD_URL="https://downloads.raspberrypi.com/raspios_lite_arm64/images/raspios_lite_arm64-2026-09-15/${IMG_FILE_BASENAME}.img.xz"

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

read -p "Hostname of new pi [newrpi]: " RPI_HOSTNAME
read -p "WIFI SSID (only used if a password is given) [WLAN]: " WIFI_SSID
read -p "WIFI password (leave empty for non-wifi setup) [] " WIFI_PASSWORD
RPI_HOSTNAME=${RPI_HOSTNAME:-newrpi}
WIFI_SSID=${WIFI_SSID:-WLAN}

mkdir -p "$TARGET_FOLDER"
cd "$TARGET_FOLDER"

if [ -e "${IMG_FILE_BASENAME}.img" ]; then
   echo "${TARGET_FOLDER}/${IMG_FILE_BASENAME}.img already exists (already customised?)."
   echo "Remove it first if you want to start over."
   exit 1
fi

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

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

echo "Unpack image"
unxz ${IMG_FILE_BASENAME}.img.xz


echo "Mounting image"
MOUNT_DIR=$(mktemp -d "$TARGET_FOLDER/mounted_image.XXXXXX")
LOOP_DEV=""
cleanup() {
   set +e
   umount "$MOUNT_DIR/boot" "$MOUNT_DIR/system" 2>/dev/null
   [ -n "$LOOP_DEV" ] && losetup -d "$LOOP_DEV"
   rmdir "$MOUNT_DIR/boot" "$MOUNT_DIR/system" "$MOUNT_DIR" 2>/dev/null
}
trap cleanup EXIT

mkdir "$MOUNT_DIR/boot" "$MOUNT_DIR/system"
LOOP_DEV=$(losetup -f --show -P ${IMG_FILE_BASENAME}.img)
mount ${LOOP_DEV}p1 "$MOUNT_DIR/boot"    # firmware partition, /boot/firmware on the Pi
mount ${LOOP_DEV}p2 "$MOUNT_DIR/system"


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


if [ -z "$WIFI_PASSWORD" ]
then
   echo "Skipping WIFI setup"
elif [ -d "$MOUNT_DIR/system/etc/NetworkManager" ] && [ -e "$MOUNT_DIR/boot/user-data" ]
then
   # Trixie (and bookworm images from 2023-10 on, when they use NetworkManager) ignore
   # wpa_supplicant.conf on the boot partition; give NetworkManager a connection profile.
   echo "Setting up initial wifi config (NetworkManager)"
   NM_DIR="$MOUNT_DIR/system/etc/NetworkManager/system-connections"
   mkdir -p "$NM_DIR"
   /bin/cat <<EOM >"$NM_DIR/wifi.nmconnection"
[connection]
id=wifi
type=wifi
autoconnect=true

[wifi]
mode=infrastructure
ssid=${WIFI_SSID}

[wifi-security]
key-mgmt=wpa-psk
psk=${WIFI_PASSWORD}

[ipv4]
method=auto

[ipv6]
method=auto
EOM
   chmod 600 "$NM_DIR/wifi.nmconnection"  # NetworkManager ignores profiles readable by others
   # regulatory domain, so the 5GHz channels and the right power limits are usable at boot
   sed -i '1 s/$/ cfg80211.ieee80211_regdom=DE/' "$MOUNT_DIR/boot/cmdline.txt"
else
   echo "Setting up initial wifi config (wpa_supplicant)"
   /bin/cat <<EOM >"$MOUNT_DIR/boot/wpa_supplicant.conf"
country=DE
ctrl_interface=DIR=/var/run/wpa_supplicant GROUP=netdev
update_config=1

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

fi


echo "Unmounting image and other cleanup"
sleep 3 # to prevent "device is busy"
cleanup
trap - EXIT
rm ${IMG_FILE_BASENAME}.img.xz.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"