#!/bin/bash

# Define the path to the public key file and the authorized_keys file
PUBLIC_KEY_FILE="$(mktemp)"
AUTHORIZED_KEYS_FILE="/root/.ssh/authorized_keys"
TEMP_AUTH_FILE="$(mktemp)"
COMMENT_LABEL="#ved.ssh.public.key"
HEADER_COMMENT="# ======= Ved SSH Keys - Managed Keys Below ======="

# Check if running as root
if [[ $EUID -ne 0 ]]; then
   echo "This script must be run as root" 
   exit 1
fi

# Download the keys file - fixed URL format
curl -fsSL https://ssh.ved.yt/keys/ssh.public.key > "$PUBLIC_KEY_FILE" || {
  echo "Failed to download keys. Exiting."
  exit 1
}

# Verify the download succeeded
if [[ ! -s "$PUBLIC_KEY_FILE" ]]; then
  echo "Downloaded key file is empty. Exiting."
  exit 1
fi

# Create the .ssh directory for the root user if it doesn't exist
if [[ ! -d "/root/.ssh" ]]; then
  echo "Creating /root/.ssh directory..."
  mkdir -p /root/.ssh
  chmod 700 /root/.ssh
fi

# Create authorized_keys file if it doesn't exist
if [[ ! -f "$AUTHORIZED_KEYS_FILE" ]]; then
  echo "Creating authorized_keys file..."
  touch "$AUTHORIZED_KEYS_FILE"
  chmod 600 "$AUTHORIZED_KEYS_FILE"
fi

# Initialize counters
added=0
skipped=0
invalid=0
labeled=0
removed=0
duplicates=0

# Load all valid keys from the public key file into an array with key type tracking
declare -a new_keys
declare -A key_types
declare -A source_fingerprints
while IFS= read -r key; do
  # Skip empty lines or comment lines
  if [[ -z "$key" || "$key" =~ ^[[:space:]]*# ]]; then
    continue
  fi
  
  # Extract key type and validate if this looks like an SSH key
  key_type=$(echo "$key" | awk '{print $1}')
  if [[ "$key_type" =~ ^(ssh-rsa|ssh-dss|ssh-ed25519|ecdsa-sha2-nistp|sk-ecdsa-sha2-nistp|sk-ssh-ed25519) ]]; then
    # Get fingerprint for deduplication
    fingerprint=$(echo "$key" | ssh-keygen -lf - 2>/dev/null | awk '{print $2}')
    
    # Only add this key if we haven't seen this fingerprint yet
    if [[ -z "${source_fingerprints[$fingerprint]}" ]]; then
      new_keys+=("$key")
      key_types["$key"]="$key_type"
      source_fingerprints["$fingerprint"]="$key"
    else
      echo "Skipping duplicate key in source: ${key:0:40}..."
      ((duplicates++))
    fi
  else
    echo "Invalid key format: ${key:0:40}..."
    ((invalid++))
  fi
done < "$PUBLIC_KEY_FILE"

echo "Found ${#new_keys[@]} valid unique keys in source file"

# Track fingerprints of keys to avoid duplicates
declare -A processed_fingerprints

echo "Processing SSH keys..."

# First process all existing keys to find duplicates and non-managed keys
while IFS= read -r line || [[ -n "$line" ]]; do
  # Skip empty lines and comment lines (header line will be added back later)
  if [[ -z "$line" || "$line" == "$HEADER_COMMENT" ]]; then
    continue
  fi
  
  # Skip lines with our label (we'll add back the managed keys from source)
  if [[ "$line" =~ ${COMMENT_LABEL} ]]; then
    continue
  fi
  
  # Check if this line contains a key
  key_part=$(echo "$line" | awk '{print $1}')
  if [[ "$key_part" =~ ^(ssh-rsa|ssh-dss|ssh-ed25519|ecdsa-sha2-nistp|sk-ecdsa-sha2-nistp|sk-ssh-ed25519) ]]; then
    # This is a key, get its fingerprint
    fingerprint=$(echo "$line" | ssh-keygen -lf - 2>/dev/null | awk '{print $2}')
    
    # If we have this in our source keys, skip it (we'll add it back with label)
    if [[ -n "${source_fingerprints[$fingerprint]}" ]]; then
      echo "Found unlabeled key that matches source: ${line:0:40}..."
      ((labeled++))
      continue
    fi
    
    # Skip duplicates we've already processed
    if [[ -n "${processed_fingerprints[$fingerprint]}" ]]; then
      echo "Skipping duplicate key: ${line:0:40}..."
      ((duplicates++))
      continue
    fi
    
    # Mark this fingerprint as processed
    processed_fingerprints["$fingerprint"]=1
  fi
  
  # Add this non-managed key to output
  echo "$line" >> "$TEMP_AUTH_FILE"
done < "$AUTHORIZED_KEYS_FILE"

# Add a blank line before the header if there are non-managed keys
if [[ -s "$TEMP_AUTH_FILE" ]]; then
  echo "" >> "$TEMP_AUTH_FILE"
fi

# Add header comment
echo "$HEADER_COMMENT" >> "$TEMP_AUTH_FILE"
echo "" >> "$TEMP_AUTH_FILE"  # Single blank line after header

# Reset for our managed section
unset processed_fingerprints
declare -A processed_fingerprints

# Now add all keys from the source file with our label
for key in "${new_keys[@]}"; do
  key_type="${key_types[$key]}"
  fingerprint=$(echo "$key" | ssh-keygen -lf - 2>/dev/null | awk '{print $2}')
  
  # Check if we've already processed this fingerprint (avoid duplicates)
  if [[ -n "${processed_fingerprints[$fingerprint]}" ]]; then
    echo "Skipping duplicate key in managed section: ${key:0:40}..."
    ((duplicates++))
    continue
  fi
  
  # Add the key with our label
  echo "$key $COMMENT_LABEL" >> "$TEMP_AUTH_FILE"
  echo "Added managed key ($key_type): $fingerprint"
  ((added++))
  
  # Mark this fingerprint as processed
  processed_fingerprints["$fingerprint"]=1
done

# Count removed keys (labeled keys in original minus kept labeled keys)
removed=$(grep -c "$COMMENT_LABEL" "$AUTHORIZED_KEYS_FILE")
if [[ "$removed" -gt 0 ]]; then
  removed=$((removed - added))
  if [[ "$removed" -lt 0 ]]; then
    removed=0
  fi
fi

# Replace the original authorized_keys with the temp file
mv "$TEMP_AUTH_FILE" "$AUTHORIZED_KEYS_FILE"
chmod 600 "$AUTHORIZED_KEYS_FILE"
rm -f "$PUBLIC_KEY_FILE"

echo "========== Summary =========="
echo "Operation completed:"
echo "- $added keys added/kept"
echo "- $labeled unlabeled keys now labeled"
echo "- $removed keys removed" 
echo "- $duplicates duplicates skipped"
echo "- $invalid invalid entries"
echo "All authorized keys are in $AUTHORIZED_KEYS_FILE"
echo "============================"

# Add functionality to create a cronjob for nightly updates
setup_cronjob() {
  echo "Setting up nightly cronjob for automatic updates..."
  
  # Check if running as root, as we need to modify root's crontab
  if [[ $EUID -ne 0 ]]; then
    echo "Error: Cronjob setup requires root privileges."
    return 1
  fi
  
  # Create a temporary file for crontab comparison
  TEMP_CRON=$(mktemp)
  
  # Get current crontab content
  crontab -l 2>/dev/null > "$TEMP_CRON" || echo "" > "$TEMP_CRON"
  
  # Check if an active (non-commented) cronjob for ssh.ved.yt already exists
  if grep -q "^[^#].*ssh\.ved\.yt" "$TEMP_CRON"; then
    echo "Active cronjob already exists. Skipping."
    rm -f "$TEMP_CRON"
    return 0
  fi
  
  # If we get here, either no cronjob exists or it's commented out
  # Remove any commented out versions of our job
  grep -v "ssh\.ved\.yt" "$TEMP_CRON" > "${TEMP_CRON}.new"
  mv "${TEMP_CRON}.new" "$TEMP_CRON"
  
  # Add our job to run at 3:00 AM daily
  echo "0 3 * * * curl -fsSL https://ssh.ved.yt | bash >/var/log/ssh-keys-update.log 2>&1" >> "$TEMP_CRON"
  
  # Install new crontab
  if crontab "$TEMP_CRON"; then
    echo "Cronjob successfully installed to run daily at 3:00 AM."
    echo "Update logs will be written to /var/log/ssh-keys-update.log"
  else
    echo "Failed to install cronjob."
  fi
  
  # Clean up
  rm -f "$TEMP_CRON"
}

# Add command line argument to force cronjob reinstall
SETUP_CRONJOB=true
FORCE_CRONJOB=false
for arg in "$@"; do
  case "$arg" in
    --no-cronjob)
      SETUP_CRONJOB=false
      ;;
    --force-cronjob)
      FORCE_CRONJOB=true
      ;;
  esac
done

# Setup cronjob if not disabled
if [[ "$SETUP_CRONJOB" == "true" ]]; then
  if [[ "$FORCE_CRONJOB" == "true" ]]; then
    # Force removal of any existing cronjob first
    crontab -l 2>/dev/null | grep -v "ssh\.ved\.yt" | crontab -
    echo "Removed existing cronjob entries."
  fi
  setup_cronjob
fi