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

# ubuntu-vm-security-update.sh
# Robust Ubuntu VM security update + unattended-upgrades setup.
# No Ubuntu Pro / paid features.
#
# Usage:
#   sudo bash ubuntu-vm-security-update.sh
#   sudo bash ubuntu-vm-security-update.sh --no-reboot
#   sudo bash ubuntu-vm-security-update.sh --no-auto-reboot
#   sudo bash ubuntu-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"
SSH_PORT_OVERRIDE=""
SCRIPT_NAME="$(basename "$0")"

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 $SCRIPT_NAME [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
  --auto-reboot-time HH:MM
                     Set unattended-upgrades auto reboot time (default: $AUTO_REBOOT_TIME)
  --ssh-port PORT    Override detected SSH port when configuring UFW
  --dry-run          Show what would be done without applying changes
  --help             Show this help

Examples:
  sudo bash $SCRIPT_NAME
  sudo bash $SCRIPT_NAME --no-reboot
  sudo bash $SCRIPT_NAME --no-auto-reboot
  sudo bash $SCRIPT_NAME --ssh-port 2222
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 "$*"
}

safe_write_file() {
  local target="$1"
  local mode="$2"
  local tmp

  tmp="$(mktemp)"
  cat > "$tmp"

  if [[ -f "$target" ]] && cmp -s "$tmp" "$target"; then
    log "No change needed for $target"
    rm -f "$tmp"
    return 0
  fi

  log "Writing $target"

  if [[ "$DRY_RUN" == "true" ]]; then
    rm -f "$tmp"
    return 0
  fi

  install -m "$mode" "$tmp" "$target"
  rm -f "$tmp"
}

have_systemd() {
  command -v systemctl >/dev/null 2>&1 && [[ -d /run/systemd/system ]]
}

ensure_service_enabled_and_restarted() {
  local service="$1"

  if ! have_systemd; then
    log "systemd not detected. Skipping enable/restart for $service"
    return 0
  fi

  run systemctl enable "$service"
  run systemctl restart "$service"
}

wait_for_fail2ban_ready() {
  local attempts=10
  local delay_seconds=1
  local i

  if [[ "$DRY_RUN" == "true" ]]; then
    log "DRY RUN: skipping fail2ban readiness check"
    return 0
  fi

  for ((i = 1; i <= attempts; i++)); do
    if fail2ban-client ping >/dev/null 2>&1; then
      log "fail2ban is ready"
      return 0
    fi
    sleep "$delay_seconds"
  done

  return 1
}

detect_ssh_port() {
  local detected_port=""
  local sshd_output=""

  if [[ -n "$SSH_PORT_OVERRIDE" ]]; then
    echo "$SSH_PORT_OVERRIDE"
    return 0
  fi

  if command -v sshd >/dev/null 2>&1; then
    if sshd_output="$(sshd -T 2>/dev/null)"; then
      detected_port="$(printf '%s\n' "$sshd_output" | awk '$1 == "port" {print $2}')"
      detected_port="$(printf '%s\n' "$detected_port" | head -n 1)"
    fi
  fi

  if [[ -z "$detected_port" ]] && [[ -f /etc/ssh/sshd_config ]]; then
    detected_port="$(awk '
      $1 ~ /^[Pp]ort$/ && $2 ~ /^[0-9]+$/ {print $2}
    ' /etc/ssh/sshd_config | tail -n 1)"
  fi

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

  echo "$detected_port"
}

validate_time_hhmm() {
  [[ "$1" =~ ^([01][0-9]|2[0-3]):[0-5][0-9]$ ]]
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    --no-reboot)
      REBOOT_NOW="false"
      shift
      ;;
    --no-auto-reboot)
      AUTO_REBOOT="false"
      shift
      ;;
    --auto-reboot-time)
      [[ $# -ge 2 ]] || die "--auto-reboot-time requires a value like 04:00"
      validate_time_hhmm "$2" || die "Invalid --auto-reboot-time value: $2"
      AUTO_REBOOT_TIME="$2"
      shift 2
      ;;
    --ssh-port)
      [[ $# -ge 2 ]] || die "--ssh-port requires a numeric port"
      [[ "$2" =~ ^[0-9]+$ ]] || die "Invalid --ssh-port value: $2"
      (( "$2" >= 1 && "$2" <= 65535 )) || die "SSH port must be between 1 and 65535"
      SSH_PORT_OVERRIDE="$2"
      shift 2
      ;;
    --dry-run)
      DRY_RUN="true"
      shift
      ;;
    --help|-h)
      usage
      exit 0
      ;;
    *)
      die "Unknown argument: $1"
      ;;
  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 $SCRIPT_NAME"
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 "Automatic reboot time: $AUTO_REBOOT_TIME"
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
run apt-get clean

log "Configuring unattended-upgrades daily security updates"

safe_write_file /etc/apt/apt.conf.d/20auto-upgrades 0644 <<'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.
safe_write_file /etc/apt/apt.conf.d/52bytelyst-unattended-upgrades 0644 <<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"
ensure_service_enabled_and_restarted unattended-upgrades

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="$(detect_ssh_port)"

log "Detected SSH port: $SSH_PORT"

run ufw allow "${SSH_PORT}/tcp" comment "SSH"
run ufw --force enable
if [[ "$DRY_RUN" == "false" ]]; then
  ufw status verbose | tee -a "$LOG_FILE"
else
  log "DRY RUN: skipping ufw status command"
fi

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

safe_write_file /etc/fail2ban/jail.d/sshd.local 0644 <<'EOF'
[sshd]
enabled = true
port = ssh
filter = sshd
backend = auto
maxretry = 5
findtime = 10m
bantime = 1h
EOF

ensure_service_enabled_and_restarted fail2ban

if [[ "$DRY_RUN" == "false" ]] && have_systemd && command -v fail2ban-client >/dev/null 2>&1; then
  wait_for_fail2ban_ready || die "fail2ban-client ping failed"
  fail2ban-client ping | tee -a "$LOG_FILE"
fi

log "Checking package integrity database availability"
if [[ "$DRY_RUN" == "false" ]]; then
  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."