#!/usr/bin/python3

from __future__ import print_function

import contextlib
import errno
import itertools
import os
import shutil
import sys
import tempfile
import time
import unittest

from http.server import HTTPServer
from http.server import SimpleHTTPRequestHandler as HTTPRequestHandler

import apt_pkg
import apt.auth

import testcommon

WHEEZY_KEYID = "8B48AD6246925553"
WHEEZY_KEYDATE = "1335553717"
WHEEZY_KEY = """-----BEGIN PGP PUBLIC KEY BLOCK-----
Version: GnuPG v1.4.12 (GNU/Linux)


DHcut3Yey8o=
=id4q
-----END PGP PUBLIC KEY BLOCK-----"""


def normalize_key(keystr):
    """Remove the Version: header from a key block"""
    lines = keystr.split("\n")
    if lines[1].startswith("Version:"):
        return lines[:1] + lines[2:]
    return lines


class TestAuthKeys(testcommon.TestCase):

    """Test handling of keys for signed repositories."""

    def setUp(self):
        testcommon.TestCase.setUp(self)

        self.tmpdir = tempfile.mkdtemp()
        self.addCleanup(shutil.rmtree, self.tmpdir)
        apt_pkg.config.set("Dir", self.tmpdir)
        apt_pkg.config.set("Dir::Bin::Apt-key", "fakeroot-apt-key")
        apt_pkg.config.set("Dir::Etc", "etc/apt/")
        trustedparts_dir = apt_pkg.config.find_dir("Dir::Etc::Trustedparts")
        confparts_dir = apt_pkg.config.find_dir("Dir::Etc::parts")
        self.assertTrue(trustedparts_dir.startswith(self.tmpdir))
        os.makedirs(trustedparts_dir)
        os.makedirs(confparts_dir)
        shutil.copy("fakeroot-apt-key", self.tmpdir)

    @contextlib.contextmanager
    def _discard_stderr(self):
        stderr_fd = sys.stderr.fileno()
        stderr_save = os.dup(stderr_fd)
        try:
            devnull = os.open("/dev/null", os.O_WRONLY)
            try:
                os.dup2(devnull, stderr_fd)
                yield
            finally:
                os.close(devnull)
        finally:
            os.dup2(stderr_save, stderr_fd)
            os.close(stderr_save)

    def testAddAndExportKey(self):
        """Add an example key."""
        apt.auth.add_key(WHEEZY_KEY)
        # Strip the headers from the keys to avoid test errors because
        # the exported key used a differenct GnuPG version than the
        # original example key
        self.assertEqual(
            normalize_key(apt.auth.export_key(WHEEZY_KEYID)), normalize_key(WHEEZY_KEY)
        )

    def testAddAndListKey(self):
        """Add an example key and test if it is correctly returned by
        list_keys()
        """
        apt.auth.add_key(WHEEZY_KEY)
        ret = apt.auth.list_keys()
        self.assertEqual(len(ret), 1)
        key = ret[0]
        self.assertEqual(
            key.name,
            "Debian Archive Automatic Signing Key (7.0/wheezy) "
            "<ftpmaster@debian.org>",
        )
        self.assertEqual(key.keyid, WHEEZY_KEYID)
        self.assertEqual(key.date, WHEEZY_KEYDATE)

    def testAddKeyFromFile(self):
        """Test adding a key from file."""
        keyfd, keyname = tempfile.mkstemp()
        self.addCleanup(os.close, keyfd)
        os.write(keyfd, WHEEZY_KEY.encode("UTF-8"))

        apt.auth.add_key_from_file(keyname)

        ret = apt.auth.list_keys()
        self.assertEqual(len(ret), 1)
        key = ret[0]
        self.assertEqual(
            key.name,
            "Debian Archive Automatic Signing Key (7.0/wheezy) "
            "<ftpmaster@debian.org>",
        )
        self.assertEqual(key.keyid, WHEEZY_KEYID)
        self.assertEqual(key.date, WHEEZY_KEYDATE)

    def test_add_key_from_keyserver_too_short(self):
        """Ensure that short keyids are not imported"""
        with self.assertRaises(apt.auth.AptKeyIDTooShortError):
            apt.auth.add_key_from_keyserver(WHEEZY_KEYID, "hkp://localhost:19191")
        with self.assertRaises(apt.auth.AptKeyIDTooShortError):
            apt.auth.add_key_from_keyserver(
                "0101010178F7FE5C3E65D8AF8B48AD624692555", "hkp://localhost:19191"
            )
        with self.assertRaises(apt.auth.AptKeyIDTooShortError):
            apt.auth.add_key_from_keyserver(
                "0x0101010178F7FE5C3E65D8AF8B48AD624692555", "hkp://localhost:19191"
            )
        with self.assertRaises(apt.auth.AptKeyIDTooShortError):
            apt.auth.add_key_from_keyserver(
                "0101 0101 78F7 FE5C 3E65 D8AF 8B48 AD62 4692 555",
                "hkp://localhost:19191",
            )

    def test_add_key_from_server_mitm(self):
        """Verify that the key fingerprint is verified after download"""
        self._start_keyserver()
        self.addCleanup(self._stop_keyserver)
        with self.assertRaises(apt.auth.AptKeyError) as cm:
            with self._discard_stderr():
                apt.auth.add_key_from_keyserver(
                    "0101010178F7FE5C3E65D8AF8B48AD6246925553",
                    "hkp://localhost:%d" % self.keyserver_port,
                )
        self.assertTrue(
            str(cm.exception).startswith(
                "recv from 'hkp://localhost:%d' failed for '%s'"
                % (self.keyserver_port, "0101010178F7FE5C3E65D8AF8B48AD6246925553")
            ),
            cm.exception,
        )

    def testAddKeyFromServer(self):
        """Install a GnuPG key from a remote server."""
        self._start_keyserver()
        self.addCleanup(self._stop_keyserver)

        with self._discard_stderr():
            apt.auth.add_key_from_keyserver(
                "0xa1bD8E9D78F7FE5C3E65D8AF8B48AD6246925553",
                "hkp://localhost:%d" % self.keyserver_port,
            )

        ret = apt.auth.list_keys()
        self.assertEqual(len(ret), 1)
        key = ret[0]
        self.assertEqual(
            key.name,
            "Debian Archive Automatic Signing Key (7.0/wheezy) "
            "<ftpmaster@debian.org>",
        )
        self.assertEqual(key.keyid, WHEEZY_KEYID)
        self.assertEqual(key.date, WHEEZY_KEYDATE)

    def _start_keyserver(self):
        """Start a fake keyserver on http://localhost:19191
        If port 19191 is unavailable, try successive ports until one is.
        Store the port actually in use in self.keyserver_port.
        Thanks pitti.
        """
        dir = tempfile.mkdtemp()
        self.addCleanup(shutil.rmtree, dir)
        os.mkdir(os.path.join(dir, "pks"))
        with open(os.path.join(dir, "pks", "lookup"), "w") as key_file:
            key_file.write(WHEEZY_KEY)

        keyserver_pipe = os.pipe()
        self.keyserver_pid = os.fork()
        if self.keyserver_pid == 0:
            os.close(keyserver_pipe[0])
            # quiesce server log
            os.dup2(os.open("/dev/null", os.O_WRONLY), sys.stderr.fileno())
            os.chdir(dir)
            for port in itertools.count(19191):
                try:
                    httpd = HTTPServer(("localhost", port), HTTPRequestHandler)
                    break
                except IOError as e:
                    if e.errno != errno.EADDRINUSE:
                        raise
            keyserver_write = os.fdopen(keyserver_pipe[1], "w")
            print(port, file=keyserver_write)
            keyserver_write.close()
            httpd.serve_forever()
            os._exit(0)

        os.close(keyserver_pipe[1])
        keyserver_read = os.fdopen(keyserver_pipe[0])
        self.keyserver_port = int(keyserver_read.readline())
        keyserver_read.close()

        # temporarily disable proxy, as gnupg does not get along with that
        # (LP #789049)
        self.orig_proxy = os.environ.get("http_proxy")
        try:
            del os.environ["http_proxy"]
        except KeyError:
            pass

        # wait a bit until server is ready
        time.sleep(0.5)

    def _stop_keyserver(self):
        """Stop fake keyserver"""
        assert self.keyserver_pid

        os.kill(self.keyserver_pid, 15)
        os.wait()

        # restore proxy
        if self.orig_proxy is not None:
            os.environ["http_proxy"] = self.orig_proxy


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