#!/bin/bash
set -e

hostname=pi
ssh_key="$HOME/.ssh/id_rsa.pub"
release=stretch
rootfs=

print_usage() {
    cat <<EOF
Usage: $0 [options] <rootfs_directory>

Create a bare debian root filesystem for the Raspberry Pi 2 B. The
password for the root account will be set to 'guest' by default.

Options:
  --hostname <hostname>  Use a given hostname for the device. Defaults to '$hostname'.
  --ssh-key <key file>	 Add the given SSH key to root's authorized_keys file. Defaults to '$ssh_key'.
  --release <release>    Create a root filesystem for the given debian release. Defaults to '$release'.
  --no-debootstrap       Don't create an initial file system with debootstrap. This may be useful for debugging purposes
EOF
}

fail() {
    echo "$1" >&2
    exit 1
}

log() {
    echo "mkrootfs: $1" >&2
}

# Process arguments
while [ $# -gt 0 ]; do
    case "$1" in
	--help|-h)
	    print_usage
	    exit 0
	    ;;	
	--hostname)
	    shift
	    [[ -z "$1" ]] && fail "no hostname specified"
	    hostname="$1"
	    ;;
	--release)
	    shift
	    [[ -z "$1" ]] && fail "no release specified"
	    release="$1"
	    ;;
	--ssh-key)
	    shift
	    [[ -z "$1" ]] && fail "no ssh-key specified"
	    ssh_key="$1"
	    ;;
	--no-debootstrap)
	    nodeboostrap=true
	    ;;
	*)
	    [[ -n "$rootfs" ]] && fail "rootfs already defined"
	    rootfs="$1"
	    ;;
    esac
    shift
done

[ -n "$rootfs" ] || fail "$(print_usage)"
[ ! -e "$rootfs" ] || fail "$rootfs already exists"
[ "$EUID" -eq 0 ] || fail "$0 must be run as root"

if [ -n "$nodebootstrap" ]; then
    echo "skipping deboostrap" >&2
else
    log "debootstrapping"
    if  [ $(dpkg -l | grep -c apt-cacher-ng) -gt 0 ]; then
	export http_proxy="http://localhost:3142"
    else
	echo "WARNING: running debootstrap will download a lot of data. It is recommended to use apt-cacher-ng." >&2
    fi

    qemu-debootstrap \
	--verbose \
	--include=wget,curl,openssl,ca-certificates \
	--arch=armhf \
	"$release" \
	"$rootfs" \
	http://httpredir.debian.org/debian

    log "debootstrap complete"
    # http_proxy conflicts with direct calls to apt-get
    #unset http_proxy
fi

# Run a command in the root file system
chroot_exec() {
    LANG=C LC_ALL=C DEBIAN_FRONTEND=noninteractive chroot "$rootfs" "$@"
}

# Unmount and remove any temporary files
cleanup() {
    log "cleaning"
    set +e
    fuser --kill -SIGTERM "$rootfs"
    sleep 2
    fuser --kill -SIGKILL --verbose "$rootfs"
    umount -l "$rootfs/proc" 2> /dev/null
    umount -l "$rootfs/sys" 2> /dev/null
    umount -l "$rootfs/dev/pts" 2> /dev/null
    umount -l "$rootfs/dev" 2> /dev/null
    trap - 0 1 2 3 6
}
trap cleanup 0 1 2 3 6

# Prepare filesystem to be chrooted into
log "creating temporary system mount points"
mount -t proc chproc "$rootfs/proc"
mount -t sysfs chsys "$rootfs/sys"
mount -t devtmpfs chdev "$rootfs/dev"
mount -t devpts chpts "$rootfs/dev/pts"

# Set up package archive
log "configuring apt archive"
cat > "$rootfs/etc/apt/sources.list" <<EOF
deb http://httpredir.debian.org/debian $release main contrib non-free
deb http://httpredir.debian.org/debian $release-updates main contrib non-free
deb http://security.debian.org $release/updates main contrib non-free
EOF
chroot_exec apt-get update

# Configure fstab
log "configuring fstab"
cat > "$rootfs/etc/fstab" <<EOF
/dev/mmcblk0p1 /boot/firmware vfat  noatime  0 2
/dev/mmcblk0p2 /              ext4  noatime  0 1
tmpfs          /tmp           tmpfs defaults 0 0
EOF

# Install proprietary 2nd-stage bootloader firmware
log "installing proprietary 2nd-stage bootloader"
mkdir -p "$rootfs/boot/firmware"
wget https://github.com/raspberrypi/firmware/raw/master/boot/bootcode.bin \
     -O "$rootfs/boot/firmware/bootcode.bin"
wget https://github.com/raspberrypi/firmware/raw/master/boot/start.elf \
     -O "$rootfs/boot/firmware/start.elf"

# Install u-boot 3rd-stage bootloader
# https://blog.night-shade.org.uk/2015/05/booting-a-raspberry-pi2-with-u-boot-and-hyp-enabled/
log "installing u-boot 3rd stage bootloader"
chroot_exec apt-get install -y u-boot-rpi
cp "$rootfs/usr/lib/u-boot/rpi_2/u-boot.bin" "$rootfs/boot/firmware/u-boot.bin"
cat <<EOF > "$rootfs/boot/firmware/config.txt"
gpu_mem=64
kernel=u-boot.bin
EOF

# Install and configure kernel flashing utility
# http://sjoerd.luon.net/posts/2015/02/debian-jessie-on-rpi2/
chroot_exec apt-get install -y flash-kernel u-boot-tools
echo "Raspberry Pi 2 Model B" > "$rootfs/etc/flash-kernel/machine"
cat <<EOF > "$rootfs/etc/flash-kernel/ubootenv.d/10-partition"
setenv distro_bootpart 2
setenv prefix /boot/
EOF
echo "Machine: Raspberry Pi 2 Model B" >> "$rootfs/etc/flash-kernel/db"
echo "Boot-Script-Path: /boot/firmware/boot.scr" >> "$rootfs/etc/flash-kernel/db"

# Install kernel
log "installing kernel"
chroot_exec apt-get install -y linux-image-armmp-lpae

# Configure kernel boot args
log "setting kernel boot arguments"
cat <<EOF > "$rootfs/etc/default/flash-kernel"
LINUX_KERNEL_CMDLINE="quiet"
LINUX_KERNEL_CMDLINE_DEFAULTS="console=ttyAMA0 root=/dev/mmcblk0p2"
EOF
chroot_exec flash-kernel

# Set up hostname configuration
log "configuring network host names"
echo "$hostname" > "$rootfs/etc/hostname"
sed -i "/^127.0.0.1/ s/\$/ $hostname/" "$rootfs/etc/hosts"
sed -i "/^::1/ s/\$/ $hostname/" "$rootfs/etc/hosts"

# Set up network interfaces
log "configuring network interfaces"
cat > "$rootfs/etc/network/interfaces.d/lo" <<EOF
auto lo
iface lo inet loopback
EOF
cat > "$rootfs/etc/network/interfaces.d/eth0" <<EOF
allow-hotplug eth0
iface eth0 inet dhcp
EOF

# Set up root account
log "creating root account"
chroot_exec sh -c "echo root:guest | chpasswd"

# Install OpenSSH
log "setting up openssh-server"
chroot_exec apt-get -y install openssh-server
if [ -f "$ssh_key" ]; then
    mkdir -p "$rootfs/root/.ssh"
    chmod 600 "$rootfs/root/.ssh"
    cat "$ssh_key" > "$rootfs/root/.ssh/authorized_keys"
    chmod 600 "$rootfs/root/.ssh/authorized_keys"
    sed -i 's/#PasswordAuthentication yes/PasswordAuthentication no/g' "$rootfs/etc/ssh/sshd_config"
    echo "SSH password login disabled. Use the key in $ssh_key to login." 1>&2
else
    sed -i 's/PermitRootLogin prohibit-password/PermitRootLogin yes/g' "$rootfs/etc/ssh/sshd_config"
    echo "Warning: SSH password login is enabled! This can be a security risk." 1>&2
fi

# Prepare first boot
log "setting up first boot script"
chroot_exec apt-get -y install parted
cat <<'EOM' > "$rootfs/etc/rc.firstboot"
#!/bin/bash
set -e
logger -t "rc.firstboot" "Starting first boot"

led_blink() {
    local led=/sys/class/leds/"$1"

    local time_on="$2"
    local time_off="$time_on"
    local blinks="$3"

    local trigger=$(cat "$led/trigger" | sed 's/.*\[\(.*\)\].*/\1/g' -)
    local brightness=$(cat "$led/brightness")

    echo "none" > "$led/trigger"
    for i in $(seq 1 $blinks); do
        echo 255 > "$led/brightness"
        sleep "$time_on"
	echo 0 > "$led/brightness"
        sleep "$time_off"
    done
    echo "$brightness" > "$led/brightness"
    echo "$trigger" > "$led/trigger"
}

led_trigger() {
    local led=/sys/class/leds/"$1"

    if [ "$2" == "heartbeat" ]; then
        modprobe ledtrig_heartbeat || true
    fi
    echo "$2" > "$led/trigger"
}

led_blink ACT 0.1 20 &
led_blink PWR 0.1 20
led_trigger PWR heartbeat
led_trigger ACT mmc0

logger -t "rc.firstboot" "Generating ssh host keys"
rm -f /etc/ssh/ssh_host_*
systemctl stop sshd
ssh-keygen -q -t rsa -N "" -f /etc/ssh/ssh_host_rsa_key
ssh-keygen -q -t dsa -N "" -f /etc/ssh/ssh_host_dsa_key
ssh-keygen -q -t ecdsa -N "" -f /etc/ssh/ssh_host_ecdsa_key
ssh-keygen -q -t ed25519 -N "" -f /etc/ssh/ssh_host_ed25519_key
systemctl start sshd

logger -t "rc.firstboot" "Expanding root file system"
fdisk /dev/mmcblk0 <<EOF || true
p
d
2
n
p
2


n
w
EOF
partprobe
resize2fs /dev/mmcblk0p2

logger -t "rc.firstboot" "Generating D-Bus machine-id"
rm -f /var/lib/dbus/machine-id 
dbus-uuidgen --ensure

logger -t "rc.firstboot" "Regenerating initramfs"
update-initramfs -u

logger -t "rc.firstboot" "Restarting core services"
systemctl daemon-reload
systemctl restart networking.service
systemctl restart systemd-networkd.service

logger -t "rc.firstboot" "First boot actions finished"
rm -f /etc/rc.firstboot
sed -i '/.*rc.firstboot/d' /etc/rc.local

led_trigger ACT heartbeat
led_trigger PWR none
EOM
chmod +x "$rootfs/etc/rc.firstboot"

cat <<EOF > "$rootfs/etc/rc.local"
#!/bin/sh -e
/etc/rc.firstboot
exit 0
EOF
chmod +x "$rootfs/etc/rc.local"

cleanup

log "done"