From bfbd76f61270ed11424ac86f250ada331b1add48 Mon Sep 17 00:00:00 2001
From: "Daniel F. Dickinson" <cshored@thecshore.com>
Date: Fri, 17 Jan 2020 22:45:03 -0500
Subject: [PATCH] Enable adding SSH pubkey and setting pubkey only

This commit add the ability to specify an SSH public key as well as the
option to disable password authentication and only allow public key
authentication for SSH.

Signed-off-by: Daniel F. Dickinson <cshored@thecshore.com>
---
 README.md                      | 12 ++++++++++++
 build.sh                       |  8 ++++++++
 stage2/01-sys-tweaks/01-run.sh | 12 ++++++++++++
 3 files changed, 32 insertions(+)

diff --git a/README.md b/README.md
index 73d4961..007ee20 100644
--- a/README.md
+++ b/README.md
@@ -125,6 +125,18 @@ The following environment variables are supported:
 
    Setting to `1` will enable ssh server for remote log in. Note that if you are using a common password such as the defaults there is a high risk of attackers taking over you Raspberry Pi.
 
+  * `PUBKEY_SSH_FIRST_USER` (Default: unset)
+
+   Setting this to a value will make that value the contents of the FIRST_USER_NAME's ~/.ssh/authorized_keys.  Obviously the value should
+   therefore be a valid authorized_keys file.  Note that this does not
+   automatically enable SSH.
+
+  * `PUBKEY_ONLY_SSH` (Default: `0`)
+
+   * Setting to `1` will disable password authentication for SSH and enable
+   public key authentication.  Note that if SSH is not enabled this will take
+   effect when SSH becomes enabled.
+
  * `STAGE_LIST` (Default: `stage*`)
 
     If set, then instead of working through the numeric stages in order, this list will be followed. For example setting to `"stage0 stage1 mystage stage2"` will run the contents of `mystage` before stage2. Note that quotes are needed around the list. An absolute or relative path can be given for stages outside the pi-gen directory.
diff --git a/build.sh b/build.sh
index 9e78728..8009181 100755
--- a/build.sh
+++ b/build.sh
@@ -170,6 +170,7 @@ export WPA_ESSID
 export WPA_PASSWORD
 export WPA_COUNTRY
 export ENABLE_SSH="${ENABLE_SSH:-0}"
+export PUBKEY_ONLY_SSH="${PUBKEY_ONLY_SSH:-0}"
 
 export LOCALE_DEFAULT="${LOCALE_DEFAULT:-en_GB.UTF-8}"
 
@@ -180,6 +181,8 @@ export TIMEZONE_DEFAULT="${TIMEZONE_DEFAULT:-Europe/London}"
 
 export GIT_HASH=${GIT_HASH:-"$(git rev-parse HEAD)"}
 
+export PUBKEY_SSH_FIRST_USER
+
 export CLEAN
 export IMG_NAME
 export APT_PROXY
@@ -225,6 +228,11 @@ if [[ -n "${WPA_PASSWORD}" && ${#WPA_PASSWORD} -lt 8 || ${#WPA_PASSWORD} -gt 63
 	exit 1
 fi
 
+if [[ "${PUBKEY_ONLY_SSH}" = "1" && -z "${PUBKEY_SSH_FIRST_USER}" ]]; then
+	echo "Must set 'PUBKEY_SSH_FIRST_USER' to a valid SSH public key if using PUBKEY_ONLY_SSH"
+	exit 1
+fi
+
 mkdir -p "${WORK_DIR}"
 log "Begin ${BASE_DIR}"
 
diff --git a/stage2/01-sys-tweaks/01-run.sh b/stage2/01-sys-tweaks/01-run.sh
index 8d28adc..c1836f9 100755
--- a/stage2/01-sys-tweaks/01-run.sh
+++ b/stage2/01-sys-tweaks/01-run.sh
@@ -11,6 +11,18 @@ install -m 644 files/console-setup   	"${ROOTFS_DIR}/etc/default/"
 
 install -m 755 files/rc.local		"${ROOTFS_DIR}/etc/"
 
+if [ -n "${PUBKEY_SSH_FIRST_USER}" ]; then
+	install -v -m 0700 -o 1000 -g 1000 -d "${ROOTFS_DIR}"/home/"${FIRST_USER_NAME}"/.ssh
+	echo "${PUBKEY_SSH_FIRST_USER}" >"${ROOTFS_DIR}"/home/"${FIRST_USER_NAME}"/.ssh/authorized_keys
+	chown 1000:1000 "${ROOTFS_DIR}"/home/"${FIRST_USER_NAME}"/.ssh/authorized_keys
+	chmod 0600 "${ROOTFS_DIR}"/home/"${FIRST_USER_NAME}"/.ssh/authorized_keys
+fi
+
+if [ "${PUBKEY_ONLY_SSH}" = "1" ]; then
+	sed -i -Ee 's/^#?[[:blank:]]*PubkeyAuthentication[[:blank:]]*no[[:blank:]]*$/PubkeyAuthentication yes/
+s/^#?[[:blank:]]*PasswordAuthentication[[:blank:]]*yes[[:blank:]]*$/PasswordAuthentication no/' "${ROOTFS_DIR}"/etc/ssh/sshd_config
+fi
+
 on_chroot << EOF
 systemctl disable hwclock.sh
 systemctl disable nfs-common