#!/usr/bin/env bash

# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2, June 1991.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.	See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/gpl-2.0.html>,
# or in pdf format at <http://www.dhampir.no/stuff/private/gpl-2.0.pdf>

# Copyright 2018 - Øyvind 'bolt' Hvidsten	 <bolt@dhampir.no>
_scriptname="create_pi_image"
set -u
set -e


# intro
cat >&2 <<-EOF

This script produces compressed Raspberry Pi (RPI) images from an sd card device
Currently, only clean Raspbian images are supported. No NOOBS.

EOF

function usage
{
cat >&2 <<-EOF

Usage: ${_scriptname} [options] <sda|sdb|...>

Options:
  -?
    Print this helpful help
  -7
    Don't use 7zip to compress of the resulting image to a .zip file (which is the default behaviour)
    Avoiding this is faster, and needs less space to do the job, as the image is saved directly with gzip compression
  -h <hostname>
    Set hostname inside the image to <hostname>
  -i <imagename>
    Set the name of the resulting image (the part before the timestamp)
    Defaults to "image"
  -k
    Wipe host SSH keys from the image
  -K
    Keep the uncompressed .img file as well as the compressed one
  -m
    Wipe machine-id from the image
  -s
    Single partition - used for vague compatibility with images for non-Raspberry flavored images
  -w <workdir>
    Override working directory (usually /tmp by default, from \$TMP)
  -y
    Assume yes to the initial sanity check question ("Do you want to use device <dev>..?")

EOF
}


# options?
assume_yes=false
ssh_key_wipe=false
machine_id_wipe=false
singlepart=false
keep_img=false
do7z=true
workdir=""
hostname=""
image=""
OPTIND=1
while getopts ":7h:i:kKmsw:y" opt; do
	case "$opt" in
		7) do7z=false ;;
		h) hostname="$OPTARG" ;;
		i) image="$OPTARG" ;;
		k) ssh_key_wipe=true ;;
		K) keep_img=true ;;
		m) machine_id_wipe=true ;;
		s) singlepart=true ;;
		y) assume_yes=true ;;
		w) workdir="$OPTARG" ;;
		?) usage; exit 0 ;;
		[?])
			echo "Unknown option -${OPTARG}" >&2
			exit 1
		;;
		:)
			echo "Option -${OPTARG} requires an argument" >&2
			exit 1
		;;
	esac
done
shift $((OPTIND-1))
if [[ "${1:-}" = "--" ]]; then shift; fi
if (( $# != 1 )); then
	usage
	exit 1
fi


# steps
color_reset="$(tput sgr0)"
color_white="${color_reset}$(tput bold)$(tput setaf 7)"
color_green="${color_reset}$(tput setaf 2)"
color_yellow="${color_reset}$(tput bold)$(tput setaf 3)"
function step { echo "${color_white}==> ${color_green}$*${color_reset}" >&2; }
function yell { echo "${color_yellow}$*${color_reset}" >&2; }


# check required tools
requirements=(
	dd
	resize2fs
	dumpe2fs
	e4defrag
	e2fsck
)
if $do7z; then
	requirements+=(	"7z" )
fi
missing=false
for req in "${requirements[@]}"; do
	if ! type "$req" >/dev/null 2>&1; then
		echo "Required command missing: $req" >&2
		missing=true
	fi
done
! $missing || exit 1


# select device
step "Interrogate device"
dev="$1"
orig_dev="$dev"
dev=$(realpath "/dev/${dev}" 2>/dev/null || :)
if [[ "$dev" = "/dev" ]] || ! [[ -e "$dev" ]]; then
	echo "No such device: $orig_dev" >&2
	exit 1
fi
echo "Selected device path: $dev" >&2
if ! $singlepart && { ! [[ -e "${dev}1" ]] || ! [[ -e "${dev}2" ]] || [[ -e "${dev}3" ]]; }; then
	echo "This does not look like a clean Raspbian block device: $dev" >&2
	echo "Tip: Do not enter sdb1, just sdb" >&2
	exit 2
elif $singlepart && { ! [[ -e "${dev}1" ]] || [[ -e "${dev}2" ]]; }; then
	echo "This does not look like a clean Single Partition block device: $dev" >&2
	echo "Tip: Do not enter sdb1, just sdb" >&2
	exit 2
elif $assume_yes; then
	echo "Not asking you to make sure $dev is correct, as '-y' was specified." >&2
else
	cat >&2 <<-EOF

		WARNING: This will alter the contents of the device.
		Make sure it's not your boot device, or anything else stupid.

	EOF
	read -r -p "Use '$dev'? [y/N] " answer
	case "${answer,,}" in
		y|yes) ;;
		*)
			echo "Abort mission." >&2
			exit 0
		;;
	esac
fi


# input and output devices
filename="${image:-image}-$(date +'%Y_%m_%d')"
echo -e "\nWill produce RPI image from: $dev"
echo -e "Output: ${filename}.zip\n"
if $singlepart; then
	boot="${dev}1"
	root="${dev}1"
else
	boot="${dev}1"
	root="${dev}2"
fi


# mount point, cache, and cleanup
tmpdir="${workdir:-"${HOME}/.config/${_scriptname}"}"
mkdir -p "$tmpdir"
mount=$(mktemp -d)
function cleanup
{
	umount "$mount" 2>/dev/null || :
	rmdir "$mount"
	[[ -n "$workdir" ]] || find "$tmpdir" -mindepth 1 -depth -delete
}
trap cleanup EXIT


# tmp should be empty
if find "$tmpdir" -mindepth 1 | grep -q "^"; then
	echo "Temporary directory \`${tmpdir}\` is not empty!" >&2
	exit 1
fi


# tune-up
step "Check root file system"
function check_root { e2fsck -f -y "$root" || e2fsck -f -y "$root"; } # try twice
check_root


# how to clean
function remove { [[ -e "$1" ]] || return 0; rm -r "$1"; echo "Removed:   ${1#$mount}"; }
function truncate { [[ -e "$1" ]] || return 0; >"$1";    echo "Truncated: ${1#$mount}"; }


# user files
step "Erase user files"
mount "$root" "$mount"
while read -r -d $'\0' file; do
	case "$file" in
		*"/.bash_history") truncate "$file" ;;
		*"/.vim") remove "$file" ;;
		*"/.nano") remove "$file" ;;
		*"/.lesshst") remove "$file" ;;
		*"/.wget-hsts") remove "$file" ;;
	esac
done < <(
	find "${mount}/root" -mindepth 1 -depth -print0
	find "${mount}/home" -mindepth 2 -depth -print0
)
umount "$mount"


# logs
step "Erase logs"
mount "$root" "$mount"
find "${mount}/var/log" -type f -print0 | while read -r -d $'\0' file; do
	if [[ "$file" =~ \.(gz|xz|tgz|old|[0-9]+)$ ]]; then
		remove "$file"
	else
		truncate "$file"
	fi
done
umount "$mount"


# tmp
step "Erase temporary files"
mount "$root" "$mount"
for dir in "${mount}/tmp" "${mount}/var/tmp" "${mount}/var/cache/apt/archives" "${mount}/var/lib/apt/lists"; do
	[[ -d "$dir" ]] || continue
	find "$dir" -mindepth 1 -depth -delete -print0 | while read -r -d $'\0' file; do
		echo "Removed: ${file#$mount}"
	done
done
umount "$mount"


# dhcp
step "Erase DHCP leases"
mount "$root" "$mount"
for dir in "${mount}/var/lib/dhcp"*; do
	[[ -d "$dir" ]] || continue
	find "$dir" -mindepth 1 -type f \( -name "*.lease" -o -name "*.leases" \) -delete -print0 |
		while read -r -d $'\0' file; do
			echo "Removed: ${file#$mount}"
		done
done
umount "$mount"


# special cases
step "Erase some more stuff"
mount "$root" "$mount"
if [[ -e "${mount}/opt/vc/src/hello_pi/hello_video/test.h264" ]]; then
	truncate "${mount}/opt/vc/src/hello_pi/hello_video/test.h264"
fi
truncate "${mount}/etc/resolv.conf"
umount "$mount"


# ssh keys
if $ssh_key_wipe; then
	step "Erase SSH server host keys"
	mount "$root" "$mount"
	find "${mount}/etc/ssh" -name "ssh_host_*key*" -type f -print0 | while read -r -d $'\0' file; do
		remove "$file"
	done
	umount "$mount"
fi


# machine id
if $machine_id_wipe; then
	step "Erase machine-id"
	mount "$root" "$mount"
	truncate "${mount}/etc/machine-id"
	remove "${mount}/var/lib/dbus/machine-id"
	umount "$mount"
fi


# reset hostname?
if [[ -n "$hostname" ]]; then
	step "Set hostname"
	mount "$root" "$mount"
	old="$(<"${mount}/etc/hostname")"
	sed -i "s/${old}/${hostname}/g" "${mount}/etc/hostname" "${mount}/etc/hosts"
	umount "$mount"
fi


# defrag
step "Defragment root file system"
mount "$root" "$mount"
e4defrag "$root" >/dev/null
umount "$mount"


# re-tune-up
step "Re-check root file system"
check_root


# resize!
step "Resize root file system"
resize2fs -M "$root"


# determine sizes
step "Resize root partition"
root_count=""
root_size=""
root_begin=""
root_sector=""
while read -r line; do
	case "${line,,}" in
		"block count: "*) root_count=${line##* } ;;
		"block size: "*) root_size=${line##* } ;;
	esac
done < <(dumpe2fs -h "$root")
while read -r line; do
	case "${line,,}" in
		"$root "*"*"*) root_begin=$(awk '{print $3}' <<<"$line") ;;
		"$root "*) root_begin=$(awk '{print $2}' <<<"$line") ;;
		"units: sectors of "*) root_sector=$(awk '{print $1}' <<< "${line##* = }") ;;
	esac
done < <(fdisk -l "${dev}")
if [[ -z "$root_count" ]] || [[ -z "$root_size" ]] || [[ -z "$root_begin" ]] || [[ -z "$root_sector" ]]; then
	echo "Unable to determine root file system size :(" >&2
	exit 1
fi
root_end=$((root_begin + (root_count*root_size/root_sector) - 1))
echo "Root FS block count:             $root_count"
echo "Root FS block size:              $root_size"
echo "Root FS begins at:               $root_begin"
echo "Root FS sector size:             $root_sector"
echo "Root FS partition should end at: $root_end"
if parted "${dev}" unit s resizepart "${root#"$dev"}" "${root_end}s" Yes | tee /dev/stderr | grep -q "Usage:"; then
	parted "${dev}" unit s resizepart "${root#"$dev"}" "${root_end}s" >&2
fi

sync


# re-tune
step "Re-check root file system"
check_root


# auto expanding
step "Re-enable auto expanding root file system"
mount "$root" "$mount"
if ! [[ -e "${mount}/usr/lib/raspi-config/init_resize.sh" ]]; then
	echo "No idea how to re-enable the auto expand on this system :(" >&2
else
	remove "${mount}/usr/local/etc/auto-expand-fs.count"
	cat >"${mount}/usr/local/sbin/auto-expand-fs" <<-"EOF"
		#!/usr/bin/env bash
		sleep 30
		root_dev=$(/bin/findmnt / -o source -n)
		store="/usr/local/etc/auto-expand-fs.count"
		>>"$store"
		count=$(<store)
		count=${count:-0}
		echo "$((++count))" >"$store"
		if (( count <= 3 )) && /sbin/resize2fs "$root_dev"; then
			/bin/rm -v "/usr/local/sbin/auto-expand-fs" "/etc/cron.d/auto-expand-fs"
			rm "$store"
			exit 0
		fi
		/sbin/shutdown -r now
	EOF
	cat >"${mount}/etc/cron.d/auto-expand-fs" <<-"EOF"
		@reboot root /bin/bash /usr/local/sbin/auto-expand-fs
	EOF
	umount "$mount"
	mount "$boot" "$mount"
	cmdline=$(<"${mount}/cmdline.txt")
	init='init=/usr/lib/raspi-config/init_resize.sh'
	if [[ "  $cmdline " != *" $init "* ]]; then
		cmdline+=" $init"
		printf '%s\n' "$cmdline" >"${mount}/cmdline.txt"
	fi
fi
umount "$mount"


# zero fill
step "Zero-fill free space"
for i in 1 2; do
	mount "${dev}$i" "$mount"
	cat /dev/zero >"${mount}/${_scriptname}-zero" 2>/dev/null || :
	rm "${mount}/${_scriptname}-zero"
	umount "$mount"
	! $singlepart || break
done
sync


# create image file
step "Create image file"
if $do7z; then
	dd if="${dev}" bs="$root_sector" count="$((root_end+1))" >"${tmpdir}/${filename}.img"

	# the user can carry on now...
	echo
	yell "We're done creating your image file now."
	yell "You may now remove the SD card and test it in an RPi, or eat it, while I compress this thing..."

	# keep the img?
	if $keep_img; then
		ln "${tmpdir}/${filename}.img" 2>/dev/null || cp "${tmpdir}/${filename}.img" ./
	fi

	# compress it
	7z a -tzip -mx=9 "${filename}.zip" "${tmpdir}/${filename}.img"
else
	dd if="${dev}" bs="$root_sector" count="$((root_end+1))" |
	{
		if $keep_img; then
			tee >"${tmpdir}/${filename}.img"
		else
			cat
		fi
	} |
	gzip >"${tmpdir}/${filename}.img.gz"

	# move to target
	if $keep_img; then
		mv "${tmpdir}/${filename}.img" ./
	fi
	mv "${tmpdir}/filename.img.gz" ./
fi


# woop!
echo
yell "Done!"



# vim: tabstop=4:softtabstop=4:shiftwidth=4:noexpandtab
