#!/usr/bin/python

from __future__ import print_function

import optparse
import datetime
import os
import re
import subprocess
import sys

import apt
from UpdateManager.Core.utils import twrap

# set locale early so that the subsequent imports have localized
# strings
import locale
try:
    locale.setlocale(locale.LC_ALL, "")
except:
    pass

from gettext import gettext as _
import gettext
gettext.textdomain("update-manager")


from HweSupportStatus.consts import (
    Messages,
    PRECISE_EOL_DATE,
    PRECISE_HWE_EOL_DATE,
    TRUSTY_DOT1_DATE,
)
    

# HWE backport names that are no longer supported
HWE_UNSUPPORTED_BACKPORTS = (
    "-lts-quantal",
    "-lts-raring",
    "-lts-saucy"
)
# from https://wiki.ubuntu.com/Kernel/LTSEnablementStack
    # version 3.11.x 3.8.x and 3.5.x are unsupported in precise
UNSUPPORTED_KERNEL_IMAGE_REGEX = r'linux-.*-3\.(11|8|5)(\.[0-9]+)?-.*'

# supported HWE stack
HWE_SUPPORTED_BACKPORT = "-lts-trusty"
SUPPORTED_KERNEL_IMAGE_REGEX = r'linux-.*-3\.13(\.[0-9]+)?-.*'


KERNEL_METAPKGS = (
    "linux-generic",
    "linux-image-generic",
    "linux-signed-generic",                                                            
    "linux-signed-image-generic",
)
XORG_METAPKGS = (
    "xserver-xorg",
    "libgl1-mesa-glx",
)
METAPKGS = KERNEL_METAPKGS + XORG_METAPKGS


class Package:
    """A lightweight apt package """
    def __init__(self, name, version, arch, foreign=False):
        self.name = name
        self.installed_version = version
        self.arch = arch
        self.foreign = foreign


def find_hwe_packages(installed_packages):
    unsupported_hwe_packages = set()
    supported_hwe_packages = set()
    for pkg in installed_packages:
        # metapackages and X are marked with the -lts-distro string
        for name in HWE_UNSUPPORTED_BACKPORTS:
            if pkg.name.endswith(name):
                unsupported_hwe_packages.add(pkg)
        # The individual backported kernels have names like 
        #   linux-image-3.11.0-17-generic
        # so we match via a regexp.
        # 
        # The linux-image-generic-lts-saucy metapkg has additional
        #  dependencies (like linux-firmware) so we can't just walk the
        # dependency chain.
        if re.match(UNSUPPORTED_KERNEL_IMAGE_REGEX, pkg.name):
            unsupported_hwe_packages.add(pkg)
        # SUPPORTED
        if pkg.name.endswith(HWE_SUPPORTED_BACKPORT):
            supported_hwe_packages.add(pkg)
        if re.match(SUPPORTED_KERNEL_IMAGE_REGEX, pkg.name):
            supported_hwe_packages.add(pkg)
    return unsupported_hwe_packages, supported_hwe_packages


def is_unsupported_hwe_kernel_running(unsupported_hwe_package):
    # kernels do not conflict with each other, so we need to check
    # what version is actually running
    running_kernel_ver = os.uname()[2]
    # the running kernel without the abi or buildver
    running_kernel_ver = running_kernel_ver.split("-")[0]
    for pkg in unsupported_hwe_package:
        if not pkg.name.startswith("linux-"):
            continue
        # we only care about the version, not abi or build
        if pkg.installed_version.startswith(running_kernel_ver):
            return True
    return False


def is_unsupported_xstack_running(unsupported_hwe_packages):
    # the HWE xstacks conflict with each other, so we can simply test
    # for existance in the installed unsupported hwe packages
    for pkg in unsupported_hwe_packages:
        for xorg_meta in XORG_METAPKGS:
            if pkg.name.startswith(xorg_meta):
                return True
    return False


def find_supported_replacement_hwe_packages(unsupported_hwe_packages, installed_packages=[]):
    unsupported_metapkg_names = set()
    replacement_names = set()
    buildstamp = subprocess.check_output(["tail", "-1", "/etc/buildstamp"])
    dell_type = subprocess.check_output("awk -F: '{print $6}' /sys/devices/virtual/dmi/id/modalias", shell=True)
    pci_id=subprocess.check_output("lspci -n | awk '{print $3}'", shell=True).split()
    special_fglrx_id=['1002:130a','1002:6900','1002:9853','1002:666f','1002:6823','1002:6606','1002:6664','1002:9850','1002:9851','1002:6665','1002:9993','1002:1315','1002:130f','1002:130a','1002:1309','1002:130b','1002:1316','1002:1318','1002:6600','1002:6604','1002:6821','1002:6901','1002:6640','1002:9648','1002:674a','1002:6660']
    special_fglrx_installed = False
    for metapkg in METAPKGS:
        for unsupported_backport in HWE_UNSUPPORTED_BACKPORTS:
            metapkg_name = metapkg + unsupported_backport
            for pkg in unsupported_hwe_packages:
                if pkg.name == metapkg_name:
                    replacement_name = metapkg + HWE_SUPPORTED_BACKPORT
                    if pkg.foreign:
                        replacement_name += ':' + pkg.arch
                    replacement_names.add(replacement_name)
                    unsupported_metapkg_names.add(metapkg_name)

    if len(set(pci_id).intersection(set(special_fglrx_id))) > 0:
        replacement_names.add("fglrx=1:14.201-0ubuntu2wataugafour2")
        replacement_names.add("fglrx-amdcccle=1:14.201-0ubuntu2wataugafour2")
        replacement_names.add("fglrx-core=1:14.201-0ubuntu2wataugafour2")
        special_fglrx_installed = True

    for pkg in installed_packages:
        if pkg.name == "oem-wireless-rtl-92ce-92se-92de-8723ae-88ee-8723be-92ee-dkms" or pkg.name == "oem-wireless-focus-rtl8188ee-1245323-dkms":
            replacement_names.add("oem-wireless-rtlwifi-1367658-dkms")
        elif pkg.name == "oem-wireless-skyray-rtl8723be-1265442-dkms":
            replacement_names.add("oem-wireless-skyray-rtl8723be-1265442-dkms-")
            replacement_names.add("oem-wireless-rtlwifi-1367658-dkms")
        elif pkg.name == "rts5229-dkms":
            replacement_names.add("rts5229-dkms-")
            replacement_names.add("oem-cardreader-rts5229-1409576-dkms")
        elif pkg.name == "bt-bcm20702a0-dkms":
            replacement_names.add("bt-bcm20702a0-trusty-dkms")
        elif pkg.name == "oem-bt-rtk8723b-dkms":
            replacement_names.add("oem-bt-rtk8723b-trusty-dkms")
        elif pkg.name == "oem-wireless-rt2x00-1099652-dkms":
            replacement_names.add("oem-wireless-rt2x00-1099652-dkms-")
            # replacement_names.add("oem-wireless-rt2x00-1099652-trusty-dkms")
            replacement_names.add("oem-wireless-mt76xx-1386997-dkms")
        elif pkg.name == "thinkpad-acpi-x1carbon-dkms":
            replacement_names.add("thinkpad-acpi-trusty-dkms")
        elif pkg.name == "oem-audio-hda-daily-lts-quantal-dkms":
            replacement_names.add("oem-audio-hda-daily-dkms")
        elif pkg.name == "oem-bluetooth-rtkbtusb-1265449-dkms":
            replacement_names.add("oem-bluetooth-rtkbtusb-1265449-dkms-")
            replacement_names.add("oem-bluetooth-rtkbtusb-1328352-dkms")
        elif pkg.name == "oem-bluetooth-mt76xx-1112206-dkms":
            replacement_names.add("oem-bluetooth-mt76xx-1112206-dkms-")
            #replacement_names.add("oem-bluetooth-mt76xx-1112206-trusty-dkms")
            replacement_names.add("oem-bluetooth-mt76xx-1386997-dkms")
        elif pkg.name == "oem-bluetooth-rt3298-1009059-dkms" or pkg.name == "rt3290bt-dkms":
            replacement_names.add("oem-bluetooth-rt3298-1009059-trusty-dkms")
        elif pkg.name == "oem-input-synaptics-1057412":
            pass
            #replacement_names.add("oem-input-synaptics-1336620") # this is a Close Source driver
        elif pkg.name == "oem-wireless-supra-1185223-dkms":
            replacement_names.add("oem-wireless-wl-1269220-dkms")
        elif pkg.name == "oem-bluetooth-btusb-1269661-dkms": # LP: #1384639
            replacement_names.add("oem-bluetooth-btusb-1269661-dkms-")
        elif pkg.name == "dw1704-resume-delay": # LP: #1397271
            replacement_names.add("dw1704-resume-delay-")
        elif pkg.name == "fglrx":
            if pkg.installed_version.startswith("2:13.352"):
                # replacement_names.add("fglrx=2:14.201-0ubuntu2wataugafour1")
                pass
            elif pkg.installed_version.startswith("2:14.201"):
                pass
                #replacement_names.add("fglrx=2:14.201-0ubuntu2wataugafour1")
            elif buildstamp == "stella-daan-precise-amd64-20140416-1\n" or buildstamp == "stella-daan-precise-amd64-20140506-1\n" or dell_type == "pnInspiron5537\n":
                #replacement_names.add("fglrx-")
                #replacement_names.add("fglrx-amdcccle-")
                replacement_names.add("fglrx=1:14.201-0ubuntu2wataugafour2")
                replacement_names.add("fglrx-amdcccle=1:14.201-0ubuntu2wataugafour2")
                replacement_names.add("fglrx-core=1:14.201-0ubuntu2wataugafour2")
                replacement_names.add("wataugafour-fglrx-helper")
            elif special_fglrx_installed == False:
                #replacement_names.add("fglrx=2:13.350.1-0ubuntu0.0.1")
                #os.system("gksu -D update-manager -- apt-get -y purge fglrx fglrx-amdcccle") # must purge first
                replacement_names.add("fglrx")
                replacement_names.add("fglrx-amdcccle")
            else:
                replacement_names.add("wataugafour-fglrx-helper")
        if pkg.name == "fglrx-13.352.1003-140402a-170450c":
                replacement_names.add("fglrx-13.352.1003-140402a-170450c-")
                replacement_names.add("fglrx-amdcccle-13.352.1003-140402a-170450c-")
                replacement_names.add("fglrx=1:14.201-0ubuntu2wataugafour2")
                replacement_names.add("fglrx-amdcccle=1:14.201-0ubuntu2wataugafour2")
                replacement_names.add("fglrx-core=1:14.201-0ubuntu2wataugafour2")
        elif pkg.name == "fglrx-13.352.1007-140428a-171604c":
                replacement_names.add("fglrx-13.352.1007-140428a-171604c-")
                replacement_names.add("fglrx-amdcccle-13.352.1007-140428a-171604c-")
                replacement_names.add("fglrx=1:14.201-0ubuntu2wataugafour2")
                replacement_names.add("fglrx-amdcccle=1:14.201-0ubuntu2wataugafour2")
                replacement_names.add("fglrx-core=1:14.201-0ubuntu2wataugafour2")
        elif pkg.name == "fglrx-13.352.1009linux-140520a-172286c":
                replacement_names.add("fglrx-13.352.1009linux-140520a-172286c-")
                replacement_names.add("fglrx-amdcccle-13.352.1009linux-140520a-172286c-")
                replacement_names.add("fglrx=1:14.201-0ubuntu2wataugafour2")
                replacement_names.add("fglrx-amdcccle=1:14.201-0ubuntu2wataugafour2")
                replacement_names.add("fglrx-core=1:14.201-0ubuntu2wataugafour2")
        elif pkg.name == "fglrx-13.352-140426a-171353e":
                replacement_names.add("fglrx-13.352-140426a-171353e-")
                replacement_names.add("fglrx-amdcccle-13.352-140426a-171353e-")
                replacement_names.add("fglrx=1:14.201-0ubuntu2wataugafour2")
                replacement_names.add("fglrx-amdcccle=1:14.201-0ubuntu2wataugafour2")
                replacement_names.add("fglrx-core=1:14.201-0ubuntu2wataugafour2")
        elif pkg.name == "fglrx-14.201.1004.1001-140722n-173681esomverille":
                replacement_names.add("fglrx-14.201.1004.1001-140722n-173681esomverille-")
                replacement_names.add("fglrx-amdcccle-14.201.1004.1001-140722n-173681esomverille-")
                replacement_names.add("fglrx=1:14.201-0ubuntu2wataugafour2")
                replacement_names.add("fglrx-amdcccle=1:14.201-0ubuntu2wataugafour2")
                replacement_names.add("fglrx-core=1:14.201-0ubuntu2wataugafour2")
        elif pkg.name == "fglrx-12.10.14.2-130409a-171042e":
                replacement_names.add("fglrx-12.10.14.2-130409a-171042e-")
                replacement_names.add("fglrx-amdcccle-12.10.14.2-130409a-171042e-")
                if special_fglrx_installed == False:
                    replacement_names.add("fglrx")
                    replacement_names.add("fglrx-amdcccle")
        elif pkg.name == "fglrx-12.104.2.9-131017a-165437c":
                replacement_names.add("fglrx-12.104.2.9-131017a-165437c-")
                replacement_names.add("fglrx-amdcccle-12.104.2.9-131017a-165437c-")
                if special_fglrx_installed == False:
                    replacement_names.add("fglrx")
                    replacement_names.add("fglrx-amdcccle")
        elif pkg.name == "fglrx-13.12.8":
                replacement_names.add("fglrx-13.12.8-")
                replacement_names.add("fglrx-amdcccle-13.12.8-")
                if special_fglrx_installed == False:
                    replacement_names.add("fglrx")
                    replacement_names.add("fglrx-amdcccle")
        elif pkg.name == "fglrx-13.201.3-131115a-165087c":
                replacement_names.add("fglrx-13.201.3-131115a-165087c-")
                replacement_names.add("fglrx-amdcccle-13.201.3-131115a-165087c-")
                if special_fglrx_installed == False:
                    replacement_names.add("fglrx")
                    replacement_names.add("fglrx-amdcccle")
        elif pkg.name == "fglrx-13.251-131206a-171983c":
                replacement_names.add("fglrx-13.251-131206a-171983c-")
                replacement_names.add("fglrx-amdcccle-13.251-131206a-171983c-")
                if special_fglrx_installed == False:
                    replacement_names.add("fglrx")
                    replacement_names.add("fglrx-amdcccle")
        elif pkg.name == "fglrx-13.25-131116a-165401e":
                replacement_names.add("fglrx-13.25-131116a-165401e-")
                replacement_names.add("fglrx-amdcccle-13.25-131116a-165401e-")
                if special_fglrx_installed == False:
                    replacement_names.add("fglrx")
                    replacement_names.add("fglrx-amdcccle")
        elif pkg.name == "fglrx-13.301.1001-131220a-167061c-somerville2":
                replacement_names.add("fglrx-13.301.1001-131220a-167061c-somerville2-")
                replacement_names.add("fglrx-amdcccle-13.301.1001-131220a-167061c-somerville2-")
                if special_fglrx_installed == False:
                    replacement_names.add("fglrx")
                    replacement_names.add("fglrx-amdcccle")
        elif pkg.name == "fglrx-13.302.1904-140603a-172653c":
                replacement_names.add("fglrx-13.302.1904-140603a-172653c-")
                replacement_names.add("fglrx-amdcccle-13.302.1904-140603a-172653c-")
                if special_fglrx_installed == False:
                    replacement_names.add("fglrx")
                    replacement_names.add("fglrx-amdcccle")
        elif pkg.name == "fglrx-13.251-131206a-166143c": # line 50 in unified DKMS list
                replacement_names.add("fglrx-13.251-131206a-166143c-")
                replacement_names.add("fglrx-amdcccle-13.251-131206a-166143c-")
                if special_fglrx_installed == False:
                    replacement_names.add("fglrx")
                    replacement_names.add("fglrx-amdcccle")
        #elif pkg.name == "bbswitch":
        #    replacement_names.add("bbswtich")
        #elif pkg.name == "oem-video-fglrx-1226451": # for Stella
        #    replacement_names.add("oem-video-fglrx-1226451-")
        #    replacement_names.add("fglrx")
        #elif pkg.name == "oem-video-fglrx-1091203": # LP: 1371741
        #    replacement_names.add("oem-video-fglrx-1091203-")
        #    replacement_names.add("fglrx")
        #elif pkg.name == "oem-video-fglrx-1128809": # LP: 1371742
        #    replacement_names.add("oem-video-fglrx-1128809-")
        #    replacement_names.add("fglrx")
        elif pkg.name == "gfx-fglrx": # LP: 1384068, Line 114-119
            replacement_names.add("gfx-fglrx-")
            if special_fglrx_installed == False:
                replacement_names.add("fglrx")
        elif pkg.name == "video-fglrx": # LP: 1384079
            replacement_names.add("video-fglrx-")
            if special_fglrx_installed == False:
                replacement_names.add("fglrx")
        #elif pkg.name.find("fglrx") != -1:
        #    replacement_names.add(" fglrx")
        elif pkg.name == "nvidia-settings":
            replacement_names.add("nvidia-settings-")
            #pass # on Mondavi-3 2D module
            #replacement_names.add("nvidia-331-updates") # not nvidia-331
        elif pkg.name == "nvidia-current-updates":
            replacement_names.add("nvidia-current-updates-")
        elif pkg.name == "nvidia-settings-updates":
            replacement_names.add("nvidia-settings-updates-")
        elif pkg.name == "nvidia-304-updates":
            replacement_names.add("nvidia-304-updates-")
            replacement_names.add("nvidia-331-updates")
        elif pkg.name == "nvidia-319":
            replacement_names.add("nvidia-319-")
            replacement_names.add("nvidia-331-updates")
        elif pkg.name == "nvidia-319-updates":
            replacement_names.add("nvidia-319-updates-")
            replacement_names.add("nvidia-331-updates")
        elif pkg.name == "nvidia-331":
            replacement_names.add("nvidia-331-")
            replacement_names.add("nvidia-331-updates")
        elif pkg.name == "nvidia-334":
            replacement_names.add("nvidia-334-")
            replacement_names.add("nvidia-340")
        elif pkg.name == "nvidia-tulum-340":
            replacement_names.add("nvidia-tulum-340-")
            replacement_names.add("nvidia-340")
        elif pkg.name == "synaptics-touchpad": # LP: 1370923
            replacement_names.add("synaptics-touchpad-")
        elif pkg.name == "oem-input-synaptics-1057412":
            replacement_names.add("oem-input-synaptics-1057412-")
        elif pkg.name == "oem-input-synapticsled-1098063-dkms":
            replacement_names.add("oem-input-synapticsled-1098063-dkms-")
        elif pkg.name == "hp-laptop-dkms-vine-touchpad-alps-acelan-1pt-dkms": # LP: 1373326
            replacement_names.add("oem-touchpad-alps-1373326-dkms")
        elif pkg.name == "oem-wireless-dw1550-1160187-dkms" or pkg.name == "wifi-brcm-dw1550-dkms" or pkg.name == "bcmwl-kernel-source" or pkg.name == "sutton-bcmwl-kernel-source": # LP: #1383175, #1383177, #1383180, #1395670
            replacement_names.add("broadcom-sta-dkms")
            #replacement_names.add("broadcom-sta-common") # this package conflict with above.
    if dell_type == "pnVostro23-3340\n":
        replacement_names.add("workaround-xhci-quirk-precise-dkms")
    elif dell_type == "pnInspiron3137\n":
        replacement_names.add("oem-atheros-bluetooth-1417476-dkms")

    replacement_names.add("wataugafour-meta")
    return unsupported_metapkg_names, replacement_names


def is_unsupported_hwe_running(unsupported_hwe_packages):
    return (is_unsupported_hwe_kernel_running(unsupported_hwe_packages) or
            is_unsupported_xstack_running(unsupported_hwe_packages))


def advice_about_hwe_status(unsupported_hwe_packages, supported_hwe_packages,
                            has_update_manager, verbose, installed_packages):
    unsupported_hwe_stack_running = is_unsupported_hwe_running(
        unsupported_hwe_packages)
    unsupported_hwe_metapkgs, supported_replacement_hwe = \
        find_supported_replacement_hwe_packages(unsupported_hwe_packages, installed_packages)
    # we need the "-p" option until 14.04.1 is released on 2014-07-15
    if datetime.date.today() < TRUSTY_DOT1_DATE:
        do_release_upgrade_option = "-p"
    else:
        do_release_upgrade_option = ""

    if unsupported_hwe_stack_running:
        if datetime.date.today() < PRECISE_HWE_EOL_DATE:
            s = Messages.HWE_SUPPORT_ENDS
        else:
            s = Messages.HWE_SUPPORT_HAS_ENDED
        if has_update_manager:
            print(s + Messages.UM_UPGRADE)
        else:
            # bug #1341320 - if no metapkg is left we need to show
            #                what is no longer supported
            if supported_replacement_hwe:
                print(s + Messages.APT_UPGRADE % (
                    do_release_upgrade_option,
                    " ".join(supported_replacement_hwe)))
            else:
                print(s + Messages.APT_SHOW_UNSUPPORTED % (
                    " ".join([pkg.name for pkg in unsupported_hwe_packages])))

    # some unsupported package installed but not running and not superseeded
    # - this is worth reporting
    elif (unsupported_hwe_packages and
          not supported_hwe_packages and
          not unsupported_hwe_stack_running):
        s = _("""
You have packages from the Hardware Enablement Stack (HWE) installed that
are going out of support on %s.
        """) % PRECISE_HWE_EOL_DATE
        if has_update_manager:
            print(s + Messages.UM_UPGRADE)
        else:
            print(s + Messages.APT_UPGRADE % (
                do_release_upgrade_option,
                " ".join(supported_replacement_hwe)))

    elif supported_hwe_packages:
        print(Messages.HWE_SUPPORTED)
    elif verbose:
        print(
            _("You are not running a system with a Hardware Enablement Stack. "
              "Your system is supported until %(month)s %(year)s.") % {
                  'month': PRECISE_EOL_DATE.strftime("%B"),
                  'year': PRECISE_EOL_DATE.year})


if __name__ == "__main__":
    parser = optparse.OptionParser(description=_("Check HWE support status"))
    parser.add_option('--quiet', action='store_true', default=False,
                        help="No output, exit code 1 on unsupported HWE "
                        "packages")
    parser.add_option('--verbose', action='store_true', default=False,
                        help="more verbose output")
    parser.add_option('--show-all-unsupported', action='store_true',
                        default=False,
                        help="Show unsupported HWE packages")
    parser.add_option('--show-replacements', action='store_true',
                      default=False,
                      help='show what packages need installing to be supported')
    # hidden, only useful for testing
    parser.add_option(
        '--disable-hwe-check-semaphore-file', 
        default="/var/lib/update-notifier/disable-hwe-eol-messages",
        help=optparse.SUPPRESS_HELP)
    options, args = parser.parse_args()

    if options.quiet:
        nullfd = os.open(os.devnull, os.O_WRONLY)
        os.dup2(nullfd, sys.stdout.fileno())

    # request from PSE to be able to disabled the hwe check via a special
    # semaphore file
    HWE_CHECK_DISABLED_FILE = options.disable_hwe_check_semaphore_file
    if os.path.exists(HWE_CHECK_DISABLED_FILE):
        if options.verbose:
            print("Forcefully disabled hwe-support-status via file %s" %
                  HWE_CHECK_DISABLED_FILE, file=sys.stderr)
        sys.exit(0)

    foreign_archs = set(subprocess.check_output(
        ['dpkg', '--print-foreign-architectures']).split())

    # do the actual check
    installed_packages = set()
    tagf = apt.apt_pkg.TagFile("/var/lib/dpkg/status")
    while tagf.step():
        if tagf.section.find("Status", "") != "install ok installed":
            continue
        pkgname = tagf.section.find("Package")
        version = tagf.section.find("Version")
        arch = tagf.section.find("Architecture")
        foreign = arch in foreign_archs
        installed_packages.add(Package(pkgname, version, arch, foreign))

    has_update_manager = "update-manager" in [
        pkg.name for pkg in installed_packages]
    unsupported_hwe_packages, supported_hwe_packages = find_hwe_packages(
        installed_packages)
    if options.show_all_unsupported:
        print(twrap(" ".join([
            pkg.foreign and pkg.name + ':' + pkg.arch or pkg.name \
                    for pkg in unsupported_hwe_packages])))
    if options.show_replacements:
        unsupported, replacements =  find_supported_replacement_hwe_packages(
            unsupported_hwe_packages, installed_packages)
        print(" ".join(replacements))
    
    if not options.show_all_unsupported and not options.show_replacements:
        advice_about_hwe_status(
            unsupported_hwe_packages, supported_hwe_packages, 
            has_update_manager, options.verbose, installed_packages)

    if is_unsupported_hwe_running(unsupported_hwe_packages):
        sys.exit(10)

    sys.exit(0)
