#!/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

if sys.version_info[0] > 2:
    from http.server import HTTPServer
    from http.server import SimpleHTTPRequestHandler as HTTPRequestHandler
    HTTPServer  # pyflakes
    HTTPRequestHandler  # pyflakes
else:
    from BaseHTTPServer import HTTPServer
    from SimpleHTTPServer 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()
