# Copyright 2021 Offensive Security
# SPDX-license-identifier: GPL-3.0-only

import unittest
from textwrap import dedent, indent
from urllib.parse import urlparse

from kali_tweaks.settings.aptrepositories import (
    Deb822Source,
    OneLineSource,
    mk_deb822_source,
    mk_one_line_source,
    parse_deb822_style_sources,
    parse_deb822_style_stanza,
    parse_one_line_style_line,
    parse_one_line_style_sources,
    print_deb822_style_default_stanza,
    print_deb822_style_stanza,
    print_one_line_style_default_line,
    print_one_line_style_line,
    update_deb822_style_sources,
    update_one_line_style_sources,
)

# Deb822-Style tests ----------------------------------------------------------


class TestAptRepositoriesMkDeb822Source(unittest.TestCase):
    def test_common_case(self):
        source = mk_deb822_source(
            "deb deb-src",
            "http://http.kali.org/kali/",
            "kali-rolling",
            "main contrib non-free non-free-firmware",
            {},
        )
        expected = Deb822Source(
            ["deb", "deb-src"],
            [urlparse("http://http.kali.org/kali/")],
            ["kali-rolling"],
            ["main", "contrib", "non-free", "non-free-firmware"],
            {},
        )
        self.assertEqual(source, expected)

    def test_options_are_copied(self):
        # mk_deb822_source must make a copy of the options
        options = {"foo": "bar"}
        source = mk_deb822_source("", "", "", "", options)
        expected = Deb822Source([], [], [], [], {"foo": "bar"})
        self.assertEqual(source, expected)
        options["foo"] = "charlie"  # modify options
        self.assertEqual(source, expected)


class TestAptRepositoriesParseDeb822StyleStanza(unittest.TestCase):
    def assert_parse(self, stanza, expected):
        source = parse_deb822_style_stanza(stanza)
        self.assertEqual(source, expected)

    def test_empty_stanza(self):
        self.assert_parse("", None)
        self.assert_parse("# this is a comment", None)
        self.assert_parse("# comment 1\n# comment 2", None)

    def test_kali_default(self):
        stanza = dedent("""\
            Types: deb
            URIs: http://http.kali.org/kali/
            Suites: kali-rolling
            Components: main contrib non-free non-free-firmware
            Signed-By: /usr/share/keyrings/kali-archive-keyring.gpg""")
        exp = mk_deb822_source(
            "deb",
            "http://http.kali.org/kali/",
            "kali-rolling",
            "main contrib non-free non-free-firmware",
            {"signed-by": "/usr/share/keyrings/kali-archive-keyring.gpg"},
        )
        self.assert_parse(stanza, exp)

    def test_embedded_gpg_public_key(self):
        gpg_block = dedent("""\
            -----BEGIN PGP PUBLIC KEY BLOCK-----
            .
            mQINBGgBJJUBEADlMTZVDCjrSXIAuYfL3VZt8OoplUdw3mSPlhIjZQmIo2sdzvAF
            EMSCQ+vWeD4VqV9tBtiVx6j8VSfyW18YHHAkvajWDRg5hPLf80wGxrtXYu+vj3Ri
            5dOMhrl9fHKIifPOoV3pFTtOk0dB9lkcmtNzjWgwOJduLbjjraE1BBKqc0uaXDCa
            [...]
            -----END PGP PUBLIC KEY BLOCK----""")
        stanza = dedent("""\
            Types: deb
            URIs: http://http.kali.org/kali/
            Suites: kali-rolling
            Components: main contrib non-free non-free-firmware
            Signed-By:
            """)
        stanza += indent(gpg_block, " ")
        exp = mk_deb822_source(
            "deb",
            "http://http.kali.org/kali/",
            "kali-rolling",
            "main contrib non-free non-free-firmware",
            {"signed-by": gpg_block},
        )
        self.assert_parse(stanza, exp)

    def test_no_signed_by(self):
        stanza = dedent("""\
            Types: deb
            URIs: http://http.kali.org/kali/
            Suites: kali-rolling
            Components: main""")
        exp = mk_deb822_source(
            "deb",
            "http://http.kali.org/kali/",
            "kali-rolling",
            "main",
            {},
        )
        self.assert_parse(stanza, exp)

    def test_stanza_disabled(self):
        stanza = dedent("""\
            Enabled: no
            Types: deb
            URIs: http://http.kali.org/kali/
            Suites: kali-rolling
            Components: main""")
        self.assert_parse(stanza, None)

    def test_suites_not_kali(self):
        stanza = dedent("""\
            Types: deb
            URIs: http://http.kali.org/kali/
            Suites: this-is-not-a-kali-suite
            Components: main""")
        self.assert_parse(stanza, None)

    def test_messy_sources(self):
        # example taken and adjusted from:
        # apt/integration/test-apt-sources-deb822
        stanza = """\
Types: deb
 deb-src
URIs:
 http://http.kali.org/kali/
 http://kali.download/kali/
Suites: kali-rolling
          kali-experimental
Components:
        main
 contrib        non-free
Architectures: amd64
 armhf     i386
"""
        exp = mk_deb822_source(
            "deb deb-src",
            "http://http.kali.org/kali/ http://kali.download/kali/",
            "kali-rolling kali-experimental",
            "main contrib non-free",
            {"architectures": "amd64\narmhf     i386"},
        )
        self.assert_parse(stanza, exp)

    def test_missing_mandatory_fields(self):
        stanza = "Enabled: yes"
        with self.assertRaises(ValueError):
            self.assert_parse(stanza, None)

    def test_malformed_stanza(self):
        stanza = "This line is missing a colon"
        with self.assertRaises(ValueError):
            self.assert_parse(stanza, None)


class TestAptRepositoriesPrintDeb822StyleStanza(unittest.TestCase):
    def assert_print(self, source, expected):
        stanza = print_deb822_style_stanza(source)
        self.assertEqual(stanza, expected)

    def test_kali_default(self):
        src = mk_deb822_source(
            "deb",
            "http://http.kali.org/kali/",
            "kali-rolling",
            "main contrib non-free non-free-firmware",
            {"signed-by": "/usr/share/keyrings/kali-archive-keyring.gpg"},
        )
        exp = dedent("""\
            Types: deb
            URIs: http://http.kali.org/kali/
            Suites: kali-rolling
            Components: main contrib non-free non-free-firmware
            Signed-By: /usr/share/keyrings/kali-archive-keyring.gpg
            """)
        self.assert_print(src, exp)


class TestAptRepositoriesPrintDeb822StyleDefaultStanza(unittest.TestCase):
    def assert_print(self, proto, mirror, extra_suites, expected):
        stanza = print_deb822_style_default_stanza(proto, mirror, extra_suites)
        self.assertEqual(stanza, expected)

    def test_it(self):
        exp = dedent("""\
            Types: deb
            URIs: http://http.kali.org/kali/
            Suites: kali-rolling
            Components: main contrib non-free non-free-firmware
            Signed-By: /usr/share/keyrings/kali-archive-keyring.gpg
            """)
        self.assert_print("http", "http.kali.org", [], exp)
        exp = dedent("""\
            Types: deb
            URIs: https://kali.download/kali/
            Suites: kali-rolling kali-bleeding-edge
            Components: main contrib non-free non-free-firmware
            Signed-By: /usr/share/keyrings/kali-archive-keyring.gpg
            """)
        self.assert_print("https", "kali.download", ["kali-bleeding-edge"], exp)


class TestAptRepositoriesParseDeb822StyleSources(unittest.TestCase):
    def assert_parse(self, content, expected):
        sources = parse_deb822_style_sources(content)
        self.assertEqual(sources, expected)

    def test_kali_default_sources(self):
        content = dedent("""\
            Types: deb
            URIs: http://http.kali.org/kali/
            Suites: kali-rolling
            Components: main contrib non-free non-free-firmware
            Signed-By: /usr/share/keyrings/kali-archive-keyring.gpg""")
        src = mk_deb822_source(
            "deb",
            "http://http.kali.org/kali/",
            "kali-rolling",
            "main contrib non-free non-free-firmware",
            {"signed-by": "/usr/share/keyrings/kali-archive-keyring.gpg"},
        )
        self.assert_parse(content, [src])

    def test_multiple_stanzas(self):
        content = dedent("""\
            # This is a comment
            Types: deb
            URIs: http://http.kali.org/kali/
            Suites: kali-rolling
            Components: main contrib non-free non-free-firmware
            Signed-By: /usr/share/keyrings/kali-archive-keyring.gpg

            Types: deb deb-src
            URIs: http://http.kali.org/kali/
            # A comment right here!
            Suites: kali-experimental
            Components: main
            """)
        src1 = mk_deb822_source(
            "deb",
            "http://http.kali.org/kali/",
            "kali-rolling",
            "main contrib non-free non-free-firmware",
            {"signed-by": "/usr/share/keyrings/kali-archive-keyring.gpg"},
        )
        src2 = mk_deb822_source(
            "deb deb-src",
            "http://http.kali.org/kali/",
            "kali-experimental",
            "main",
            {},
        )
        self.assert_parse(content, [src1, src2])


class TestAptRepositoriesUpdateDeb822StyleSources(unittest.TestCase):
    def assert_update(self, content, expected, **kwargs):
        res = update_deb822_style_sources(content, **kwargs)
        self.assertEqual(res, expected)

    @staticmethod
    def mk_stanza(
        types="deb",
        proto="http",
        mirror="http.kali.org",
        suites="kali-rolling",
        components="main contrib non-free non-free-firmware",
    ):
        stanza = dedent("""\
            Types: {types}
            URIs: {proto}://{mirror}/kali/
            Suites: {suites}
            Components: {components}
            Signed-By: /usr/share/keyrings/kali-archive-keyring.gpg
            """)
        return stanza.format(
            types=types,
            proto=proto,
            mirror=mirror,
            suites=suites,
            components=components,
        )

    def test_not_modified(self):
        old = self.mk_stanza()
        self.assert_update(old, (old, False))

    def test_protocol(self):
        old = self.mk_stanza()
        new = self.mk_stanza(proto="https")
        self.assert_update(old, (old, False), protocol="http")
        self.assert_update(old, (new, True), protocol="https")

    def test_mirror(self):
        old = self.mk_stanza()
        new = self.mk_stanza(mirror="kali.download")
        self.assert_update(old, (old, False), mirror="http.kali.org")
        self.assert_update(old, (new, True), mirror="kali.download")

    def test_protocol_plus_mirror(self):
        old = self.mk_stanza()
        new = self.mk_stanza(proto="https", mirror="kali.download")
        self.assert_update(old, (new, True), protocol="https", mirror="kali.download")

    def test_add_suite(self):
        old = self.mk_stanza(suites="kali-rolling")
        new = self.mk_stanza(suites="kali-rolling kali-experimental")
        # add suite that was not there
        self.assert_update(old, (new, True), extra_suites=["kali-experimental"])
        # add suite that was already there
        self.assert_update(new, (new, False), extra_suites=["kali-experimental"])

    def test_remove_suite(self):
        old = self.mk_stanza(suites="kali-rolling kali-experimental")
        new = self.mk_stanza(suites="kali-rolling")
        # remove suite that was there
        self.assert_update(old, (new, True), extra_suites=[])
        # remove suite that was not listed
        self.assert_update(new, (new, False), extra_suites=[])

    def test_types_components_are_preserved(self):
        old = self.mk_stanza(types="deb deb-src", components="main")
        new = self.mk_stanza(types="deb deb-src", components="main", proto="https")
        self.assert_update(old, (new, True), protocol="https")

    def test_kali_dev(self):
        old = self.mk_stanza(suites="kali-dev")
        new = self.mk_stanza(suites="kali-dev kali-bleeding-edge")
        self.assert_update(old, (new, True), extra_suites=["kali-bleeding-edge"])

    def test_signed_by_added_if_missing(self):
        old = dedent("""\
            Types: deb
            URIs: http://http.kali.org/kali/
            Suites: kali-rolling
            Components: main
            """)
        new = dedent("""\
            Types: deb
            URIs: https://http.kali.org/kali/
            Suites: kali-rolling
            Components: main
            Signed-By: /usr/share/keyrings/kali-archive-keyring.gpg
            """)
        self.assert_update(old, (new, True), protocol="https")

    def test_multiple_stanzas(self):
        # Only the stanza with a kali suite is preserved
        old = dedent("""\
            # A non-kali suite in the first stanza
            Types: deb
            URIs: http://thisisnotkali.com
            Suites: this-is-not-kali
            Components: comp1 comp2

            # The kali stanza is there
            Types: deb
            URIs: http://http.kali.org/kali/
            Suites: kali-rolling
            Components: main
            Signed-By: /usr/share/keyrings/kali-archive-keyring.gpg
            """)
        new = self.mk_stanza(components="main", proto="https")
        self.assert_update(old, (new, True), protocol="https")


# One-Line-Style tests --------------------------------------------------------


class TestAptRepositoriesMkOneLineSource(unittest.TestCase):
    def test_common_cases(self):
        url = "http://http.kali.org/kali"
        components = "main contrib non-free non-free-firmware"
        source = mk_one_line_source("deb", "", url, "kali-rolling", components)
        url = urlparse(url)
        components = ["main", "contrib", "non-free", "non-free-firmware"]
        expected = OneLineSource("deb", [], url, "kali-rolling", components)
        self.assertEqual(source, expected)


class TestAptRepositoriesParseOneLineStyleLine(unittest.TestCase):
    def assert_parse(self, line, expected):
        source = parse_one_line_style_line(line)
        self.assertEqual(source, expected)

    def test_common_cases(self):
        components = "main contrib non-free non-free-firmware"
        for mirror in ["http.kali.org", "kali.download"]:
            url = f"http://{mirror}/kali"
            src = f"deb {url} kali-rolling {components}"
            exp = mk_one_line_source("deb", "", url, "kali-rolling", components)
            self.assert_parse(src, exp)

    def test_deb_src(self):
        url = "http://http.kali.org/kali"
        src = f"deb-src {url} kali-dev main"
        exp = mk_one_line_source("deb-src", "", url, "kali-dev", "main")
        self.assert_parse(src, exp)

    def test_with_options(self):
        url = "http://http.kali.org/kali"
        # empty options array
        src = f"deb [] {url} kali-dev main"
        exp = mk_one_line_source("deb", "", url, "kali-dev", "main")
        self.assert_parse(src, exp)
        # one option, no space around brackets
        src = f"deb [opt1=foo] {url} kali-dev main"
        exp = mk_one_line_source("deb", "opt1=foo", url, "kali-dev", "main")
        self.assert_parse(src, exp)
        # two options, spaces around brackets
        src = f"deb [ opt1=foo opt2=bar ] {url} kali-dev main"
        exp = mk_one_line_source("deb", "opt1=foo opt2=bar", url, "kali-dev", "main")
        self.assert_parse(src, exp)

    def test_third_party_mirror(self):
        url = "http://third.party.mirror.com/Linux/kali"
        src = f"deb {url} kali-dev main"
        exp = mk_one_line_source("deb", "", url, "kali-dev", "main")
        self.assert_parse(src, exp)

    def test_cdrom(self):
        url = "cdrom://http.kali.org/kali"
        src = f"deb {url} kali-dev main"
        exp = mk_one_line_source("deb", "", url, "kali-dev", "main")
        self.assert_parse(src, exp)

    def test_invalid_suite(self):
        url = "http://http.kali.org/kali"
        src = f"deb {url} this-is-not-kali main"
        exp = mk_one_line_source("deb", "", url, "this-is-not-kali", "main")
        self.assert_parse(src, exp)

    def test_components(self):
        components = "main contrib non-free"
        url = "http://http.kali.org/kali"
        src = f"deb {url} kali-dev {components}"
        exp = mk_one_line_source("deb", "", url, "kali-dev", components)
        self.assert_parse(src, exp)
        components = "main contrib non-free non-free-firmware"
        src = f"deb {url} kali-dev {components}"
        exp = mk_one_line_source("deb", "", url, "kali-dev", components)
        self.assert_parse(src, exp)

    def test_empty_line(self):
        self.assert_parse("", None)
        self.assert_parse("\n\n\n", None)

    def test_commented_line(self):
        src = "#deb http://http.kali.org/kali kali-dev main"
        self.assert_parse(src, None)

    def test_malformed_lines(self):
        # invalid type
        src = "foobar http://http.kali.org/kali kali-dev main"
        with self.assertRaises(ValueError):
            self.assert_parse(src, None)
        # nothing after options
        src = "deb []"
        with self.assertRaises(ValueError):
            self.assert_parse(src, None)
        # opening bracket for options, and nothing else
        src = "deb [ http://http.kali.org/kali kali-dev main"
        with self.assertRaises(ValueError):
            self.assert_parse(src, None)
        # options missing the closing bracket
        src = "deb [opt1=foo http://http.kali.org/kali kali-dev main"
        with self.assertRaises(ValueError):
            self.assert_parse(src, None)
        # no component
        src = "deb http://http.kali.org/kali kali-dev"
        with self.assertRaises(ValueError):
            self.assert_parse(src, None)


class TestAptRepositoriesPrintOneLineStyleLine(unittest.TestCase):
    def assert_print(self, source, expected):
        line = print_one_line_style_line(source)
        self.assertEqual(line, expected)

    def test_common_cases(self):
        components = "main contrib non-free non-free-firmware"
        for mirror in ["http.kali.org", "kali.download"]:
            url = f"http://{mirror}/kali"
            src = mk_one_line_source("deb", "", url, "kali-rolling", components)
            exp = f"deb {url} kali-rolling {components}"
            self.assert_print(src, exp)

    def test_with_options(self):
        url = "http://http.kali.org/kali"
        src = mk_one_line_source("deb", "opt=foo", url, "kali-rolling", "main")
        exp = f"deb [ opt=foo ] {url} kali-rolling main"
        self.assert_print(src, exp)


class TestAptRepositoriesPrintOneLineStyleDefaultLine(unittest.TestCase):
    def assert_print(self, proto, mirror, suite, expected):
        line = print_one_line_style_default_line(proto, mirror, suite)
        self.assertEqual(line, expected)

    def test_it(self):
        components = "main contrib non-free non-free-firmware"
        exp = f"deb http://http.kali.org/kali kali-experimental {components}"
        self.assert_print("http", "http.kali.org", "kali-experimental", exp)
        exp = f"deb https://kali.download/kali kali-bleeding-edge {components}"
        self.assert_print("https", "kali.download", "kali-bleeding-edge", exp)


class TestAptRepositoriesParseOneLineStyleSources(unittest.TestCase):
    def assert_parse(self, content, expected):
        sources = parse_one_line_style_sources(content)
        self.assertEqual(sources, expected)

    def test_one_source(self):
        url = "http://http.kali.org/kali"
        components = "main contrib non-free non-free-firmware"
        content = f"deb {url} kali-rolling {components}"
        exp = [mk_one_line_source("deb", "", url, "kali-rolling", components)]
        self.assert_parse(content, exp)

    def test_deb_src(self):
        url = "http://http.kali.org/kali"
        content = f"deb-src {url} kali-rolling main"
        self.assert_parse(content, [])

    def test_invalid_suite(self):
        url = "http://http.kali.org/kali"
        content = f"deb {url} this-is-not-kali main"
        self.assert_parse(content, [])

    def test_multiple_sources_with_comments(self):
        url = "http://kali.download/kali"
        components = "main contrib non-free non-free-firmware"
        content = dedent(f"""
            ## kali rolling, with sources ##
            deb     {url} kali-rolling {components}
            deb-src {url} kali-rolling {components}
            ## kali exp, no sources ##
            deb     {url} kali-experimental {components}
            """)
        exp = [
            mk_one_line_source("deb", "", url, "kali-rolling", components),
            mk_one_line_source("deb", "", url, "kali-experimental", components),
        ]
        self.assert_parse(content, exp)


class TestAptRepositoriesUpdateOneLineStyleSources(unittest.TestCase):
    def assert_update(self, content, expected, **kwargs):
        res = update_one_line_style_sources(content, **kwargs)
        self.assertEqual(res, expected)

    def test_not_modified(self):
        old = "deb http://http.kali.org/kali kali-dev main"
        self.assert_update(old, (old, False))

    def test_protocol(self):
        old = "deb http://http.kali.org/kali kali-dev main"
        new = "deb https://http.kali.org/kali kali-dev main\n"
        self.assert_update(old, (old, False), protocol="http")
        self.assert_update(old, (new, True), protocol="https")

    def test_mirror(self):
        old = "deb http://http.kali.org/kali kali-dev main"
        new = "deb http://kali.download/kali kali-dev main\n"
        self.assert_update(old, (old, False), mirror="http.kali.org")
        self.assert_update(old, (new, True), mirror="kali.download")

    def test_protocol_plus_mirror(self):
        old = "deb http://http.kali.org/kali kali-dev main"
        new = "deb https://kali.download/kali kali-dev main\n"
        self.assert_update(old, (new, True), protocol="https", mirror="kali.download")

    def test_remove_suite(self):
        old = "deb http://http.kali.org/kali kali-dev main"
        new = ""
        self.assert_update(old, (old, False), remove_suites=["foo"])
        self.assert_update(old, (new, True), remove_suites=["kali-dev"])

    def test_non_kali_source(self):
        old = "deb http://http.kali.org/kali not-kali main"
        self.assert_update(old, (old, False), protocol="https")

    def test_all_at_once(self):
        old = dedent("""
            # This is a bit surprising
            deb http://foo.bar/abcd not-kali hence not modified
            deb http://foo.bar/whatever kali-dev this is a kali source
            # This is just the usual
            deb http://http.kali.org/kali kali-rolling main contrib non-free non-free-firmware
            deb-src http://http.kali.org/kali kali-rolling main contrib
            deb http://http.kali.org/kali kali-bleeding-edge main contrib non-free non-free-firmware
            deb-src http://http.kali.org/kali kali-bleeding-edge main contrib
            """)
        new = dedent("""
            # This is a bit surprising
            deb http://foo.bar/abcd not-kali hence not modified
            deb https://kali.download/kali kali-dev this is a kali source
            # This is just the usual
            deb https://kali.download/kali kali-rolling main contrib non-free non-free-firmware
            deb-src https://kali.download/kali kali-rolling main contrib
            """)
        self.assert_update(
            old,
            (new, True),
            protocol="https",
            mirror="kali.download",
            remove_suites=["kali-bleeding-edge"],
        )


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