#!/bin/bash

# --- ANSI Color Codes ---
RED='\033[0;31m'        # Error/Danger
GREEN='\033[0;32m'      # Success/OK
YELLOW='\033[0;33m'     # Warning/Info
BLUE='\033[0;34m'       # Section Headers
PURPLE='\033[0;35m'     # Prompts
CYAN='\033[0;36m'       # General Info
NC='\033[0m'            # No Color (Reset)
BOLD='\033[1m'          # Bold text
UNDERLINE='\033[4m'     # Underline text

# --- Global Configuration ---
SFTP_GROUP="ASFTPLIMITED"
CHROOT_ROOT_DIR="/var/sftp"
SSHD_CONFIG_FILE="/etc/ssh/sshd_config"
PAM_SSHD_FILE="/etc/pam.d/sshd"
PAM_COMMON_AUTH_FILE="/etc/pam.d/common-auth"

# --- Helper Function: Check Last Command Status ---
check_status() {
	local exit_code="$?"
	local success_msg="$1"
	local failure_msg="$2"
	if [ "$exit_code" -eq 0 ]; then
		echo -e "${GREEN}✓ ${success_msg}${NC}"
		return 0
	else
		echo -e "${RED}X ${failure_msg} (Exit Code: $exit_code)${NC}" >&2
		return 1
	fi
}

# --- Initial Setup: Create Chroot Root Directory & Set Permissions ---
sudo mkdir -p "${CHROOT_ROOT_DIR}"
sudo chown root:root /
sudo chmod 755 /
sudo chown root:root /var
sudo chmod 755 /var
sudo chown root:root "${CHROOT_ROOT_DIR}"
sudo chmod 755 "${CHROOT_ROOT_DIR}"

# --- Helper Function: Get User Input with Validation ---
get_user_input() {
	local prompt_msg="$1"
	local validation_pattern="$2"
	local error_msg="${3:-${RED}Invalid input. Please try again.${NC}}"
	local input_val=""

	while true; do
		# Send the prompt to standard error so it's displayed on the terminal
		# even when stdout is being captured.
		echo -e "\n${PURPLE}${prompt_msg}${NC}" >&2 # <--- CHANGE IS HERE
		read -r input_val

		if [ -z "$input_val" ]; then
			echo -e "${RED}Input cannot be empty. Please try again.${NC}" >&2
		elif [ -n "$validation_pattern" ] && ! echo "$input_val" | grep -Eq "$validation_pattern"; then
			echo -e "$error_msg" >&2
		else
			# Only the final, valid input should go to stdout for capture
			echo "$input_val"
			return 0
		fi
	done
}

# --- Helper Function: Get Password Securely with Confirmation ---
get_password_securely() {
	local entity_name="$1"
	local password=""
	local confirm_password=""

	while true; do
		echo -e "\n${PURPLE}Enter password for '$entity_name': ${NC}"
		read -rs password
		echo "" # Newline after silent input
		echo -e "${PURPLE}Confirm password: ${NC}"
		read -rs confirm_password
		echo "" # Newline after silent input

		if [ -z "$password" ]; then
			echo -e "${RED}Password cannot be empty. Please enter a password.${NC}" >&2
		elif [ "$password" = "$confirm_password" ]; then
			echo "$password"
			return 0
		else
			echo -e "${RED}Passwords do not match. Please try again.${NC}" >&2
		fi
	done
}

# --- Helper Function: Yes/No Prompt ---
confirm_action() {
	local prompt_msg="$1"
	local choice=""

	while true; do
		echo -e "\n${PURPLE}${prompt_msg}${NC}"
		read -r choice
		choice=$(echo "$choice" | tr '[:upper:]' '[:lower:]')
		if [ "$choice" = "y" ] || [ -z "$choice" ]; then
			return 0
		elif [ "$choice" = "n" ]; then
			return 1
		else
			echo -e "${RED}Invalid choice. Please enter 'Y' or 'n'.${NC}" >&2
		fi
	done
}

# --- Helper Function: Get list of SFTP Routes ---
# Returns an array of SFTP route usernames (e.g., "user1 user2")
_list_sftp_routes_array() {
	local -a routes=()
	local members
	members=$(getent group "${SFTP_GROUP}" 2>/dev/null | cut -d: -f4)

	local IFS=$'\n' # Set IFS to newline for read
	for username in $(echo "$members" | tr ',' '\n'); do
		if [ -n "$username" ]; then
			local user_shell
			user_shell=$(getent passwd "$username" | cut -d: -f7 2>/dev/null)
			if [ "$user_shell" = "/bin/false" ]; then
				routes+=("$username")
			fi
		fi
	done
	echo "${routes[@]}"
}

display_routes()
{
	# Display existing SFTP routes every time
	current_routes=($(_list_sftp_routes_array))
	echo -e "Existing ASFTP Routes:"
	if [ "${#current_routes[@]}" -eq 0 ]; then
		echo -e "${YELLOW}No ASFTP routes found.${NC}"
	else
		for route in "${current_routes[@]}"; do
			echo -e "  - ${route}"
		done
	fi
	echo "" # Add a newline for spacing
}

# --- Helper Function: Creates SFTP_GROUP ---
create_group() {
	if grep -q "^${SFTP_GROUP}:" /etc/group > /dev/null 2>&1; then
		echo -e "Group '${SFTP_GROUP}' already exists. Skipping group creation."
	else
		echo -e "${CYAN}Creating group '${SFTP_GROUP}'...${NC}"
		sudo groupadd "${SFTP_GROUP}"
		check_status "Group '${SFTP_GROUP}' created successfully." "Failed to create group '${SFTP_GROUP}'. Check permissions." || return 1
	fi
	echo ""
	return 0
}

# --- Function: Add SFTP Route ---
add_route() {
	clear 
	echo -e "${BLUE}--- Add New ASFTP Route ---${NC}"
	display_routes
	if ! create_group; then
		echo -e "${RED}Aborting route creation due to group creation failure.${NC}"
		return 1
	fi

	local username
	username=$(get_user_input "Enter the ASFTP route name to create (e.g., clientA_uploads): " "^[a-zA-Z0-9_.-]+$" "${RED}The route name can only contain letters, numbers, underscores, periods, and hyphens.${NC}")
	if [ -z "$username" ]; then
		echo -e "${RED}Invalid route name provided. Aborting.${NC}"
		return 1
	fi

	echo -e "${CYAN}Attempting to create ASFTP route '$username'...\n${NC}"

	if id "$username" > /dev/null 2>&1; then
		echo -e "${YELLOW}Warning: Route '$username' already exists. Skipping user creation.${NC}"
		return 0
	else
		sudo useradd -r -s /bin/false -M "$username"
		check_status "Route '$username' created successfully." "Failed to create route '$username'. Check permissions or try a different route name." || return 1

		echo ""
		if confirm_action "Do you want to set a password for '$username'? [Y/n]: "; then
			local password
			password=$(get_password_securely "$username")
			if [ -n "$password" ]; then
				echo -e "${CYAN}Setting password for '$username'...${NC}"
				echo "$username:$password" | sudo chpasswd
				check_status "Password for '$username' set successfully." "Failed to set password for '$username'."
			else
				echo -e "${YELLOW}Password setting cancelled or failed. Setting password to empty.${NC}"
				sudo passwd -d "$username"
				check_status "Password for '$username' set to empty." "Failed to clear password for '$username'."
			fi
		else
			sudo passwd -d "$username"
			check_status "Password for '$username' set to empty." "Failed to clear password for '$username'."
			echo -e "${CYAN}Note: For key-based authentication, a password isn't needed. To set one manually later, run 'sudo passwd $username'.${NC}"
		fi

		echo -e "${CYAN}Adding route '$username' to group '${SFTP_GROUP}'...${NC}"
		if groups "$username" | grep -q "\b${SFTP_GROUP}\b" > /dev/null 2>&1; then
			echo -e "${YELLOW}Warning: Route '$username' is already a member of group '${SFTP_GROUP}'. Skipping group assignment.${NC}"
		else
			sudo usermod -aG "${SFTP_GROUP}" "$username"
			check_status "Route '$username' added to group '${SFTP_GROUP}'." "Failed to add route '$username' to group '${SFTP_GROUP}'."
		fi
		echo ""

		local USER_CHROOT_DIR="${CHROOT_ROOT_DIR}/$username"

		echo -e "${CYAN}Creating chroot directory for '$username' at '${USER_CHROOT_DIR}'...${NC}"
		sudo mkdir -p "${USER_CHROOT_DIR}"
		check_status "Directory '${USER_CHROOT_DIR}' created." "Failed to create chroot directory for '$username'."

		sudo chown root:root "${USER_CHROOT_DIR}"
		check_status "Owner for '${USER_CHROOT_DIR}' set to root:root." "Failed to set owner for '${USER_CHROOT_DIR}'."
		sudo chmod 755 "${USER_CHROOT_DIR}"

		check_status "Permissions for ${USER_CHROOT_DIR} set to 755." "Failed to set permissions for ${USER_CHROOT_DIR}."

		echo -e "${BOLD}${GREEN}SFTP route '$username' setup complete!${NC}"
	fi
	echo ""
}

# --- Function: Modify an existing ASFTP Route ---
modify_route() {
	clear
	echo -e "${BLUE}--- Modify Existing ASFTP Route ---${NC}"
	display_routes

	local username
	username=$(get_user_input "Enter the name of the ASFTP route to modify (e.g., clientA_uploads): " "^[a-zA-Z0-9_.-]+$" "${RED}Invalid route name. Please use only letters, numbers, underscores, periods, and hyphens.${NC}")

	if [ -z "$username" ]; then
		echo -e "${YELLOW}No route name provided. Returning to main menu.${NC}"
		return 1
	fi

    # Validate that the entered username exists as an SFTP route
    local -a all_sftp_routes=($(_list_sftp_routes_array))
    local found="false"
    for existing_route in "${all_sftp_routes[@]}"; do
	    if [ "$username" = "$existing_route" ]; then
		    found="true"
		    break
	    fi
    done

    if [ "$found" = "false" ]; then
	    echo -e "${RED}Error: ASFTP route '$username' not found. Please ensure the name is correct.${NC}"
	    return 1
    fi

    echo -e "${CYAN}Modifying route '$username'. What would you like to do?${NC}"
    local options=("Rename Route" "Change Password" "Add SSH Public Key" "Return to Main Menu")
    local choice_num

    while true; do
	    local i=1
	    for opt in "${options[@]}"; do
		    echo "  ${i}) ${opt}"
		    i=$((i+1))
	    done
	    echo -e "\n${PURPLE}Select an option for '$username': ${NC}"
	    read -r choice_num

	    case "$choice_num" in
		    1) # Rename Route
			    local newname
			    newname=$(get_user_input "Enter the new name of '$username' to modify (e.g., clientA_uploads): " "^[a-zA-Z0-9_.-]+$" "${RED}Invalid route name. Please use only letters, numbers, underscores, periods, and hyphens.${NC}")
			    
			    local USER_CHROOT_DIR="${CHROOT_ROOT_DIR}/$username"
			    local NEW_USER_CHROOT_DIR="${CHROOT_ROOT_DIR}/$newname"

			    echo -e "${CYAN}Moving '${USER_CHROOT_DIR}' chroot directory for '$username' to '${NEW_USER_CHROOT_DIR}'...${NC}"
			    sudo usermod -l "$newname" -m -d "$NEW_USER_CHROOT" "$username"
			    sudo groupmod -n "$newname" "$username"
			    sudo usermod -aG "${SFTP_GROUP}" "$newname"
			    sudo gpasswd -d "$username" "${SFTP_GROUP}"
			    sudo chmod 755 "${NEW_USER_CHROOT_DIR}"
			    echo -e "${BOLD}${GREEN}SFTP route '$newname' setup complete!${NC}"
			    break
			    ;;
		    2) # Change Password
			    if confirm_action "Do you want to set a new password for '$username'? [Y/n] (n will remove any prevously set password) : "; then
				    local password
				    password=$(get_password_securely "$username")
				    if [ -n "$password" ]; then
					    echo -e "${CYAN}Setting new password for '$username'...${NC}"
					    echo "$username:$password" | sudo chpasswd
					    check_status "Password for '$username' updated successfully." "Failed to set new password for '$username'."
				    else
					    sudo passwd -d "$username"
					    echo -e "${YELLOW}Password removed.${NC}"
				    fi
			    fi
			    break
			    ;;
		    3) # Add SSH Public Key
			    # Call add_ssh_public_key with the confirmed username
			    add_ssh_public_key "$username"
			    break
			    ;;
		    4) # Return to Main Menu
			    break
			    ;;
		    *)
			    echo -e "${RED}Invalid option $choice_num. Please choose a number from the list.${NC}" >&2
			    ;;
	    esac
	    echo ""
	    clear
    done
    echo ""
}

# --- Function: Add SSH Public Key to an SFTP Route ---
add_ssh_public_key() {
	local username="$1" # Get username from argument if provided by modify_route
	# If called directly from main menu without a pre-selected username
	if [ -z "$username" ]; then
		echo -e "${BLUE}--- Add SSH Public Key ---${NC}"
		username=$(get_user_input "Enter the name of the ASFTP route to add the SSH key to (e.g., clientA_uploads): " "^[a-zA-Z0-9_.-]+$" "${RED}Invalid route name. Please use only letters, numbers, underscores, periods, and hyphens.${NC}")

		if [ -z "$username" ]; then
			echo -e "${YELLOW}No route name provided. Returning to main menu.${NC}"
			return 1
		fi

		local -a all_sftp_routes=($(_list_sftp_routes_array))
		local found="false"
		for existing_route in "${all_sftp_routes[@]}"; do
			if [ "$username" = "$existing_route" ]; then
				found="true"
				break
			fi
		done

		if [ "$found" = "false" ]; then
			echo -e "${RED}Error: ASFTP route '$username' not found. Please ensure the name is correct.${NC}"
			return 1
		fi
	else
		echo -e "${BLUE}--- Add SSH Public Key for Route '$username' ---${NC}"
	fi

	local USER_CHROOT_DIR="${CHROOT_ROOT_DIR}/$username"
	local SSH_DIR="${USER_CHROOT_DIR}/.ssh"
	local AUTHORIZED_KEYS_FILE="${SSH_DIR}/authorized_keys"

	echo -e "${CYAN}Creating .ssh directory for '$username' at '${SSH_DIR}'...${NC}"
	sudo mkdir -p "${SSH_DIR}"
	check_status ".ssh directory ensured." "Failed to create .ssh directory." || return 1

	sudo chmod 700 "${SSH_DIR}"
	check_status "Permissions for .ssh set to 700." "Failed to set .ssh directory permissions." || return 1
	sudo chown "$username":"${SFTP_GROUP}" "${SSH_DIR}"

	echo -e "${CYAN}Please paste the SSH public key content (e.g., ssh-rsa AAAA...) or type 'FILE' to provide a path.${NC}"
	local key_input
	key_input=$(get_user_input "Public key content or 'FILE': ")

	local public_key_content=""
	if [ "$key_input" = "FILE" ]; then
		local key_path
		key_path=$(get_user_input "Enter the path to the public key file (e.g., /home/youruser/id_rsa.pub): " "" "${RED}Invalid file path.${NC}")
		if [ -z "$key_path" ]; then
			echo -e "${RED}No valid key file path provided. Aborting.${NC}"
			return 1
		fi
		if [ ! -f "$key_path" ]; then
			echo -e "${RED}Error: Key file '$key_path' not found. Aborting.${NC}"
			return 1
		fi
		public_key_content=$(cat "$key_path" 2>/dev/null)
		if [ -z "$public_key_content" ]; then
			echo -e "${RED}Error: Key file '$key_path' is empty or unreadable. Aborting.${NC}"
			return 1
		fi
	else
		public_key_content="$key_input"
		if [ -z "$public_key_content" ]; then
			echo -e "${RED}No public key content provided. Aborting.${NC}"
			return 1
		fi
	fi

	echo -e "${CYAN}Adding public key to '${AUTHORIZED_KEYS_FILE}'...${NC}"
	echo "$public_key_content" | sudo tee -a "${AUTHORIZED_KEYS_FILE}" > /dev/null
	check_status "Public key added." "Failed to add public key." || return 1

	sudo chmod 600 "${AUTHORIZED_KEYS_FILE}"
	check_status "Permissions for authorized_keys set to 600." "Failed to set authorized_keys permissions." || return 1
	sudo chown "$username":"${SFTP_GROUP}" "${AUTHORIZED_KEYS_FILE}"

	echo -e "${BOLD}${GREEN}SSH Public Key successfully added for '$username'!${NC}"
	echo -e "${CYAN}The user can now connect via SFTP using this key.${NC}"
	echo ""
}

# --- Function: Add Writable Directory to an ASFTP Route ---
add_writable_dir_to_route() {
	clear
	echo -e "${BLUE}--- Add Writable Directory to ASFTP Route ---${NC}"
	display_routes

	local username
	username=$(get_user_input "Enter the name of the ASFTP route to add a writable directory to (e.g., clientA_uploads): " "^[a-zA-Z0-9_.-]+$" "${RED}Invalid route name. Please use only letters, numbers, underscores, periods, and hyphens.${NC}")

	if [ -z "$username" ]; then
		echo -e "${YELLOW}No route name provided. Returning to main menu.${NC}"
		return 1
	fi

	local -a all_sftp_routes=($(_list_sftp_routes_array))
	local found="false"
	for existing_route in "${all_sftp_routes[@]}"; do
		if [ "$username" = "$existing_route" ]; then
			found="true"
			break
		fi
	done

	if [ "$found" = "false" ]; then
		echo -e "${RED}Error: ASFTP route '$username' not found. Please ensure the name is correct.${NC}"
		return 1
	fi

	local USER_CHROOT_DIR="${CHROOT_ROOT_DIR}/$username"

	local foldername
	foldername=$(get_user_input "Enter the name of the directory to create.${NC}")

	local UPLOADS_DIR="${USER_CHROOT_DIR}/$foldername"

	echo -e "${CYAN}Checking/creating writable directory '${UPLOADS_DIR}' for route '$username'...${NC}"
	sudo mkdir -p "${UPLOADS_DIR}"
	check_status "Writable directory '${UPLOADS_DIR}' ensured." "Failed to create writable directory '${UPLOADS_DIR}'. Check permissions." || return 1

	sudo chown "$username":"${SFTP_GROUP}" "${UPLOADS_DIR}"
	check_status "Owner for '${UPLOADS_DIR}' set to $username:${SFTP_GROUP}." "Failed to set owner for '${UPLOADS_DIR}'." || return 1
	sudo chmod 775 "${UPLOADS_DIR}"
	check_status "Permissions for '${UPLOADS_DIR}' set to 775." "Failed to set permissions for '${UPLOADS_DIR}'." || return 1

	echo -e "${CYAN}Route '$username' can now upload files to '${UPLOADS_DIR}' inside their SFTP chroot.${NC}"
	echo ""
}

# --- Function: Remove an ASFTP Route ---
remove_route() {
	clear
	echo -e "${BLUE}--- Remove ASFTP Route ---${NC}"
	display_routes

	local username
	username=$(get_user_input "Enter the name of the ASFTP route to remove (e.g., clientA_uploads): " "^[a-zA-Z0-9_.-]+$" "${RED}Invalid route name. Please use only letters, numbers, underscores, periods, and hyphens.${NC}")

	if [ -z "$username" ]; then
		echo -e "${YELLOW}No route name provided. Returning to main menu.${NC}"
		return 1
	fi

	local -a all_sftp_routes=($(_list_sftp_routes_array))
	local found="false"
	for existing_route in "${all_sftp_routes[@]}"; do
		if [ "$username" = "$existing_route" ]; then
			found="true"
			break
		fi
	done

	if [ "$found" = "false" ]; then
		echo -e "${RED}Error: ASFTP route '$username' not found. Please ensure the name is correct.${NC}"
		return 1
	fi

	if confirm_action "${RED}WARNING: This will permanently delete route '$username' and its data. Are you sure? [y/N]: ${NC}"; then
		echo -e "${CYAN}Proceeding with removal of route '$username'...\n${NC}"

		sudo userdel "$username"
		check_status "User '$username' removed from system." "Failed to remove user '$username'. Manual intervention may be required."

		local USER_CHROOT_DIR="${CHROOT_ROOT_DIR}/$username"
		if [ -d "${USER_CHROOT_DIR}" ]; then
			echo -e "${CYAN}Removing chroot directory '${USER_CHROOT_DIR}'...${NC}"
			sudo rm -rf "${USER_CHROOT_DIR}"
			check_status "Chroot directory '${USER_CHROOT_DIR}' removed." "Failed to remove chroot directory '${USER_CHROOT_DIR}'. Manual intervention may be required."
		else
			echo -e "${YELLOW}Warning: Chroot directory '${USER_CHROOT_DIR}' not found. Skipping directory removal.${NC}"
		fi

		echo -e "${BOLD}${GREEN}SFTP route '$username' removed successfully.${NC}"
	else
		echo -e "${YELLOW}Route removal cancelled.${NC}"
	fi
	echo ""
}


# --- Function: Setup SSHD and PAM for SFTP ---
setup() {
	clear
	echo -e "${BLUE}--- Initializing OpenSSH Server Configuration ---${NC}"

	echo -e "${CYAN}1. Backing up OpenSSH configuration...${NC}"
	local SSHD_CONFIG_BAK="${SSHD_CONFIG_FILE}.bak"
	if [ ! -f "${SSHD_CONFIG_BAK}" ]; then
		sudo cp "${SSHD_CONFIG_FILE}" "${SSHD_CONFIG_BAK}"
		check_status "Backed up ${SSHD_CONFIG_FILE} to ${SSHD_CONFIG_BAK}" "Failed to backup ${SSHD_CONFIG_FILE}."
	else
		echo -e "${YELLOW}Warning: Backup of ${SSHD_CONFIG_FILE} already exists at ${SSHD_CONFIG_BAK}. Skipping.${NC}"
	fi
	echo ""

	echo -e "${CYAN}2. Configuring 'PermitEmptyPasswords yes'...${NC}"
	if grep -q "^PermitEmptyPasswords yes" "${SSHD_CONFIG_FILE}" > /dev/null 2>&1; then
		echo -e "${GREEN}✓ PermitEmptyPasswords is already 'yes'.${NC}"
	else
		if grep -q "^#PermitEmptyPasswords" "${SSHD_CONFIG_FILE}" > /dev/null 2>&1; then
			sudo sed -i.bak 's/^#PermitEmptyPasswords.*/PermitEmptyPasswords yes/' "${SSHD_CONFIG_FILE}"
			check_status "Set PermitEmptyPasswords to 'yes'." "Failed to set PermitEmptyPasswords."
		elif grep -q "^PermitEmptyPasswords" "${SSHD_CONFIG_FILE}" > /dev/null 2>&1; then
			sudo sed -i.bak 's/^PermitEmptyPasswords.*/PermitEmptyPasswords yes/' "${SSHD_CONFIG_FILE}"
			check_status "Updated PermitEmptyPasswords to 'yes'." "Failed to update PermitEmptyPasswords."
		else
			echo -e "PermitEmptyPasswords yes" | sudo tee -a "${SSHD_CONFIG_FILE}" > /dev/null
			check_status "Added PermitEmptyPasswords yes." "Failed to add PermitEmptyPasswords."
		fi
	fi
	echo ""

	echo -e "${CYAN}3. Configuring 'Subsystem sftp internal-sftp'...${NC}"
	if grep -q "^Subsystem sftp internal-sftp" "${SSHD_CONFIG_FILE}" > /dev/null 2>&1; then
		echo -e "${GREEN}✓ Subsystem sftp internal-sftp is already configured.${NC}"
	else
		if grep -q "^Subsystem sftp" "${SSHD_CONFIG_FILE}" > /dev/null 2>&1; then
			sudo sed -i.bak 's/^Subsystem sftp.*/Subsystem sftp internal-sftp/' "${SSHD_CONFIG_FILE}"
			check_status "Updated Subsystem sftp to internal-sftp." "Failed to update Subsystem sftp."
		else
			echo -e "Subsystem sftp internal-sftp" | sudo tee -a "${SSHD_CONFIG_FILE}" > /dev/null
			check_status "Added Subsystem sftp internal-sftp." "Failed to add Subsystem sftp."
		fi
	fi
	echo ""

	echo -e "${CYAN}4. Adding/Updating 'Match Group ${SFTP_GROUP}' block...${NC}"
	local MATCH_BLOCK="
	Match Group ${SFTP_GROUP}
	ChrootDirectory ${CHROOT_ROOT_DIR}/%u
	ForceCommand internal-sftp
	AllowTcpForwarding no
	X11Forwarding no
	PermitTunnel no
	PasswordAuthentication yes
	"
	sudo sed -i.bak "/^Match Group ${SFTP_GROUP}/,/^[[:space:]]*PasswordAuthentication yes/d" "${SSHD_CONFIG_FILE}" 2>/dev/null || true
	echo "${MATCH_BLOCK}" | sudo tee -a "${SSHD_CONFIG_FILE}" > /dev/null
	check_status "'Match Group ${SFTP_GROUP}' block added/updated in ${SSHD_CONFIG_FILE}." "Failed to add/update 'Match Group' block."
	echo ""

	echo -e "${CYAN}5. Modifying PAM configuration for empty passwords...${NC}"

	local PAM_SSHD_BAK="${PAM_SSHD_FILE}.bak"
	if [ ! -f "${PAM_SSHD_BAK}" ]; then
		sudo cp "${PAM_SSHD_FILE}" "${PAM_SSHD_BAK}"
		check_status "Backed up ${PAM_SSHD_FILE} to ${PAM_SSHD_BAK}" "Failed to backup ${PAM_SSHD_FILE}."
	else
		echo -e "${YELLOW}Warning: Backup of ${PAM_SSHD_FILE} already exists at ${PAM_SSHD_BAK}. Skipping.${NC}"
	fi

	if ! grep -q "^@include common-auth" "${PAM_SSHD_FILE}" > /dev/null 2>&1; then
		echo -e "${CYAN}Adding '@include common-auth' to ${PAM_SSHD_FILE}.${NC}"
		if grep -q "^#%PAM-1.0" "${PAM_SSHD_FILE}" > /dev/null 2>&1; then
			sudo sed -i.bak '/^#%PAM-1.0/a @include common-auth' "${PAM_SSHD_FILE}"
		else
			echo "@include common-auth" | sudo tee "${PAM_SSHD_FILE}.tmp" > /dev/null
			cat "${PAM_SSHD_FILE}" | sudo tee -a "${PAM_SSHD_FILE}.tmp" > /dev/null
			sudo mv "${PAM_SSHD_FILE}.tmp" "${PAM_SSHD_FILE}"
		fi
		sudo sed -i.bak '/^auth \[success=1 default=ignore\] pam_unix.so nullok/d' "${PAM_SSHD_FILE}"
		sudo sed -i.bak '/^auth requisite pam_deny.so/d' "${PAM_SSHD_FILE}"
		sudo sed -i.bak '/^auth required pam_permit.so/d' "${PAM_SSHD_FILE}"
		check_status "'@include common-auth' added to ${PAM_SSHD_FILE}." "Failed to add '@include common-auth'."
	else
		echo -e "${GREEN}✓ '@include common-auth' is already present in ${PAM_SSHD_FILE}.${NC}"
	fi
	echo ""

	local PAM_COMMON_AUTH_BAK="${PAM_COMMON_AUTH_FILE}.bak"
	if [ ! -f "${PAM_COMMON_AUTH_BAK}" ]; then
		sudo cp "${PAM_COMMON_AUTH_FILE}" "${PAM_COMMON_AUTH_BAK}"
		check_status "Backed up ${PAM_COMMON_AUTH_FILE} to ${PAM_COMMON_AUTH_BAK}" "Failed to backup ${PAM_COMMON_AUTH_FILE}."
	else
		echo -e "${YELLOW}Warning: Backup of ${PAM_COMMON_AUTH_FILE} already exists at ${PAM_COMMON_AUTH_BAK}. Skipping.${NC}"
	fi

	echo -e "${CYAN}Modifying 'pam_unix.so' to allow 'nullok' in ${PAM_COMMON_AUTH_FILE}...${NC}"
	if grep -q "pam_unix.so nullok" "${PAM_COMMON_AUTH_FILE}" > /dev/null 2>&1; then
		echo -e "${GREEN}✓ 'pam_unix.so nullok' already configured in ${PAM_COMMON_AUTH_FILE}.${NC}"
	elif grep -q "pam_unix.so nullok_secure" "${PAM_COMMON_AUTH_FILE}" > /dev/null 2>&1; then
		sudo sed -i.bak 's/pam_unix.so nullok_secure/pam_unix.so nullok/' "${PAM_COMMON_AUTH_FILE}"
		check_status "Changed 'pam_unix.so nullok_secure' to 'pam_unix.so nullok' in ${PAM_COMMON_AUTH_FILE}." "Failed to change 'pam_unix.so nullok_secure'."
	elif grep -q "pam_unix.so" "${PAM_COMMON_AUTH_FILE}" > /dev/null 2>&1; then
		sudo sed -i.bak 's/\(auth.*pam_unix.so\)/\1 nullok/' "${PAM_COMMON_AUTH_FILE}"
		check_status "Added 'nullok' to 'pam_unix.so' in ${PAM_COMMON_AUTH_FILE}." "Failed to add 'nullok' to 'pam_unix.so'."
	else
		echo -e "${YELLOW}Warning: 'pam_unix.so' line not found in ${PAM_COMMON_AUTH_FILE}. Manual check might be required.${NC}"
	fi
	echo ""

	echo -e "${BLUE}--- Restarting SSH Service ---${NC}"
	sudo systemctl restart ssh
	check_status "SSH service restarted successfully." "Failed to restart SSH service. Please check 'systemctl status ssh'." || return 1
	echo ""

	echo -e "${BOLD}${GREEN}SFTP Setup Completed!${NC}"
	echo -e "${CYAN}Please test your SFTP connection.${NC}"
	echo -e "${CYAN}Example: sftp your_route_name@YOUR_SERVER_IP${NC}"
	echo -e "${CYAN}If no password was set, or if using SSH key, just press ENTER when prompted for password.${NC}"
	echo -e "${CYAN}You should be chrooted to the route's specific directory.${NC}"
	echo ""
	return 0
}


# --- Main Menu Loop ---
while true; do
	clear
	echo -e "=== Welcome to the ASFTP Management Utility ==="

	display_routes 

    # Main menu options
    echo "1) Initialize ASFTP for the first time (SSH and PAM configuration)"
    echo "2) Add a new ASFTP Route (User/Chroot Setup)"
    echo "3) Modify an existing ASFTP Route (Change Password/Add SSH Key)"
    echo "4) Add a writable directory to an ASFTP Route (e.g., 'uploads' folder)"
    echo "5) Remove an ASFTP Route"
    echo "6) Quit"
    echo -e "${PURPLE}What would you like to do? Select an option: ${NC}"
    read -r main_choice

    case "$main_choice" in
	    1) setup ;;
	    2) add_route ;;
	    3) modify_route ;; # Now prompts for name
	    4) add_writable_dir_to_route ;; # Now prompts for name
	    5) remove_route ;; # Now prompts for name
	    6)
		    echo -e "${BLUE}Exiting ASFTP Management Utility. Goodbye!${NC}"
		    exit 0
		    ;;
	    *)
		    echo -e "${RED}Invalid option $main_choice. Please choose a number from the list.${NC}" >&2
		    ;;
    esac
    echo "" # Add a newline after each action
done
