From acd3eff7617971d5efd90a270fcfa05155a3bfcd Mon Sep 17 00:00:00 2001
From: "Brock A. Martin" <bamartin125@gmail.com>
Date: Mon, 14 Dec 2020 10:41:18 -0600
Subject: [PATCH] Adds ability to set custom bootstrap url, and apt source list
 file/dir

---
 README.md                         |  24 ++++++
 build-docker.sh                   | 132 +++++++++++++++++++++---------
 build.sh                          |  20 +++++
 stage0/00-configure-apt/00-run.sh |  32 +++++++-
 stage0/prerun.sh                  |   2 +-
 5 files changed, 166 insertions(+), 44 deletions(-)

diff --git a/README.md b/README.md
index 8b7f875af5..b6d063f629 100644
--- a/README.md
+++ b/README.md
@@ -41,6 +41,30 @@ The following environment variables are supported:
    The release version to build images against. Valid values are jessie, stretch
    buster, bullseye, and testing.
 
+ * `BOOTSTRAP_URL` (Default: http://raspbian.raspberrypi.org/raspbian/)
+
+   The apt source to use as the source for the stage0 Debian bootstrap process.
+
+   Public mirrors may be found at https://www.raspbian.org/RaspbianMirrors
+
+ * `CUSTOM_LIST` (Default: unset)
+
+   Set this variable to a path to a file to install as the
+   `/etc/apt/sources.list` file. The path may be absolute or relative to the
+   location of the `config` file.
+
+   If this variable is set, the standard list files as provided with `pi-gen`
+   are ignored.
+
+ * `CUSTOM_LIST_DIR` (Default: unset)
+
+   Set this variable to a path to a directory holding list files to be
+   installed to the `/etc/apt/sources.list.d` directory. The path may be
+   absolute or relative to the location of the `config` file.
+
+   If this variable is set, the standard list files as provided with `pi-gen`
+   are ignored.
+
  * `APT_PROXY` (Default: unset)
 
    If you require the use of an apt proxy, set it here.  This proxy setting
diff --git a/build-docker.sh b/build-docker.sh
index 350f722bde..fe24a79a55 100755
--- a/build-docker.sh
+++ b/build-docker.sh
@@ -3,7 +3,7 @@ DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
 
 BUILD_OPTS="$*"
 
-DOCKER="docker"
+DOCKER=${DOCKER:-docker}
 
 if ! ${DOCKER} ps >/dev/null 2>&1; then
 	DOCKER="sudo docker"
@@ -15,10 +15,8 @@ if ! ${DOCKER} ps >/dev/null; then
 fi
 
 CONFIG_FILE=""
-if [ -f "${DIR}/config" ]; then
-	CONFIG_FILE="${DIR}/config"
-fi
 
+# Arguments passed on command line have highest priority (others are fallbacks)
 while getopts "c:" flag
 do
 	case "${flag}" in
@@ -30,19 +28,20 @@ do
 	esac
 done
 
-# Ensure that the configuration file is an absolute path
-if test -x /usr/bin/realpath; then
-	CONFIG_FILE=$(realpath -s "$CONFIG_FILE" || realpath "$CONFIG_FILE")
+if [ -z "${CONFIG_FILE}" ]; then # config file not yet defined
+	if [ -f "${DIR}/config" ]; then # guess location relative to this script
+		CONFIG_FILE="${DIR}/config"
+	fi
 fi
 
-# Ensure that the confguration file is present
+# Ensure that the configuration file is present
 if test -z "${CONFIG_FILE}"; then
 	echo "Configuration file need to be present in '${DIR}/config' or path passed as parameter"
 	exit 1
-else
-	# shellcheck disable=SC1090
-	source ${CONFIG_FILE}
 fi
+CONFIG_FILE_ORG_DIR="$(dirname "${CONFIG_FILE}")"
+
+source "${CONFIG_FILE}"
 
 CONTAINER_NAME=${CONTAINER_NAME:-pigen_work}
 CONTINUE=${CONTINUE:-0}
@@ -51,7 +50,7 @@ PRESERVE_CONTAINER=${PRESERVE_CONTAINER:-0}
 if [ -z "${IMG_NAME}" ]; then
 	echo "IMG_NAME not set in 'config'" 1>&2
 	echo 1>&2
-exit 1
+	exit 1
 fi
 
 # Ensure the Git Hash is recorded before entering the docker container
@@ -66,46 +65,101 @@ fi
 if [ "${CONTAINER_EXISTS}" != "" ] && [ "${CONTINUE}" != "1" ]; then
 	echo "Container ${CONTAINER_NAME} already exists and you did not specify CONTINUE=1. Aborting."
 	echo "You can delete the existing container like this:"
-	echo "  ${DOCKER} rm -v ${CONTAINER_NAME}"
+	echo "	${DOCKER} rm -v ${CONTAINER_NAME}"
 	exit 1
 fi
 
 # Modify original build-options to allow config file to be mounted in the docker container
-BUILD_OPTS="$(echo "${BUILD_OPTS:-}" | sed -E 's@\-c\s?([^ ]+)@-c /config@')"
+BUILD_OPTS="$(echo "${BUILD_OPTS:-}" | sed -E 's@\-c\s+?([^ ]+)@@')"
 
 # Check the arch of the machine we're running on. If it's 64-bit, use a 32-bit base image instead
 case "$(uname -m)" in
-  x86_64|aarch64)
-    BASE_IMAGE=i386/debian:buster
-    ;;
-  *)
-    BASE_IMAGE=debian:buster
-    ;;
+	x86_64|aarch64)
+		BASE_IMAGE=i386/debian:buster
+		;;
+	*)
+		BASE_IMAGE=debian:buster
+		;;
 esac
+
+# Build the pi-gen image
 ${DOCKER} build --build-arg BASE_IMAGE=${BASE_IMAGE} -t pi-gen "${DIR}"
 
+# Create the pi-gen container
 if [ "${CONTAINER_EXISTS}" != "" ]; then
-	trap 'echo "got CTRL+C... please wait 5s" && ${DOCKER} stop -t 5 ${CONTAINER_NAME}_cont' SIGINT SIGTERM
-	time ${DOCKER} run --rm --privileged \
-		--volume "${CONFIG_FILE}":/config:ro \
-		-e "GIT_HASH=${GIT_HASH}" \
-		--volumes-from="${CONTAINER_NAME}" --name "${CONTAINER_NAME}_cont" \
-		pi-gen \
-		bash -e -o pipefail -c "dpkg-reconfigure qemu-user-static &&
-	cd /pi-gen; ./build.sh ${BUILD_OPTS} &&
-	rsync -av work/*/build.log deploy/" &
-	wait "$!"
+	CONTAINER_ID=$(
+		${DOCKER} create \
+			--rm \
+			--name "${CONTAINER_NAME}_cont" \
+			--privileged \
+			-e "GIT_HASH=${GIT_HASH}" \
+			--volumes-from="${CONTAINER_NAME}" \
+			pi-gen \
+			bash -e -o pipefail -c "dpkg-reconfigure qemu-user-static &&
+		cd /pi-gen; ./build.sh ${BUILD_OPTS} &&
+		rsync -av work/*/build.log deploy/"
+	)
 else
-	trap 'echo "got CTRL+C... please wait 5s" && ${DOCKER} stop -t 5 ${CONTAINER_NAME}' SIGINT SIGTERM
-	time ${DOCKER} run --name "${CONTAINER_NAME}" --privileged \
-		--volume "${CONFIG_FILE}":/config:ro \
-		-e "GIT_HASH=${GIT_HASH}" \
-		pi-gen \
-		bash -e -o pipefail -c "dpkg-reconfigure qemu-user-static &&
-	cd /pi-gen; ./build.sh ${BUILD_OPTS} &&
-	rsync -av work/*/build.log deploy/" &
-	wait "$!"
+	CONTAINER_ID=$(
+		${DOCKER} create \
+			--name "${CONTAINER_NAME}" \
+			--privileged \
+			-e "GIT_HASH=${GIT_HASH}" \
+			pi-gen \
+			bash -e -o pipefail -c "dpkg-reconfigure qemu-user-static &&
+		cd /pi-gen; ./build.sh ${BUILD_OPTS} &&
+		rsync -av work/*/build.log deploy/"
+	)
 fi
+
+# Create a temporary working dir for file tweaks prior to copying into container
+PIGEN_TMP_DIR="$(mktemp -d -p "" pi-gen.XXXXXX)" || { echo "Failed to create temp dir"; exit 1; }
+
+finish() {
+	rm -rf "$PIGEN_TMP_DIR"
+}
+
+trap finish EXIT
+
+cp "${CONFIG_FILE}" "${PIGEN_TMP_DIR}"/config
+
+OPTIONAL_EXT_CONFIGS=(
+CUSTOM_LIST
+CUSTOM_LIST_DIR
+)
+
+# Add optional config files to target area
+pushd "${CONFIG_FILE_ORG_DIR}" >/dev/null || { echo "Unable to cd to ${CONFIG_FILE_ORG_DIR}" 1>&2; exit 1; }
+for ext_config_item in ${OPTIONAL_EXT_CONFIGS[@]}; do
+	# Skip undefined ext configs
+	if [ -z ${!ext_config_item+x} ]; then
+		continue
+	fi
+
+	declare "${ext_config_item}"="${!ext_config_item//$'\r'}" # remove any trailing carriage returns
+
+	if [ ! -e ${!ext_config_item} ]; then
+		echo "The target of config item $ext_config_item (${!ext_config_item}) does not exist" 1>&2
+		exit 1
+	fi
+
+	target_config_path=/pi-gen/"${ext_config_item,,}"
+
+	# Tweak config file path to ext config
+	sed -i -E 's@('"$ext_config_item"=').*@\1'"$target_config_path"'@' "${PIGEN_TMP_DIR}"/config
+
+	# Copy file into container
+	${DOCKER} cp "${!ext_config_item}" "$CONTAINER_ID":"${target_config_path}"
+done
+popd >/dev/null
+
+${DOCKER} cp "${PIGEN_TMP_DIR}"/config "$CONTAINER_ID":/pi-gen/config
+
+# Start a pi-gen container
+trap 'echo "got CTRL+C... please wait 5s" && ${DOCKER} stop -t 5 ${CONTAINER_ID}' SIGINT SIGTERM
+time ${DOCKER} start -a "$CONTAINER_ID" &
+wait "$!"
+
 echo "copying results from deploy/"
 ${DOCKER} cp "${CONTAINER_NAME}":/pi-gen/deploy .
 ls -lah deploy
diff --git a/build.sh b/build.sh
index 6f720b1f50..86370ff71f 100755
--- a/build.sh
+++ b/build.sh
@@ -125,9 +125,12 @@ fi
 BASE_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
 export BASE_DIR
 
+CONFIG_FILE_DIR="$(realpath ".")"
+
 if [ -f config ]; then
 	# shellcheck disable=SC1091
 	source config
+	CONFIG_FILE_DIR="$(realpath "$(dirname config)")"
 fi
 
 while getopts "c:" flag
@@ -137,6 +140,7 @@ do
 			EXTRA_CONFIG="$OPTARG"
 			# shellcheck disable=SC1090
 			source "$EXTRA_CONFIG"
+			CONFIG_FILE_DIR="$(realpath "$(dirname "$EXTRA_CONFIG")")"
 			;;
 		*)
 			;;
@@ -156,6 +160,7 @@ export IMG_DATE="${IMG_DATE:-"$(date +%Y-%m-%d)"}"
 export IMG_FILENAME="${IMG_FILENAME:-"${IMG_DATE}-${IMG_NAME}"}"
 export ZIP_FILENAME="${ZIP_FILENAME:-"image_${IMG_DATE}-${IMG_NAME}"}"
 
+export CONFIG_FILE_DIR
 export SCRIPT_DIR="${BASE_DIR}/scripts"
 export WORK_DIR="${WORK_DIR:-"${BASE_DIR}/work/${IMG_DATE}-${IMG_NAME}"}"
 export DEPLOY_DIR=${DEPLOY_DIR:-"${BASE_DIR}/deploy"}
@@ -164,6 +169,21 @@ export LOG_FILE="${WORK_DIR}/build.log"
 
 export TARGET_HOSTNAME=${TARGET_HOSTNAME:-raspberrypi}
 
+export BOOTSTRAP_URL=${BOOTSTRAP_URL:-"http://raspbian.raspberrypi.org/raspbian/"}
+
+pushd "${CONFIG_FILE_DIR}" >/dev/null # allow relative paths to config file
+CUSTOM_LIST=${CUSTOM_LIST:-}
+CUSTOM_LIST_DIR=${CUSTOM_LIST_DIR:-}
+if [ -n "$CUSTOM_LIST" ]; then
+	CUSTOM_LIST="$(realpath "$CUSTOM_LIST")"
+fi
+if [ -n "$CUSTOM_LIST_DIR" ]; then
+	CUSTOM_LIST_DIR="$(realpath "$CUSTOM_LIST_DIR")"
+fi
+export CUSTOM_LIST
+export CUSTOM_LIST_DIR
+popd >/dev/null
+
 export FIRST_USER_NAME=${FIRST_USER_NAME:-pi}
 export FIRST_USER_PASS=${FIRST_USER_PASS:-raspberry}
 export RELEASE=${RELEASE:-buster}
diff --git a/stage0/00-configure-apt/00-run.sh b/stage0/00-configure-apt/00-run.sh
index d10a49f174..9db8363a1b 100755
--- a/stage0/00-configure-apt/00-run.sh
+++ b/stage0/00-configure-apt/00-run.sh
@@ -1,9 +1,33 @@
 #!/bin/bash -e
 
-install -m 644 files/sources.list "${ROOTFS_DIR}/etc/apt/"
-install -m 644 files/raspi.list "${ROOTFS_DIR}/etc/apt/sources.list.d/"
-sed -i "s/RELEASE/${RELEASE}/g" "${ROOTFS_DIR}/etc/apt/sources.list"
-sed -i "s/RELEASE/${RELEASE}/g" "${ROOTFS_DIR}/etc/apt/sources.list.d/raspi.list"
+# if either CUSTOM_LIST or CUSTOM_LIST_DIR is set, then install sources from there
+if [ -n "$CUSTOM_LIST" -o -n "$CUSTOM_LIST_DIR" ]; then
+    if [ -n "$CUSTOM_LIST" ]; then
+        if [ -f "$CUSTOM_LIST" ]; then
+            install -m 644 "$CUSTOM_LIST" "${ROOTFS_DIR}/etc/apt/sources.list" || exit 1
+        else
+            echo "$CUSTOM_LIST cannot be found"; exit 1
+        fi
+    fi
+    if [ -n "$CUSTOM_LIST_DIR" ]; then
+        if [ -d "$CUSTOM_LIST_DIR" ]; then
+            install -m 644 "$CUSTOM_LIST_DIR"/* "${ROOTFS_DIR}/etc/apt/sources.list.d/" || exit 1
+        else
+            echo "$CUSTOM_LIST_DIR cannot be found"; exit 1
+        fi
+    fi
+
+# otherwise, use the standard sources as provided by pi-gen
+else
+    install -m 644 files/sources.list "${ROOTFS_DIR}/etc/apt/"
+    install -m 644 files/raspi.list "${ROOTFS_DIR}/etc/apt/sources.list.d/"
+fi
+
+# replace 'RELEASE' with "$RELEASE" in all .list files
+if [ -f "${ROOTFS_DIR}/etc/apt/sources.list" ]; then
+    sed -i "s/RELEASE/${RELEASE}/g" "${ROOTFS_DIR}/etc/apt/sources.list"
+fi
+find "${ROOTFS_DIR}/etc/apt/sources.list.d/" -type f -exec sed -i "s/RELEASE/${RELEASE}/g" {} \;
 
 if [ -n "$APT_PROXY" ]; then
 	install -m 644 files/51cache "${ROOTFS_DIR}/etc/apt/apt.conf.d/51cache"
diff --git a/stage0/prerun.sh b/stage0/prerun.sh
index 3b257836b8..7537cd1d94 100755
--- a/stage0/prerun.sh
+++ b/stage0/prerun.sh
@@ -1,5 +1,5 @@
 #!/bin/bash -e
 
 if [ ! -d "${ROOTFS_DIR}" ]; then
-	bootstrap ${RELEASE} "${ROOTFS_DIR}" http://raspbian.raspberrypi.org/raspbian/
+	bootstrap ${RELEASE} "${ROOTFS_DIR}" "${BOOTSTRAP_URL}"
 fi