#! /usr/bin/python3

# Wrapper script to run Ubiquity as root using the appropriate privilege
# escalation method for the frontend.

from __future__ import print_function

import os
import sys

sys.path.insert(0, '/usr/lib/ubiquity')

from ubiquity import osextras


def main():
    newargv = []
    frontend = None
    toexec = []

    if os.getenv('GTK_MODULES', False):
        os.environ['GTK_MODULES'] += os.pathsep + 'overlay-scrollbar'
    else:
        os.environ['GTK_MODULES'] = 'overlay-scrollbar'

    i = 1
    while i < len(sys.argv):
        if not sys.argv[i].startswith('-'):
            frontend = sys.argv[i]
        elif sys.argv[i] == '--autopilot':
            os.environ['GTK_MODULES'] += os.pathsep + 'autopilot'
        newargv.append(sys.argv[i])
        i += 1

    if os.getuid() == 0:
        # no privilege escalation required
        inner = ['/usr/lib/ubiquity/bin/ubiquity'] + newargv

        # Make sure ibus works
        if (os.getenv("GTK_IM_MODULE") == "ibus" and
                os.getenv("XDG_SESSION_COOKIE")):
            ibus_path = os.path.expanduser("~/.config/ibus/bus")
            if os.path.exists(ibus_path):
                # Guess the path to the ibus config file
                ibus_filename = "%s-unix-0" % (
                    os.getenv("XDG_SESSION_COOKIE").split("-")[0])
                ibus_config = os.path.join(ibus_path, ibus_filename)
                with open(ibus_config, "r") as fp:
                    for line in fp:
                        fields = line.strip().split('=', 1)
                        # If we get KEY=VALUE, export it in the environment
                        if len(fields) == 2:
                            os.environ[fields[0]] = fields[1]

        # Ensure the OOM killer doesn't nom on us.
        with open('/proc/%d/oom_score_adj' % os.getpid(), 'w') as fp:
            fp.write('-1000')

        # Disable storage automounting
        udisks2 = '/usr/lib/udisks2/udisks2-inhibit'
        if os.path.isfile(udisks2) and os.access(udisks2, os.X_OK):
            toexec += [udisks2] + inner
        elif osextras.find_on_path('udisks'):
            toexec += ['udisks', '--inhibit', '--'] + inner
        elif (osextras.find_on_path('hal-lock') and
              not os.system('pgrep hald>/dev/null')):
            toexec += ['hal-lock',
                       '--interface', 'org.freedesktop.Hal.Device.Storage',
                       '--exclusive', '--run', ' '.join(inner)]
        else:
            toexec += inner
    else:
        if frontend is None:
            # Try to detect which frontend will be used by looking for a
            # frontend module.
            import imp
            import ubiquity.frontend
            frontend_names = ['gtk_ui', 'kde_ui']
            for f in frontend_names:
                try:
                    imp.find_module(f, ubiquity.frontend.__path__)
                    frontend = f
                    break
                except ImportError:
                    pass

        toexec = ['pkexec', sys.argv[0]]
        toexec.extend(newargv)

    if 'UBIQUITY_WRAPPER_DEBUG' in os.environ:
        print(toexec, file=sys.stderr)
    os.execvp(toexec[0], toexec)
    sys.exit(127)


if __name__ == '__main__':
    main()