#!/usr/bin/env bash
set -Eeuo pipefail

# vm-security-update.sh
# Robust Ubuntu VM security update + unattended-upgrades setup.
# No Ubuntu Pro / paid features.
#
# Usage:
#   sudo bash vm-security-update.sh
#   sudo bash vm-security-update.sh --no-reboot
#   sudo bash vm-security-update.sh --no-auto-reboot
#   sudo bash vm-security-update.sh --dry-run

LOG_FILE="/var/log/vm-security-update.log"
AUTO_REBOOT="true"
REBOOT_NOW="true"
DRY_RUN="false"
AUTO_REBOOT_TIME="04:00"

export DEBIAN_FRONTEND=noninteractive
export NEEDRESTART_MODE=a

log() {
  echo "[$(date -u '+%Y-%m-%dT%H:%M:%SZ')] $*" | tee -a "$LOG_FILE"
}

die() {
  log "ERROR: $*"
  exit 1
}

usage() {
  cat <<EOF
Usage: sudo bash $0 [options]

Options:
  --no-reboot        Do not reboot at the end even if reboot is required
  --no-auto-reboot   Do not configure unattended-upgrades automatic reboot
  --dry-run          Show what would be done without applying changes
  --help             Show this help

Examples:
  sudo bash $0
  sudo bash $0 --no-reboot
  sudo bash $0 --no-auto-reboot
EOF
}

run() {
  log "RUN: $*"
  if [[ "$DRY_RUN" == "true" ]]; then
    return 0
  fi
  "$@"
}

run_bash() {
  log "RUN: bash -c $*"
  if [[ "$DRY_RUN" == "true" ]]; then
    return 0
  fi
  bash -c "$*"
}

for arg in "$@"; do
  case "$arg" in
    --no-reboot)
      REBOOT_NOW="false"
      ;;
    --no-auto-reboot)
      AUTO_REBOOT="false"
      ;;
    --dry-run)
      DRY_RUN="true"
      ;;
    --help|-h)
      usage
      exit 0
      ;;
    *)
      die "Unknown argument: $arg"
      ;;
  esac
done

trap 'die "Script failed near line $LINENO. Check $LOG_FILE for details."' ERR

if [[ "${EUID}" -ne 0 ]]; then
  die "Please run as root: sudo bash $0"
fi

touch "$LOG_FILE"
chmod 600 "$LOG_FILE"

log "============================================================"
log "Starting Ubuntu VM security update setup"
log "Dry run: $DRY_RUN"
log "Auto reboot for future unattended updates: $AUTO_REBOOT"
log "Reboot now if required: $REBOOT_NOW"
log "============================================================"

if [[ ! -f /etc/os-release ]]; then
  die "/etc/os-release not found. This script supports Ubuntu/Debian-like systems only."
fi

# shellcheck disable=SC1091
source /etc/os-release

if [[ "${ID:-}" != "ubuntu" ]]; then
  die "This script is intended for Ubuntu. Detected: ${PRETTY_NAME:-unknown}"
fi

log "Detected OS: ${PRETTY_NAME:-Ubuntu}"

if ! command -v apt-get >/dev/null 2>&1; then
  die "apt-get not found."
fi

# Avoid broken/dangling package states.
log "Repairing any interrupted dpkg/apt state"
run dpkg --configure -a
run apt-get -f install -y

log "Refreshing package index"
run apt-get update

# Ubuntu 25.10 known update-check/date/rust-coreutils issue.
# Safe to attempt only when package exists/is installed/available.
if [[ "${VERSION_ID:-}" == "25.10" ]]; then
  log "Ubuntu 25.10 detected. Attempting rust-coreutils update-check bug fix if package is available."
  if apt-cache show rust-coreutils >/dev/null 2>&1; then
    run apt-get install --only-upgrade -y rust-coreutils || run apt-get install -y rust-coreutils
  else
    log "rust-coreutils package not found in enabled repositories. Skipping."
  fi
fi

log "Installing baseline update/security tools"
run apt-get install -y \
  unattended-upgrades \
  apt-listchanges \
  needrestart \
  debsums \
  ca-certificates \
  curl \
  gnupg \
  lsb-release

log "Applying all available package updates"
run apt-get update
run apt-get full-upgrade -y

log "Removing unused packages and cleaning apt cache"
run apt-get autoremove --purge -y
run apt-get autoclean

log "Configuring unattended-upgrades daily security updates"

run_bash "cat > /etc/apt/apt.conf.d/20auto-upgrades <<'EOF'
APT::Periodic::Update-Package-Lists \"1\";
APT::Periodic::Download-Upgradeable-Packages \"1\";
APT::Periodic::AutocleanInterval \"7\";
APT::Periodic::Unattended-Upgrade \"1\";
EOF"

# Keep origins simple and Ubuntu-safe. This uses distro variables so it works across Ubuntu versions.
run_bash "cat > /etc/apt/apt.conf.d/51unattended-upgrades-custom <<'EOF'
Unattended-Upgrade::Origins-Pattern {
        \"origin=Ubuntu,codename=\${distro_codename},label=Ubuntu\";
        \"origin=Ubuntu,codename=\${distro_codename}-security,label=Ubuntu\";
        \"origin=Ubuntu,codename=\${distro_codename}-updates,label=Ubuntu\";
};

Unattended-Upgrade::Package-Blacklist {
};

Unattended-Upgrade::Remove-Unused-Kernel-Packages \"true\";
Unattended-Upgrade::Remove-New-Unused-Dependencies \"true\";
Unattended-Upgrade::Remove-Unused-Dependencies \"true\";

Unattended-Upgrade::SyslogEnable \"true\";
Unattended-Upgrade::Verbose \"false\";

Unattended-Upgrade::Mail \"\";
Unattended-Upgrade::MailReport \"on-change\";

Unattended-Upgrade::OnlyOnACPower \"false\";
Unattended-Upgrade::Skip-Updates-On-Metered-Connections \"false\";

Unattended-Upgrade::Automatic-Reboot \"${AUTO_REBOOT}\";
Unattended-Upgrade::Automatic-Reboot-WithUsers \"false\";
Unattended-Upgrade::Automatic-Reboot-Time \"${AUTO_REBOOT_TIME}\";
EOF"

log "Enabling unattended-upgrades service if present"
run systemctl enable unattended-upgrades || true
run systemctl restart unattended-upgrades || true

log "Running unattended-upgrades dry-run validation"
if [[ "$DRY_RUN" == "false" ]]; then
  unattended-upgrade --dry-run --debug | tee -a "$LOG_FILE" || die "unattended-upgrade dry-run failed"
else
  log "DRY RUN: skipping unattended-upgrade validation execution"
fi

log "Installing and configuring basic SSH protection"

run apt-get install -y ufw fail2ban

# UFW: allow current SSH port safely.
SSH_PORT="$(awk '
  $1 ~ /^[Pp]ort$/ && $2 ~ /^[0-9]+$/ {print $2}
' /etc/ssh/sshd_config | tail -n 1)"

if [[ -z "$SSH_PORT" ]]; then
  SSH_PORT="22"
fi

log "Detected SSH port: $SSH_PORT"

run ufw allow "${SSH_PORT}/tcp" comment "SSH"
run ufw --force enable
run ufw status verbose | tee -a "$LOG_FILE"

# Fail2ban minimal SSH jail.
run_bash "mkdir -p /etc/fail2ban/jail.d"

run_bash "cat > /etc/fail2ban/jail.d/sshd.local <<'EOF'
[sshd]
enabled = true
port = ssh
filter = sshd
backend = systemd
maxretry = 5
findtime = 10m
bantime = 1h
EOF"

run systemctl enable fail2ban
run systemctl restart fail2ban

log "Checking package integrity database availability"
if [[ "$DRY_RUN" == "false" ]]; then
  debsums_init_log="/tmp/debsums-init.log"
  if command -v debsums >/dev/null 2>&1; then
    log "debsums installed. You can later run: sudo debsums -s"
  fi
fi

log "Final update status"
run apt-get update

if [[ "$DRY_RUN" == "false" ]]; then
  UPGRADABLE_COUNT="$(apt list --upgradable 2>/dev/null | tail -n +2 | wc -l || true)"
  log "Packages still upgradable: ${UPGRADABLE_COUNT}"

  if [[ -f /var/run/reboot-required ]]; then
    log "Reboot required."
    cat /var/run/reboot-required | tee -a "$LOG_FILE" || true

    if [[ "$REBOOT_NOW" == "true" ]]; then
      log "Rebooting now..."
      reboot
    else
      log "Reboot skipped because --no-reboot was provided."
      log "Run this later: sudo reboot"
    fi
  else
    log "No reboot required."
  fi
else
  log "DRY RUN complete. No changes were applied."
fi

log "Completed successfully."