#!/usr/bin/python3

import sys, os, apt
import subprocess
import filecmp
import platform
import subprocess
import tempfile

ORIGIN = "Linux Mint 19.3 'Tricia'"
ORIGIN_CODENAME = "tricia"
ORIGIN_BASE_CODENAME = "bionic"

DESTINATION = "Linux Mint 20 'Ulyana'"
DESTINATION_CODENAME = "ulyana"
DESTINATION_BASE_CODENAME = "focal"

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

CHECK_ABSENT = []

CHECK_PRESENT = ["default-jre", "os-prober"]
CHECK_UP_TO_DATE = ["mintupgrade", "apt", "dpkg", "linuxmint-keyring", "ubuntu-keyring", "mintsystem"]

BACKUP_APT_SOURCES = os.path.expanduser("~/Upgrade-Backup-%s/APT/" % ORIGIN_CODENAME)
BACKUP_FSTAB = os.path.expanduser("~/Upgrade-Backup-%s/fstab" % ORIGIN_CODENAME)
BACKUP_CRYPTTAB = os.path.expanduser("~/Upgrade-Backup-%s/crypttab" % ORIGIN_CODENAME)

PACKAGES_PRE_REMOVALS = []

PACKAGES_REMOVALS = [
    "tomboy",
    "libxplayer-plparser18",
    "xplayer-common",
    "gksu",
    "memtest86+",
    "*hwe-18.04*",
    "linux-hwe*",
    "python3-tinycss" #
]

PACKAGES_ADDITIONS = [
    "neofetch",
    "ffmpegthumbnailer",
    "amd64-microcode",
    "intel-microcode",
    "celluloid",
    "drawing",
    "gnote",
    "grub2-theme-mint",
    "adwaita-icon-theme-full", # 19.3->20
    "warpinator", #
    "alsa-topology-conf", #
    "alsa-ucm-conf", #
    "mesa-vdpau-drivers", #
    "mesa-vulkan-drivers", #
    "cryptsetup-initramfs",
    "cryptsetup-run"
]

IMPORTANT_PACKAGES = [
    "mint-meta-cinnamon",
    "mint-meta-mate",
    "mint-meta-xfce",
    "xreader",
    "xed",
    "xviewer",
    "pix",
    "mintsystem",
    "metacity",
    "nemo-preview",
    "mintdrivers",
    "mintupdate",
    "mintsources",
    "mintinstall"
]

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'

# fstab columns
FSTAB_ID = 0
FSTAB_MNTPT = 1
FSTAB_TYPE = 2
FSTAB_OPTS = 3

CRYPTTAB_NAME = 0
CRYPTTAB_PATH = 1
CRYPTTAB_PW = 2
CRYPTTAB_OPTS = 3

def commented_out(line):
    return line.strip().startswith("#")

def comment_out(line):
    if not commented_out(line):
        return "# %s" % line
    else:
        return line

def uncomment(line):
    if not commented_out(line):
        return line
    else:
        return line.replace("#", "").lstrip()

def write_tabfile(path, lines):
    tmp_filename = None

    with tempfile.NamedTemporaryFile(mode="w", delete=False) as f:
        tmp_filename = f.name
        for line in lines:
            f.write(line)

    subprocess.run(["sudo", "mv", tmp_filename, path], check=True)

class MintUpgrade():

    def __init__(self, mode):
        self.reversible = True
        self.command = mode

        self.encrypted_swap = False
        self.write_tabs = False
        self.make_swapfile = False
        self.crypttab_name = None
        self.existing_swap_path = None
        self.new_fstab = []
        self.new_crypttab = []

        # 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

        self.check_ac()

    def check_ac(self):
        status = subprocess.getoutput("acpi -a")
        if "off-line" in status:
            messages = ["You should connect this computer to a power source before attempting to upgrade."]
            self.continue_yes_no(messages, False)
            os.system("clear")

    def restore_sources(self):
        self.reversible = False
        self.progress("Restoring your backed up APT sources")
        print("")
        if os.path.exists(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):
        messages = []
        messages.append("Executing '%s'. This will perform the following:" % self.command)
        messages.append("")
        messages.append("1 - Your system will be checked to ensure upgrading is possible.")
        messages.append("")
        messages.append("2 - Your repositories will be switched to point to %s and" % DESTINATION)
        messages.append("    any 3rd party repositories will be removed. A backup of your APT sources")
        messages.append("    will be written to %s." % BACKUP_APT_SOURCES)
        if self.command == "prepare":
            messages.append("")
            messages.append("Your sources will NOT be restored to %s at the end of this command." % ORIGIN)
        elif self.command == "check":
            messages.append("")
            messages.append("3 - The upgrade will be simulated so impacted packages can be evaluated.")
            messages.append("")
            messages.append("Your sources will be restored to %s at the end of this command." % ORIGIN)
        elif self.command == "download":
            messages.append("")
            messages.append("3 - The packages and updates to perform the upgrade will be downloaded. You")
            messages.append("    should run 'check' prior to this to ensure there are no conflicts, otherwise")
            messages.append("    this step could fail.")
            messages.append("")
            messages.append("Your sources will be restored to %s at the end of this command." % ORIGIN)
        elif self.command == "upgrade":
            messages.append("")
            messages.append("3 - The upgrade will be simulated so impacted packages can be evaluated.")
            messages.append("")
            messages.append("4 - The required packages will be downloaded. There will be a final confirmation")
            messages.append("    when this is completed before moving to the next step.")
            messages.append("")
            messages.append("5 - The upgrade will be performed.  It is very important not to interrupt this")
            messages.append("    step. You may be required to interact or re-authenticate during the upgrade.")

        self.continue_yes_no(messages, False)

        # 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))

        bits, linkage = platform.architecture()
        if "32" in bits:
            self.fail("You are on a 32-bit system. It cannot be upgraded to %s." % 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.")

        if self.command == "upgrade":
            self.progress("Checking for encrypted swap")

            # We'll do all the figuring out now, so we can discuss with the user,
            # then at the end of the upgrade we'll write changes.

            blkdev_name = None
            blkdev_real_id = None
            existing_swap_in_fstab = False

            # Check crypttab to get /dev/mapper/id
            with open("/etc/crypttab", "r") as f:
                for line in f:
                    if not commented_out(line):
                        entry = line.split()
                        if "swap" in entry[CRYPTTAB_OPTS]:
                            self.encrypted_swap = True
                            self.crypttab_name = blkdev_name = entry[CRYPTTAB_NAME]
                            self.existing_swap_path = blkdev_real_id = entry[CRYPTTAB_PATH]
                            line = comment_out(line)

                    self.new_crypttab.append(line)

            if blkdev_name != None:
                with open("/etc/fstab", "r") as f:
                    # We're looking for a) the /dev/mapper name from crypttab and
                    # b) a line with an id that matches the crypttab entry's path -
                    # if we find one we can uncomment it (this happens if it was a
                    # partition and not a file).

                    for line in f:
                        commented = False
                        if commented_out(line):
                            commented = True
                            orig_line = line
                            line = uncomment(line)

                        entry = line.split()
                        if len(entry) < 6:
                            self.new_fstab.append(orig_line)
                            continue

                        if entry[FSTAB_ID] == ("/dev/mapper/%s" % blkdev_name) and not commented:
                            # found the /dev/mapper entry from crypttab
                            commented = True
                        elif entry[FSTAB_ID] == blkdev_real_id and "swap" in entry[FSTAB_TYPE]:
                            # This is either the original entry for the swap partition that was
                            # found during installation, which was commented out when encryption
                            # was setup, or it's an uncommented entry for a swapfile.  Either
                            # way, we leave it, enabled. Any # was already removed above.
                            existing_swap_in_fstab = True
                            commented = False

                        line = comment_out(line) if commented else line
                        self.new_fstab.append(line)

            messages = []
            if self.encrypted_swap:
                messages.append("You appear to be using an encrypted swap space. This is no longer supported in")
                messages.append("%s and will be disabled." % DESTINATION)

                if existing_swap_in_fstab:
                    messages.append("")
                    messages.append("The same swap space will be used, but will now be unencrypted.")
                    self.continue_yes_no(messages, restore=False)
                else:
                    messages.append("")
                    messages.append("No usable swap space detected. A swap file of the default size (2gb)")
                    messages.append("can be created instead.")
                    print ("")
                    for message in messages:
                        print ("%s    %s%s" % (bcolors.WARNING, message, bcolors.ENDC))
                    answer = None
                    while (answer not in ["y", "yes", "n", "no"]):
                        print ("")
                        answer = input("%s    Create a swap file? [y/n]:%s " % (bcolors.OKGREEN, bcolors.ENDC)).lower()
                        if answer in ["y", "yes"]:
                            print ("")
                            self.make_swapfile = True
                            self.new_fstab.append("\n")
                            self.new_fstab.append("# Added by mintupgrade\n")
                            self.new_fstab.append(
                                "/swapfile                                 none            swap    sw              0       0\n"
                            )

                self.write_tabs = True

        if not self.points_to_destination:
            # Check packages

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

            self.progress("Checking packages")
            cache = apt.Cache()
            need_to_remove = []
            for pkg_name in CHECK_ABSENT:
                if pkg_name in cache.keys():
                    pkg = cache[pkg_name]
                    if pkg.is_installed:
                        need_to_remove.append(pkg.name)
            if len(need_to_remove) > 0:
                self.fail("The following packages create issues with this upgrade, please remove them:\n\n    %s\n" % ", ".join(need_to_remove))
            need_to_install = []
            for pkg_name in CHECK_PRESENT:
                if pkg_name in cache.keys():
                    pkg = cache[pkg_name]
                    if not pkg.is_installed:
                        need_to_install.append(pkg.name)
            if len(need_to_install) > 0:
                self.fail("The following packages are required for a smooth install, please install them:\n\n    %s\n" % ", ".join(need_to_install))

            # # Check that we don't have any Debian Multimedia packages
            # self.progress("Checking deb-multimedia.org repository")
            # output = subprocess.check_output("apt policy", shell=True).decode("UTF-8")
            # if "deb-multimedia.org" in output:
            #     self.fail("The Unofficial Debian Multimedia repository (https://deb-multimedia.org) is no longer used in LMDE 4. Please remove it from your APT configuration (usually in /etc/apt/sources.list.d/official-package-repositories.list).")

            # self.progress("Checking deb-multimedia.org packages")
            # dmo_packages = []
            # for pkg in cache:
            #     if pkg.is_installed:
            #         if "-dmo" in pkg.installed.version:
            #             dmo_packages.append(pkg.name)
            #         if pkg.installed.source_name.endswith("-dmo"):
            #             dmo_packages.append(pkg.name)
            # if len(dmo_packages) > 0:
            #     self.fail("The following packages come from the Unofficial Debian Multimedia repository: \n\n     %s.\n\nLaunch mintsources and go to the \"Maintenance tab\".\nFirst, use \"Downgrade Foreign Packages\" to downgrade all the packages.\nThen, use \"Remove Foreign Packages\" to remove any remaining packages with '-dmo' in their version number.\n" % ", ".join(dmo_packages))

            # 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")
            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)
        print("")
        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.try_command(2, 'DEBIAN_PRIORITY=critical sudo apt-get update', [])

        if self.command == "prepare":
            self.warn("Command '%s' completed successfully" % self.command)
            self.warn_inconsistent_sources()

    def check(self):
        messages = []
        messages.append("APT will now calculate the package changes necessary to upgrade to %s.\n" % DESTINATION)
        messages.append("If conflicts are detected and APT is unable to perform the upgrade, take note of the packages")
        messages.append("causing the issue, remove them, and restart the upgrade.\n")
        messages.append("Pay close attention to what appears on the screen, particularly the packages being REMOVED. Take")
        messages.append("note of any you may wish to reinstall after the upgrade.")
        self.continue_press_enter(messages)
        self.progress("Simulating an upgrade")
        print("")
        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:\n\n    %s\n" % ", ".join(sorted(incorrect_removals)))
        if len(kept_packages) > 0:
            self.warn("The following packages will be kept back during the upgrade:\n\n    %s\n\n"
                      "    This might or might not indicate a problem. Check the APT output above to decide\n"
                      "    whether to continue with the upgrade. If this OK you might need to update them\n"
                      "    manually after the upgrade." % ", ".join(sorted(kept_packages)))
        if self.command == "check":
            self.warn("Command '%s' completed successfully" % self.command)

    def download(self):
        self.progress("Downloading upgrade packages")
        messages = []
        if self.command == "download":
            messages.append("APT will now download the package updates necessary for the upgrade to %s." % DESTINATION)
            self.continue_yes_no(messages)
        print("")
        self.check_command("DEBIAN_PRIORITY=critical sudo apt-get dist-upgrade --download-only --yes", "Failed to download packages for the upgrade.")

        if self.command == "download":
            self.warn("Command '%s' completed successfully" % self.command)
            self.warn_inconsistent_sources()

    def upgrade(self):
        self.progress("Disabling screensaver and power management")
        os.system("killall cinnamon-screensaver")
        os.system("killall mate-screensaver")
        os.system("killall light-locker")

        current_desktop = os.environ["XDG_CURRENT_DESKTOP"].lower()
        current_desktop = current_desktop.replace("x-", "") # X-Cinnamon

        if current_desktop == "xfce":
            self.warn("Could not inhibit session.  You should disable power management for the duration\n"
                      "    of this upgrade, and refrain from logging out or switching users.")
        else:
            subprocess.Popen(["mintupgrade-inhibit-power", str(os.getpid())])

        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.

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

        self.progress("Performing upgrade")
        print("")
        self.reversible = False

        # 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)

        # Enable APT recommends
        self.progress("Re-enabling APT recommends")
        if os.path.exists("/etc/apt/apt.conf.d/00recommends"):
            os.system("sudo rm -f /etc/apt/apt.conf.d/00recommends")

        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("Running autoclean to remove unused packages")
        self.check_command("sudo apt-get --purge autoremove --yes", "Failed to autoremove unused packages.")

        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")

        # Re-enable the plymouth theme (specific to LMDE)
        # os.system("sudo plymouth-set-default-theme mint-logo")
        # os.system("sudo update-initramfs -u")

        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\n"
                      "    upgrader restored your original /etc/fstab and saved the modified file in \n"
                      "    %s.upgraded." % BACKUP_FSTAB)

        if self.encrypted_swap:
            self.progress("Disabling encrypted swap partition")
            print("")

            try:
                subprocess.run(["sudo", "swapoff", "-a"], check=True)
                subprocess.run(["sudo", "cryptsetup", "remove", self.crypttab_name], check=False)
                subprocess.run(["cp", "/etc/crypttab", BACKUP_CRYPTTAB], check=True)

                write_tabfile("/etc/crypttab", self.new_crypttab)
                write_tabfile("/etc/fstab", self.new_fstab)

                if self.make_swapfile:
                    subprocess.run(["sudo", "fallocate", "-l", "2G", "/swapfile"], check=True)
                    subprocess.run(["sudo", "chmod", "600", "/swapfile"], check=True)
                    subprocess.run(["sudo", "mkswap", "/swapfile"], check=True)
                else:
                    print("Running mkswap to reinitialize the swap space. This will fail if you're using")
                    print("a swap partition, and that's ok.\n")
                    subprocess.run(["sudo", "mkswap", self.existing_swap_path], check=False)

                subprocess.run(["sudo", "update-initramfs", "-u"])
                subprocess.run(["sudo", "update-grub"])
            except Exception as e:
                messages = []
                messages.append("Something went wrong trying to update your swap configuration: %s" % e)
                # FIXME: Do we print all the steps required??
                self.critical_warn(messages)
                os.system("sudo cp %s /etc/fstab" % 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..." % (i+1))
            if (i+1) < num_times:
                self.progress("Retrying...")
            for fallback_command in fallback_commands:
                self.progress ("Running fallback command '%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 ("------------------------------------------------")
        if self.reversible:
            self.restore_sources()
        sys.exit(1)

    def warn_inconsistent_sources(self):
        messages = []
        messages.append("Your system is now in an inconsistent state. You must run 'mintupgrade restore-sources'")
        messages.append("if you wish to abort the upgrade process.")
        self.critical_warn(messages)

    def critical_warn(self, messages):
        print("")
        print("%s!!  WARNING:%s\n" % (bcolors.WARNING, bcolors.ENDC))
        for line in messages:
            print ("%s    %s%s" % (bcolors.WARNING, line, bcolors.ENDC))
        print("")

    def continue_press_enter(self, messages):
        print ("")
        for message in messages:
            print ("%s    %s%s" % (bcolors.WARNING, message, bcolors.ENDC))
        print("")
        answer = input("%s    Press Enter to continue%s" % (bcolors.OKGREEN, bcolors.ENDC))

    def continue_yes_no(self, messages, restore=True):
        print ("")
        for message in messages:
            print ("%s    %s%s" % (bcolors.WARNING, message, bcolors.ENDC))
        answer = None
        while (answer not in ["y", "yes", "n", "no"]):
            print ("")
            answer = input("%s    Do you want to continue? [y/n]:%s " % (bcolors.OKGREEN, bcolors.ENDC)).lower()
        if answer in ["n", "no"]:
            print ("")
            print ("    Exiting...")
            if self.reversible and restore:
                self.restore_sources()
            else:
                print ("")
            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))
        print ("")

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. You should run this first." % DESTINATION)
    print ("  upgrade           - upgrades to %s, performing all necessary steps." % DESTINATION)
    print ("")
    print ("  prepare           - prepares the upgrade to %s" % DESTINATION)
    print ("  download          - downloads the packages for the upgrade to %s" % DESTINATION)
    print ("  restore-sources   - restores the backed up APT sources (only use this")
    print ("                      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(command)

    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()
