From c1a7dae113ce943abe8bf5d5ab36fb7aac829fd2 Mon Sep 17 00:00:00 2001
From: "Daniel F. Dickinson" <cshored@danielfdickinson.ca>
Date: Tue, 8 Sep 2020 16:16:07 -0400
Subject: [PATCH] Enable adding SSH pubkey and setting pubkey only (#380)

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 29d8166..e7f224a 100644
--- a/README.md
+++ b/README.md
@@ -130,6 +130,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 a8247ab..6f720b1 100755
--- a/build.sh
+++ b/build.sh
@@ -171,6 +171,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}"
 
@@ -181,6 +182,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
@@ -226,6 +229,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