switchlibglx

switchlibglx

#!/usr/bin/python
#
#       switchlibglx
#
#       Copyright 2011 Canonical Ltd.
#       Author: Alberto Milone <alberto.milone@canonical.com>
#
#       Script to switch between AMD and Intel graphics driver libraries.
#
#       Usage:
#           switchlibglx   amd|intel|query
#           amd:   switches to AMD's version of libglx.so
#           intel: switches to the open-source version of libglx.so
#           query: checks which version is currently active and writes
#                  "amd", "intel" or "unknown" to the standard output
#
#       Permission is hereby granted, free of charge, to any person
#       obtaining a copy of this software and associated documentation
#       files (the "Software"), to deal in the Software without
#       restriction, including without limitation the rights to use,
#       copy, modify, merge, publish, distribute, sublicense, and/or sell
#       copies of the Software, and to permit persons to whom the
#       Software is furnished to do so, subject to the following
#       conditions:
#
#       The above copyright notice and this permission notice shall be
#       included in all copies or substantial portions of the Software.
#
#       THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
#       EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
#       OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
#       NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
#       HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
#       WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
#       FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
#       OTHER DEALINGS IN THE SOFTWARE.


import optparse
import os
import sys
import re
import subprocess
from subprocess import Popen, PIPE, CalledProcessError

class Alternatives:

    def __init__(self, master_link):
        self._open_drivers_alternative = 'mesa/ld.so.conf'
        self._command = 'update-alternatives'
        self._master_link = master_link

        # Make sure that the PATH environment variable is set
        if not os.environ.get('PATH'):
            os.environ['PATH'] = '/sbin:/usr/sbin:/bin:/usr/bin'

    def _find_process(self, process_name):
        ps = subprocess.Popen("ps -eaf | grep " + process_name + " | grep -v grep",
                              shell=True, stdout=subprocess.PIPE)
        output = ps.stdout.read()
        ps.stdout.close()
        ps.wait()
        m = re.search('^(\w+)\s+(\d+)\s+.*', output, re.M)
        if m:
            try:
                output = int(m.group(2))
            except ValueError:
                output = -1
        else:
            output = -1
        return output

    def _is_process_running(self, process_name):
        process_id = self._find_process(process_name)
        if process_id < 0:
            return False
        try:
            os.kill(process_id, 0)
            return True
        except OSError:
            return False

    def list_alternatives(self):
        '''Get the list of alternatives for the master link'''
        dev_null = open('/dev/null', 'w')
        alternatives = []
        p1 = Popen([self._command, '--list', self._master_link], stdout=PIPE, stderr=dev_null)
        p = p1.communicate()[0]
        dev_null.close()
        c = p.split('\n')
        for line in c:
            line.strip() and alternatives.append(line.strip())

        return alternatives

    def get_current_alternative(self):
        '''Get the alternative in use'''
        dev_null = open('/dev/null', 'w')
        current_alternative = None
        p1 = Popen([self._command, '--query', self._master_link], stdout=PIPE, stderr=dev_null)
        p = p1.communicate()[0]
        dev_null.close()
        c = p.split('\n')
        for line in c:
            if line.strip().startswith('Value:'):
                return line.replace('Value:', '').strip()
        return None

    def get_alternative_by_name(self, name):
        '''Get the alternative link by providing the driver name'''
        alternatives = self.list_alternatives()

        for alternative in alternatives:
            if alternative.__contains__(name):
                return alternative

        return None

    def get_open_drivers_alternative(self):
        '''Get the alternative link for open drivers'''
        return self.get_alternative_by_name(self._open_drivers_alternative)

    def update_gmenu(self):
        '''Trigger gmenu so that the icons will show up in the menu'''
        try:
            subprocess.check_call(['dpkg-trigger', '--by-package=fakepackage', 'gmenucache'], stdout=open(os.devnull, 'w'))
            # Let's not call dpkg if dpkg is already running
            if not self._is_process_running('dpkg'):
                subprocess.check_call(['dpkg', '--configure', '-a'], stdout=open(os.devnull, 'w'))
        except (OSError, CalledProcessError):
            pass

    def set_alternative(self, path):
        '''Tries to set an alternative and returns the boolean exit status'''
        try:
            subprocess.check_call([self._command, '--set', self._master_link, path], stdout=open(os.devnull, 'w'))
            self.ldconfig()
        except CalledProcessError:
            return False

        self.update_gmenu()

        return True

    def ldconfig(self):
        '''Call ldconfig'''
        try:
            subprocess.check_call(['ldconfig'], stdout=open(os.devnull, 'w'))
        except CalledProcessError:
            return False
        return True


class Switcher(object):

    def __init__(self):
        self._supported_architectures = {'i386': 'i386', 'amd64': 'x86_64'}
        main_arch = self._get_architecture()
        other_arch = self._supported_architectures.values()[
                          int(not self._supported_architectures
                          .values().index(main_arch))]
        main_alternative_name = self._get_alternative_name_from_arch(main_arch)
        other_alternative_name = self._get_alternative_name_from_arch(other_arch)

        # We have 2 alternatives, one for each architecture
        self._gl_switcher = Alternatives(main_alternative_name)
        self._gl_switcher_other = Alternatives(other_alternative_name)

    def _get_architecture(self):
        dev_null = open('/dev/null', 'w')
        p1 = Popen(['dpkg', '--print-architecture'], stdout=PIPE, stderr=dev_null)
        p = p1.communicate()[0]
        dev_null.close()
        architecture = p.strip()
        return self._supported_architectures.get(architecture)

    def _get_alternative_name_from_arch(self, architecture):
        alternative = '%s-linux-gnu_gl_conf' % architecture
        return alternative

    def _simplify_x_alternative_name(self, alternative):
        return alternative.split('/')[-1]

    def _simplify_gl_alternative_name(self, alternative):
        return alternative.split('/')[-2]

    def _get_gl_alternatives_list(self):
        raw_alternatives = self._gl_switcher.list_alternatives()
        return  [ self._simplify_gl_alternative_name(x) for x in raw_alternatives ]

    def _update_initramfs(self):
        subprocess.call(['update-initramfs', '-u'])
        # This may not be necessary
        subprocess.call(['update-initramfs', '-u', '-k', os.uname()[2]])

    def _get_current_alternative(self):
        # This is a list as the 2nd item could be another alternative
        # we might be looking for e.g. an xorg.conf alternative
        alternatives = [None, None]
        raw_gl_alternative = self._gl_switcher.get_current_alternative()
        raw_gl_alternative_other = self._gl_switcher_other.get_current_alternative()

        if (raw_gl_alternative != None):
            alternatives[0] = self._simplify_gl_alternative_name(raw_gl_alternative)
        if (raw_gl_alternative_other != None):
            alternatives[1] = self._simplify_gl_alternative_name(raw_gl_alternative_other)

        return alternatives

    def enable_alternative(self, alternative):
        # Make sure that the alternative exists
        gl_alternative = self._gl_switcher.get_alternative_by_name(alternative)
        gl_alternative_other = self._gl_switcher_other.get_alternative_by_name(alternative)

        # See if they are null
        # Abort if gl_alternative is null
        if gl_alternative and gl_alternative_other:
            success = (self._gl_switcher.set_alternative(gl_alternative) and
                       self._gl_switcher_other.set_alternative(gl_alternative_other))
            return success
        else:
            sys.stderr.write("Error: no alternative can be found for %s\n" % alternative)

        return False

    def print_current_alternative(self):
        alternatives = self._get_current_alternative()

        try:
            alternative = str(alternatives[0])
            alternative_other = str(alternatives[1])
        except IndexError:
            # No alternative exists
            return False

        if alternative == alternative_other == 'fglrx':
            print 'amd'
        elif alternative == alternative_other == 'pxpress':
            print 'intel'
        else:
            print 'unknown'

        return True

def check_root():
    if not os.geteuid() == 0:
        sys.stderr.write("This operation requires root privileges\n")
        exit(1)

def handle_alternatives_error(mode):
    sys.stderr.write('Error: %s mode can\'t be enabled\n' % mode)
    exit(1)

def handle_query_error():
    sys.stderr.write("Error: no alternative can be found\n")
    exit(1)

def usage():
    sys.stderr.write("Usage: %s amd|intel|query\n" % (sys.argv[0]))

if __name__ == '__main__':
    try:
        arg = sys.argv[1]
    except IndexError:
        arg = None

    check_root()

    if len(sys.argv[1:]) != 1:
        usage()
        exit(1)

    switcher = Switcher()

    if arg == "amd":
        if not switcher.enable_alternative('fglrx'):
            handle_alternatives_error(arg)
    elif arg == "intel":
        if not switcher.enable_alternative('pxpress'):
            handle_alternatives_error(arg)
    elif arg == "query":
        if not switcher.print_current_alternative():
            handle_query_error()
    else:
        usage()
        sys.exit(1)

    exit(0)