sabayon 14 patch for nvidia drivers

#!/usr/bin/python
# -*- coding: utf-8 -*-

import os
import errno
import subprocess
try:
    from subprocess import getoutput
except ImportError:
    from commands import getoutput
import shutil
import sys

# Variables
xorgfile = "/etc/X11/xorg.conf"
lspci = '/usr/sbin/lspci'
nvidia_settings = "/usr/share/applications/nvidia-settings.desktop"

device_id_prefix = "SabayonVga"
nvidia_option_prefix = "--nvidia-opt--"
screen_layout_sections = []
device_sections = []
xorg_conf_structure = """
Section "Module"
    SubSection  "extmod"
       Option    "omit xfree86-dga"
    EndSubSection
    Load    "i2c"
    Load    "ddc"
    Load    "vbe"
    Load    "dri"
    Load    "glx"
    Load    "synaptics"
EndSection

Section "ServerFlags"
     Option    "AllowMouseOpenFail"    "true"
EndSection

Section "Monitor"
    Identifier    "Generic Monitor"
    VertRefresh    43 - 60
    HorizSync      28 - 80
EndSection

__device_section__

__screen_section__

Section "DRI"
    Mode 0666
EndSection

Section "ServerLayout"
    Identifier   "Main Layout"
    __screen_layout_section__
EndSection

Section "Extensions"
   #Option "Composite" "Enable"
EndSection
"""

screen_sections = []
screen_section = """
Section "Screen"

    Identifier    "Screen __screen_id__"
    Device        "%s__screen_id__"
    Monitor       "Generic Monitor"
    %sOption       "AddARGBGLXVisuals" "true"
    %sOption       "RegistryDwords" "EnableBrightnessControl=1"

    DefaultDepth 24

    SubSection "Display"
        Depth       8
        ViewPort    0 0
        #Modes      "1024x768" "800x600" "640x480"
    EndSubsection

    SubSection "Display"
        Depth           16
        ViewPort        0 0
        #Modes      "1024x768" "800x600" "640x480"
    EndSubsection

    SubSection "Display"
        Depth           24
        ViewPort        0 0
        #Modes      "1024x768" "800x600" "640x480"
    EndSubsection

EndSection
""" % (device_id_prefix, nvidia_option_prefix,
       nvidia_option_prefix,)

# cmdlines
options = sys.argv[1:]
dryrun = False
noproprietary = False
nvidia_forcefail = False
nvidia_disablelegacy = False
legacy = False
livecd = False
steps = []
forced_xdriver = ''
current_arch = os.uname()[4]
nomodeset = False
noefi = False

fglrx_supported = sorted(getoutput(
        "modinfo fglrx | grep alias | grep pci | "
        "cut -d':' -f 3 | cut -d'*' -f 1 | "
        "sed 's/.*1002d//' | sed 's/^0000//' | sed 's/sv$//'"
        ).lower().split())

nvidia_71xx_supported = ['0020', '0028', '0029', '002c', '002d', '00a0',
    '0100', '0101', '0103', '0150', '0151', '0152', '0153']
nvidia_96xx_supported = ['0110', '0111', '0112', '0113', '0170', '0171',
    '0172', '0173', '0174', '0175', '0176', '0177', '0178', '0179', '017a',
    '017c', '017d', '0181', '0182', '0183', '0185', '0188', '018a', '018b',
    '018c', '01a0', '01f0', '0200', '0201', '0202', '0203', '0250', '0251',
    '0253', '0258', '0259', '025b', '0280', '0281', '0282', '0286', '0288',
    '0289', '028c']
nvidia_173xx_supported = ['00fa', '00fb', '00fc', '00fd', '00fe', '0301',
    '0302', '0308', '0309', '0311', '0312', '0314', '031a', '031b', '031c',
    '0320', '0321', '0322', '0323', '0324', '0325', '0326', '0327', '0328',
    '032a', '032b', '032c', '032d', '0330', '0331', '0332', '0333', '0334',
    '0338', '033f', '0341', '0342', '0343', '0344', '0347', '0348', '034c',
    '034e']
# Taken from here:
# http://www.nvidia.com/object/IO_32667.html
nvidia_304xx_supported = ['0040', '0041', '0042', '0043', '0044', '0045',
    '0046', '0047', '0048', '004e', '0090', '0091', '0092', '0093', '0095',
    '0098', '0099', '009d', '00c0', '00c1', '00c2', '00c3', '00c8', '00c9',
    '00cc', '00cd', '00ce', '00f1', '00f2', '00f3', '00f4', '00f5', '00f6',
    '00f8', '00f9', '0140', '0141', '0142', '0143', '0144', '0145', '0146',
    '0147', '0148', '0149', '014a', '014c', '014d', '014e', '014f', '0160',
    '0161', '0162', '0163', '0164', '0165', '0166', '0167', '0168', '0169',
    '016a', '01d0', '01d1', '01d2', '01d3', '01d6', '01d7', '01d8', '01da',
    '01db', '01dc', '01dd', '01de', '01df', '0211', '0212', '0215', '0218',
    '0221', '0222', '0240', '0241', '0242', '0244', '0245', '0247', '0290',
    '0291', '0292', '0293', '0294', '0295', '0297', '0298', '0299', '029a',
    '029b', '029c', '029d', '029e', '029f', '02e0', '02e1', '02e2', '02e3',
    '02e4', '038b', '0390', '0391', '0392', '0393', '0394', '0395', '0397',
    '0398', '0399', '039c', '039e', '03d0', '03d1', '03d2', '03d5', '03d6',
    '0531', '0533', '053a', '053b', '053e', '07e0', '07e1', '07e2', '07e3',
    '07e5']
savage_supported = ['8a20', '8a21', '8a22', '9102', '8c10', '8c11', '8c12',
    '8c13', '8c22', '8c24', '8c26', '8c2a', '8c2b', '8c2c', '8c2d', '8c2e',
    '8c2f', '8a25', '8a26', '8d01', '8d02', '8d03', '8d04']
unichrome_supported = ['3108', '3118', '3157', '3343', '3344', '7205']

lspci_output = ''
for option in options:
    if option == "--dry-run":
        dryrun = True
    elif option.startswith('--with-lspci=') and len(option.split("=")) >= 2:
        option = option.split("=")[1:]
        option = "=".join(option)
        if option.startswith('"'):
            option = option[1:]
        if option.startswith("'"):
            option = option[1:]
        if option.endswith("'"):
            option = option[:len(option)-1]
        if option.endswith('"'):
            option = option[:len(option)-1]
        lspci_output = option
    elif option.startswith('--forced-xdriver=') and len(option.split("=")) == 2:
        forced_xdriver = option.split("=")[1]

if not lspci_output:
    lspci_output = getoutput(lspci+' -mm -n')

# parse cmdline
with open("/proc/cmdline","r") as f:
    cmdline = f.readline().split()
for cmd in cmdline:
    if cmd == "noproprietary":
        noproprietary = True
    elif cmd == "nomodeset":
        nomodeset = True
    elif cmd == "nvidia=forcefail":
        nvidia_forcefail = True
    elif cmd == "nvidia=disablelegacy":
        nvidia_disablelegacy = True
    elif cmd == "legacy":
        legacy = True
    elif cmd == "cdroot":
        livecd = True
    elif cmd == "noefi":
        noefi = True
    elif cmd.startswith("xdriver=") and (len(cmd.split("=")) == 2):
        if not forced_xdriver:
            forced_xdriver = cmd.split("=")[1] # --forced-xdriver= owns

def openrc_running():
    return os.path.isfile("/run/openrc/softlevel")

def systemd_running():
    return os.path.isdir("/run/systemd/system")

def remove_proprietary_opengl(bumblebee):
    if not dryrun:
        if not bumblebee:
            os.system("""
            mount -t tmpfs none /usr/lib/opengl/ati &> /dev/null
            mount -t tmpfs none /usr/lib/opengl/nvidia &> /dev/null
            sed -i '/LIBGL_DRIVERS_PATH/ s/.*//' /etc/profile.env
            """)
            fix_possible_opengl_misconfiguration('xorg-x11')
        else:
            print("Bumblebee enabled, not deactivating proprietary drivers")
    else:
        print("I was about to remove proprietary OpenGL libraries")

def get_kernel_version():
    try:
        return int(os.uname()[2].replace(".", "")[:3])
    except (ValueError, TypeError) as err:
        print("get_kernel_version: ouch: %s" % (err,))
        return None

def setup_radeon_kms():
    # Starting from kernel 3.6, we have CONFIG_DRM_RADEON_KMS=y
    kver = get_kernel_version()
    if kver is None:
        kver = 360 # assume new kernel
    if not dryrun and kver < 360:
        os.system("""
        modprobe -r radeon &> /dev/null
        modprobe radeon modeset=1 && touch /tmp/.radeon.kms
        """)
    else:
        print("I was about to modprobe radeon modeset=1")

def generate_fglrx_steps(videocard, cardnumber, total_cards, bus_id):
    print("AMD!")
    print("total supported AMD cards: %s" % (len(fglrx_supported),))
    print("supported list:", fglrx_supported)
    supported = card_id in fglrx_supported
    if supported:
        print("fglrx driver supports this card")
    # check if nomodeset is enabled for >=3.6.0 kernel
    kver = get_kernel_version()
    if kver is None:
        kver = 360 # assume new kernel
    if not nomodeset and kver >= 360:
        print("however, nomodeset is not set, though KMS is active,"
              " defaulting to OSS driver")
        supported = False

    if supported:
        if noproprietary:
            steps.append((drop_kernel_mod, "fglrx",))
            steps.append((setup_radeon_kms,))
        else:
            steps.append((fix_possible_opengl_misconfiguration,
                          "ati"))
            steps.append((copy_ati_settings_on_desktop,))
            steps.append((opengl_activate, "ati"))
            steps.append((set_xorg_device, "fglrx",
                          cardnumber, total_cards, bus_id,))
    else:
        # video card not supported by fglrx
        print("using OSS 'ati' drivers")
        generate_generic_steps()
        # This works for Mach64, Rage128
        # Radeon and in future RadeonHD driver
        steps.append((drop_kernel_mod, "fglrx",))
        steps.append((setup_radeon_kms,))

def check_if_driver_is_available(xdriver):
    drv_path = "/usr/lib/xorg/modules/drivers/" + xdriver + "_drv.so"
    if os.path.isfile(drv_path):
        print("check_if_driver_is_available for " + xdriver + ": available")
        return True
    print("check_if_driver_is_available for " + xdriver + ": not available")
    return False

def check_if_proprietary_driver_system_is_healthy(kernelmod):
    rc = subprocess.call(["modprobe", kernelmod])
    if rc == 0:
        if kernelmod == "nvidia":
            if os.path.exists("/usr/lib/opengl/nvidia/lib"):
                print("check_if_proprietary_driver_system_is_healthy:"
                      " nvidia healthy")
                return True
            print("check_if_proprietary_driver_system_is_healthy:"
                  " nvidia NOT healthy")
            return False
        elif kernelmod == "fglrx":
            kver = get_kernel_version()
            if kver is None:
                kver = 360 # assume new kernel
            if not nomodeset and kver >= 360:
                print("check_if_proprietary_driver_system_is_healthy:"
                      " fglrx (ati) NOT healthy, 'nomodeset' boot argument"
                      " is mising")
                return False
            if os.path.exists("/usr/lib/opengl/ati/lib"):
                print("check_if_proprietary_driver_system_is_healthy:"
                      " fglrx (ati) healthy")
                return True
            print("check_if_proprietary_driver_system_is_healthy:"
                  " fglrx (ati) NOT healthy")
            return False
    return False

def deploy_nvidia_xxxxxx_drivers(ver):
    if dryrun:
        print("I was about to run deploy_nvidia_xxxxxx_drivers"
              ", ver: %s" % (ver,))
        return False

    drivers_dir = "/install-data/drivers"
    # are they available ? we're on livecd...
    if not os.path.isdir(drivers_dir):
        print("drivers_dir not available")
        return False

    packages = os.listdir(drivers_dir)
    _packages = []
    for pkg in packages:
        if not pkg.endswith(".tbz2"):
            continue
        if pkg.startswith("x11-drivers:nvidia-drivers-" + ver):
            _packages.append(pkg)
        elif pkg.startswith("x11-drivers:nvidia-userspace-" + ver):
            _packages.append(pkg)

    packages = [os.path.join(drivers_dir, x) for x in _packages]
    if not packages:
        return False

    rc = subprocess.call(["/usr/bin/equo", "install", "--nodeps"] + packages)
    if rc:
        return False

    # try to check driver status now
    return check_if_proprietary_driver_system_is_healthy("nvidia")

efivars_loaded = False
def is_efi():
    """
    Return whether the system boots from EFI
    """
    global efivars_loaded

    if noefi:
        return False

    if not efivars_loaded:
        subprocess.call(["modprobe", "efivars"])
        efivars_loaded = True

    return os.path.exists("/sys/firmware/efi")

def get_vesa_driver():
    """
    Return either "vesa" or "fbdev" as the fallback
    vesa-like X driver.
    """
    if is_efi():
        # vesa does not work
        return "fbdev"
    return "vesa"

def set_xorg_device(xdriver, cardnum, total_cards, bus_id):
    if (xdriver not in ("nvidia", "fglrx",)) and \
            (not check_if_driver_is_available(xdriver)):
        xdriver = get_vesa_driver() # fallback to vesa
    bus_id_mark = "#"
    if total_cards > 1:
        bus_id_mark = ""

    device_sections.append("""
Section "Device"

    Identifier  "%s%s"
    Driver      "%s"
    %sBusID    "%s"
    #Option "RenderAccel" "on"
    #Option "XAANoOffscreenPixmaps"
    #Option "BusType" "PCI"
    #Option "ColorTiling" "on"
    #Option "EnablePageFlip" "on"
    # UseEvents is causing segmentation faults with
    # NVIDIA 6xxx, 7xxx and >=275.xx.xx drivers
    #Option "UseEvents" "True"
    Option "LogoPath" "/usr/share/backgrounds/sabayonlinux-nvidia.png"

EndSection
    """ % (device_id_prefix, cardnum, xdriver, bus_id_mark, bus_id,))

    my_screen_section = screen_section.replace("__screen_id__", str(cardnum))
    # setup Option AddARGBVisuals
    # especially needed for legacy nvidia drivers, but works
    # on all of them
    if xdriver == "nvidia":
        my_screen_section = my_screen_section.replace(nvidia_option_prefix, "")
    else:
        my_screen_section = my_screen_section.replace(nvidia_option_prefix, "#")
    screen_sections.append(my_screen_section)
    screen_layout_sections.append('Screen %s    "Screen %s"' % (
            cardnum, cardnum,))

def opengl_activate(profile, force=False):
    if not dryrun:
        if not force:
            current = opengl_show()
            if current == profile:
                print("OpenGL profile is already set to: " + profile)
                return
        subprocess.call(["eselect", "opengl", "set", profile])
    else:
        print("I was about to set opengl subsystem to: " + profile)

def opengl_show():
    return getoutput("eselect opengl show").split("\n")[0].strip()

def fix_possible_opengl_misconfiguration(profile):
    # get current subsystem
    current = opengl_show()
    if not dryrun:
        if (profile in ("ati","nvidia","xorg-x11")) and (profile != current):
            if profile == "ati" or profile == "nvidia":
                subprocess.call(["umount", "/usr/lib/opengl/" + profile])
                subprocess.call(["umount", "/usr/lib/opengl/" + profile])
                opengl_activate(profile)
    else:
        print("I was about to fix OpenGL subsystem to: " + \
                  profile + " while the current implementation is: " + \
                  current)

def copy_nvidia_settings_on_desktop():
    homes = []
    if os.path.isfile(nvidia_settings):
        _homes = os.listdir("/home")
        homes += [x for x in os.listdir("/home") \
                      if os.path.isdir("/home/" + x + "/Desktop")]

    for home in homes:
        try:

            full_home = os.path.join("/home", home)
            st = os.stat(full_home)
            dest_path = "/home/" + home + "/Desktop/" + \
                os.path.basename(nvidia_settings)
            shutil.copy2(nvidia_settings, dest_path)
            os.chmod(dest_path, 0o755)
            os.chown(dest_path, st.st_uid, st.st_gid)

            if os.path.isdir("/etc/skel/Desktop"):
                dest_path = os.path.join(
                    "/etc/skel/Desktop",
                    os.path.basename(nvidia_settings))
                shutil.copy2(nvidia_settings, dest_path)
                os.chmod(dest_path, 0o755)

        except Exception:
            pass

def copy_ati_settings_on_desktop():
    desktop_files = getoutput(
        'equo query files ati-drivers --quiet | grep ".desktop"').split("\n")
    desktop_files = [x for x in desktop_files if os.path.isfile(x)]
    print("copy_ati_settings_on_desktop: found files: "+str(desktop_files))

    for ati_settings in desktop_files:
        homes = os.listdir("/home")
        homes = [x for x in homes if os.path.isdir("/home/" + x + "/Desktop")]
        for home in homes:
            try:
                full_home = os.path.join("/home", home)
                st = os.stat(full_home)
                dest_path = "/home/" + home + "/Desktop/" + \
                    os.path.basename(ati_settings)
                shutil.copy2(ati_settings, dest_path)
                os.chmod(dest_path, 0o755)
                os.chown(dest_path, st.st_uid, st.st_gid)

                if os.path.isdir("/etc/skel/Desktop"):
                    dest_path = os.path.join(
                        "/etc/skel/Desktop",
                        os.path.basename(ati_settings))
                    shutil.copy2(ati_settings, dest_path)
                    os.chmod(dest_path, 0o755)
            except Exception:
                pass

def setup_nvidia_drivers(card_id):
    drv_string = ''
    done_legacy = False

    drivers_map = (
        ("304", nvidia_304xx_supported,),
        ("173", nvidia_173xx_supported,),
        ("96", nvidia_173xx_supported,),
        ("71", nvidia_173xx_supported,),
        )

    if not nvidia_disablelegacy:
        for ver, lst in drivers_map:
            if card_id not in lst:
                continue
            print("NVIDIA %s driver selected" % (ver,))
            drv_string = ver
            if livecd:
                rc = deploy_nvidia_xxxxxx_drivers(ver)
                if rc:
                    print("NVIDIA %s deployed correctly" % (ver,))
                    done_legacy = True
                    break

    if not done_legacy:
        drv_string = '[latest]'
        print("latest and greatest NVIDIA driver selected or unsupported")

    healthy = check_if_proprietary_driver_system_is_healthy("nvidia")
    if healthy:
        print("NVIDIA proprietary driver %s is loaded" % (drv_string,))

        if done_legacy:
            try:
                os.makedirs("/lib/nvidia/legacy")
            except OSError  as err:
                if err.errno != errno.EEXIST:
                    raise
            with open("/lib/nvidia/legacy/running", "w") as f:
                f.write("%s" % (drv_string,))

    return done_legacy, healthy

def generate_nvidia_bumblebee_steps(v3dcard, company_id, card_id):
    done_legacy, healthy = setup_nvidia_drivers(card_id)
    if not healthy:
        print("NVIDIA drivers couldn't be loaded, cannot enable bumblebee")
        return

    if dryrun:
        print("Was about to start bumblebee")
        return

    if not livecd:
        print("LiveCD mode off, not starting bumblebee service")
        return

    # This is used by our Installer
    with open("/tmp/.bumblebee.enabled", "w") as f:
        pass

    if openrc_running():
        os.system("/etc/init.d/bumblebee start")
    elif systemd_running():
        os.system("/usr/bin/systemctl start bumblebeed")

def generate_nvidia_steps(videocard, cardnumber, total_cards, bus_id):
    comp_id, card_id = extract_pci_ids(videocard)
    done_legacy, healthy = setup_nvidia_drivers(card_id)

    if healthy:
        if done_legacy:
            # then activate nvidia opengl subsystem after resetting it
            steps.append((opengl_activate, "xorg-x11"))
            steps.append((opengl_activate, "nvidia"))

            steps.append((set_xorg_device, "nvidia",
                          cardnumber, total_cards, bus_id,))
            steps.append((fix_possible_opengl_misconfiguration, "nvidia"))
            steps.append((copy_nvidia_settings_on_desktop,))

        else:

            steps.append((fix_possible_opengl_misconfiguration, "nvidia"))
            steps.append((copy_nvidia_settings_on_desktop,))
            steps.append((opengl_activate, "nvidia"))
            steps.append((set_xorg_device, "nvidia",
                          cardnumber, total_cards, bus_id,))
    else:
        print("NVIDIA drivers couldn't be loaded, switchting to nv driver")
        steps.append((opengl_activate, "xorg-x11"))

def generate_generic_steps():
    steps.append((remove_proprietary_opengl, bb_enabled))
    steps.append((opengl_activate, "xorg-x11",))

def drop_kernel_mod(kmod):
    return subprocess.call(["modprobe", "-r", kmod])

def extract_pci_ids(videocard_str):
    videocard_split = [x.strip() for x in videocard_str.strip().split('"') \
                           if x.strip()]
    try:
        card_id = videocard_split[3].split()[-1].lower().strip("[]")
    except IndexError:
        card_id = None

    try:
        company_id = videocard_split[2].split()[-1].lower().strip("[]")
    except IndexError:
        company_id = None

    return company_id, card_id

def extract_vga_cards(lspci_list):
    cards = []
    for item in lspci_list:
        try:
            class_type = item.split()[1].strip('"')
            if class_type == "0300":
                cards.append(item)
        except IndexError:
            continue
    return cards

def extract_3d_cards(lspci_list):
    # bumblebee support
    cards = []
    for item in lspci_list:
        try:
            class_type = item.split()[1].strip('"')
            if class_type == "0302":
                cards.append(item)
        except IndexError:
            continue
    return cards


# Create videocards list
lspci_out_split = lspci_output.split("\n")
videocards = extract_vga_cards(lspci_out_split)
v3dcards = extract_3d_cards(lspci_out_split)
# Run the program
cardnumber = -1

total_cards = len(videocards)
forced_monitor_modes = False
steps = []
bb_enabled = False
write_config = False

for v3dcard in v3dcards:

    company_id, card_id = extract_pci_ids(v3dcard)

    if company_id == "10de":
        print("NVIDIA Optimus 3D Acceleration detected, enabling bumblebee")
        generate_nvidia_bumblebee_steps(v3dcard, company_id, card_id)
        bb_enabled = True

for videocard in videocards:

    # setup card number
    cardnumber += 1
    print("Card Number: " + str(cardnumber))
    try:
        bus_id = "PCI:%s" % (
                videocard.split()[0].split(".", 1)[0]
            )
    except (IndexError,ValueError,TypeError,):
        bus_id = None

    if forced_xdriver:
        print("You have chosen to force the X driver: " + forced_xdriver)
        if forced_xdriver == "fglrx":
            if check_if_proprietary_driver_system_is_healthy("fglrx") \
                    or noproprietary:
                steps.append((opengl_activate, "xorg-x11"))
                forced_xdriver = "ati"
                steps.append((drop_kernel_mod, "fglrx",))
            else:
                steps.append((fix_possible_opengl_misconfiguration, "ati"))
                steps.append((copy_ati_settings_on_desktop,))
                steps.append((opengl_activate, "ati"))

        elif forced_xdriver == "nvidia" and (not noproprietary):
            generate_nvidia_steps(videocard, cardnumber, total_cards, bus_id)
        elif forced_xdriver == "vesa":
            forced_monitor_modes = True
        else:
            generate_generic_steps()
            steps.append((set_xorg_device, forced_xdriver,
                          cardnumber, total_cards, bus_id,))
        write_config = True

    else:
        company_id, card_id = extract_pci_ids(videocard)
        print("[%s] company_id: %s | card_id: %s" % (
                cardnumber, company_id, card_id,))

        if company_id == "10de": # NVIDIA
            if noproprietary:
                steps.append((set_xorg_device, "nv",
                              cardnumber, total_cards, bus_id,))
            else:
                generate_nvidia_steps(
                    videocard, cardnumber, total_cards, bus_id)
                print("NVIDIA!")
            write_config = True

        elif company_id == "1002":
            generate_fglrx_steps(
                videocard, cardnumber, total_cards, bus_id)
            write_config = True

        else:
            generate_generic_steps()
            print("GPU will be automatically detected by X.Org and udevd")


# now create the file
for args in steps:
    func, args = args[0], args[1:]
    func(*args)

if write_config:
    config = xorg_conf_structure.replace(
        '__device_section__',
        '\n\n'.join(device_sections))
    config = config.replace(
        '__screen_section__',
        '\n\n'.join(screen_sections))
    config = config.replace(
        '__screen_layout_section__',
        '\n    '.join(screen_layout_sections))
    if forced_monitor_modes:
        config = config.replace('#Modes', 'Modes')

    if not dryrun:
        with open(xorgfile, "w") as f:
            f.write(config)
            f.flush()
else:
    try:
        os.remove(xorgfile)
    except (OSError, IOError):
        pass

raise SystemExit(0)