#
# Copyright © 2023 Red Hat, Inc
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License as
# published by the Free Software Foundation; either version 2 of the
# License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA
# 02111-1307, USA.
#
#   Olivier Fourdan <ofourdan@redhat.com>
#

""" Common routines for xwayland modules. """

import os
import subprocess
from secrets import token_hex
from signal import signal, SIGTERM, SIG_IGN
from tempfile import mkstemp, mkdtemp


class Xwayland:

    """
    Common routines for Xwayland.
    """


    def __call__(self):
        return self


    def __init__(self):
        self.tempdir = None
        self.xauthfile = None
        try:
            os.setpgrp()
        except Exception as error:
            print(f'Warning, failed to create process group: {error}\n')


    def run_command(self, command):
        """ Helper function to spawn a command. """
        try:
            with subprocess.Popen(command) as proc:
                proc.wait()
                return proc.returncode
        except Exception as error:
            print(f'Failed to run the command: {error}\n')
            return -1
        except KeyboardInterrupt:
            return 0


    def spawn_xwayland(self, xserver_args = [], command = []):
        """ Helper function to spawn Xwayland. """
        command.extend(['Xwayland', '-noreset', '-displayfd', '1'])
        command.extend(xserver_args)
        try:
            proc = subprocess.Popen(command, stdout = subprocess.PIPE)
            outs = proc.stdout.readline()
            disp_num = outs.decode().rstrip()
            if not disp_num:
                print('Xwayland failed to start!\n')
                return -1
            os.environ['DISPLAY'] = f':{disp_num}'
            return proc.pid
        except Exception as error:
            print(f'Failed to start Xwayland: {error}\n')
            return -1


    def setup_xauth(self, xauthfile = None):
        """ Creates the Xauthority file if needed and set the envvar. """
        if xauthfile is None:
            _, self.xauthfile = mkstemp(prefix = 'Xauthority.', dir = self.tempdir)
        else:
            self.xauthfile = xauthfile
        os.environ['XAUTHORITY'] = self.xauthfile


    def cleanup_xauth(self):
        """ Removes the Xauthority file if it's ours. """
        if os.path.dirname(self.xauthfile) != self.tempdir:
            return
        try:
            os.unlink(self.xauthfile)
        except Exception as error:
            print(f'Error removing xauth entry: {error}\n')


    def __generate_cookie(self):
        """ Generates a secure random key to use a cookie with Xauthority. """
        return token_hex(16)


    def create_xauth_entry(self, proto = '.'):
        """ Uses xauth to add an entry to Xauthority. """
        xauth_command_line = [ 'xauth', 'source', '-' ]
        entry = f'add {os.environ["DISPLAY"]} {proto} {self.__generate_cookie()}\n'
        try:
            with subprocess.Popen(xauth_command_line, stdin = subprocess.PIPE) as proc:
                proc.communicate(entry.encode())
                if proc.returncode != 0:
                    print('Failed to add xauth entry!\n')
                return proc.returncode
        except Exception as error:
            print(f'Failed to execute xauth: {error}\n - Is xauth installed?\n')
            return -1


    def remove_xauth_entry(self):
        """ Uses xauth to remove the entry in Xauthority. """
        xauth_command_line = ['xauth', 'remove', os.environ['DISPLAY']]
        try:
            with subprocess.Popen(xauth_command_line) as proc:
                proc.wait()
        except Exception as error:
            print(f'Failed to run xauth to remove the entry in Xauthority: {error}\n')


    def setup_tempdir(self):
        """ Creates a temporary directory. """
        self.tempdir = mkdtemp()


    def cleanup_tempdir(self):
        """ Removes the temporary directory if empty. """
        if self.tempdir:
            try:
                os.rmdir(self.tempdir)
            except OSError as error:
                print(f'Error removing directory {self.tempdir}: {error}\n')


    def cleanup(self):
        """ Clean-up function. """
        handler = signal(SIGTERM, SIG_IGN)
        try:
            os.killpg(os.getpgid(0), SIGTERM)
        except Exception as error:
            print(f'Warning, failed to kill process group: {error}\n')
        signal(SIGTERM, handler)
        self.remove_xauth_entry()
        self.cleanup_xauth()
        self.cleanup_tempdir()
