# Copyright 2022 Google Inc. All Rights Reserved.
# Modifications Copyright 2022-2023 Aerleon Project Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Unittest for Nftables rendering module."""

from absl import logging
from absl.testing import absltest, parameterized

from aerleon.lib import nacaddr, naming, nftables, policy
from tests.regression_utils import capture


class DictObj:
    """Helper class to use a dictionary of dictionaries to form an object.

    We can then specifically test using it.
    """

    def __init__(self, in_dict: dict):
        assert isinstance(in_dict, dict)
        for key, val in in_dict.items():
            if isinstance(val, (list, tuple)):
                setattr(self, key, [DictObj(x) if isinstance(x, dict) else x for x in val])
            else:
                setattr(self, key, DictObj(val) if isinstance(val, dict) else val)


# "logging" is not a token.
SUPPORTED_TOKENS = frozenset(
    {
        'action',
        'comment',
        'destination_address',
        'destination_address_exclude',
        'destination_port',
        'expiration',
        'icmp_type',
        'name',  # obj attribute, not token
        'option',
        'protocol',
        'platform',
        'platform_exclude',
        'source_address',
        'source_address_exclude',
        'source_port',
        'translated',  # obj attribute, not token
        'stateless_reply',
    }
)

SUPPORTED_SUB_TOKENS = {
    'action': {'accept', 'deny'},
    'option': {'established', 'tcp-established'},
    'icmp_type': {
        'alternate-address',
        'certification-path-advertisement',
        'certification-path-solicitation',
        'conversion-error',
        'destination-unreachable',
        'echo-reply',
        'echo-request',
        'mobile-redirect',
        'home-agent-address-discovery-reply',
        'home-agent-address-discovery-request',
        'icmp-node-information-query',
        'icmp-node-information-response',
        'information-request',
        'inverse-neighbor-discovery-advertisement',
        'inverse-neighbor-discovery-solicitation',
        'mask-reply',
        'mask-request',
        'information-reply',
        'mobile-prefix-advertisement',
        'mobile-prefix-solicitation',
        'multicast-listener-done',
        'multicast-listener-query',
        'multicast-listener-report',
        'multicast-router-advertisement',
        'multicast-router-solicitation',
        'multicast-router-termination',
        'neighbor-advertisement',
        'neighbor-solicit',
        'packet-too-big',
        'parameter-problem',
        'redirect',
        'redirect-message',
        'router-advertisement',
        'router-renumbering',
        'router-solicit',
        'router-solicitation',
        'source-quench',
        'time-exceeded',
        'timestamp-reply',
        'timestamp-request',
        'unreachable',
        'version-2-multicast-listener-report',
    },
}

HEADER_TEMPLATE = """
header {
  target:: nftables %s
}
"""

HEAD_OVERRIDE_DEFAULT_ACTION = """
header {
  target:: nftables inet output ACCEPT
}
"""

HEADER_COMMENT = """
header {
  comment:: "Noverbose + custom priority policy example"
  target:: nftables inet output ACCEPT
}
"""

# TODO(gfm): Noverbose testing once Term handling is added.
HEADER_NOVERBOSE = """
header {
  target:: nftables mixed output noverbose
}
"""

GOOD_HEADER_1 = """
header {
  target:: nftables inet6 INPUT
}
"""

GOOD_HEADER_2 = """
header {
  target:: nftables mixed output accept
}
"""

ESTABLISHED_OPTION_TERM = """
term established-term {
  protocol:: udp
  option:: established
  action:: accept
}
"""

TCP_ESTABLISHED_OPTION_TERM = """
term tcp-established-term {
  protocol:: tcp
  option:: tcp-established
  action:: accept
}
"""

ICMP_TERM = """
term good-icmp {
  protocol:: icmp
  action:: accept
}
"""

GOOD_TERM_1 = """
term good-term-1 {
  action:: accept
}
"""

IPV6_TERM_2 = """
term inet6-icmp {
  action:: deny
}
"""

EXCLUDE = {'ip6': [nacaddr.IP('::/3'), nacaddr.IP('::/0')]}

# Print a info message when a term is set to expire in that many weeks.
# This is normally passed from command line.
EXP_INFO = 2


def IPhelper(addresses):
    """Helper for string to nacaddr.IP conversion for parametized tests."""
    normalized = []
    if not addresses:
        # if empty list of addresses.
        return addresses
    else:
        for addr in addresses:
            normalized.append(nacaddr.IP(addr))
        return normalized


class NftablesTest(parameterized.TestCase):
    def setUp(self):
        super().setUp()
        self.naming = naming.Naming()
        self.dummyterm = nftables.Term('', '', '')

    @parameterized.parameters(('ip protocol tcp', ' ip protocol tcp'), ('', ''))
    def testAdd(self, statement, expected_output):
        result = nftables.Add(statement)
        self.assertEqual(result, expected_output)

    @parameterized.parameters((2, 'chain acl_name', '  chain acl_name'))
    def testTabSpacer(self, num_spaces, statement, expected_output):
        result = nftables.TabSpacer(num_spaces, statement)
        self.assertEqual(result, expected_output)

    @parameterized.parameters(
        (
            'inet',
            ['200.1.1.3/32', '9782:b30a:e5c6:1aa4:29ff:e57c:44a0:1b84'],
            ['200.1.1.3/32', '2606:4700:4700::1111'],
            [
                'ip saddr 200.1.1.3/32 ip daddr 200.1.1.3/32',
                'ip6 saddr 9782:b30a:e5c6:1aa4:29ff:e57c:44a0:1b84/128 ip6 daddr 2606:4700:4700::1111/128',
            ],
        ),
        (
            'inet',
            ['9782:b30a:e5c6:1aa4:29ff:e57c:44a0:1b84'],
            ['200.1.1.3/32', '2606:4700:4700::1111'],
            [
                'ip6 saddr 9782:b30a:e5c6:1aa4:29ff:e57c:44a0:1b84/128 ip6 daddr 2606:4700:4700::1111/128'
            ],
        ),
        (
            'inet',
            ['200.1.1.3/32', '9782:b30a:e5c6:1aa4:29ff:e57c:44a0:1b84'],
            ['200.1.1.3/32'],
            [
                'ip saddr 200.1.1.3/32 ip daddr 200.1.1.3/32',
            ],
        ),
        (
            'ip',
            ['200.1.1.3/32', '9782:b30a:e5c6:1aa4:29ff:e57c:44a0:1b84'],
            ['200.1.1.3/32', '2606:4700:4700::1111'],
            ['ip saddr 200.1.1.3/32 ip daddr 200.1.1.3/32'],
        ),
        (
            'ip6',
            ['8.8.8.8', '9782:b30a:e5c6:1aa4:29ff:e57c:44a0:1b84'],
            ['200.1.1.3/32', '2606:4700:4700::1111'],
            [
                'ip6 saddr 9782:b30a:e5c6:1aa4:29ff:e57c:44a0:1b84/128 ip6 daddr 2606:4700:4700::1111/128'
            ],
        ),
        (
            'inet',
            [],
            ['200.1.1.3/32', '2606:4700:4700::1111'],
            ['ip daddr 200.1.1.3/32', 'ip6 daddr 2606:4700:4700::1111/128'],
        ),
        (
            'inet',
            ['9782:b30a:e5c6:1aa4:29ff:e57c:44a0:1b84'],
            [],
            ['ip6 saddr 9782:b30a:e5c6:1aa4:29ff:e57c:44a0:1b84/128'],
        ),
        ('inet', [], [], []),
    )
    def test_AddrStatement(self, af, src_addr, dst_addr, expected):
        # Necessary object format.
        src_obj = IPhelper(src_addr)
        dst_obj = IPhelper(dst_addr)
        result = self.dummyterm._AddrStatement(af, src_obj, dst_obj)
        self.assertEqual(result, expected)

    @parameterized.parameters(
        (
            ['nd-router-advert', 'nd-neighbor-solicit', 'nd-neighbor-advert'],
            '{ nd-router-advert, nd-neighbor-solicit, nd-neighbor-advert }',
        ),
        (['200.1.1.3/32'], '200.1.1.3/32'),
        (['1.1.1.1', '8.8.8.8'], '{ 1.1.1.1, 8.8.8.8 }'),
        (['tcp', 'udp', 'icmp'], '{ tcp, udp, icmp }'),
        (['80', '443'], '{ 80, 443 }'),
        ('53', '53'),
    )
    def testCreateAnonymousSet(self, input_data, expected):
        result = self.dummyterm.CreateAnonymousSet(input_data)
        self.assertEqual(result, expected)

    @parameterized.parameters(
        (
            [
                'ip6 saddr 2606:4700:4700::1111/128 ip6 daddr { 2001:4860:4860::8844/128, 2001:4860:4860::8888/128'
            ],
            ['tcp sport 80 tcp dport 80'],
            'ct state { ESTABLISHED, RELATED } log prefix "combo_cnt_log_established" counter',
            'accept',
            [
                'ip6 saddr 2606:4700:4700::1111/128 ip6 daddr { 2001:4860:4860::8844/128, 2001:4860:4860::8888/128 tcp sport 80 tcp dport 80 ct state { ESTABLISHED, RELATED } log prefix "combo_cnt_log_established" counter accept'
            ],
        ),
    )
    def testGroupExpressions(self, address_expr, porst_proto_expr, opt, verdict, expected_output):
        result = self.dummyterm.GroupExpressions(address_expr, porst_proto_expr, opt, verdict)
        self.assertEqual(result, expected_output)

    def testDuplicateTerm(self):
        pol = policy.ParsePolicy(GOOD_HEADER_1 + GOOD_TERM_1 + GOOD_TERM_1, self.naming)
        with self.assertRaises(nftables.TermError):
            nftables.Nftables.__init__(nftables.Nftables.__new__(nftables.Nftables), pol, EXP_INFO)

    @parameterized.parameters(([(80, 80)], '80'), ([(1024, 65535)], '1024-65535'), ([], ''))
    def testGroup(self, data, expected_output):
        """Test _Group function we use in Ports."""
        result = self.dummyterm._Group(data)
        self.assertEqual(result, expected_output)

    @parameterized.parameters(
        (
            'inet',
            ['tcp'],
            [(3199, 3199)],
            [(80, 80)],
            [],
            ['tcp sport 3199 tcp dport 80', 'tcp sport 3199 tcp dport 80'],
        ),
        (
            'inet',
            ['tcp'],
            [(3199, 3199)],
            [],
            [],
            ['tcp sport 3199', 'tcp sport 3199'],
        ),
        (
            'inet',
            ['tcp'],
            [],
            [(80, 80)],
            [],
            ['tcp dport 80', 'tcp dport 80'],
        ),
        ('inet', ['tcp'], [], [], [], ['ip protocol tcp', 'meta l4proto tcp']),
        ('ip6', ['tcp'], [], [], [], ['meta l4proto tcp']),
        (
            'inet',
            ['tcp', 'udp'],
            [],
            [],
            [],
            ['ip protocol { tcp, udp }', 'meta l4proto { tcp, udp }'],
        ),
    )
    def testPortsAndProtocols(self, af, proto, src_p, dst_p, icmp_type, expected):
        result = self.dummyterm.PortsAndProtocols(af, proto, src_p, dst_p, icmp_type)
        self.assertEqual(result, expected)

    @parameterized.parameters(
        'chain_name input 0 inet extraneous_target_option',
        'ip6 OUTPUT 300 400' 'mixed input',  # pylint: disable=implicit-str-concat
        'ip forwarding',
        'ip7 0 spaghetti',
        'ip6 prerouting',
        'chain_name',
        '',
    )
    def testBadHeader(self, case):
        logging.info('Testing bad header case %s.', case)
        header = HEADER_TEMPLATE % case
        pol = policy.ParsePolicy(header + GOOD_TERM_1, self.naming)
        with self.assertRaises(nftables.HeaderError):
            nftables.Nftables.__init__(nftables.Nftables.__new__(nftables.Nftables), pol, EXP_INFO)

    @parameterized.parameters((HEADER_NOVERBOSE, False), (HEADER_COMMENT, True))
    def testVerboseHeader(self, header_to_use, expected_output):
        pol = policy.ParsePolicy(header_to_use + GOOD_TERM_1, self.naming)
        data = nftables.Nftables(pol, EXP_INFO)
        for _, _, _, _, _, _, verbose, _ in data.nftables_policies:
            result = verbose
        self.assertEqual(result, expected_output)

    @capture.stdout
    def testGoodHeader(self):
        nftables.Nftables(policy.ParsePolicy(GOOD_HEADER_1 + GOOD_TERM_1, self.naming), EXP_INFO)
        nft = str(
            nftables.Nftables(
                policy.ParsePolicy(
                    GOOD_HEADER_1 + GOOD_TERM_1 + GOOD_HEADER_2 + IPV6_TERM_2, self.naming
                ),
                EXP_INFO,
            )
        )
        self.assertIn('type filter hook input', nft)
        print(nft)

    @capture.stdout
    def testStatefulFirewall(self):
        nftables.Nftables(policy.ParsePolicy(GOOD_HEADER_1 + GOOD_TERM_1, self.naming), EXP_INFO)
        nft = str(
            nftables.Nftables(
                policy.ParsePolicy(
                    GOOD_HEADER_1 + GOOD_TERM_1 + GOOD_HEADER_2 + IPV6_TERM_2, self.naming
                ),
                EXP_INFO,
            )
        )
        self.assertIn('ct state established,related accept', nft)
        print(nft)

    @capture.stdout
    def testOverridePolicyHeader(self):
        expected_output = 'accept'

        pol = policy.ParsePolicy(HEAD_OVERRIDE_DEFAULT_ACTION + GOOD_TERM_1, self.naming)
        data = nftables.Nftables(pol, EXP_INFO)
        for _, _, _, _, _, default_policy, _, _ in data.nftables_policies:
            result = default_policy
        self.assertEqual(result, expected_output)
        print(data)

    @parameterized.parameters(
        (['127.0.0.1', '8.8.8.8'], {'ip': ['127.0.0.1/32', '8.8.8.8/32']}),
        (['0.0.0.0/8', '2001:db8::/32'], {'ip': ['0.0.0.0/8'], 'ip6': ['2001:db8::/32']}),
    )
    def testAddressClassifier(self, addr_to_classify, expected_output):
        result = nftables.Term._AddressClassifier(self, IPhelper(addr_to_classify))
        self.assertEqual(result, expected_output)

    @parameterized.parameters(
        ('ip6', ['multicast-listener-query'], ['mld-listener-query']),
        (
            'ip6',
            ['echo-request', 'multicast-listener-query'],
            ['echo-request', 'mld-listener-query'],
        ),
        (
            'ip6',
            ['router-solicit', 'multicast-listener-done', 'router-advertisement'],
            ['nd-router-solicit', 'mld-listener-done', 'nd-router-advert'],
        ),
        ('ip4', ['echo-request', 'echo-reply'], ['echo-request', 'echo-reply']),
    )
    def testMapICMPtypes(self, af, icmp_types, expected_output):
        result = self.dummyterm.MapICMPtypes(af, icmp_types)
        self.assertEqual(result, expected_output)

    @parameterized.parameters(
        (
            {
                'name': 'tcp_established',
                'option': ['tcp-established', 'established'],
                'counter': None,
                'logging': [],
                'protocol': ['tcp', 'icmp'],
                'action': ['deny'],
            },
            '',
        ),
        (
            {
                'name': 'dont_render_tcp_established',
                'option': ['tcp-established', 'established'],
                'counter': None,
                'logging': [],
                'protocol': ['icmp'],
                'action': ['accept'],
            },
            'ct state new',
        ),
        (
            {
                'name': 'blank_option_donothing',
                'option': [],
                'counter': None,
                'logging': [],
                'protocol': ['icmp'],
                'action': ['accept'],
            },
            'ct state new',
        ),
        (
            {
                'name': 'syslog',
                'option': [],
                'counter': None,
                'logging': ['syslog'],
                'protocol': ['tcp'],
                'action': ['accept'],
            },
            'ct state new log prefix "syslog"',
        ),
        (
            {
                'name': 'logging_disabled',
                'option': [],
                'counter': None,
                'logging': ['disable'],
                'protocol': ['tcp'],
                'action': ['accept'],
            },
            'ct state new',
        ),
        (
            {
                'name': 'combo_logging_tcp_established',
                'option': ['tcp-established'],
                'counter': None,
                'logging': ['true'],
                'protocol': ['tcp'],
                'action': ['accept'],
            },
            'ct state new log prefix "combo_logging_tcp_established"',
        ),
        (
            {
                'name': 'combo_cnt_log_established',
                'option': ['tcp-established'],
                'counter': 'whatever-name-you-want',
                'logging': ['true'],
                'protocol': ['tcp'],
                'action': ['deny'],
            },
            'log prefix "combo_cnt_log_established" counter',
        ),
    )
    def testOptionsHandler(self, term_dict, expected_output):
        term = DictObj(term_dict)
        result = self.dummyterm._OptionsHandler(term)
        self.assertEqual(result, expected_output)

    def testBuildTokens(self):
        pol1 = nftables.Nftables(
            policy.ParsePolicy(GOOD_HEADER_1 + GOOD_TERM_1, self.naming), EXP_INFO
        )
        st, sst = pol1._BuildTokens()
        self.assertEqual(st, SUPPORTED_TOKENS)
        self.assertEqual(sst, SUPPORTED_SUB_TOKENS)

    @parameterized.parameters(
        (
            ESTABLISHED_OPTION_TERM,
            'WARNING: Term established-term is a established term and will not be rendered.',
        ),
        (
            TCP_ESTABLISHED_OPTION_TERM,
            'WARNING: Term tcp-established-term is a tcp-established term and will not be rendered.',
        ),
    )
    def testSkippedTerm(self, termdata, messagetxt):
        with self.assertLogs() as ctx:
            # run a policy object expected to be skipped and logged.
            nft = nftables.Nftables(
                policy.ParsePolicy(GOOD_HEADER_1 + termdata, self.naming), EXP_INFO
            )
        # self.assertEqual(len(ctx.records), 2)
        record = ctx.records[1]
        self.assertEqual(record.message, messagetxt)


if __name__ == '__main__':
    absltest.main()
