#!/usr/bin/python3

import apt

from functools import partial
import unittest
import sys
import os

from aptsources.distro import get_distro
from aptsources.sourceslist import (SourceEntry, Deb822SourceEntry)
from contextlib import contextmanager
from http.client import HTTPException
from launchpadlib.launchpad import Launchpad
from lazr.restfulclient.errors import Unauthorized
from urllib.request import urlopen
from urllib.error import URLError
from unittest.mock import (patch, Mock)

sys.path.insert(0, "..")

from softwareproperties.sourceslist import SourcesListShortcutHandler
from softwareproperties.uri import URIShortcutHandler
from softwareproperties.cloudarchive import CloudArchiveShortcutHandler
from softwareproperties.ppa import PPAShortcutHandler
from softwareproperties.shortcuthandler import InvalidShortcutException, ShortcutException
from softwareproperties.shortcuts import shortcut_handler


DISTRO = get_distro()
CODENAME = DISTRO.codename

# These must match the ppa used in the VALID_PPAS
PPA_LINE = f"""Types: deb
URIs: https://ppa.launchpadcontent.net/ddstreet/ppa/ubuntu/
Suites: {CODENAME}
Components: main
"""
PPA_FILEBASE = "ddstreet-ubuntu-ppa"
PPA_SOURCEFILE = f"{PPA_FILEBASE}-{CODENAME}.sources"
PPA_NETRCFILE = f"{PPA_FILEBASE}.conf"

PRIVATE_PPA_PASSWORD = "thisisnotarealpassword"
PRIVATE_PPA_LINE = f"""Types: deb
URIs: https://private-ppa.launchpadcontent.net/ddstreet/ppa/ubuntu/
Suites: {CODENAME}
Components: main
"""
PRIVATE_PPA_NETRCCONTENT = f"machine private-ppa.launchpadcontent.net/ddstreet/ppa/ubuntu/ login ddstreet password {PRIVATE_PPA_PASSWORD}"
PRIVATE_PPA_SUBSCRIPTION_URL = f"https://ddstreet:{PRIVATE_PPA_PASSWORD}@private-ppa.launchpadcontent.net/ddstreet/ppa/ubuntu/"

# These must match the uca used in VALID_UCAS
UCA_CANAME = "train"
UCA_ARCHIVE = "http://ubuntu-cloud.archive.canonical.com/ubuntu"
UCA_LINE = f"deb {UCA_ARCHIVE} {CODENAME}-updates/{UCA_CANAME} main"
UCA_LINE_PROPOSED = f"deb {UCA_ARCHIVE} {CODENAME}-proposed/{UCA_CANAME} main"
UCA_FILEBASE = f"cloudarchive-{UCA_CANAME}"
UCA_SOURCEFILE = f"{UCA_FILEBASE}.list"

# This must match the VALID_URIS
URI = "http://fake.mirror.private.com/ubuntu"
URI_FILEBASE = "archive_uri-http_fake_mirror_private_com_ubuntu"
URI_SOURCEFILE = f"{URI_FILEBASE}-{CODENAME}.list"

VALID_LINES = [f"deb {URI} {CODENAME} main"]
VALID_URIS = [URI]
VALID_PPAS = ["ppa:ddstreet", "ppa:~ddstreet", "ppa:ddstreet/ppa", "ppa:~ddstreet/ppa", "ppa:ddstreet/ubuntu/ppa", "ppa:~ddstreet/ubuntu/ppa"]
VALID_UCAS = [f"cloud-archive:{UCA_CANAME}", f"cloud-archive:{UCA_CANAME}-updates", f"cloud-archive:{UCA_CANAME}-proposed", f"uca:{UCA_CANAME}", f"uca:{UCA_CANAME}-updates", f"uca:{UCA_CANAME}-proposed"]
VALID_ALL = VALID_LINES + VALID_URIS + VALID_PPAS + VALID_UCAS

INVALID_LINES = ["xxx invalid deb line"]
INVALID_URIS = ["invalid"]
INVALID_PPAS = ["ppainvalid:ddstreet", "ppa:ddstreet/ubuntu/ppa/invalid"]
INVALID_UCAS = [f"cloud-invalid:{UCA_CANAME}", "cloud-archive:"]
INVALID_ALL = INVALID_LINES + INVALID_URIS + INVALID_PPAS + INVALID_UCAS

def has_network():
    try:
        with urlopen("https://launchpad.net/"):
            pass
    except (URLError, HTTPException):
        return False
    return True

def mock_login_with(*args, **kwargs):
    has_subscription = kwargs.pop("has_subscription", True)

    _lp = Launchpad.login_anonymously(*args, **kwargs)
    lp = Mock(wraps=_lp)

    lp.me = Mock()
    lp.me.name = 'ddstreet'

    def mock_getPPAByName(_team, name):
        _ppa = _team.getPPAByName(name=name)
        ppa = Mock(wraps=_ppa)
        ppa.signing_key_fingerprint = _ppa.signing_key_fingerprint
        ppa.private = True
        return ppa

    def mock_people(teamname):
        _team = _lp.people(teamname)
        team = Mock(wraps=_team)
        team.getPPAByName = lambda name: mock_getPPAByName(_team, name)
        team.getArchiveSubscriptionURL = lambda archive: PRIVATE_PPA_SUBSCRIPTION_URL
        if has_subscription:
            team.getArchiveSubscriptionURL = lambda archive: PRIVATE_PPA_SUBSCRIPTION_URL
        else:
            def raise_unauthorized(archive):
                raise Unauthorized(Mock(), Mock())

            team.getArchiveSubscriptionURL = raise_unauthorized

        return team

    lp.people = mock_people
    return lp


class ShortcutsTestcase(unittest.TestCase):
    enable_source = False

    @classmethod
    def setUpClass(cls):
        for k in apt.apt_pkg.config.keys():
            apt.apt_pkg.config.clear(k)
        apt.apt_pkg.init()

    def setUp(self):
        # avoid polution from other tests
        apt.apt_pkg.config.set("Dir", "/")
        apt.apt_pkg.config.set("Dir::Etc", "etc/apt")
        apt.apt_pkg.config.set("Dir::Etc::sourcelist", "sources.list")
        apt.apt_pkg.config.set("Dir::Etc::sourceparts", "sources.list.d")
        apt.apt_pkg.config.set("Dir::Etc::trustedparts", "trusted.gpg.d")
        apt.apt_pkg.config.set("Dir::Etc::netrcparts", "auth.conf.d")

    def create_handler(self, line, handler, *args, **kwargs):
        return handler(line, *args, enable_source=self.enable_source, **kwargs)

    def create_handlers(self, line, handler, *args, **kwargs):
        handlers = handler if isinstance(handler, list) else [handler]
        # note, always appends shortcut_handler
        return [self.create_handler(line, handler, *args, **kwargs)
                for handler in handlers + [shortcut_handler]]

    def check_shortcut(self, shortcut, line, sourcefile=None, trustedfile=None, netrcfile=None,
                       sourceparts=apt.apt_pkg.config.find_dir("Dir::Etc::sourceparts"),
                       trustedparts=apt.apt_pkg.config.find_dir("Dir::Etc::trustedparts"),
                       netrcparts=apt.apt_pkg.config.find_dir("Dir::Etc::netrcparts"),
                       trustedcontent=False, netrccontent=None):
        self.assertEqual(shortcut.SourceEntry().line, line)

        self.assertEqual(shortcut.sourceparts_path, sourceparts)
        if sourcefile:
            self.assertEqual(shortcut.SourceEntry().file, os.path.join(sourceparts, sourcefile))
            self.assertEqual(shortcut.sourceparts_filename, sourcefile)
            self.assertEqual(shortcut.sourceparts_file, os.path.join(sourceparts, sourcefile))

        if shortcut.deb822:
            binentry = Deb822SourceEntry(line, '')
            binentry.types = [DISTRO.binary_type]
        else:
            binentry = SourceEntry(line)
            binentry.type = DISTRO.binary_type

        self.assertEqual(shortcut.SourceEntry(shortcut.binary_type), binentry)

        if shortcut.deb822:
            srcentry = Deb822SourceEntry(line, '')
            srcentry.types = [DISTRO.source_type]
        else:
            srcentry = SourceEntry(line)
            srcentry.type = DISTRO.source_type

        srcentry.set_enabled(self.enable_source)
        self.assertEqual(shortcut.SourceEntry(shortcut.source_type), srcentry)

        self.assertEqual(shortcut.trustedparts_path, trustedparts)
        if trustedfile:
            self.assertEqual(shortcut.trustedparts_filename, trustedfile)
            self.assertEqual(shortcut.trustedparts_file, os.path.join(trustedparts, trustedfile))

        # Checking the actual gpg key content is too much work.
        if trustedcontent:
            self.assertIsNotNone(shortcut.trustedparts_content)
        else:
            self.assertIsNone(shortcut.trustedparts_content)

        self.assertEqual(shortcut.netrcparts_path, netrcparts)
        if netrcfile:
            self.assertEqual(shortcut.netrcparts_filename, netrcfile)
            self.assertEqual(shortcut.netrcparts_file, os.path.join(netrcparts, netrcfile))

        self.assertEqual(shortcut.netrcparts_content, netrccontent)

    def test_shortcut_sourceslist(self):
        for line in VALID_LINES:
            for shortcut in self.create_handlers(line, SourcesListShortcutHandler):
                self.check_shortcut(shortcut, line)

    def test_shortcut_uri(self):
        for uri in VALID_URIS:
            line = f"deb {uri} {CODENAME} main"
            for shortcut in self.create_handlers(uri, URIShortcutHandler):
                self.check_shortcut(shortcut, line, sourcefile=URI_SOURCEFILE)

    @unittest.skipUnless(has_network(), "requires network")
    def test_shortcut_ppa(self):
        for ppa in VALID_PPAS:
            for shortcut in self.create_handlers(ppa, PPAShortcutHandler):
                self.check_shortcut(shortcut, PPA_LINE,
                                    sourcefile=PPA_SOURCEFILE,
                                    netrcfile=PPA_NETRCFILE,
                                    trustedcontent=True)

    @unittest.skipUnless(has_network(), "requires network")
    def test_shortcut_private_ppa(self):
        # this is the same tests as the public ppa, but login=True will use the mocked lp instance
        # this *does not* actually test/verify this works with a real private ppa; that must be done manually
        with patch('launchpadlib.launchpad.Launchpad.login_with', new=mock_login_with):
            for ppa in VALID_PPAS:
                for shortcut in self.create_handlers(ppa, PPAShortcutHandler, login=True):
                    self.check_shortcut(shortcut, PRIVATE_PPA_LINE,
                                        sourcefile=PPA_SOURCEFILE,
                                        netrcfile=PPA_NETRCFILE,
                                        trustedcontent=True,
                                        netrccontent=PRIVATE_PPA_NETRCCONTENT)

    @unittest.skipUnless(has_network(), "requires network")
    def test_shortcut_private_ppa_no_subscription(self):
        # this is the same tests as the public ppa, but login=True will use the mocked lp instance
        # this *does not* actually test/verify this works with a real private ppa; that must be done manually
        with patch('launchpadlib.launchpad.Launchpad.login_with', new=partial(mock_login_with, has_subscription=False)):
            for ppa in VALID_PPAS:
                with self.assertRaises(ShortcutException):
                    self.create_handler(ppa, PPAShortcutHandler, login=True)

    @contextmanager
    def ca_allow_codename(self, codename):
        key = "CA_ALLOW_CODENAME"
        orig = os.environ.get(key, None)
        try:
            os.environ[key] = codename
            yield
        finally:
            if orig:
                os.environ[key] = orig
            else:
                os.environ.pop(key, None)

    def test_shortcut_cloudarchive(self):
        for uca in VALID_UCAS:
            line = UCA_LINE_PROPOSED if 'proposed' in uca else UCA_LINE
            with self.ca_allow_codename(CODENAME):
                for shortcut in self.create_handlers(uca, CloudArchiveShortcutHandler):
                    self.check_shortcut(shortcut, line, sourcefile=UCA_SOURCEFILE)

    def check_invalid_shortcut(self, handler, shortcut):
        msg = "'%s' should have rejected '%s'" % (handler, shortcut)
        with self.assertRaises(InvalidShortcutException, msg=msg):
            self.create_handler(shortcut, handler)

    def test_shortcut_invalid(self):
        for s in INVALID_ALL + VALID_URIS + VALID_PPAS + VALID_UCAS:
            self.check_invalid_shortcut(SourcesListShortcutHandler, s)
        for s in INVALID_ALL + VALID_LINES + VALID_PPAS + VALID_UCAS:
            self.check_invalid_shortcut(URIShortcutHandler, s)
        for s in INVALID_ALL + VALID_LINES + VALID_URIS + VALID_UCAS:
            self.check_invalid_shortcut(PPAShortcutHandler, s)
        for s in INVALID_ALL + VALID_LINES + VALID_URIS + VALID_PPAS:
            self.check_invalid_shortcut(CloudArchiveShortcutHandler, s)
        for s in INVALID_ALL:
            self.check_invalid_shortcut(shortcut_handler, s)


class EnableSourceShortcutsTestcase(ShortcutsTestcase):
    enable_source = True


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