#!/usr/bin/python3

import dbusmock
from gi.repository import Gio, GLib

import os
import subprocess
import threading
import unittest

import apt

from test.test_base import TestBase, MockOptions
import unattended_upgrade

apt.apt_pkg.config.set("APT::Architecture", "amd64")

APT_CONF = """
Unattended-Upgrade::Keep-Debs-After-Install "true";
Unattended-Upgrade::Allowed-Origins {{
    "Ubuntu:lucid-security";
}};
Unattended-Upgrade::Skip-Updates-On-Metered-Connections "false";
Unattended-Upgrade::Postpone-For-Days "{days}";
Unattended-Upgrade::Postpone-Wait-Time "2";
"""


class PostponeClient(threading.Thread):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.on_about_to_start = None
        self.on_started = None
        self.on_finished = None
        self.on_canceled = None
        self.received = []
        self.error = None

        self.proxy = None
        self.loop = None

    def call_postpone(self):
        def on_res(proxy, res):
            try:
                proxy.call_finish(res)
            except GLib.GError as e:
                self.error = e

        if self.proxy:
            self.proxy.call(
                "Postpone", None,
                Gio.DBusCallFlags.ALLOW_INTERACTIVE_AUTHORIZATION,
                -1, None, on_res)

    def quit(self):
        if self.loop:
            self.loop.quit()

    def start(self):
        self.started_sem = threading.Semaphore(0)
        threading.Thread.start(self)
        self.started_sem.acquire(timeout=10)

    def run(self):
        self.thread_context = GLib.MainContext()
        self.thread_context.push_thread_default()
        self.loop = GLib.MainLoop(self.thread_context)
        self.proxy = Gio.DBusProxy.new_for_bus_sync(
            Gio.BusType.SYSTEM,
            Gio.DBusProxyFlags.NONE,
            None,
            "com.ubuntu.UnattendedUpgrade",
            "/com/ubuntu/UnattendedUpgrade/Pending",
            "com.ubuntu.UnattendedUpgrade.Pending",
            None)

        def on_signal(proxy, sender, signal, params):
            self.received.append(signal)
            if signal == "AboutToStart":
                if self.on_about_to_start:
                    self.on_about_to_start()
            elif signal == "Started":
                if self.on_started:
                    self.on_started()
            elif signal == "Canceled":
                if self.on_canceled:
                    self.on_canceled()
            elif signal == "Finished":
                if self.on_finished:
                    self.on_finished()

        handler_id = self.proxy.connect('g-signal', on_signal)

        source = GLib.idle_source_new()
        source.set_callback(self.started_sem.release)
        source.attach(self.thread_context)

        self.loop.run()

        self.proxy.disconnect(handler_id)
        del self.proxy


class TestPostpone(TestBase, dbusmock.DBusTestCase):
    @classmethod
    def setUpClass(cls):
        TestBase.setUpClass()
        cls.start_system_bus()
        cls.dbus_con = cls.get_dbus(system_bus=True)

    def setUp(self):
        TestBase.setUp(self)
        self.rootdir = self.make_fake_aptroot(
            template=os.path.join(self.testdir, "root.postpone"),
            fake_pkgs=[
                ("test-package", "1.0.test.pkg", {}),
            ]
        )
        self.apt_conf = os.path.join(self.rootdir, "etc", "apt", "apt.conf")
        self.stamp_file = os.path.join(self.rootdir, ".postpone-stamp")
        unattended_upgrade.POSTPONE_STAMP_FILE = self.stamp_file
        unattended_upgrade.DBUS_DIR = os.path.join(
            self.rootdir, "usr/share/dbus-1")

        # FIXME: make this more elegant
        # fake on_ac_power
        os.environ["PATH"] = (os.path.join(self.rootdir, "usr", "bin") + ":"
                              + os.environ["PATH"])

        # clean log
        self.log = os.path.join(
            self.rootdir, "var", "log", "unattended-upgrades",
            "unattended-upgrades.log")
        if not os.path.exists(os.path.dirname(self.log)):
            os.makedirs(os.path.dirname(self.log))
        with open(self.log, "w"):
            pass
        # clean cache
        subprocess.check_call(
            ["rm", "-f",
             os.path.join(self.rootdir, "var", "cache", "apt", "*.bin")])

        # clear stamp file
        os.makedirs(os.path.dirname(self.stamp_file), exist_ok=True)
        unattended_upgrade.reset_stamp_file(
            unattended_upgrade.POSTPONE_STAMP_FILE)

    def tearDown(self):
        try:
            os.remove(os.path.join(self.rootdir, "etc", "apt", "apt.conf"))
        except OSError:
            pass

    def test_postpone_days(self):
        apt.apt_pkg.config.set("Unattended-Upgrade::Postpone-For-Days", "0")
        self.assertEqual(unattended_upgrade.get_remaining_postpone_days(), 0)

        apt.apt_pkg.config.set("Unattended-Upgrade::Postpone-For-Days", "1")
        self.assertEqual(unattended_upgrade.get_remaining_postpone_days(), 1)

        subprocess.check_call(["touch", "-d", "3 days ago", self.stamp_file])
        apt.apt_pkg.config.set("Unattended-Upgrade::Postpone-For-Days", "1")
        self.assertEqual(unattended_upgrade.get_remaining_postpone_days(), 0)

        subprocess.check_call(["touch", "-d", "2 days ago", self.stamp_file])
        apt.apt_pkg.config.set("Unattended-Upgrade::Postpone-For-Days", "5")
        self.assertEqual(unattended_upgrade.get_remaining_postpone_days(), 3)

    def test_postpone_reset(self):
        subprocess.check_call(["touch", "-d", "3 days ago", self.stamp_file])
        apt.apt_pkg.config.set("Unattended-Upgrade::Postpone-For-Days", "2")
        self.assertEqual(unattended_upgrade.get_remaining_postpone_days(), 0)

        with open(self.apt_conf, "w") as fp:
            fp.write(APT_CONF.format(days=2))
        unattended_upgrade.main(
            MockOptions(), rootdir=self.rootdir)

        self.assertEqual(unattended_upgrade.get_remaining_postpone_days(), 2)

    def test_no_wait_for_postpone(self):
        with open(self.apt_conf, "w") as fp:
            fp.write(APT_CONF.format(days=0))
        unattended_upgrade.main(
            MockOptions(), rootdir=self.rootdir)

        pkgs = unattended_upgrade.UnattendedUpgradesCache(rootdir=self.rootdir)
        self.assertFalse(pkgs["test-package"].is_upgradable)

        with open(self.log) as f:
            needle = "Waiting for a postpone request..."
            haystack = f.read()
            self.assertFalse(
                needle in haystack,
                "Should not have found '%s' in '%s'" % (needle, haystack))

    def test_wait_for_postpone(self):
        with open(self.apt_conf, "w") as fp:
            fp.write(APT_CONF.format(days=1))
        unattended_upgrade.main(
            MockOptions(), rootdir=self.rootdir)

        pkgs = unattended_upgrade.UnattendedUpgradesCache(rootdir=self.rootdir)
        self.assertFalse(pkgs["test-package"].is_upgradable)

        with open(self.log) as f:
            needle = "Waiting for a postpone request..."
            haystack = f.read()
            self.assertTrue(
                needle in haystack,
                "Can not find '%s' in '%s'" % (needle, haystack))

    def test_all_signals(self):
        client = PostponeClient()
        client.start()

        with open(self.apt_conf, "w") as fp:
            fp.write(APT_CONF.format(days=1))
        unattended_upgrade.main(
            MockOptions(), rootdir=self.rootdir)

        client.quit()
        client.join()
        self.assertEqual(
            client.received, ['AboutToStart', 'Started', 'Finished'])

    def test_signals_no_postpone(self):
        client = PostponeClient()
        client.start()

        with open(self.apt_conf, "w") as fp:
            fp.write(APT_CONF.format(days=0))
        unattended_upgrade.main(
            MockOptions(), rootdir=self.rootdir)

        client.quit()
        client.join()
        self.assertEqual(client.received, ['Started', 'Finished'])

    def test_do_postpone(self):
        # simulate polkit with dbusmock
        (p_mock, obj_polkit) = self.spawn_server_template(
            'polkitd', None, stdout=subprocess.PIPE)
        obj_polkit.SetAllowed(
            ["com.ubuntu.UnattendedUpgrade.Pending.Postpone"])

        client = PostponeClient()
        client.on_about_to_start = client.call_postpone
        client.start()

        with open(self.apt_conf, "w") as fp:
            fp.write(APT_CONF.format(days=3))
        unattended_upgrade.main(MockOptions(), rootdir=self.rootdir)

        client.quit()
        client.join()

        p_mock.stdout.close()
        p_mock.terminate()
        p_mock.wait()

        self.assertIsNone(client.error)
        self.assertEqual(client.received, ["AboutToStart", "Canceled"])
        self.assertTrue(os.path.isfile(unattended_upgrade.POSTPONE_STAMP_FILE))
        self.assertEqual(unattended_upgrade.get_remaining_postpone_days(), 3)

        pkgs = unattended_upgrade.UnattendedUpgradesCache(rootdir=self.rootdir)
        self.assertTrue(pkgs["test-package"].is_upgradable)

    def test_postpone_stamp_persists(self):
        apt.apt_pkg.config.set("Unattended-Upgrade::Postpone-For-Days", "3")
        subprocess.check_call(["touch", "-d", "1 days ago", self.stamp_file])
        self.assertEqual(unattended_upgrade.get_remaining_postpone_days(), 2)

        # simulate polkit with dbusmock
        (p_mock, obj_polkit) = self.spawn_server_template(
            'polkitd', None, stdout=subprocess.PIPE)
        obj_polkit.SetAllowed(
            ["com.ubuntu.UnattendedUpgrade.Pending.Postpone"])

        client = PostponeClient()
        client.on_about_to_start = client.call_postpone
        client.start()

        with open(self.apt_conf, "w") as fp:
            fp.write(APT_CONF.format(days=3))
        unattended_upgrade.main(MockOptions(), rootdir=self.rootdir)

        client.quit()
        client.join()

        p_mock.stdout.close()
        p_mock.terminate()
        p_mock.wait()

        self.assertIsNone(client.error)
        self.assertEqual(client.received, ["AboutToStart", "Canceled"])
        self.assertEqual(unattended_upgrade.get_remaining_postpone_days(), 2)

    def test_postpone_too_late(self):
        # simulate polkit with dbusmock
        (p_mock, obj_polkit) = self.spawn_server_template(
            'polkitd', None, stdout=subprocess.PIPE)
        obj_polkit.SetAllowed(
            ["com.ubuntu.UnattendedUpgrade.Pending.Postpone"])

        client = PostponeClient()
        client.on_started = client.call_postpone
        client.start()

        with open(self.apt_conf, "w") as fp:
            fp.write(APT_CONF.format(days=3))
        unattended_upgrade.main(MockOptions(), rootdir=self.rootdir)

        client.quit()
        client.join()

        p_mock.stdout.close()
        p_mock.terminate()
        p_mock.wait()

        self.assertTrue(isinstance(client.error, GLib.GError))
        self.assertTrue(
            client.error.matches(Gio.DBusError.quark(), Gio.DBusError.FAILED))
        self.assertEqual(
            client.received, ["AboutToStart", "Started", "Finished"])
        self.assertFalse(
            os.path.isfile(unattended_upgrade.POSTPONE_STAMP_FILE))

        pkgs = unattended_upgrade.UnattendedUpgradesCache(rootdir=self.rootdir)
        self.assertFalse(pkgs["test-package"].is_upgradable)


if __name__ == "__main__":
    # do not setup logging in here or the test will break
    unittest.main()
