#!/usr/bin/python3

import hashlib
import http.server as SimpleHTTPServer
import os
import random
import shutil
import socketserver
import sys
import time
import tempfile
import unittest

from mock import patch
from multiprocessing import Process

sys.path.insert(0, "data")
sys.path.insert(1, "../data")
import package_data_downloader  # nopep8


class TestHttpd(Process):
    def __init__(self, port, servdir):
        super(TestHttpd, self).__init__()
        self.port = port
        self.servdir = servdir

    def run(self):
        os.chdir(self.servdir)
        Handler = SimpleHTTPServer.SimpleHTTPRequestHandler
        httpd = socketserver.TCPServer(("localhost", self.port), Handler)
        httpd.serve_forever()


class DownloadTestCase(unittest.TestCase):

    def setUp(self):
        self.tempdir = tempfile.mkdtemp()
        self._setup_canary_file()
        self.PORT = random.randint(1025, 60000)
        self.p = TestHttpd(self.PORT, self.tempdir)
        self.p.start()

    def tearDown(self):
        self.p.terminate()
        shutil.rmtree(self.tempdir)

    def _setup_canary_file(self):
        self.canary_text = "meep"
        self.canary_file = os.path.join(self.tempdir, "canary-file.txt")
        with open(self.canary_file, "w") as f:
            f.write(self.canary_text)
        self.canary_hashsum = hashlib.sha256(
            self.canary_text.encode('utf-8')).hexdigest()

    def test_download_file(self):
        package_data_downloader.STAMPDIR = self.tempdir
        os.makedirs(os.path.join(self.tempdir, "partial"))
        uri = "http://localhost:%s/%s" % (
            self.PORT, os.path.basename(self.canary_file))
        destfile = package_data_downloader.download_file(
            uri, self.canary_hashsum)
        self.assertNotEqual(destfile, None)
        self.assertEqual(hashlib.sha256(
            open(destfile).read().encode('utf-8')).hexdigest(),
            self.canary_hashsum)


class ProcessDownloadRequestsTestCase(unittest.TestCase):

    def setUp(self):
        self.tempdir = tempfile.mkdtemp()
        self._setup_canary_file()
        self.PORT = random.randint(1025, 60000)
        self.p = TestHttpd(self.PORT, self.tempdir)
        self.p.start()
        self.tempdir = tempfile.mkdtemp()
        # stampdir
        stampdir = os.path.join(self.tempdir, "stampdir")
        os.makedirs(stampdir)
        package_data_downloader.STAMPDIR = stampdir
        os.makedirs(os.path.join(stampdir, "partial"))
        # datadir
        datadir = os.path.join(self.tempdir, "datadir")
        os.makedirs(datadir)
        package_data_downloader.DATADIR = datadir
        # override the location of the notifier files
        notifierdir = os.path.join(self.tempdir, "notifierdir")
        os.makedirs(notifierdir)
        package_data_downloader.NOTIFIER_FILE = \
            os.path.join(notifierdir, "data-downloads-failed")
        package_data_downloader.NOTIFIER_PERMANENT_FILE = \
            package_data_downloader.NOTIFIER_FILE + "-permanently"

    def tearDown(self):
        self.p.terminate()
        shutil.rmtree(self.tempdir)

    def _setup_canary_file(self):
        self.canary_text = "meep"
        self.canary_file = os.path.join(self.tempdir, "canary-file.txt")
        with open(self.canary_file, "w") as f:
            f.write(self.canary_text)
        self.canary_hashsum = hashlib.sha256(
            self.canary_text.encode('utf-8')).hexdigest()

    def _setup_hook_file(self, filename, script="/bin/true"):
        with open(os.path.join(package_data_downloader.DATADIR,
                  filename), "w") as foo:
            foo.write("Url: http://localhost:%s/%s\n" %
                      (self.PORT, os.path.basename(self.canary_file)))
            foo.write("Sha256: %s\n" % self.canary_hashsum)
            foo.write("\n")
            foo.write("Script: %s" % script)

    def test_hookfile_older_than_stampfile(self):
        # hook file older than stamp file nothing happens
        # create the hook file
        hookfile = "older-hookfile"
        self._setup_hook_file(hookfile)
        time.sleep(0.01)
        # create the stamp file
        stampfile = os.path.join(package_data_downloader.STAMPDIR,
                                 hookfile)
        with open(stampfile, "w"):
            pass
        orig_stamp_time = os.stat(stampfile).st_mtime
        package_data_downloader.process_download_requests()
        new_stamp_time = os.stat(stampfile).st_mtime
        self.assertEqual(orig_stamp_time, new_stamp_time)
        # confirm failure files not created
        self.assertFalse(os.path.exists(
            os.path.join(package_data_downloader.STAMPDIR,
                         "%s.permanent-failure" % hookfile)))
        self.assertFalse(os.path.exists(package_data_downloader.NOTIFIER_FILE))

    def test_stampfile_older_than_hookfile(self):
        # hook file newer than stampfile, download, run script, update
        # stampfile
        hookfile = "older-stampfile"
        stampfile = os.path.join(package_data_downloader.STAMPDIR,
                                 hookfile)
        with open(stampfile, "w"):
            pass
        orig_stamp_date = os.stat(stampfile).st_mtime
        time.sleep(0.01)
        # create the hook file
        self._setup_hook_file(hookfile)
        package_data_downloader.process_download_requests()
        # the download succeeded so the stamp file is touched
        new_stamp_date = os.stat(stampfile).st_mtime
        self.assertGreater(new_stamp_date, orig_stamp_date)
        # confirm failure files not created
        self.assertFalse(os.path.exists(
            os.path.join(package_data_downloader.STAMPDIR,
                         "%s.permanent-failure" % hookfile)))
        self.assertFalse(os.path.exists(package_data_downloader.NOTIFIER_FILE))

    def test_only_retry_failure(self):
        # two hook files but only has failed
        hookfile = "success"
        stampfile = os.path.join(package_data_downloader.STAMPDIR,
                                 hookfile)
        with open(stampfile, "w"):
            pass
        orig_stamp_date = os.stat(stampfile).st_mtime
        time.sleep(0.01)
        # create the hook file
        self._setup_hook_file(hookfile)
        fail_hookfile = "failure"
        # the stampfile indicates a failure
        fail_stampfile = os.path.join(package_data_downloader.STAMPDIR,
                                      "%s.failed" % fail_hookfile)
        with open(fail_stampfile, "w"):
            pass
        time.sleep(0.01)
        # create the hook file
        self._setup_hook_file(fail_hookfile)
        # create an empty notifier file
        with open(package_data_downloader.NOTIFIER_FILE, "w"):
            pass
        package_data_downloader.process_download_requests()
        new_stamp_date = os.stat(stampfile).st_mtime
        # if the downloaded succeeded it shouldn't be done again
        self.assertEqual(new_stamp_date, orig_stamp_date)

    def test_hookfile_script_fails(self):
        # script in the hook file fails, failure recorded as permanent
        hookfile = "script-failure"
        stampfile = os.path.join(package_data_downloader.STAMPDIR,
                                 hookfile)
        with open(stampfile, "w"):
            pass
        time.sleep(0.01)
        # create the hook file
        self._setup_hook_file(hookfile, "/bin/false")
        package_data_downloader.process_download_requests()
        # script failures are considered permanent
        self.assertTrue(os.path.exists(
            os.path.join(package_data_downloader.STAMPDIR,
                         "%s.permanent-failure" % hookfile)))

    def test_stampfile_and_notifierfile(self):
        # notifier file removed on successful download
        hookfile = "stampfile_and_notifierfile"
        stampfile = os.path.join(package_data_downloader.STAMPDIR,
                                 hookfile)
        with open(stampfile, "w"):
            pass
        time.sleep(0.01)
        self._setup_hook_file(hookfile)
        # create an empty notifier file
        with open(package_data_downloader.NOTIFIER_FILE, "w"):
            pass
        package_data_downloader.process_download_requests()
        self.assertFalse(os.path.exists(package_data_downloader.NOTIFIER_FILE))

    def test_stampfile_notifierfile_and_notifierpermanent(self):
        # both notifier files removed on successful download
        hookfile = "stampfile_notifierfile_and_notifierpermanentfile"
        stampfile = os.path.join(package_data_downloader.STAMPDIR,
                                 hookfile)
        with open(stampfile, "w"):
            pass
        time.sleep(0.01)
        self._setup_hook_file(hookfile)
        # create empty notifier files
        with open(package_data_downloader.NOTIFIER_FILE, "w"):
            pass
        with open(package_data_downloader.NOTIFIER_PERMANENT_FILE, "w"):
            pass
        package_data_downloader.process_download_requests()
        self.assertFalse(os.path.exists(package_data_downloader.NOTIFIER_FILE))
        self.assertFalse(os.path.exists(
            package_data_downloader.NOTIFIER_PERMANENT_FILE))

    def test_stampfile_and_notifierpermanent(self):
        # permanent notifier file removed on successful download
        hookfile = "stampfile_and_permanentnotifierfile"
        stampfile = os.path.join(package_data_downloader.STAMPDIR,
                                 hookfile)
        with open(stampfile, "w"):
            pass
        time.sleep(0.01)
        self._setup_hook_file(hookfile)
        # create an empty permanent notifier file
        with open(package_data_downloader.NOTIFIER_PERMANENT_FILE, "w"):
            pass
        package_data_downloader.process_download_requests()
        self.assertFalse(os.path.exists(
            package_data_downloader.NOTIFIER_PERMANENT_FILE))

    def test_hookfile_download_failure(self):
        # can't download the file, non-permanent failure created
        hookfile = "download-failure"
        stampfile = os.path.join(package_data_downloader.STAMPDIR,
                                 hookfile)
        with open(stampfile, "w"):
            pass
        time.sleep(0.01)
        # overwrite canary file to create a failure
        self.canary_file = "not-there.txt"
        # create the hook file
        self._setup_hook_file(hookfile)
        package_data_downloader.process_download_requests()
        self.assertTrue(os.path.exists(
            os.path.join(package_data_downloader.STAMPDIR,
                         "%s.failed" % hookfile)))
        self.assertFalse(os.path.exists(
            os.path.join(package_data_downloader.STAMPDIR,
                         "%s.permanent-failure" % hookfile)))

    def test_hookfile_notifierfile_download_failure(self):
        # can't download the file, notifier file stays around
        hookfile = "download-failure-with-notifierfile"
        # overwrite canary file to create a failure
        self.canary_file = "not-there.txt"
        # create the hook file
        self._setup_hook_file(hookfile)
        # create an empty notifier file
        with open(package_data_downloader.NOTIFIER_FILE, "w"):
            pass
        package_data_downloader.process_download_requests()
        # because there was a failure it shouldn't be removed
        self.assertTrue(os.path.exists(
            package_data_downloader.NOTIFIER_FILE))


class PackageDataDownloaderTestCase(unittest.TestCase):

    def setUp(self):
        self.tmpdir = tempfile.mkdtemp()
        # stampdir
        stampdir = os.path.join(self.tmpdir, "stampdir")
        os.makedirs(stampdir)
        package_data_downloader.STAMPDIR = stampdir
        # datadir
        datadir = os.path.join(self.tmpdir, "datadir")
        os.makedirs(datadir)
        package_data_downloader.DATADIR = datadir

    def tearDown(self):
        shutil.rmtree(self.tmpdir)

    def test_wrong_template_translations(self):
        package_data_downloader.NOTIFIER_SOURCE_FILE = \
            'data/package-data-downloads-failed.in'
        package_data_downloader.NOTIFIER_FILE = \
            self.tmpdir + "/data-downloads-failed"
        package_data_downloader.NOTIFIER_PERMANENT_SOURCE_FILE = \
            'data/package-data-downloads-failed-permanently.in'
        package_data_downloader.NOTIFIER_PERMANENT_FILE = \
            package_data_downloader.NOTIFIER_FILE + '-permanently'
        package_data_downloader.trigger_update_notifier([], False)
        package_data_downloader.trigger_update_notifier([], True)

    def test_permanently_failed(self):
        # create a bunch of files using the provided mechanism
        test_files = ["foo.permanent-failure", "bar.failure", "baz"]
        for f in test_files:
            package_data_downloader.create_or_update_stampfile(
                os.path.join(package_data_downloader.STAMPDIR, f))
        self.assertEqual(
            sorted(os.listdir(package_data_downloader.STAMPDIR)),
            sorted(test_files))
        # test hook_is_permanently_failed()
        self.assertTrue(
            package_data_downloader.hook_is_permanently_failed("foo"))
        self.assertFalse(
            package_data_downloader.hook_is_permanently_failed("bar"))
        self.assertFalse(
            package_data_downloader.hook_is_permanently_failed("baz"))
        # existing_permanent_failures()
        self.assertEqual(
            package_data_downloader.existing_permanent_failures(),
            ["foo"])

    def test_hook_aged_out(self):
        # test to see if a hook has failed for more than 3 days
        # create a failure file
        with open(os.path.join(package_data_downloader.STAMPDIR,
                               "aged.failed"), "w"):
            pass
        from datetime import datetime, timedelta
        # patch datetime so we think its the future
        with patch('package_data_downloader.datetime') as mock_datetime:
            thefuture = (datetime.now() + timedelta(days=3))
            mock_datetime.now.return_value = thefuture
            mock_datetime.fromtimestamp.side_effect = \
                lambda *args, **kw: datetime.fromtimestamp(*args, **kw)
            self.assertEqual(package_data_downloader.datetime.now(),
                             thefuture)
            self.assertTrue(package_data_downloader.hook_aged_out("aged"))

    def test_mark_hook_failed(self):
        # prepare
        package_data_downloader.create_or_update_stampfile(
            os.path.join(package_data_downloader.STAMPDIR, "foo"))
        # temp failure
        package_data_downloader.mark_hook_failed("foo")
        self.assertEqual(os.listdir(package_data_downloader.STAMPDIR),
                         ["foo.failed"])
        self.assertFalse(
            package_data_downloader.hook_is_permanently_failed("foo"))
        # permanent
        package_data_downloader.mark_hook_failed("foo", permanent=True)
        self.assertEqual(os.listdir(package_data_downloader.STAMPDIR),
                         ["foo.permanent-failure"])
        self.assertTrue(
            package_data_downloader.hook_is_permanently_failed("foo"))

    def test_get_hook_file_names(self):
        # test that .dpkg-* is ignored
        test_datadir_files = ["foo.dpkg-new", "bar"]
        for name in test_datadir_files:
            with open(os.path.join(package_data_downloader.DATADIR, name),
                      "w"):
                pass
        self.assertEqual(package_data_downloader.get_hook_file_names(),
                         ["bar"])

    def test_trigger_update_notifier(self):
        # write to tmpdir
        package_data_downloader.NOTIFIER_FILE = os.path.join(
            self.tmpdir, "data-downloads-failed")
        # point to local repo file
        package_data_downloader.NOTIFIER_SOURCE_FILE = \
            "data/package-data-downloads-failed.in"
        package_data_downloader.trigger_update_notifier(
            failures=["foo", "bar"], permanent=False)
        data = open(package_data_downloader.NOTIFIER_FILE).read()
        self.assertEqual(data, """Priority: High
_Name: Failure to download extra data files
Terminal: True
Command: pkexec /usr/lib/update-notifier/package-data-downloader
_Description:
 The following packages requested additional data downloads after package
 installation, but the data could not be downloaded or could not be
 processed.
 .
 .
   foo, bar
 .
 .
 The download will be attempted again later, or you can try the download
 again now.  Running this command requires an active Internet connection.
""")


if __name__ == "__main__":
    unittest.main()