#!/usr/bin/python3
# -*- coding: utf-8 -*-
from __future__ import absolute_import, division, print_function, unicode_literals
import argparse
import sys
import os
import glob
import subprocess
import time
import shutil
import re
try:
    import configparser
except ImportError:
    import ConfigParser as configparser

from dogtail.sessions import Script

SCRIPT_DESCRIPTION = """
Unlike the original headless script this will make use of a Display Manager
(DM) to handle starting the X server and user session. It's motivated by
changes related to systemd that disallows running a GNOME session from an
environment spawned by 'su'. The original headless will not work in these
cases on systemd systems.

This version uses the AutoLogin feature of the DM, so that when its
service is started the session will login for the user immediately. It then
grabs the new session’s environment and runs the target script in there.

Works on distros where 'service gdm/kdm/lightdm start/stop' takes effect,
and on systemd systems using systemd-logind.
"""

PRESERVE_ENVS = ["PYTHONPATH", "TEST"]


def getSessionEnvironment(sessionBinary):
    def isSessionProcess(file_name):
        session_binary_path = ("/usr/bin/plasma-desktop"
                               if sessionBinary.split("/")[-1] == "startkde"
                               else sessionBinary)
        try:
            if os.path.realpath(file_name + "exe") != session_binary_path:
                return False
        except OSError:
            return False
        pid = file_name.split("/")[2]
        if pid == "self" or pid == str(os.getpid()):
            return False
        return True

    def get_environment_dictionary(file_name):
        try:
            data = open(file_name, "r").read()
        except IOError as e:
            print("dogtail exception:", e)
            return {}
        items = data.split("\x00")
        env = {}
        for item in items:
            if "=" not in item:
                continue
            k, v = item.split("=", 1)
            env[k] = v
        return env

    def copy_environment_variables(env_dict):
        for var in PRESERVE_ENVS:
            if var in os.environ:
                env_dict[var] = os.environ[var]
        return env_dict

    env = None
    for path in glob.glob("/proc/*/"):
        if not isSessionProcess(path):
            continue
        env = get_environment_dictionary(path + "environ")

    if not env:
        raise RuntimeError("Can't find our environment!")
    return copy_environment_variables(env)


def execCodeWithEnv(code, env=None):
    with open("/tmp/execcode.dogtail", "w") as f:
        f.write(code)
    subprocess.Popen(
        "python /tmp/execcode.dogtail".split(),
        env=(os.environ if env is None else env)
    ).wait()


### qecore headless code to quickfix EL9 #####################
ENVIRONMENT_DICTIONARY = {}
ENVIRONMENT_VARIABLES_TO_PRESERVE = ["PYTHONPATH", "TEST", "TERM"]


def run(command, verbose=False):
    try:
        out = subprocess.check_output(
            command, shell=True, stderr=subprocess.STDOUT, encoding="utf-8"
        )
        return out if not verbose else (out, 0, None)
    except subprocess.CalledProcessError as e:
        return e.output if not verbose else (e.output, e.returncode, e)


def _get_environment_as_dictionary(path_to_process):
    def is_session_process(p):
        pid = p.split("/")[2]
        return pid != "self" and pid != str(os.getpid())

    if not is_session_process(path_to_process):
        return False

    file_name = path_to_process + "environ"
    try:
        data = open(file_name, "r").read()
    except IOError:
        return False

    for item in data.split("\x00"):
        if "=" in item:
            k, v = item.split("=", 1)
            ENVIRONMENT_DICTIONARY[k] = v

    for var in ENVIRONMENT_VARIABLES_TO_PRESERVE:
        if var in os.environ:
            ENVIRONMENT_DICTIONARY[var] = os.environ[var]

    return True


def get_environment_dictionary(session_binary):
    def get_gs_pid():
        out = run("ps o pid,command,user").split("\n")
        for line in out:
            if "gnome-session" in line and os.getenv("USER", "") in line:
                return int(line.strip().split()[0])

    pid = get_gs_pid()
    if pid:
        _get_environment_as_dictionary(f"/proc/{pid}/")
        print("headless: Setting environment variable TERM as 'xterm'")
        ENVIRONMENT_DICTIONARY["TERM"] = "xterm"
        return ENVIRONMENT_DICTIONARY

    status = run("systemctl is-active gdm").strip()
    print(
        f"headless: Display Manager Status is '{status}'\n"
        "headless: Can't find our environment!\n"
        "headless: Stopping Display Manager"
    )
    subprocess.Popen("sudo systemctl stop gdm".split()).wait()
    sys.exit(1)


###################################


class DisplayManagerSession:
    gdm_config     = "/etc/gdm/custom.conf"
    kdm_config     = "/etc/kde/kdm/kdmrc"
    lightdm_config = "/etc/lightdm/lightdm.conf"

    gdm_options = {
        "section": "daemon",
        "enable":  "AutomaticLoginEnable",
        "user":    "AutomaticLogin"
    }
    kdm_options = {
        "section": "X-:0-Core",
        "enable":  "AutoLoginEnable",
        "user":    "AutoLoginUser"
    }
    lightdm_options = {
        "section": "Seat:*",
        "user":    "autologin-user",
        "timeout": "autologin-user-timeout",
        "session": "autologin-session"
    }

    def __init__(self, display_manager="gdm", session="gnome",
                 session_binary="gnome-shell", user=None):
        self.display_manager = display_manager
        self.session         = session
        self.session_binary  = session_binary
        self.user            = user or os.getenv("USER", "test")
        self.accountfile     = f"/var/lib/AccountsService/users/{self.user}"

        if display_manager == "gdm":
            self.config = self.gdm_config
            self.options = self.gdm_options
        elif display_manager == "kdm":
            self.config = self.kdm_config
            self.options = self.kdm_options
        elif display_manager == "lightdm":
            self.config = self.lightdm_config
            self.options = self.lightdm_options
        else:
            raise ValueError(f"Unsupported display manager: {display_manager}")

        self.tmp_file = "/tmp/" + os.path.basename(self.config)

    def isProcessRunning(self, proc_name):
        ps = subprocess.Popen(["ps", "axw"], stdout=subprocess.PIPE)
        for line in ps.stdout:
            if "headless" in line.decode("utf-8"):
                continue
            if re.search(proc_name, line.decode("utf-8")):
                return True
        return False

    def waitForProcess(self, proc_name, invert=False, max_cycles=None):
        """
        Wait until a process appears (invert=False) or disappears (invert=True).

        :param proc_name: substring of the process name to wait for
        :param invert: if False, wait until process **is** running; if True, wait until it's gone
        :param max_cycles: optional int, maximum number of 1-second polls before timing out;
                           None means wait indefinitely
        :return: True if the desired condition was met within the limit, False on timeout
        """
        cycles = 0
        # keep looping while condition is not yet met
        while self.isProcessRunning(proc_name) is invert:
            # if we've reached the cycle limit, bail out
            if max_cycles is not None and cycles >= max_cycles:
                return False
            time.sleep(1)
            cycles += 1
        return True


    def setup(self, restore=False, force_xorg=False):
        # copy and edit config
        shutil.copy(self.config, self.tmp_file)
        if hasattr(configparser, "SafeConfigParser"):
            cfg = configparser.SafeConfigParser()
        else:
            cfg = configparser.ConfigParser()
        cfg.optionxform = str
        cfg.read(self.tmp_file)

        section = self.options["section"]
        if not cfg.has_section(section):
            cfg.add_section(section)

        if not restore:
            if self.display_manager in ("gdm", "kdm"):
                cfg.set(section, self.options["enable"], "true")
                cfg.set(section, self.options["user"], self.user)
                if force_xorg and self.display_manager == "gdm":
                    cfg.set(section, "WaylandEnable", "false")
            else:  # lightdm
                cfg.set(section, self.options["user"], self.user)
                cfg.set(section, self.options["timeout"], "0")
                if self.session:
                    cfg.set(section, self.options["session"], self.session)
        else:
            if self.display_manager in ("gdm", "kdm"):
                cfg.remove_option(section, self.options["enable"])
                cfg.remove_option(section, self.options["user"])
            else:
                for key in ("user", "timeout", "session"):
                    cfg.remove_option(section, self.options[key])

        with open(self.tmp_file, "w") as out:
            cfg.write(out)
        subprocess.Popen(f"sudo mv -f {self.tmp_file} {self.config}", shell=True).wait()

        # post-setup for GDM/KDM only
        if not restore and self.display_manager in ("gdm", "kdm"):
            if "kwin" in self.session_binary:
                try:
                    os.makedirs(os.path.expanduser("~/.kde/env/"), exist_ok=True)
                except Exception:
                    pass
                subprocess.Popen(
                    "echo 'export QT_ACCESSIBILITY=1' > ~/.kde/env/qt-at-spi.sh",
                    shell=True
                ).wait()

            if self.display_manager == "gdm":
                need_restart = False
                tmp_acc = f"/tmp/{self.user}_headless"
                if os.path.isfile(self.accountfile):
                    subprocess.Popen(
                        f"cp -f {self.accountfile} {tmp_acc}", shell=True
                    ).wait()

                if hasattr(configparser, "SafeConfigParser"):
                    acc_cfg = configparser.SafeConfigParser()
                else:
                    acc_cfg = configparser.ConfigParser()
                acc_cfg.optionxform = str
                acc_cfg.read(tmp_acc)

                try:
                    saved = acc_cfg.get("User", "Session")
                    if self.session is None:
                        if "kde" in saved and "gnome-shell" in self.session_binary:
                            self.session_binary = "/usr/bin/kwin"
                        elif "gnome" in saved and "kwin" in self.session_binary:
                            self.session_binary = "/usr/bin/gnome-shell"
                    elif saved != self.session:
                        acc_cfg.set("User", "Session", self.session)
                        need_restart = True
                except (configparser.NoOptionError, configparser.NoSectionError):
                    if self.session is not None:
                        if not acc_cfg.has_section("User"):
                            acc_cfg.add_section("User")
                        acc_cfg.set("User", "Session", self.session)
                        acc_cfg.set("User", "SystemAccount", "false")
                        need_restart = True

                if need_restart:
                    with open(tmp_acc, "w") as o:
                        acc_cfg.write(o)
                    subprocess.Popen(
                        f"sudo mv -f {tmp_acc} {self.accountfile}", shell=True
                    ).wait()
                    time.sleep(1)
                    subprocess.Popen(
                        "sudo systemctl restart accounts-daemon", shell=True
                    ).wait()
                    time.sleep(1)
                    subprocess.Popen(
                        "sudo systemctl restart systemd-logind", shell=True
                    ).wait()
                    time.sleep(1)
                else:
                    subprocess.Popen(f"sudo rm -f {tmp_acc}", shell=True).wait()

            elif self.display_manager == "kdm":
                if self.session is not None:
                    subprocess.Popen(
                        f"echo '[Desktop]\nSession={self.session}' > /home/{self.user}/.dmrc",
                        shell=True
                    ).wait()

    def start(self, restart=False):
        subprocess.Popen(f"sudo systemctl stop {self.display_manager}".split())
        time.sleep(0.5)

        # kill any pre-existing seat0 session for this user
        if os.system(
            f"sudo loginctl | grep {self.user} | grep seat0"
        ) == 0:
            subprocess.Popen(
                f"sudo loginctl kill-session --signal=9 $(loginctl | grep {self.user} | grep seat0 | awk '{{print $1}}')",
                shell=True
            ).wait()
            time.sleep(5)

        subprocess.Popen(f"sudo systemctl start {self.display_manager}".split())
        self.waitForProcess(os.path.basename(self.session_binary))

        # give DM-specific delay
        if self.display_manager == "kdm":
            time.sleep(10)
        elif self.display_manager == "gdm":
            time.sleep(4)

    def set_accessibility_to(self, enable):
        subprocess.Popen(
            f"/usr/bin/gsettings set org.gnome.desktop.interface toolkit-accessibility {'true' if enable else 'false'}",
            shell=True,
            env=os.environ
        )

    def stop(self):
        subprocess.Popen(
            f"sudo systemctl stop {self.display_manager}".split()
        ).wait()
        self.waitForProcess(self.display_manager, invert=True)
        time.sleep(3)

        # if session processes linger, force-kill them
        if self.isProcessRunning(os.path.basename(self.session_binary)):
            subprocess.Popen(
                f"sudo loginctl terminate-session $(loginctl | grep {self.user} | grep seat0 | awk '{{print $1}}')",
                shell=True
            ).wait()
            time.sleep(5)

        if any(self.isProcessRunning(n) for n in ("Xorg", "gnome-shell", "Xwayland")):
            subprocess.Popen(
                f"sudo loginctl kill-session --signal=9 $(loginctl | grep {self.user} | grep seat0 | awk '{{print $1}}')",
                shell=True
            ).wait()
            time.sleep(3)


def parse():
    parser = argparse.ArgumentParser(
        prog="$ dogtail-run-headless-next",
        description=SCRIPT_DESCRIPTION,
        formatter_class=argparse.RawDescriptionHelpFormatter
    )

    parser.add_argument("script",
                        help="Command to execute in the headless session")

    parser.add_argument("--session",
                        required=False,
                        help="\n".join((
                            "What session to use (e.g. 'gnome', 'kde-plasma', 'mate').",
                            "Comes from /usr/share/xsessions/ or /usr/share/wayland-sessions/."
                        )))

    parser.add_argument("--session-binary",
                        required=False,
                        help="Full path to an in-session binary (e.g. '/usr/bin/gnome-shell').")

    parser.add_argument("--dm",
                        required=False,
                        help="\n".join((
                            "What display manager to use for spawning session.",
                            "Supported: 'gdm' (default), 'kdm', or 'lightdm'."
                        )))

    parser.add_argument("--restart",
                        action="store_true",
                        help="Restart previously running display manager session before script execution.")

    parser.add_argument("--dont-start",
                        action="store_true",
                        help="Use already running session (doesn't have to be under Display Manager).")

    parser.add_argument("--dont-kill",
                        action="store_true",
                        help="Do not kill session when script exits.")

    parser.add_argument("--disable-a11y",
                        action="store_true",
                        help="Disable accessibility technologies on exit.")

    parser.add_argument("--debug",
                        action="store_true",
                        help="Enable DOGTAIL_DEBUG in the session environment.")

    parser.add_argument("--force-xorg",
                        action="store_true",
                        help="Force Xorg session by setting WaylandEnable=false in GDM config.")

    return parser.parse_args()


def main():
    args = parse()
    script_command_list = args.script.split()

    # default session / binary logic
    if args.dm == "lightdm":
        args.session = args.session or "mate"
        args.session_binary = args.session_binary or "/usr/bin/mate-session"
    elif args.session is None or "gnome" in args.session:
        args.session_binary = args.session_binary or "/usr/bin/gnome-session"
    elif "kde" in args.session:
        args.session_binary = args.session_binary or "/usr/bin/kwin"
        if args.session == "kde":
            args.session = "kde-plasma"
    else:
        if not args.session_binary:
            print("dogtail-run-headless-next: Need to specify --session-binary for '%s'." %
                  args.session)
            sys.exit(-1)

    # pick display manager
    if args.dm in (None, "gdm"):
        display_manager_name = "gdm"
    elif args.dm == "kdm":
        display_manager_name = "kdm"
    elif args.dm == "lightdm":
        display_manager_name = "lightdm"
    else:
        print("dogtail-run-headless-next: Unknown display manager '%s'!" % args.dm)
        sys.exit(-1)

    print("dogtail-run-headless-next: Using display manager:", display_manager_name)
    display_manager_session = DisplayManagerSession(
        display_manager_name,
        args.session,
        args.session_binary,
        user=os.getenv("USER")
    )

    if not args.dont_start:
        display_manager_session.setup(force_xorg=args.force_xorg)
        display_manager_session.start(restart=args.restart)

    print("dogtail-run-headless-next: Binding to session binary:",
          display_manager_session.session_binary)

    # grab the session environment
    try:
        os.environ = getSessionEnvironment(
            display_manager_session.session_binary
        )
    except RuntimeError:
        print("Could not get environment from initial session binary",
              display_manager_session.session_binary)
        if display_manager_session.session_binary == "/usr/bin/gnome-shell":
            try:
                binary = "/usr/libexec/gvfsd"
                try:
                    with open('/etc/redhat-release', 'r') as f:
                        release = f.read()
                    if "Linux release 9" in release:
                        binary = "/usr/bin/gnome-session"
                except Exception:
                    pass
                print("Try fallback binary:", binary)
                os.environ = getSessionEnvironment(binary)
                print("dogtail-run-headless-next: Trying workaround with", binary)
            except RuntimeError:
                print("dogtail-run-headless-next: Fallback failed for gvfsd")
                display_manager_session.stop()
                display_manager_session.setup(restore=True)
                sys.exit(1)
        else:
            display_manager_session.stop()
            display_manager_session.setup(restore=True)
            sys.exit(1)

    if args.debug:
        os.environ["DOGTAIL_DEBUG"] = "true"

    if display_manager_name == "gdm":
        display_manager_session.set_accessibility_to(True)

    script = Script(script_command_list)

    if os.path.isfile('/tmp/headless_enable_fatal_criticals'):
        os.environ['G_DEBUG'] = 'fatal-criticals'
    if os.path.isfile('/tmp/headless_enable_fatal_warnings'):
        os.environ['G_DEBUG'] = 'fatal-warnings'

    pid = script.start()
    print("dogtail-run-headless-next: Started script with PID", pid)
    exit_code = script.wait()
    print("dogtail-run-headless-next: Script finished with return code", exit_code)

    if args.disable_a11y:
        display_manager_session.set_accessibility_to(False)

    if not args.dont_kill:
        display_manager_session.stop()
        display_manager_session.setup(restore=True)

    sys.exit(exit_code)


if __name__ == "__main__":
    main()