From 226d1e292c4fa96ba18213d8028c3e4753b8edd0 Mon Sep 17 00:00:00 2001
From: Nick Hepler <nick@nickhepler.com>
Date: Fri, 11 Oct 2024 17:01:20 -0400
Subject: [PATCH] Fix issue with sshd_config for PasswordAuthentication,
 AddressFamily.

---
 setup.sh | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/setup.sh b/setup.sh
index d4bf6df..8d46908 100644
--- a/setup.sh
+++ b/setup.sh
@@ -100,10 +100,18 @@ if [ -f "$sshd_config" ]; then
     sed -i 's/^PermitRootLogin .*/PermitRootLogin no/' "$sshd_config" || echo "PermitRootLogin no" >> "$sshd_config"
 
     # Set PasswordAuthentication to no
-    sed -i 's/^PasswordAuthentication .*/PasswordAuthentication no/' "$sshd_config" || echo "PasswordAuthentication no" >> "$sshd_config"
+    if grep -q '^#PasswordAuthentication' "$sshd_config"; then
+        sed -i 's/^#PasswordAuthentication.*/PasswordAuthentication no/' "$sshd_config"
+    else
+        sed -i 's/^PasswordAuthentication .*/PasswordAuthentication no/' "$sshd_config" || echo "PasswordAuthentication no" >> "$sshd_config"
+    fi
 
-    # Ensure AddressFamily inet is not commented out
-    sed -i 's/^#\?\(AddressFamily inet\)/\1/' "$sshd_config"
+    # Ensure AddressFamily inet is set
+    if grep -q '^#AddressFamily' "$sshd_config"; then
+        sed -i 's/^#AddressFamily.*/AddressFamily inet/' "$sshd_config"
+    else
+        sed -i 's/^AddressFamily .*/AddressFamily inet/' "$sshd_config" || echo "AddressFamily inet" >> "$sshd_config"
+    fi
 else
     echo "sshd_config file not found. Exiting."
     exit 1
@@ -124,4 +132,4 @@ systemctl restart sshd
 # Prompt for Docker installation
 prompt_for_docker_install
 
-echo "System update complete! Timezone set to $timezone, hostname set to $hostname, /etc/hosts updated, and sshd_config modified."
\ No newline at end of file
+echo "System update complete! Timezone set to $timezone, hostname set to $hostname, /etc/hosts updated, and sshd_config modified."