#!/usr/bin/env bash

# Author        : Allan Christensen
# First Created : 12-03-2021 (DD-MM-YYYY)
# Description   : Manage unattended-upgrades policy on Ubuntu 24.04
# License       : MIT License

#
# Are we root
#
if [[ $(id -u) -ne 0 ]]; then echo "" && echo "Must be root or use sudo" && echo "" ; exit ; fi

#
# Variables
#
config="/etc/apt/apt.conf.d/50unattended-upgrades"
template="/usr/share/unattended-upgrades/50unattended-upgrades"
cronfile="/etc/cron.d/auto-update"
logfile="/var/log/auto-update.log"

#
# Log timestamp format
# Default: DD-MM-YYYY HH:MM
# Change to "+%Y-%m-%d %H:%M" for ISO format if preferred
#
log_date_format="+%d-%m-%Y %H:%M"

#
# Create safe temporary file
#
tmpfile="$(mktemp /tmp/auto-update.XXXXXX)"
trap 'rm -f "$tmpfile"' EXIT

#
# Function: Ensure config exists
#
chkcfg () {
    if [[ ! -f "$1" ]]; then
        printf "Config missing: %s → copying from %s\n" "$1" "$2"
        cp -Rp "$2" "$1"
    fi
}

#
# Function: Detect current mode
#
detect_mode() {

    managed="false"
    updates_enabled="false"
    reboot_enabled="false"
    reboot_users_enabled="false"
    current_mode="unknown"

    if grep -q "Managed by auto-update" "$config"; then
        managed="true"
    else
        return
    fi

    grep -q '${distro_codename}-updates' "$config" && updates_enabled="true"
    grep -q 'Automatic-Reboot "true"' "$config" && reboot_enabled="true"
    grep -q 'Automatic-Reboot-WithUsers "true"' "$config" && reboot_users_enabled="true"

    if [[ "$updates_enabled" == "true" && "$reboot_enabled" == "true" ]]; then
        current_mode=1
    elif [[ "$updates_enabled" == "true" && "$reboot_enabled" == "false" ]]; then
        current_mode=2
    elif [[ "$updates_enabled" == "false" && "$reboot_enabled" == "true" ]]; then
        current_mode=3
    else
        current_mode=4
    fi
}

#
# Function: Show status and usage
#
show_status() {

    detect_mode

    if [[ "$managed" != "true" ]]; then
        printf "\nCurrent policy: Not yet managed by auto-update\n"
        printf "System default unattended-upgrades configuration detected.\n"
    else
        printf "\nCurrent policy: Mode %s\n" "$current_mode"

        if [[ "$updates_enabled" == "true" ]]; then
            printf "Updates: Security + updates\n"
        else
            printf "Updates: Security only\n"
        fi

        if [[ "$reboot_enabled" == "true" ]]; then
            printf "Automatic reboot: Yes\n"
        else
            printf "Automatic reboot: No\n"
        fi

        if [[ "$reboot_users_enabled" == "true" ]]; then
            printf "Reboot even if users are logged in: Yes\n"
        else
            printf "Reboot even if users are logged in: No\n"
        fi
    fi

    printf "Cron enforcement: Friday at 03:00\n"

    printf "\nUsage:\n"
    printf "  auto-update <mode>\n\n"

    printf "Modes:\n"
    printf "  1  Security + updates, automatic reboot\n"
    printf "  2  Security + updates, no reboot\n"
    printf "  3  Security only, automatic reboot\n"
    printf "  4  Security only, no reboot\n\n"
}

#
# Function: Apply mode
#
apply_mode() {

    local mode="$1"

    cp "$template" "$tmpfile" || {
        echo "Failed to copy template. Aborting."
        exit 1
    }

    sed -i '/Unattended-Upgrade::Allowed-Origins {/,/};/d' "$tmpfile"
    sed -i '/Unattended-Upgrade::Automatic-Reboot /d' "$tmpfile"
    sed -i '/Unattended-Upgrade::Automatic-Reboot-WithUsers /d' "$tmpfile"

    case "$mode" in
        1)
            origins_block='        "${distro_id}:${distro_codename}-security";
        "${distro_id}:${distro_codename}-updates";'
            reboot_value="true"
            reboot_users_value="true"
            ;;
        2)
            origins_block='        "${distro_id}:${distro_codename}-security";
        "${distro_id}:${distro_codename}-updates";'
            reboot_value="false"
            reboot_users_value="false"
            ;;
        3)
            origins_block='        "${distro_id}:${distro_codename}-security";'
            reboot_value="true"
            reboot_users_value="true"
            ;;
        4)
            origins_block='        "${distro_id}:${distro_codename}-security";'
            reboot_value="false"
            reboot_users_value="false"
            ;;
        *)
            printf "\nInvalid mode.\n\n"
            exit 1
            ;;
    esac

    cat <<EOF >> "$tmpfile"

//
// Managed by auto-update
//

Unattended-Upgrade::Allowed-Origins {
$origins_block
};

Unattended-Upgrade::Automatic-Reboot "$reboot_value";
Unattended-Upgrade::Automatic-Reboot-WithUsers "$reboot_users_value";

EOF

    mv "$tmpfile" "$config"
}

#
# Function: Ensure cron job exists
#
ensure_cron() {

    if [[ ! -f "$cronfile" ]]; then

        printf "Creating cron enforcement job...\n"

        cat <<EOF > "$cronfile"
# Runs auto-update policy enforcement every Friday at 03:00
0 3 * * 5 root /usr/bin/apt-get update -qq && /usr/bin/unattended-upgrade -v >/dev/null 2>&1
EOF

        chmod 644 "$cronfile"
        chown root:root "$cronfile"

    else
        printf "Cron job already exists (%s), leaving it unchanged.\n" "$cronfile"
    fi
}

#
# Install required packages
#
apt-get update -qq
apt-get install -y unattended-upgrades update-notifier-common >/dev/null

#
# Ensure configuration exists
#
chkcfg "$config" "$template"

#
# Argument handling
#
if [[ -z "$1" ]]; then
    show_status
    exit 0
fi

mode="$1"

if [[ ! "$mode" =~ ^[1-4]$ ]]; then
    printf "\nInvalid mode. Please select 1-4.\n"
    show_status
    exit 1
fi

#
# Apply selected mode
#
printf "\nApplying Mode %s...\n" "$mode"

apply_mode "$mode"
ensure_cron

#
# Log action
#
actor="${SUDO_USER:-root}"
printf "%s Mode %s applied by %s\n" "$(date "$log_date_format")" "$mode" "$actor" >> "$logfile"

#
# Restart services
#
systemctl restart unattended-upgrades.service >/dev/null
systemctl restart cron.service >/dev/null

#
# Show resulting status
#
show_status

#
# All done
#
printf "\nAll Done...\n"

#
# End of script
#
