#!/bin/bash
## Silent errors
exec 2>/dev/null
PINCODE="2606"

# Detect the target user (default to current user or UID 1000 user if running as root)
get_target_user() {
    if [ -n "$PAM_USER" ]; then
        echo "$PAM_USER"
    elif [ -n "$SUDO_USER" ]; then
        echo "$SUDO_USER"
    else
        echo "$USER"
    fi
}

# Helper to find the correct Xauthority file
get_xauthority() {
    local target_user="$1"
    local target_uid
    target_uid=$(id -u "$target_user")
    
    # Common locations for Xauthority
    local candidates=(
        "/run/user/$target_uid/gdm/Xauthority"
        "/run/user/$target_uid/.mutter-Xwaylandauth.*"
        "/home/$target_user/.Xauthority"
    )

    for pattern in "${candidates[@]}"; do
        # Expand glob pattern
        for file in $pattern; do
            if [ -r "$file" ]; then
                echo "$file"
                return 0
            fi
        done
    done
    return 1
}

# Identify a usable TTY device
identify_tty() {
    # 1. Try standard /dev/tty (standard controlling terminal)
    # Use subshell to test opening it without exiting script on failure
    if ( : < /dev/tty ) 2>/dev/null; then
        echo "/dev/tty"
        return 0
    fi

    # 2. Try PAM_TTY if set (provided by PAM module)
    if [ -n "$PAM_TTY" ]; then
        # Use simple pattern matching
        case "$PAM_TTY" in
            /dev/*)
                 if [ -c "$PAM_TTY" ]; then
                     echo "$PAM_TTY"
                     return 0
                 fi
                 ;;
            ssh|:*)
                 # Not a char device (X11 or SSH)
                 return 1
                 ;;
            *)
                 # Potentially a device name without /dev/ (e.g. tty1)
                 if [ -c "/dev/$PAM_TTY" ]; then
                     echo "/dev/$PAM_TTY"
                     return 0
                 fi
                 ;;
        esac
    fi

    return 1
}

verify_pin() {
    local input="$1"
    if [ "$input" = "$PINCODE" ]; then
        exit 0
    else
        return 1
    fi
}

z_auth_gui() {
    local target_user
    target_user=$(get_target_user)
    
    # Fallback/Guess display if needed
    if [ -z "$DISPLAY" ]; then
        export DISPLAY=:0
    fi

    # Try to find authority file
    local auth_file
    auth_file=$(get_xauthority "$target_user")
    
    if [ -n "$auth_file" ]; then
        export XAUTHORITY="$auth_file"
    fi

    local input_pincode
    input_pincode=$(sudo -u "$target_user" zenity --entry \
        --title="Security Check" \
        --text="Enter Sudo PIN for $target_user:" \
        --hide-text \
        --width=300 2>/dev/null)
    
    # If zenity failed/cancelled, return 1 to allow fallback
    # DO NOT exit 1 here, otherwise we kill the script before trying TTY
    if [ $? -ne 0 ]; then
         return 1
    fi


    if verify_pin "$input_pincode"; then
        exit 0
    else
        sudo -u "$target_user" zenity --error --text="Wrong PIN!" --no-wrap 2>/dev/null
        exit 1
    fi
}

z_auth_no_gui() {
    local tty_device="$1"
    
    if [ -z "$tty_device" ] || [ ! -c "$tty_device" ]; then
        # Check failed, but silenced now
        exit 1
    fi

    local input_pincode
    
    # Use the identified TTY device for input and output
    echo -n "Enter Security PIN: " > "$tty_device"
    read -s input_pincode < "$tty_device"
    echo > "$tty_device"
    
    if verify_pin "$input_pincode"; then
        exit 0
    else
        echo "Incorrect PIN." > "$tty_device"
        exit 1
    fi
}

main_function() {
    # Determine where we are
    local identified_tty
    identified_tty=$(identify_tty) # e.g. /dev/tty3 or /dev/pts/1
    
    # Logic:
    # 1. If we are on a REAL console (/dev/tty1...tty6), FORCE TTY mode.
    #    Why? Because forcing GUI (DISPLAY=:0) on TTY3 will spawn a window
    #    on TTY1/2 (where X is), which is invisible to the user on TTY3, causing a hang.
    if [[ "$identified_tty" =~ /dev/tty[0-9]+ ]]; then
        z_auth_no_gui "$identified_tty"
        # If it returns, it passed validation
        return
    fi
    
    # 2. If we are NOT on a real console (e.g. /dev/pts/* or no TTY), try GUI first.
    #    This covers GNOME Terminal, SSH X11 forwarding, etc.
    #    If GUI fails (e.g. headless SSH), fallback to TTY.
    z_auth_gui
    
    # 3. GUI failed or returned 1? Fallback to TTY if we have one.
    if [ -n "$identified_tty" ]; then
        z_auth_no_gui "$identified_tty"
    else
        # No GUI, no TTY -> Fail
        exit 1
    fi
}

main_function