#!/bin/bash

# Ubuntu Security Hardening Script
# Author: Aloke Majumder
# GitHub: https://github.com/alokemajumder
# License: MIT License

# DISCLAIMER:
# This script is provided "AS IS" without warranty of any kind, express or implied. The author expressly disclaims any and all warranties, 
# express or implied, including any warranties as to the usability, suitability or effectiveness of any methods or measures this script 
# attempts to apply. By using this script, you agree that the author shall not be held liable for any damages resulting from the use of this script.

# Define vulnerable versions and other constants
VULNERABLE_VERSIONS=("5.6.0" "5.6.1")
STABLE_VERSION="5.4.6"
STABLE_VERSION_URL="https://excellmedia.dl.sourceforge.net/project/lzmautils/xz-5.4.6.tar.bz2"

# Detect OS and package manager
detect_os_and_package_manager() {
    if [[ -f /etc/debian_version ]] || grep -qi ubuntu /etc/os-release || grep -qi kali /etc/os-release; then
        PACKAGE_MANAGER="apt-get"
        VERSION_COMMAND="apt-cache policy xz-utils | grep 'Installed:' | awk '{print \$2}'"
    elif [[ -f /etc/redhat-release ]] || grep -qi centos /etc/os-release || grep -qi fedora /etc/os-release || grep -qi rhel /etc/os-release || grep -qi rocky /etc/os-release; then
        if command -v dnf >/dev/null 2>&1; then
            PACKAGE_MANAGER="dnf"
        else
            PACKAGE_MANAGER="yum"
        fi
        VERSION_COMMAND="$PACKAGE_MANAGER list installed xz | grep xz | awk '{print \$2}' | cut -d ':' -f 2"
    elif grep -qi opensuse /etc/os-release; then
        PACKAGE_MANAGER="zypper"
        VERSION_COMMAND="zypper info xz | grep 'Version:' | awk '{print \$3}'"
    else
        echo "Unsupported distribution."
        exit 1
    fi
}

# Get xz version without executing it
get_xz_version() {
    detect_os_and_package_manager
    # Execute the version command within the detected package manager's context
    eval "$VERSION_COMMAND"
}

# Upgrade xz to the latest version available in the repositories
upgrade_xz() {
    echo "Attempting to upgrade xz to the latest version using $PACKAGE_MANAGER..."
    case $PACKAGE_MANAGER in
        apt-get|apt)
            sudo $PACKAGE_MANAGER update && sudo $PACKAGE_MANAGER install -y xz-utils
            ;;
        dnf|yum)
            sudo $PACKAGE_MANAGER update -y xz
            ;;
        zypper)
            sudo $PACKAGE_MANAGER refresh && sudo $PACKAGE_MANAGER update -y xz
            ;;
    esac
    # Check the exit status of the last command
    if [ $? -eq 0 ]; then
        echo "xz has been successfully upgraded to the latest version."
        return 0
    else
        echo "Failed to upgrade xz to the latest version."
        return 1
    fi
}

# Install the stable version of xz if upgrade fails or if a vulnerable version is installed
install_stable_xz() {
    echo "Downloading and installing the stable xz version from $STABLE_VERSION_URL..."
    wget $STABLE_VERSION_URL -O "xz-$STABLE_VERSION.tar.bz2" && \
    tar -xjf "xz-$STABLE_VERSION.tar.bz2" && \
    cd "xz-$STABLE_VERSION" && \
    ./configure && make && sudo make install

    if [ $? -eq 0 ]; then
        echo "Successfully installed xz version $STABLE_VERSION."
    else
        echo "Failed to install the stable version of xz. Please consider manual installation."
    fi
}

# Main logic
CURRENT_VERSION=$(get_xz_version)

if [[ " ${VULNERABLE_VERSIONS[*]} " =~ " ${CURRENT_VERSION} " ]]; then
    echo "Detected vulnerable xz version: $CURRENT_VERSION. Attempting to upgrade..."
    if ! upgrade_xz; then
        echo "Upgrade unsuccessful. Attempting to install a non-vulnerable version..."
        install_stable_xz
    fi
    echo "It is recommended to reboot your system to apply changes. Would you like to reboot now? (yes/no)"
    read answer
    if [[ "$answer" == "yes" ]]; then
        echo "Rebooting..."
        sudo reboot
    else
        echo "Please ensure to reboot your system later."
    fi
else
    echo "xz version is not identified as vulnerable or xz is not installed."
fi
