#!/bin/ksh

set -eu

# 1.1.4
version=master

red="\\033[01;31m"
yellow="\\033[01;33m"
green="\\033[01;32m"
bold="\\033[01;39m"
white="\\033[0m"

COLOR=true

tmp_template="snap.XXXXXXXXXX"

sets=""
xsets=""
bsds=""

set -A sets 'comp' 'game' 'man' 'base'
set -A xsets 'xbase' 'xfont' 'xserv' 'xshare'
base_sig='SHA256.sig'
build_info='BUILDINFO'

spath=$(dirname -- "$(command -v -- "$0")")
sname="${0##*/}"

usage() {
	cat <<EOF

snap options:

  -S do not check signatures.
  -C only remove partial downloads.
  -c specify location of config file (default is ~/.snaprc)
  -e just extract sets in DST.
  -d just download sets to DST, verify and exit.
  -m <machine> use <machine> instead of what 'machine' returns.
  -V <setversion> used to force snap to use <setversion> for sets (example: -V 5.3). Note: this will only append 53 to sets, ie base53.tgz.
  -r run sysmerge after extracting {x}sets. (May dump core if the snapshots have introduced ABI changes. Not recommended.)
  -x do not extract x11 sets.
  -M specify a mirror to use (example: " -M cdn.openbsd.org")
  -I [full path to SHA256.sig file] verify integrity of snap.
  -s [signify pub key] pub key to do verification with.
  -i interactive with colors.
  -n force using bsd.mp as bsd.
  -k only install kernels and exit.
  -K install only the sets without the kernel.
  -B do not backup current kernel.
  -u check for update to snap script.
  -U download new snap script (will replace currently installed version).
  -b device to install bootstrap to.
  -R reboot after installation.
  -W exit on warn
  -h help.

 Examples:

   To upgrade to the latest snapshots:

     $ doas snap

   To update to the latest snapshot using an explicit mirror
   region:

     $ doas snap -M cdn.openbsd.org

   To update to the snapshot without updating xsets:

     $ doas snap -x

   When a new beta is cut, the system version jumps from X.Y to X.Z.
   When this happens, snap will need to be told what the new version
   is:

     $ doas snap -V 6.1

 Example ~/.snaprc

   INTERACTIVE:true
   DST:/tmp/upgrade
   MERGE:true
   MIRROR:cdn.openbsd.org
   NO_X11:true

EOF
	exit 0
}

get_conf_var() {
	RET=''
	if [ -e $CONF_FILE ]; then
		RET=$( grep $1 $CONF_FILE | awk -F : '{if ($1 !~ /^#/) {print $2}}' )
	fi

	if [ "${RET}X" == "X" ]; then
		return 1
	else
		echo $RET
	fi
}

msg() {
	if [ $INTERACTIVE == true ]; then
		echo "${green}${1}${white}"
	fi
}

warn() {
	if [[ $INTERACTIVE == true ]] || [[ $WEXIT == true ]]; then
		echo "${yellow}${1}${white}"
		if [ $WEXIT == true ]; then
			exit 1
		fi
	fi
}

error() {
	if [[ $INTERACTIVE == true ]]; then
		>&2 echo "${red}${1}${white}"
	else
		>&2 echo "${sname}: $1"
	fi

	if [[ $2 == true ]]; then
		rollback
	fi

	exit 1
}

check_update() {
	if [ "${version}" != "master" ]; then
		R="https://api.github.com/repos/qbit/snap/releases/latest"
		LATEST=$(/usr/bin/ftp -o - $R | \
			awk -F , '{for(i=1;i<NF;i++){if(match($i, /tag_name/)){print $i}}}' | \
			cut -d: -f 2 | \
			tr -d \")
		rversion=$(echo $LATEST | tr -d \.)
		lversion=$(echo $version | tr -d \.)
		if [ $rversion -gt $lversion ]; then
			msg "${white}Update ${green}${LATEST}${white} available for ${bold}snap${white}!"

			if [ $INS_UPDATE == true ]; then
				msg "installing update.."
				install_update $LATEST
			fi
		else
			msg "snap-${version} ${white}is up to date!"
		fi
		else
			if [ $INS_UPDATE == true ]; then
				msg "installing update.."
				install_update $version
			fi
	fi
}

install_update() {
	sver=$1
	tmp_file=$(mktemp)
	/usr/bin/ftp $FTP_OPTS -o "$tmp_file" \
		"https://raw.githubusercontent.com/qbit/snap/${sver}/snap"
	cp "$tmp_file" "$spath/$sname"
	msg "Installed new version (${sver}) of snap - please re-run!"
	exit 0
}

backup() {
	FAIL=0
	cp /bsd /obsd || FAIL=1
	cp /bsd.rd /obsd.rd || FAIL=1
	cp /sbin/reboot /sbin/oreboot || FAIL=1

	if [ -e /bsd.sp ]; then
		cp /bsd.sp /obsd.sp || FAIL=1
	fi

	if [ $FAIL == 1 ]; then
		error "Failed to backup" false
	else
		msg "Backed up the following:
	/bsd => /obsd
	/bsd.rd => /obsd.rd
	/sbin/reboot => /sbin/oreboot"
	fi
}

rollback() {
	FAIL=0
	cp /obsd /bsd || FAIL=1
	cp /obsd.rd /bsd.rd || FAIL=1

	if [ -e /obsd.sp ]; then
		cp /obsd.sp /bsd.sp || FAIL=1
	fi

	cp /sbin/oreboot /sbin/reboot || FAIL=1

	if [ $FAIL == 1 ]; then
		error "Failed to rollback" false
	else
		msg "Restored the old files for the following:
	/bsd => /obsd
	/bsd.rd => /obsd.rd
	/sbin/reboot => /sbin/oreboot"
	fi

}

check_integ() {
	key=${key:-/etc/signify/snap.pub}
	file=${sname}

	# signify doesn't like leading ./'s
	#file=${file:S/\.\///}

	if [ ! -f "${key}" ]; then
		error "No public key (${key}).\nSee https://github.com/qbit/snap for more info!" false
	fi

	if [ ! -f "${INTEG_SIG_FILE}" ]; then
		tmp_file=$(mktemp)
		/usr/bin/ftp $FTP_OPTS -o "$tmp_file" \
			"https://raw.githubusercontent.com/qbit/snap/${version}/SHA256.sig"
		INTEG_SIG_FILE=${tmp_file}
	fi

	(
		# we need to be in the same directory as snap to verify, as that is what
		# is in SHA256.sig
		cd $spath && \
			signify -C -p "${key}" -x "${INTEG_SIG_FILE}" "${file}"
	)
}

verisigs() {
	KEY=${KEY:-/etc/signify/openbsd-${SETVER}-base.pub}
	VALID=true

	if [ -f "$KEY" ]; then
		for i in "$@"; do
			signify -V -e -p ${KEY} -x SHA256.sig -m - | sha256 -C - ${i} \
				|| VALID=false
		done

		if [ $VALID == false ]; then
			error "Invalid signature found! They are after you!" true
		fi
	else
		error "No pub key found for this release! (${KEY})" false
	fi
}

update_kernel() {
	FAIL=0
	if [ $SKIP_SIGN == false ]; then
		verisigs "bsd*"
	fi
	cp ${KERNEL} /bsd || FAIL=1
	cp ${RD} /bsd.rd || FAIL=1

	(umask 077; sha256 /bsd > /var/db/kernel.SHA256)

	if [ "${KERNEL}" == "bsd.mp" ]; then
		cp bsd /bsd.sp || FAIL=1
	fi

	if [ $FAIL == 1 ]; then
		error "Failed to copy new kernel" false
	else
		msg "Set primary kernel to ${KERNEL}:
	${KERNEL} => /bsd"
	fi
}

fetch() {
	DF=$(echo $1 | awk -F/ '{print $NF}')
	TDF="${DF}.out"
	R=0

	# this check may cause signature issues.. if old files exist in
	# the DEST directory.
	if [ ! -e $DF ]; then
		su -s/bin/sh _pkgfetch -c "/usr/bin/ftp $FTP_OPTS -o $TDF $1"
		R=$?

		# move the tmp file to actual file name so we can use -C
		mv "$TDF" "$DF"
		chown root:wheel "$DF"
	fi

	return $R
}

extract() {
	ftp -D Extracting -Vmo - "file://${1}" | tar -C / -xzphf - \
		|| error "Failed to extract ${1}" false
}

CONF_FILE="/etc/snap.conf"
if [ -e ~/.snaprc ]; then
	CONF_FILE=~/.snaprc
fi

COLOR=$(get_conf_var 'COLOR' || echo 'true')
SKIP_SIGN=false
USE_BUILDINFO=true
CPUS=$(sysctl -n hw.ncpu)
INTERACTIVE=$(get_conf_var 'INTERACTIVE' || echo 'false')
DST=$(get_conf_var 'DST' || mktemp -d -t ${tmp_template})
EXTRACT_ONLY=$(get_conf_var 'EXTRACT_ONLY' || echo 'false')
KERNEL_ONLY=false
SETS_ONLY=false
FTP_OPTS=$(get_conf_var 'FTP_OPTS' || echo " -V ")
MACHINE=$(machine)
MERGE=$(get_conf_var 'MERGE' || echo 'false')
NO_X11=$(get_conf_var 'NO_X11' || echo 'false')
SETVER=$(uname -r | tr -d \.)
CHK_INTEG=false
CHK_UPDATE=$(get_conf_var 'CHK_UPDATE' || echo 'false')
INS_UPDATE=$(get_conf_var 'INS_UPDATE' || echo 'false')
INSTBOOT=$(get_conf_var 'INSTBOOT' || echo 'false')
REBOOT=$(get_conf_var 'REBOOT' || echo 'false')
AFTER=$(get_conf_var 'AFTER' || echo 'false')
DOWNLOAD_ONLY=false
WEXIT=$(get_conf_var 'WEXIT' || echo 'false')
CLEAN_ONLY=false

if [ $COLOR == false ]; then
	green=$white
	red=$white
	yellow=$white
	bold=$white
fi

MIRROR=$(get_conf_var 'MIRROR' || \
	awk -F/ 'match($3, /[a-z]/) {print $3}' /etc/installurl 2> /dev/null || \
	echo "cdn.openbsd.org")

while getopts "b:BCc:dD:ehiIkKm:M:nrRSs:uUV:Wx" arg; do
	case $arg in
		b)
			INSTBOOT=$OPTARG
			;;
		B)
			NO_KBACKUPS=true
			;;
		C)	CLEAN_ONLY=true
			;;
		c)
			CONF_FILE=$OPTARG
			;;
		D)
			DST=$OPTARG
			;;
		d)
			DOWNLOAD_ONLY=true
			;;
		e)
			EXTRACT_ONLY=true
			;;
		h)
			usage
			;;
		i)
			INTERACTIVE=true
			;;
		I)
			CHK_INTEG=true
			shift $((${OPTIND}-1))
			INTEG_SIG_FILE=$*
			INTEG_SIG_FILE=${INTEG_SIG_FILE:-SHA256.sig}
			OPTIND=1
			;;
		k)
			KERNEL_ONLY=true
			;;
		K)
			SETS_ONLY=true
			;;
		m)
			MACHINE=$OPTARG
			;;
		M)
			MIRROR=$OPTARG
			;;
		n)
			FORCE_MP=true
			;;
		r)
			MERGE=true
			;;
		R)
			REBOOT=true
			;;
		S)
			SKIP_SIGN=true
			;;
		s)
			KEY=$OPTARG
			;;
		u)
			CHK_UPDATE=true
			;;
		U)
			CHK_UPDATE=true
			INS_UPDATE=true
			;;
		V)
			SETVER=$(echo $OPTARG | tr -d \.)
			;;
		W)
			WEXIT=true
			;;
		x)
			NO_X11=true
			;;
		*)
			exit 1
	esac
done

if [ $CLEAN_ONLY == true ]; then
	qtmp=$(echo $tmp_template | sed 's/X/\?/g')

	msg "${white}Cleaning: ${green}${DST}"
	# Only remove files from DST, someone might have set it to /
	rm -fv $DST/*.{tgz,rd,mp,sig}
	rm -fv $DST/{BUILDINFO,bsd}

	msg "${white}Purging: ${green}/tmp/${qtmp}"
	for d in /tmp/${qtmp}; do
		# Check if the dir exists before rm'ing
		# - this prevents echoing of a ???? dir that doesn't exist
		[ -d "${d}" ] && rm -rfv "$d"
	done
	exit 0
fi

if [ $KERNEL_ONLY == true ] && [ $SETS_ONLY == true ]; then
	echo 'The options -k and -K are mutually exclusive.'
	exit 1
fi

if [ $CHK_INTEG == true ]; then
	check_integ
	exit 0
fi

if [ $CHK_UPDATE == true ]; then
	check_update
	exit 0
fi

[[ $(id -u) -ne 0 ]] && error "need root privileges" false

mkdir -p -- "$DST" || exit 1
chown -R root:_pkgfetch "$DST"
chmod -R g+rwx "$DST"

case "${MIRROR}" in
	http://* | ftp://* | https://*)
		URL="${MIRROR}/pub/OpenBSD/snapshots/${MACHINE}"
		;;
	*)
		URL="http://${MIRROR}/pub/OpenBSD/snapshots/${MACHINE}"
		;;
esac

if [ ! $EXTRACT_ONLY ]; then
	msg "${white}Fetching from: ${green}${URL}"
fi

(
	cd -- "$DST" || exit 1

	# first element should be bsd, second should be mp for given kernel names.
	if [[ "${MACHINE}" == armv7 ]] || [[ "${MACHINE}" == loongson ]]; then
		# Currently there is no bsd.mp
		set -A bsds "bsd" "" "bsd.rd"
	else
		set -A bsds 'bsd' 'bsd.mp' 'bsd.rd'
	fi

	RD=${bsds[2]}

	if [ $SKIP_SIGN == false ]; then
		fetch "${URL}/${base_sig}" || error "Can't fetch signature file!" false
	fi

	fetch "${URL}/${build_info}" || USE_BUILDINFO=false

	if [ -e ~/.last_snap ]; then
		last_snap=$(cat ~/.last_snap)
		msg "last snap: ${white}${last_snap}"
		if [ $USE_BUILDINFO ]; then
			current_snap=$(awk -F- '{print $2}' "$build_info" | sed 's/^ //')
			if [ "${last_snap}" == "$current_snap" ]; then
				warn "No new snaps available, mirror has: ${current_snap}!"
			fi
		fi
	fi

	if [ $EXTRACT_ONLY == false ]; then
		if [ $SETS_ONLY == false ]; then
			msg "Fetching bsds"
			for bsd in "${bsds[@]}"; do
				fetch "${URL}/${bsd}" || error "Can't find bsds at ${URL}" false
			done

			if [ "${CPUS}" == "1" ] && [ "${FORCE_MP:=false}" != true ]; then
				msg "${white}Using ${green}bsd.."
				KERNEL=${bsds[0]}
			else
				msg "${white}Using ${green}bsd.mp.."
				KERNEL=${bsds[1]}
			fi

			if [ "${NO_KBACKUPS:=false}" == false ]; then
				if [ $DOWNLOAD_ONLY == false ]; then
					backup
				fi
			fi

			if [ $DOWNLOAD_ONLY == false ]; then
				update_kernel
			fi

			if [ $KERNEL_ONLY == true ]; then
				exit 0
			fi
		fi # SETS_ONLY

		msg "Fetching sets"
		for set in "${sets[@]}"; do
			fetch "${URL}/${set}${SETVER}.tgz" || \
				error "Perhaps you need to specify -V to set version. Example 5.2" true
		done

		if [ "${NO_X11}" == "false" ]; then
			msg "Fetching xsets"
			for set in "${xsets[@]}"; do
				fetch "${URL}/${set}${SETVER}.tgz" || \
					error "Perhaps you need to specify -V to set version. Example -V 5.2" true
			done
		fi
	fi

	if [ $SKIP_SIGN == false ]; then
		verisigs "*.tgz"
	fi

	if [ $DOWNLOAD_ONLY == true ]; then
		exit 0
	fi

	msg "Extracting sets"
	for set in "${sets[@]}"; do
		extract "${DST}/${set}${SETVER}.tgz"

		if [ "${set}" == "man" ] && [ "${NO_X11}" == "false" ]; then
			msg "Extracting xsets ${white}will continue with sets after. ${green}"

			for xset in "${xsets[@]}"; do
				extract "${DST}/${xset}${SETVER}.tgz"
			done
		fi
	done

	if [ $MERGE == true ]; then
		msg "Running sysmerge"
		sysmerge || error "Failed to sysmerge!" false
	else
		echo "/usr/sbin/sysmerge -b" >>/etc/rc.sysmerge
		chmod +x /etc/rc.sysmerge
		echo "Don't forget to run sysmerge!"
	fi

	echo -n "Relinking to create unique kernel..."
	/usr/libexec/reorder_kernel && echo "done." || echo "failed."

	if [ $INSTBOOT != false ]; then
		msg "Installing bootstrap on ${INSTBOOT}"
		installboot -v $INSTBOOT || \
			error "Something bad happened - check your boot disk!" false
	fi

	if [ $USE_BUILDINFO ]; then
		awk -F- '{print $2}' "$build_info" | sed 's/^ //' > ~/.last_snap
	else
		date > ~/.last_snap
	fi

	if [ "$AFTER" != false ]; then
		cp $AFTER /etc/rc.firsttime
		chmod +x /etc/rc.firsttime
	else
		echo 'cd /dev && sh MAKEDEV all' >>/etc/rc.firsttime
		echo "/usr/sbin/fw_update -v" >>/etc/rc.firsttime
		chmod +x /etc/rc.firsttime
	fi

	if [ $REBOOT == true ]; then
		msg "Rebooting"
		/sbin/oreboot || error "Something really bad happened - Can't reboot!" false
	fi
)