#!/usr/bin/python3

import sys, os, apt
import subprocess
import filecmp

ORIGIN = "Linux Mint 18.3 'Sylvia'"
ORIGIN_CODENAME = "sylvia"
ORIGIN_BASE_CODENAME = "xenial"

DESTINATION = "Linux Mint 19 'Tara'"
DESTINATION_CODENAME = "tara"
DESTINATION_BASE_CODENAME = "bionic"

SUPPORTED_EDITIONS = ["cinnamon", "mate", "xfce"]

CHECK_ABSENT = ["gnome-media"]
CHECK_PRESENT = ["default-jre"]
CHECK_UP_TO_DATE = ["mintupgrade", "apt", "dpkg", "linuxmint-keyring", "debian-archive-keyring", "ubuntu-archive-keyring"]

BACKUP_APT_SOURCES = os.path.expanduser("~/Upgrade-Backup/APT/")
BACKUP_FSTAB = os.path.expanduser("~/Upgrade-Backup/fstab")

PACKAGES_PRE_REMOVALS = []
PACKAGES_REMOVALS = ["vino", "banshee", "gnome-user-share"]
PACKAGES_ADDITIONS = ["rhythmbox", "libblockdev-crypto2"]

IMPORTANT_PACKAGES = ["cinnamon", "xplayer", "xreader", "xed", "mintsystem", "metacity", "nemo", "nemo-preview"]

class bcolors:
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'

class MintUpgrade():

    def __init__(self):
        # Check the Mint info file
        if not os.path.exists("/etc/linuxmint/info"):
            self.fail("Missing file '/etc/linuxmint/info'.")

        # Check the edition
        self.mint_codename = 'unknown'
        self.mint_edition = 'unknown'
        self.mint_meta = 'unknown'
        with open("/etc/linuxmint/info", "r") as info:
            for line in info:
                line = line.strip()
                if "EDITION=" in line:
                    self.mint_edition = line.split('=')[1].replace('"', '').split()[0]
                    self.mint_meta = "mint-meta-%s" % self.mint_edition.lower()
                if "CODENAME=" in line:
                    self.mint_codename = line.split('=')[1].replace('"', '').split()[0]
        self.points_to_destination = False
        if os.path.exists("/etc/apt/sources.list.d/official-package-repositories.list"):
            with open("/etc/apt/sources.list.d/official-package-repositories.list") as sources:
                for line in sources:
                    if DESTINATION_CODENAME in line:
                        self.points_to_destination = True
                        break

    def restore_sources(self):
        self.progress("Restoring your backed up APT sources")
        if not os.path.exists(BACKUP_APT_SOURCES):
            self.fail("Missing backup %s" % BACKUP_APT_SOURCES)
        self.check_command("sudo mkdir -p /etc/apt/sources.list.d", "Failed to restore APT sources")
        self.check_command("sudo rm -rf /etc/apt/sources.list.d/*", "Failed to restore APT sources")
        self.check_command("sudo cp -R %s/* /etc/apt/" % BACKUP_APT_SOURCES, "Failed to restore APT sources")
        self.check_command("sudo rm -rf '%s'" % BACKUP_APT_SOURCES, "Failed to restore APT sources")

    def prepare(self):
        # Check codename
        self.progress("Checking your Linux Mint codename")
        if self.mint_codename != ORIGIN_CODENAME and self.mint_codename != DESTINATION_CODENAME:
            self.fail("Your version of Linux Mint is '%s'. Only %s can be upgraded to %s." % (self.mint_codename.capitalize(), ORIGIN, DESTINATION))

        # Check edition
        self.progress("Checking your Linux Mint edition")
        if self.mint_edition.lower() not in SUPPORTED_EDITIONS:
            self.fail("Your edition of Linux Mint is '%s'. It cannot be upgraded to %s." % (self.mint_edition, DESTINATION))

        # Check for timeshift configuration
        self.progress("Checking your Timeshift configuration")
        if not os.path.exists("/etc/timeshift.json"):
            self.fail("Please set up system snapshots. If anything goes wrong with the upgrade, snapshots will allow you to restore your operating system. Install and configure Timeshift, and create a snapshot before proceeding with the upgrade.")

        # Check display manager
        self.progress("Checking your Display Manager")
        dm_file = "/etc/X11/default-display-manager"
        if os.path.exists(dm_file):
            with open(dm_file) as dm_handle:
                if "mdm" in dm_handle.read():
                    self.fail("MDM is no longer supported in %s, please switch to LightDM." % DESTINATION)

        if not self.points_to_destination:
            # Check packages
            self.progress("Checking packages")
            cache = apt.Cache()
            for pkg in CHECK_ABSENT:
                if pkg in cache:
                    pkg = cache[pkg]
                    if pkg.is_installed:
                        self.fail("Please remove %s. It is known to create issues with this upgrade." % pkg.name)
            for pkg in CHECK_PRESENT:
                if pkg in cache:
                    pkg = cache[pkg]
                    if not pkg.is_installed:
                        self.fail("Please install %s. It is required for a smooth upgrade." % pkg.name)

            self.progress("Updating cache")
            os.system("DEBIAN_PRIORITY=critical sudo apt-get update")
            cache = apt.Cache()

            # Check that we're up to date
            self.progress("Checking if Linux Mint is up to date")
            for pkg in CHECK_UP_TO_DATE:
                if pkg in cache:
                    pkg = cache[pkg]
                    if pkg.is_installed and pkg.installed.version != pkg.candidate.version:
                        self.fail("Your operating system is not up to date. Please apply available updates and reboot the computer.")

        # Switch to the destination APT sources
        if not os.path.exists(BACKUP_APT_SOURCES):
            self.progress("Backing up your APT sources")
            messages = []
            messages.append("Your repositories will now be switched to point to %s." % DESTINATION)
            messages.append("Any 3rd party repositories or PPA will be removed.")
            messages.append("A backup of your APT sources will be written to %s." % BACKUP_APT_SOURCES)
            self.continue_yes_no(messages)
            os.system("mkdir -p %s" % BACKUP_APT_SOURCES)
            os.system("cp -R /etc/apt/sources.* %s/" % BACKUP_APT_SOURCES)
        self.progress("Setting up the repositories for %s" % DESTINATION)
        if os.path.exists("/etc/apt/sources.list"):
            self.check_command("sudo truncate --size 0 /etc/apt/sources.list", "Failed to configure APT sources")
        self.check_command("sudo mkdir -p /etc/apt/sources.list.d", "Failed to configure APT sources")
        self.check_command("sudo rm -rf /etc/apt/sources.list.d/*", "Failed to configure APT sources")
        self.check_command("sudo cp /usr/share/linuxmint/mintupgrade/apt_destination_sources /etc/apt/sources.list.d/official-package-repositories.list", "Failed to configure APT sources")
        self.check_command("DEBIAN_PRIORITY=critical sudo apt-get update", "Failed to configure APT sources")

    def check(self):
        self.progress("Simulating an upgrade")
        messages = []
        messages.append("APT will now calculate the package changes necessary to upgrade to %s." % DESTINATION)
        messages.append("If conflicts are detected and APT is unable to perform the upgrade, take note of the packages causing the issue, remove them, and re-install them after the upgrade.")
        messages.append("Pay close attention to what appears on the screen, and review the list of packages being REMOVED during the upgrade.")
        messages.append("Take note of the packages being removed, so you can eventually reinstall them after the upgrade.")
        self.continue_yes_no(messages)
        os.system('DEBIAN_PRIORITY=critical sudo apt-get dist-upgrade -o Dpkg::Options::="--force-confnew" -o Dpkg::Options::="--force-overwrite" --assume-no')

        cache = apt.Cache()
        cache.upgrade(True)
        changes = cache.get_changes()
        incorrect_removals = []
        kept_packages = []
        for pkg in changes:
            if pkg.is_installed:
                if pkg.marked_keep:
                    kept_packages.append(pkg.name)
                elif pkg.marked_delete and pkg.name in IMPORTANT_PACKAGES:
                    incorrect_removals.append(pkg.name)
        if len(incorrect_removals) > 0:
            self.restore_sources()
            self.fail("Performing the upgrade would remove the following important packages: %s." % ", ".join(sorted(incorrect_removals)))
        if len(kept_packages) > 0:
            self.warn("The following packages will be kept back during the upgrade: %s. \n\nThis might or might not indicate a problem. Check the APT output above to decide whether to continue with the upgrade. If this OK you might need to update them manually after the upgrade." % ", ".join(sorted(kept_packages)))
        
    def download(self):
        self.progress("Downloading upgrade packages")
        messages = []
        messages.append("APT will now download the package updates necessary for the upgrade to %s." % DESTINATION)
        self.continue_yes_no(messages)
        self.check_command("DEBIAN_PRIORITY=critical sudo apt-get dist-upgrade --download-only --yes", "Failed to download packages for the upgrade.")

    def upgrade(self):
        if self.mint_edition.lower() == "cinnamon":
            self.progress("Disabling the Cinnamon screensaver")
            os.system("killall cinnamon-screensaver")
        elif self.mint_edition.lower() == "mate":
            self.progress("Disabling the MATE screensaver")
            os.system("killall mate-screensaver")

        self.progress("Saving /etc/fstab")
        os.system("cp /etc/fstab %s" % BACKUP_FSTAB)

        self.progress("Removing blacklisted packages")
        for removal in PACKAGES_PRE_REMOVALS:
            os.system('sudo apt-get remove --yes %s' % removal) # The return code indicates a failure if some packages were not found, so ignore it.

        self.progress("Performing upgrade")
        messages = []
        messages.append("APT will perform the upgrade to %s." % DESTINATION)
        messages.append("This operation is non-reversible.")
        messages.append("Make sure you made backups, you tested %s in live mode and you performed your favorite superstitious tricks before proceeding." % DESTINATION)
        self.continue_yes_no(messages)

        # Disable mintsystem during the upgrade
        os.system("sudo crudini --set /etc/linuxmint/mintSystem.conf global enabled False")

        fallback_commands = []
        fallback_commands.append("sudo dpkg --configure -a")
        fallback_commands.append("sudo apt-get install -fyq")

        result = self.try_command(5, 'DEBIAN_FRONTEND=noninteractive DEBIAN_PRIORITY=critical sudo apt-get upgrade -fyq -o Dpkg::Options::="--force-confnew" -o Dpkg::Options::="--force-overwrite"', fallback_commands)
        if not result:
            self.progress("An issue was detected during the upgrade, running the upgrade in manual mode.")
            self.check_command('sudo apt-get upgrade -o Dpkg::Options::="--force-confnew" -o Dpkg::Options::="--force-overwrite"', "Failed to upgrade some of the packages. Please review the error message, use APT to fix the situation and try again.")

        result = self.try_command(5, 'DEBIAN_FRONTEND=noninteractive DEBIAN_PRIORITY=critical sudo apt-get dist-upgrade -fyq -o Dpkg::Options::="--force-confnew" -o Dpkg::Options::="--force-overwrite"', fallback_commands)
        if not result:
            self.progress("An issue was detected during the upgrade, running dist-upgrade in manual mode.")
            self.check_command('sudo apt-get dist-upgrade -o Dpkg::Options::="--force-confnew" -o Dpkg::Options::="--force-overwrite"', "Failed to dist-upgrade some of the packages. Please review the error message, use APT to fix the situation and try again.")

        self.progress("Re-installing the meta-package for your edition of Linux Mint")
        self.check_command('sudo apt-get install --yes %s' % self.mint_meta, "Failed to install %s" % self.mint_meta)

        self.progress("Re-installing the multimedia codecs")
        self.check_command('sudo apt-get install --yes mint-meta-codecs', "Failed to install mint-meta-codecs")

        self.progress("Installing new packages")
        self.check_command('sudo apt-get install --yes %s' % " ".join(PACKAGES_ADDITIONS), "Failed to install additional packages.")

        self.progress("Removing obsolete packages")
        for removal in PACKAGES_REMOVALS:
            os.system('sudo apt-get purge --yes %s' % removal) # The return code indicates a failure if some packages were not found, so ignore it.

        self.progress("Performing system adjustments")
        os.system("sudo rm -f /etc/systemd/logind.conf")
        os.system("apt install --reinstall -o Dpkg::Options::=\"--force-confmiss\" systemd")
        os.system("sudo rm -f /etc/polkit-1/localauthority/50-local.d/com.ubuntu.enable-hibernate.pkla")
        if os.path.exists("/usr/share/ubuntu-system-adjustments/systemd/adjust-grub-title"):
            os.system("sudo /usr/share/ubuntu-system-adjustments/systemd/adjust-grub-title")
        elif os.path.exists("/usr/share/debian-system-adjustments/systemd/adjust-grub-title"):
            os.system("sudo /usr/share/debian-system-adjustments/systemd/adjust-grub-title")

        # Re-enable mintsystem
        os.system("sudo crudini --set /etc/linuxmint/mintSystem.conf global enabled True")
        os.system("sudo /usr/lib/linuxmint/mintsystem/mint-adjust.py")

        # Restore /etc/fstab if it was changed
        if not filecmp.cmp('/etc/fstab', BACKUP_FSTAB):
            os.system("cp /etc/fstab %s.upgraded" % BACKUP_FSTAB)
            os.system("sudo cp %s /etc/fstab" % BACKUP_FSTAB)
            self.warn("A package modified /etc/fstab during the upgrade. To ensure a successful boot, the upgrader restored your original /etc/fstab and saved the modified file in %d.upgraded." % BACKUP_FSTAB)

        self.progress("The upgrade is finished. Reboot the computer with \"sudo reboot\" when ready.")

    def check_command(self, command, message):
        ret = os.system(command)
        if ret != 0:
            self.fail(message)

    def try_command(self, num_times, command, fallback_commands):
        success = False
        for i in range(num_times):
            ret = os.system(command)
            if ret == 0:
                return True
            self.progress("Error detected on try #%d, running fallback commands" % (i+1))
            for fallback_command in fallback_commands:
                self.progress ("Running '%s'" % fallback_command)
                os.system(fallback_command)

    def fail(self, message):
        print ("")
        print ("------------------------------------------------")
        print ("%s!!  ERROR: %s%s" % (bcolors.FAIL, message, bcolors.ENDC))
        print ("!!  Exiting.")
        print ("------------------------------------------------")
        print ("")
        sys.exit(1)

    def continue_yes_no(self, messages):
        print ("")
        print ("-------------------------------------------------")
        for message in messages:
            print ("%s    %s%s" % (bcolors.WARNING, message, bcolors.ENDC))
        print ("")
        answer = None
        while (answer not in ["y", "yes", "n", "no"]):
            answer = input("%s    Do you want to continue? [y/n]:%s " % (bcolors.OKGREEN, bcolors.ENDC)).lower()
        if answer in ["n", "no"]:
            print ("Exiting.")
            sys.exit(0)

    def progress(self, message):
        print ("")
        print ("%s  + %s...%s" % (bcolors.HEADER, message, bcolors.ENDC))

    def warn(self, message):
        print ("")
        print ("%s  + %s%s" % (bcolors.WARNING, message, bcolors.ENDC))

def usage():
    print ("")
    print ("%sUsage:%s mintupgrade command" % (bcolors.HEADER, bcolors.ENDC))
    print ("")
    print ("%sCommands:%s" % (bcolors.HEADER, bcolors.ENDC))
    print ("  help                   - prints this usage note")
    print ("  check                  - checks the upgrade to %s" % DESTINATION)
    print ("  prepare                - prepares the upgrade to %s" % DESTINATION)
    print ("  download               - downloads the packages for the upgrade to %s" % DESTINATION)
    print ("  upgrade                - upgrades to %s" % DESTINATION)
    print ("  restore-sources        - restores the backed up APT sources (only use this command if you're still running %s)" % ORIGIN)
    print ("")
    sys.exit(0)

if __name__ == '__main__':

    if os.getuid() == 0:
        print ("")
        print ("Please don't run this command as root or with elevated privileges.")
        print ("")
        sys.exit(1)

    os.system("clear")

    if len(sys.argv) != 2:
        usage()
    command = sys.argv[1]
    if command == "help":
        usage()

    upgrader = MintUpgrade()

    if command == "restore-sources":
        upgrader.restore_sources()
    elif command == "check":
        upgrader.prepare()
        upgrader.check()
        upgrader.restore_sources()
    elif command == "prepare":
        upgrader.prepare()
    elif command == "download":
        upgrader.prepare()
        upgrader.download()
    elif command == "upgrade":
        upgrader.prepare()
        upgrader.download()
        upgrader.upgrade()
    else:
        usage()
