From 25bd635c4e58fcc866d6041c435333484e8954c7 Mon Sep 17 00:00:00 2001
From: VetheonGames <ceo@pixelatedstudios.net>
Date: Fri, 24 May 2024 18:31:37 -0600
Subject: [PATCH] Fix a bug with public keys being added to root. Made public
 key process clearer to the user

---
 InitMate.sh | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/InitMate.sh b/InitMate.sh
index 6a24535..25b5d39 100644
--- a/InitMate.sh
+++ b/InitMate.sh
@@ -1,6 +1,7 @@
 #!/bin/bash
 
 LOG_FILE="$(dirname "$0")/setup.log"
+ORIGINAL_USER=$(logname)
 
 # ASCII Art Banner
 cat << "EOF"
@@ -101,13 +102,15 @@ secure_ssh() {
     sed -i "s/#PasswordAuthentication yes/PasswordAuthentication no/" /etc/ssh/sshd_config
     sed -i "s/#PubkeyAuthentication yes/PubkeyAuthentication yes/" /etc/ssh/sshd_config
 
-    log "Please enter your public key:"
+    log "Please enter your public key (this is for your current user.):"
     read -r PUBLIC_KEY
 
-    mkdir -p ~/.ssh
-    echo "$PUBLIC_KEY" > ~/.ssh/authorized_keys
-    chmod 600 ~/.ssh/authorized_keys
-    chmod 700 ~/.ssh
+    sudo -u "$ORIGINAL_USER" bash -c "
+        mkdir -p ~/.ssh
+        echo '$PUBLIC_KEY' > ~/.ssh/authorized_keys
+        chmod 600 ~/.ssh/authorized_keys
+        chmod 700 ~/.ssh
+    "
 
     systemctl restart sshd
 }
@@ -199,7 +202,7 @@ add_users() {
             if prompt_yes_no "Do you want to add a public key for $USERNAME?"; then
                 log "Please enter the public key:"
                 read -r PUBLIC_KEY
-                su - "$USERNAME" -c "mkdir -p ~/.ssh && echo '$PUBLIC_KEY' >> ~/.ssh/authorized_keys && chmod 600 ~/.ssh/authorized_keys && chmod 700 ~/.ssh"
+                sudo -u "$USERNAME" bash -c "mkdir -p ~/.ssh && echo '$PUBLIC_KEY' >> ~/.ssh/authorized_keys && chmod 600 ~/.ssh/authorized_keys && chmod 700 ~/.ssh"
                 log "Public key added for $USERNAME."
             fi