#!/usr/bin/python3
# autopkgtest check: Boot with systemd and check different dm configurations
# (C) 2014 Canonical Ltd.
# Author: Didier Roche <didrocks@ubuntu.com>

from contextlib import suppress
import fileinput
import os
import subprocess
import shutil
import sys
import unittest
from time import sleep

DDM_CONFIG_PATH = "/etc/X11/default-display-manager"
SYSTEMD_DEFAULT_DM_PATH = "/etc/systemd/system/display-manager.service"
SYSTEMD_SYSTEM_UNIT_DIR = "/lib/systemd/system"
LIGHTDM_SYSTEMD_UNIT_PATH = os.path.join(SYSTEMD_SYSTEM_UNIT_DIR, 'lightdm.service')


class DisplayManagersTest(unittest.TestCase):
    '''Check that multiple dm configurations are handled'''

    def setUp(self):
        super().setUp()
        with suppress(FileNotFoundError):
            os.remove(SYSTEMD_DEFAULT_DM_PATH)
        with suppress(FileNotFoundError):
            os.remove(DDM_CONFIG_PATH)
        subprocess.check_call('apt-get install -y --reinstall lightdm 2>&1', shell=True)
        # Remove all Conditional ExecStartPre= as we want to check systemd logic, not unit
        for line in fileinput.input([LIGHTDM_SYSTEMD_UNIT_PATH], inplace=True):
            if not line.startswith('ExecStartPre='):
                print(line)
        self.files_to_clean = []

    def tearDown(self):
        for f in self.files_to_clean:
            os.remove(f)
        super().tearDown()

    def test_one_systemd(self):
        '''one systemd dm is started'''

        self.reload_state()

        self.assertTrue(self.is_active_unit('lightdm'))

    def test_multiple_systemd(self):
        '''only default systemd dm is started'''

        self.create_systemd_dm_unit("systemddm")
        self.reload_state()

        self.assertTrue(self.is_active_unit('lightdm'))
        self.assertFalse(self.is_active_unit('systemddm'))

    def test_multiple_systemd_ddmconfig_match(self):
        '''display-manager symlink respect ddm config and starts right unit'''

        # lightdm was the default
        self.create_systemd_dm_unit("systemddm", make_ddm_default="systemddm")
        self.reload_state()

        self.assertFalse(self.is_active_unit('lightdm'))
        self.assertTrue(self.is_active_unit('systemddm'))

    def test_multiple_systemd_ddmconfig_match_no_symlink(self):
        '''create a display-manager symlink to matching systemd unit ddm config'''

        # lightdm was the default
        self.create_systemd_dm_unit("systemddm", make_ddm_default="systemddm")
        os.remove(SYSTEMD_DEFAULT_DM_PATH)
        self.reload_state()

        self.assertFalse(self.is_active_unit('lightdm'))
        self.assertTrue(self.is_active_unit('systemddm'))

    def test_one_systemd_no_ddmconfig(self):
        '''without any ddm config, the default systemd unit via symlink is still the default'''

        os.remove(DDM_CONFIG_PATH)
        self.reload_state()

        self.assertTrue(self.is_active_unit('lightdm'))

    def test_one_systemd_masked_symlink_with_ddmconfig(self):
        '''a masked symlink will be updated to match systemd ddmconfig unit'''

        os.remove(SYSTEMD_DEFAULT_DM_PATH)
        os.symlink("/dev/null", SYSTEMD_DEFAULT_DM_PATH)
        self.reload_state()

        self.assertTrue(self.is_active_unit('lightdm'))

    def test_one_systemd_masked_symlink_no_ddmconfig(self):
        '''without any ddm config, a masked symlinked will stayed masked'''

        os.remove(DDM_CONFIG_PATH)
        os.remove(SYSTEMD_DEFAULT_DM_PATH)
        os.symlink("/dev/null", SYSTEMD_DEFAULT_DM_PATH)
        self.reload_state()

        self.assertFalse(self.is_active_unit('lightdm'))

    def test_multiple_systemd_wrong_ddmconfig(self):
        '''ddm config matches no systemd unit, don't start any of them'''

        self.create_systemd_dm_unit("systemddm", make_ddm_default="systemddm_doesnt_match")
        self.reload_state()

        self.assertFalse(self.is_active_unit('lightdm'))
        self.assertFalse(self.is_active_unit('systemddm'))

    def test_one_init(self):
        '''one init dm is started'''

        # fake removing lightdm (or we shoud remove all processes under lightdm)
        os.remove(DDM_CONFIG_PATH)
        os.remove(SYSTEMD_DEFAULT_DM_PATH)
        os.remove(LIGHTDM_SYSTEMD_UNIT_PATH)
        self.create_init_dm("initdm", make_ddm_default="initdm")
        self.reload_state()

        self.assertTrue(self.is_active_unit('initdm'))

    def test_multiple_init(self):
        '''all init dms are enabled, regardless of ddm'''

        # this enable to keep previous init behavior,
        # especially when they don't support ddm config like nodm
        os.remove(DDM_CONFIG_PATH)
        os.remove(SYSTEMD_DEFAULT_DM_PATH)
        os.remove(LIGHTDM_SYSTEMD_UNIT_PATH)
        self.create_init_dm("initdm", make_ddm_default="initdm")
        self.create_init_dm("otherinitdm")
        self.reload_state()

        self.assertTrue(self.is_active_unit('initdm'))
        self.assertTrue(self.is_active_unit('otherinitdm'))

    def test_multiple_init_no_ddm(self):
        '''all init dms are enabled, without any ddm file'''

        os.remove(DDM_CONFIG_PATH)
        os.remove(SYSTEMD_DEFAULT_DM_PATH)
        os.remove(LIGHTDM_SYSTEMD_UNIT_PATH)
        self.create_init_dm("initdm")
        self.create_init_dm("otherinitdm")
        self.reload_state()

        self.assertTrue(self.is_active_unit('initdm'))
        self.assertTrue(self.is_active_unit('otherinitdm'))

    def test_systemd_matches_ddm_and_init(self):
        '''default ddm config systemd is enabled as well as all inits'''

        # lightdm is the default
        self.create_init_dm("initdm")
        self.reload_state()

        self.assertTrue(self.is_active_unit('lightdm'))
        self.assertTrue(self.is_active_unit('initdm'))

    def test_systemd_and_init_matches_ddm(self):
        '''default ddm init prevents systemd units to start'''

        self.create_init_dm("initdm", make_ddm_default="initdm")
        self.reload_state()

        self.assertFalse(self.is_active_unit('lightdm'))
        self.assertTrue(self.is_active_unit('initdm'))

    def test_no_ddmconfig_multiple_systemd_and_init(self):
        '''no ddm config let default systemd and all init dms enabled'''

        os.remove(DDM_CONFIG_PATH)
        self.create_systemd_dm_unit("systemddm")
        self.create_init_dm("initdm")
        self.reload_state()

        self.assertTrue(self.is_active_unit('lightdm'))
        self.assertTrue(self.is_active_unit('initdm'))
        self.assertFalse(self.is_active_unit('systemddm'))

    def test_no_ddmconfig_no_default_systemd_and_init(self):
        '''no ddm config default systemd unit will only have init dms enabled'''

        os.remove(DDM_CONFIG_PATH)
        os.remove(SYSTEMD_DEFAULT_DM_PATH)
        self.create_systemd_dm_unit("systemddm")
        self.create_init_dm("initdm")
        self.reload_state()

        self.assertFalse(self.is_active_unit('lightdm'))
        self.assertTrue(self.is_active_unit('initdm'))
        self.assertFalse(self.is_active_unit('systemddm'))

    # NOTE: I think init shouldn't start in that case
    def test_one_systemd_one_init_masked_symlink_with_ddmconfig(self):
        '''a masked symlink will be updated to match systemd ddmconfig systemd unit and init is started'''

        os.remove(SYSTEMD_DEFAULT_DM_PATH)
        os.symlink("/dev/null", SYSTEMD_DEFAULT_DM_PATH)
        self.create_init_dm("initdm")
        self.reload_state()

        self.assertTrue(self.is_active_unit('lightdm'))
        self.assertTrue(self.is_active_unit('initdm'))

    # NOTE: I think init shouldn't start in that case
    def test_one_systemd_one_init_masked_symlink_no_ddmconfig(self):
        '''without any ddm config, a masked symlinked will stayed masked, but init will be started'''

        os.remove(DDM_CONFIG_PATH)
        os.remove(SYSTEMD_DEFAULT_DM_PATH)
        os.symlink("/dev/null", SYSTEMD_DEFAULT_DM_PATH)
        self.create_init_dm("initdm")
        self.reload_state()

        self.assertFalse(self.is_active_unit('lightdm'))
        self.assertTrue(self.is_active_unit('initdm'))

    # Helper methods

    def create_systemd_dm_unit(self, name, make_ddm_default=None):
        dest_unit = "{}.service".format(os.path.join(SYSTEMD_SYSTEM_UNIT_DIR, name))
        shutil.copy(LIGHTDM_SYSTEMD_UNIT_PATH, dest_unit)
        # remove BusName to avoid conflicts
        for line in fileinput.input([LIGHTDM_SYSTEMD_UNIT_PATH], inplace=True):
            if not line.startswith('BusName='):
                print(line)
        self.files_to_clean.append(dest_unit)

        if make_ddm_default:
            self.make_ddm_default(make_ddm_default)

    def create_init_dm(self, name, make_ddm_default=None):
        init_script = "/etc/init.d/{}".format(name)
        with open(init_script, 'w') as f:
            f.write('''#!/bin/sh
### BEGIN INIT INFO
# Provides:          {service}
# Required-Start:    $local_fs $remote_fs dbus
# Required-Stop:     $local_fs $remote_fs dbus
# Should-Start:      $named
# Should-Stop:       $named
# Default-Start:     2 3 4 5
# Default-Stop:      0 1 6
# Short-Description: Start {service}
### END INIT INFO
exit 0'''.format(service=name))
            os.fchmod(f.fileno(), 0o755)
        self.files_to_clean.append(init_script)

        rc2_link = "/etc/rc2.d/S05{}".format(name)
        os.symlink("../init.d/{}".format(name), rc2_link)
        self.files_to_clean.append(rc2_link)

        insserv_script = "/etc/insserv.conf.d/{}".format(name)
        with open(insserv_script, 'w') as f:
            f.write("$x-display-manager {}".format(name))
        self.files_to_clean.append(insserv_script)
        if make_ddm_default:
            self.make_ddm_default(make_ddm_default)

    def make_ddm_default(self, binary_name):
        with open(DDM_CONFIG_PATH, 'w') as f:
            f.write("/usr/bin/{}".format(binary_name))

    def is_active_unit(self, unit):
        '''Check that given unit is active'''

        if subprocess.call(['systemctl', '-q', 'is-active', unit]) != 0:
            return False
        return True

    def reload_state(self):
        subprocess.check_call(['systemctl', 'daemon-reload'])
        subprocess.check_call(['systemctl', 'default'])
        sleep(2)  # a more robust way would be to loop over remaining jobs to process


def boot_with_systemd():
    '''Reboot with systemd as init

    In case something else is currently running in the testbed
    '''
    if subprocess.call(['systemctl', 'status'], stdout=subprocess.PIPE,
                       stderr=subprocess.PIPE) != 0:
        print('Installing systemd-sysv and rebooting...')
        subprocess.check_call('apt-get -y install systemd-sysv 2>&1',
                              shell=True)
        subprocess.check_call(['autopkgtest-reboot', 'boot-systemd'])


if __name__ == '__main__':
    if not os.getenv('ADT_REBOOT_MARK'):
        boot_with_systemd()

    unittest.main(testRunner=unittest.TextTestRunner(stream=sys.stdout,
                                                     verbosity=2))
