#!/bin/bash
set -euo pipefail

BANNER=$(cat << "EOF"
           :^7J5GB##&&##GPY?~:            
        ^75B&@@@@@@&&&@@@@@@@#GJ~:        
       5&@@@&B5?7~^^^^^~!7YP#@@@@#!       
       Y##P7^               :~JB#B!       
        ::                      :         
  7PP?:          :^~!!~^:          :?PP7  
 :B@@B:        !5B&@@@@&B5!        :#@@B: 
  :!!:       ^G@@@&BPPB@@@@G^       :!!:  
            :B@@@5^    ^5@@@B:            
:7J7:       !@@@#       :&@@@~       :?J7:
J@@@5       :#@@@Y:    :Y@@@B:       5@@@J
!@@@&^       ~B@@@&G55G&@@@B~       ~&@@@~
 5@@@G:       :7P#@@@@@@#P7:       :B@@@Y 
 :P@@@B~         :~!77!~:         ~B@@@P  
   Y@@@&Y^                      ^5@@@@J   
    !G@@@&P7^                ^7P&@@@G~    
      !P&@@@&B?     ::     ?B&@@@&P!      
        ^75#&&Y   :P&&5:   5&&B57^        
           :^^    :P&&5:    ^^:           
                    ^^                    

                [nullpoint]
EOF
)

TANG_SERVERS=(
#    "https://tang1.example.com your-thumbprint-1"
#    "https://tang2.example.com your-thumbprint-2"
)
TPM_ENABLED=true
TPM_PCR_BANK="sha256"
TPM_PCR_IDS="0,1,2,3,4,5,6,7,8"
ALMA_USER="null"
ENABLE_MOTD=true
# REQUIRED: Set your SSH public key here - installation will fail without it!
SSH_KEY="ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIOkoTn2NreAXMriOUqzyj3YoFW6jMo9B5B+3R5k8yrMi dodox@dodox-ProArt"

########################################################
# Config End
########################################################

set -euo pipefail

echo -e "\n$BANNER"
echo -e "\n[+] Starting post-installation configuration..."

# Check for SSH key
if [ -z "${SSH_KEY:-}" ]; then
    echo "ERROR: No SSH key configured!"
    echo "You must set SSH_KEY variable at the top of this script."
    exit 1
fi

# Check for TPM
echo "[+] Checking for TPM..."
if [ ! -d "/sys/class/tpm/tpm0" ]; then
    echo "WARNING: No TPM detected!"
    TPM_ENABLED=false
else
    echo "TPM detected."
    TPM_ENABLED=true
fi

# Install basic packages first
echo "[+] Installing basic packages..."
dnf install -y epel-release || exit 1
dnf config-manager --set-enabled crb || exit 1
dnf install -y zsh git wget curl || exit 1

# Create user and add SSH key
echo "[+] Creating user ${ALMA_USER}..."
useradd -m -G wheel -s /bin/zsh ${ALMA_USER} || exit 1
mkdir -p /home/${ALMA_USER}/.ssh
echo "${SSH_KEY}" > /home/${ALMA_USER}/.ssh/authorized_keys
chmod 700 /home/${ALMA_USER}/.ssh
chmod 600 /home/${ALMA_USER}/.ssh/authorized_keys
chown -R ${ALMA_USER}:${ALMA_USER} /home/${ALMA_USER}/.ssh

# Install oh-my-zsh and powerlevel10k
echo "[+] Installing oh-my-zsh and powerlevel10k..."
# Download and run oh-my-zsh installer as the user with proper environment
su - ${ALMA_USER} -c 'RUNZSH=no CHSH=no bash -c "$(wget -O- https://raw.githubusercontent.com/ohmyzsh/ohmyzsh/master/tools/install.sh)"' || echo "WARNING: oh-my-zsh installation failed"
# Clone powerlevel10k theme
su - ${ALMA_USER} -c 'git clone --depth=1 https://github.com/romkatv/powerlevel10k.git ~/.oh-my-zsh/custom/themes/powerlevel10k' || echo "WARNING: powerlevel10k installation failed"

# Install dotfiles from git repo
echo "[+] Installing dotfiles..."
su - ${ALMA_USER} -c 'cd && git clone https://git.dominik-roth.eu/dodox/nullpoint.git /tmp/nullpoint-dotfiles && cp /tmp/nullpoint-dotfiles/dotfiles/.* . 2>/dev/null || true && rm -rf /tmp/nullpoint-dotfiles' || echo "WARNING: dotfiles installation failed"

# Set up MOTD
if [ "$ENABLE_MOTD" = true ]; then
    echo "[+] Setting up MOTD..."
    cat > /etc/motd << MOTD
$BANNER
MOTD
fi

# Install additional packages
echo "[+] Installing additional packages..."
dnf install -y \
  clevis clevis-luks tpm2-tools tpm2-tss \
  tmux neovim python3-pip \
  tree gcc make autoconf automake tar bzip2 || exit 1

# Build and install tinyssh from source since AlmaLinux doesn't package it
echo "[+] Installing tinyssh from source..."
cd /tmp || exit 1
wget -q https://github.com/janmojzis/tinyssh/archive/refs/tags/20250126.tar.gz || exit 1
tar xf 20250126.tar.gz || exit 1
cd tinyssh-20250126 || exit 1
make || exit 1
make install PREFIX=/usr/local || exit 1
ln -sf /usr/local/bin/tinysshd /usr/bin/tinysshd
ln -sf /usr/local/bin/tinyssh-keyconvert /usr/bin/tinyssh-keyconvert

# Install lsd and bat
echo "[+] Installing lsd and bat..."
# Install using fixed versions that should work
LSD_VERSION="1.0.0"
BAT_VERSION="0.24.0"

# Download and install lsd
if curl -sL "https://github.com/lsd-rs/lsd/releases/download/v${LSD_VERSION}/lsd-v${LSD_VERSION}-x86_64-unknown-linux-musl.tar.gz" | tar xz -C /tmp; then
    mv /tmp/lsd-*/lsd /usr/local/bin/ 2>/dev/null || true
    chmod +x /usr/local/bin/lsd 2>/dev/null || true
fi

# Download and install bat
if curl -sL "https://github.com/sharkdp/bat/releases/download/v${BAT_VERSION}/bat-v${BAT_VERSION}-x86_64-unknown-linux-musl.tar.gz" | tar xz -C /tmp; then
    mv /tmp/bat-*/bat /usr/local/bin/ 2>/dev/null || true
    chmod +x /usr/local/bin/bat 2>/dev/null || true
fi

# Create batman script for fancy man pages
cat > /usr/local/bin/batman << 'BATMAN'
#!/bin/bash
export MANPAGER="sh -c 'col -bx | bat -l man -p'"
export MANROFFOPT="-c"
man "$@"
BATMAN
chmod +x /usr/local/bin/batman

# Create .tmp directory for user
mkdir -p /home/${ALMA_USER}/.tmp
chown ${ALMA_USER}:${ALMA_USER} /home/${ALMA_USER}/.tmp

# Configure Clevis for automatic unlock
if [ ${#TANG_SERVERS[@]} -gt 0 ] || [ "$TPM_ENABLED" = true ]; then
    echo "[+] Configuring Clevis for automatic unlock..."
    
    # Find LUKS devices
    LUKS_DEVICES=$(lsblk -o NAME,FSTYPE -nr | grep crypto_LUKS | cut -d' ' -f1)
    
    for device in $LUKS_DEVICES; do
        DEVICE_PATH="/dev/${device}"
        echo "Configuring Clevis for ${DEVICE_PATH}..."
        
        if [ "$TPM_ENABLED" = true ] && [ ${#TANG_SERVERS[@]} -eq 0 ]; then
            # TPM only
            clevis luks bind -d "$DEVICE_PATH" tpm2 "{\"pcr_bank\":\"$TPM_PCR_BANK\",\"pcr_ids\":\"$TPM_PCR_IDS\"}"
        elif [ "$TPM_ENABLED" = false ] && [ ${#TANG_SERVERS[@]} -gt 0 ]; then
            # Tang only
            for server in "${TANG_SERVERS[@]}"; do
                read -r url thumbprint <<< "$server"
                clevis luks bind -d "$DEVICE_PATH" tang "{\"url\":\"$url\",\"thp\":\"$thumbprint\"}"
            done
        elif [ "$TPM_ENABLED" = true ] && [ ${#TANG_SERVERS[@]} -gt 0 ]; then
            # Both TPM and Tang (require both)
            CONFIG="{\"t\":2,\"pins\":{"
            CONFIG+="\"tpm2\":{\"pcr_bank\":\"$TPM_PCR_BANK\",\"pcr_ids\":\"$TPM_PCR_IDS\"},"
            CONFIG+="\"tang\":{\"t\":1,\"tang\":["
            for server in "${TANG_SERVERS[@]}"; do
                read -r url thumbprint <<< "$server"
                CONFIG+="{\"url\":\"$url\",\"thp\":\"$thumbprint\"},"
            done
            CONFIG="${CONFIG%,}]}}}"
            clevis luks bind -d "$DEVICE_PATH" sss "$CONFIG"
        fi
    done
fi

# Enable Clevis for early boot
echo "[+] Enabling Clevis for early boot..."
systemctl enable clevis-luks-askpass.path || echo "WARNING: clevis-luks-askpass.path not found, may be enabled by default in AlmaLinux 9.1+"

# Configure tinyssh for remote unlock
echo "[+] Configuring tinyssh for remote unlock..."

# Create tinyssh dracut module
mkdir -p /usr/lib/dracut/modules.d/60tinyssh

cat > /usr/lib/dracut/modules.d/60tinyssh/module-setup.sh << 'EOF'
#!/bin/bash
check() {
    require_binaries tinysshd tinyssh-keyconvert || return 1
    return 0
}

depends() {
    echo network
    return 0
}

install() {
    inst_multiple tinysshd tinyssh-keyconvert ssh-keygen
    mkdir -p "$initdir/etc/tinyssh"
    
    # Copy authorized keys
    [ -f /etc/tinyssh/authorized_keys ] && inst /etc/tinyssh/authorized_keys /etc/tinyssh/authorized_keys
    
    # Copy host keys
    inst /etc/tinyssh/sshkeydir /etc/tinyssh/ 2>/dev/null || true
    
    # Install startup script
    inst_hook cmdline 60 "$moddir/tinyssh-start.sh"
    
    # Install unlock script
    inst_simple "$moddir/unlock-luks.sh" /bin/unlock-luks
    chmod +x "$initdir/bin/unlock-luks"
}
EOF

cat > /usr/lib/dracut/modules.d/60tinyssh/tinyssh-start.sh << 'EOF'
#!/bin/bash
info "Starting tinyssh SSH server on port 2222..."
[ -d /etc/tinyssh ] || mkdir -p /etc/tinyssh
[ -d /etc/tinyssh/sshkeydir ] || {
    mkdir -p /etc/tinyssh/sshkeydir
    ssh-keygen -A
    for key in /etc/ssh/ssh_host_*_key; do
        [ -f "$key" ] && tinyssh-keyconvert "$key" /etc/tinyssh/sshkeydir/
    done
}
tinysshd -p 2222 /etc/tinyssh/sshkeydir &
EOF

cat > /usr/lib/dracut/modules.d/60tinyssh/unlock-luks.sh << 'EOF'
#!/bin/bash
echo "Available LUKS devices:"
ls /dev/mapper/luks-* 2>/dev/null
echo "Run: systemctl start systemd-cryptsetup@<device>.service"
echo "Or use: cryptsetup luksOpen /dev/md/<X> <name>"
echo "Then: exit"
/bin/bash
EOF

chmod +x /usr/lib/dracut/modules.d/60tinyssh/*.sh

# Setup tinyssh
mkdir -p /etc/tinyssh/sshkeydir
echo "${SSH_KEY}" > /etc/tinyssh/authorized_keys
chmod 600 /etc/tinyssh/authorized_keys

# Generate host keys using OpenSSH then convert to tinyssh format
ssh-keygen -A
for key in /etc/ssh/ssh_host_*_key; do
    [ -f "$key" ] && tinyssh-keyconvert "$key" /etc/tinyssh/sshkeydir/
done

# Configure dracut
cat > /etc/dracut.conf.d/99-tinyssh.conf << 'EOF'
add_dracutmodules+=" network tinyssh "
install_items+=" /etc/tinyssh/authorized_keys /etc/tinyssh/sshkeydir "
EOF

# Regenerate initramfs
echo "[+] Regenerating initramfs..."
dracut -f --regenerate-all

# Enable required services
echo "[+] Enabling services..."
# systemctl enable stratisd  # Not needed without Stratis
systemctl enable sshd

# Disable root login
echo "[+] Securing SSH..."
sed -i 's/^#PermitRootLogin.*/PermitRootLogin no/' /etc/ssh/sshd_config

# Set SELinux to enforcing
echo "[+] Setting SELinux to enforcing..."
sed -i 's/^SELINUX=.*/SELINUX=enforcing/' /etc/selinux/config

echo "✅ Post-installation complete!"
echo ""
echo "IMPORTANT: The LUKS passphrase is set in install.conf"
echo "Save it securely for recovery purposes."
echo ""
echo "After reboot:"
echo "- SSH to port 2222 for remote unlock: ssh -p 2222 root@<server-ip>"
echo "- Run 'unlock-luks' and follow the instructions to unlock LUKS"
echo "- Once unlocked, SSH to port 22 as user '${ALMA_USER}'"
echo "- LUKS passphrase: [see installer output]"