#! /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 subprocess
import sys

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

from ubiquity import osextras


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

    i = 1
    while i < len(sys.argv):
        if sys.argv[i] == '--desktop':
            desktop = sys.argv[i + 1]
            i += 2
            # strip option and argument from newargv
            continue
        elif not sys.argv[i].startswith('-'):
            frontend = sys.argv[i]
        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

        if frontend == 'gtk_ui':
            toexec = ['gksudo', '--preserve-env']
            if desktop:
                toexec += ['--desktop', desktop]
            toexec.append('--')
        elif frontend == 'kde_ui':
            # We need to pass through xauth information.  kdesudo won't do,
            # though, as it can't be made to tell sudo to preserve the
            # environment.
            import tempfile
            xauthority = tempfile.NamedTemporaryFile(prefix='ubiquity-',
                                                     suffix='-xauth')
            os.chmod(xauthority.name, 0o600)  # root can still read it
            xauth_extract = subprocess.Popen(
                ['xauth', 'extract', '-', os.environ['DISPLAY']],
                stdout=subprocess.PIPE)
            xauth_merge = subprocess.Popen(
                ['xauth', '-f', xauthority.name, 'merge', '-'],
                stdin=xauth_extract.stdout)
            xauth_merge.wait()
            xauth_extract.stdout.close()
            xauth_extract.wait()
            toexec = ['sudo', '-E', 'XAUTHORITY=%s' % xauthority.name]
        else:
            toexec = ['sudo', '-E']

        # re-exec ourselves first time round, to cope with broken argument
        # handling in kdesu
        toexec.append(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()
