# Copyright 2013 Google Inc. All Rights Reserved.
#
# 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 Ipset rendering module."""

from absl.testing import absltest
from unittest import mock

from capirca.lib import ipset
from capirca.lib import nacaddr
from capirca.lib import naming
from capirca.lib import policy


GOOD_HEADER_1 = """
header {
  comment:: "this is a test acl"
  target:: ipset OUTPUT DROP
}
"""

GOOD_HEADER_2 = """
header {
  comment:: "this is a test acl"
  target:: ipset OUTPUT DROP exists
}
"""

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

GOOD_TERM_2 = """
term good-term-2 {
  destination-address:: EXTERNAL
  action:: accept
}
"""

GOOD_TERM_3 = """
term good-term-3 {
  source-address:: INTERNAL
  destination-address:: EXTERNAL
  action:: accept
}
"""

GOOD_TERM_4 = """
term good-term-4 {
  source-address:: INTERNAL
  destination-address:: EXTERNAL
  policer:: batman
  action:: accept
}
"""

SUPPORTED_TOKENS = {
    'action',
    'comment',
    'counter',
    'destination_address',
    'destination_address_exclude',
    'destination_interface',
    'destination_port',
    'destination_prefix',
    'expiration',
    'fragment_offset',
    'icmp_code',
    'icmp_type',
    'stateless_reply',
    'logging',
    'log_limit',
    'name',
    'option',
    'owner',
    'packet_length',
    'platform',
    'platform_exclude',
    'protocol',
    'routing_instance',
    'source_address',
    'source_address_exclude',
    'source_interface',
    'source_port',
    'source_prefix',
    'translated',
    'verbatim',
}

SUPPORTED_SUB_TOKENS = {
    'action': {'accept', 'deny', 'reject', 'next', 'reject-with-tcp-rst'},
    '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',
    },
    'option': {'established',
               'first-fragment',
               'initial',
               'sample',
               'tcp-established',
               'tcp-initial',
               'syn',
               'ack',
               'fin',
               'rst',
               'urg',
               'psh',
               'all',
               'none'}
}

# 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


class IpsetTest(absltest.TestCase):

  def setUp(self):
    super().setUp()
    self.naming = mock.create_autospec(naming.Naming)

  def testMarkers(self):
    self.naming.GetNetAddr.return_value = [nacaddr.IPv4('10.0.0.0/8')]

    acl = ipset.Ipset(policy.ParsePolicy(GOOD_HEADER_1 + GOOD_TERM_1,
                                         self.naming), EXP_INFO)
    result = str(acl)
    self.assertIn('# begin:ipset-rules', result)
    self.assertIn('# end:ipset-rules', result)

    self.naming.GetNetAddr.assert_called_once_with('INTERNAL')

  def testGenerateSetName(self):
    # iptables superclass currently limits term name length to 26 characters,
    # but that could change
    policy_term = mock.MagicMock()
    policy_term.name = 'filter_name'
    policy_term.protocol = ['tcp']
    term = ipset.Term(policy_term, 'filter_name', False, None)
    self.assertEqual(term._GenerateSetName('good-term-1', 'src'),
                     'good-term-1-src')
    self.assertEqual(term._GenerateSetName('good-but-way-too-long-term-name',
                                           'src'),
                     'good-but-way-too-long-term--src')
    term = ipset.Term(policy_term, 'filter_name', False, None, 'inet6')
    self.assertEqual(term._GenerateSetName('good-term-1', 'src'),
                     'good-term-1-src-v6')
    self.assertEqual(term._GenerateSetName('good-but-way-too-long-term-name',
                                           'src'),
                     'good-but-way-too-long-te-src-v6')

  def testOneSourceAddress(self):
    self.naming.GetNetAddr.return_value = [nacaddr.IPv4('10.0.0.0/8')]

    acl = ipset.Ipset(policy.ParsePolicy(GOOD_HEADER_1 + GOOD_TERM_1,
                                         self.naming), EXP_INFO)
    result = str(acl)
    self.assertIn('-s 10.0.0.0/8', result)
    self.assertNotIn('-m set --match-set good-term-3-src src', result)

    self.naming.GetNetAddr.assert_called_once_with('INTERNAL')

  def testOneDestinationAddress(self):
    self.naming.GetNetAddr.return_value = [nacaddr.IPv4('172.16.0.0/12')]

    acl = ipset.Ipset(policy.ParsePolicy(GOOD_HEADER_1 + GOOD_TERM_2,
                                         self.naming), EXP_INFO)
    result = str(acl)
    self.assertIn('-d 172.16.0.0/12', result)
    self.assertNotIn('-m set --match-set good-term-2-dst dst', result)

    self.naming.GetNetAddr.assert_called_once_with('EXTERNAL')

  def testOneSourceAndDestinationAddress(self):
    self.naming.GetNetAddr.side_effect = [
        [nacaddr.IPv4('10.0.0.0/8')],
        [nacaddr.IPv4('172.16.0.0/12')]]

    acl = ipset.Ipset(policy.ParsePolicy(GOOD_HEADER_1 + GOOD_TERM_3,
                                         self.naming), EXP_INFO)
    result = str(acl)
    self.assertIn('-s 10.0.0.0/8', result)
    self.assertIn('-d 172.16.0.0/12', result)
    self.assertNotIn('-m set --match-set good-term-3-src src', result)
    self.assertNotIn('-m set --match-set good-term-3-dst dst', result)

    self.naming.GetNetAddr.assert_has_calls([
        mock.call('INTERNAL'),
        mock.call('EXTERNAL')])

  def testManySourceAddresses(self):
    self.naming.GetNetAddr.return_value = [
        nacaddr.IPv4('10.0.0.0/24'), nacaddr.IPv4('10.1.0.0/24')]

    acl = ipset.Ipset(policy.ParsePolicy(GOOD_HEADER_1 + GOOD_TERM_1,
                                         self.naming), EXP_INFO)
    result = str(acl)
    self.assertIn('create good-term-1-src hash:net family inet hashsize'
                  ' 4 maxelem 4', result)
    self.assertIn('add good-term-1-src 10.0.0.0/24', result)
    self.assertIn('add good-term-1-src 10.1.0.0/24', result)
    self.assertIn('-m set --match-set good-term-1-src src', result)
    self.assertNotIn('-s ', result)
    self.assertNotIn('-exist', result)

    self.naming.GetNetAddr.assert_called_once_with('INTERNAL')

  def testManyDestinationAddresses(self):
    self.naming.GetNetAddr.return_value = [
        nacaddr.IPv4('172.16.0.0/24'), nacaddr.IPv4('172.17.0.0/24')]

    acl = ipset.Ipset(policy.ParsePolicy(GOOD_HEADER_1 + GOOD_TERM_2,
                                         self.naming), EXP_INFO)
    result = str(acl)
    self.assertIn('create good-term-2-dst hash:net family inet hashsize '
                  '4 maxelem 4', result)
    self.assertIn('add good-term-2-dst 172.16.0.0/24', result)
    self.assertIn('add good-term-2-dst 172.17.0.0/24', result)
    self.assertIn('-m set --match-set good-term-2-dst dst', result)
    self.assertNotIn('-s ', result)
    self.assertNotIn('-exist', result)

    self.naming.GetNetAddr.assert_called_once_with('EXTERNAL')

  def testManySourceAndDestinationAddresses(self):
    self.naming.GetNetAddr.side_effect = [
        [nacaddr.IPv4('10.0.0.0/24'), nacaddr.IPv4('10.1.0.0/24')],
        [nacaddr.IPv4('172.16.0.0/24'), nacaddr.IPv4('172.17.0.0/24')]]

    acl = ipset.Ipset(policy.ParsePolicy(GOOD_HEADER_1 + GOOD_TERM_3,
                                         self.naming), EXP_INFO)
    result = str(acl)
    self.assertIn('create good-term-3-src hash:net family inet hashsize '
                  '4 maxelem 4', result)
    self.assertIn('create good-term-3-dst hash:net family inet hashsize '
                  '4 maxelem 4', result)
    self.assertIn('add good-term-3-src 10.0.0.0/24', result)
    self.assertIn('add good-term-3-src 10.1.0.0/24', result)
    self.assertIn('add good-term-3-dst 172.16.0.0/24', result)
    self.assertIn('add good-term-3-dst 172.17.0.0/24', result)
    self.assertIn('-m set --match-set good-term-3-src src', result)
    self.assertIn('-m set --match-set good-term-3-dst dst', result)
    self.assertNotIn('-s ', result)
    self.assertNotIn('-d ', result)

    self.naming.GetNetAddr.assert_has_calls([
        mock.call('INTERNAL'),
        mock.call('EXTERNAL')])

  def testBuildTokens(self):
    pol1 = ipset.Ipset(policy.ParsePolicy(GOOD_HEADER_1 + GOOD_TERM_3,
                                          self.naming), EXP_INFO)
    st, sst = pol1._BuildTokens()
    self.assertEqual(st, SUPPORTED_TOKENS)
    self.assertEqual(sst, SUPPORTED_SUB_TOKENS)

  def testBuildWarningTokens(self):
    pol1 = ipset.Ipset(policy.ParsePolicy(GOOD_HEADER_1 + GOOD_TERM_4,
                                          self.naming), EXP_INFO)
    st, sst = pol1._BuildTokens()
    self.assertEqual(st, SUPPORTED_TOKENS)
    self.assertEqual(sst, SUPPORTED_SUB_TOKENS)

  def testAddsExistsOption(self):
    self.naming.GetNetAddr.return_value = [
        nacaddr.IPv4('10.0.0.0/24'), nacaddr.IPv4('10.1.0.0/24')]
    acl = ipset.Ipset(policy.ParsePolicy(GOOD_HEADER_2 + GOOD_TERM_1,
                                         self.naming), EXP_INFO)
    self.assertIn('create -exist', str(acl))
    self.assertIn('add -exist', str(acl))


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