# Copyright 2014 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 Nftables rendering module."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import datetime
import unittest

from absl import logging
from capirca.lib import aclgenerator
from capirca.lib import nacaddr
from capirca.lib import nftables
from capirca.lib import policy

import mock


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

GOOD_HEADER_1 = """
header {
  target:: nftables chain_name input 0 inet
}
"""

GOOD_HEADER_2 = """
header {
  target:: nftables chain_name input 0 inet6
}
"""

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

IPV6_TERM_1 = """
term inet6-icmp {
  protocol:: icmpv6
  icmp-type:: destination-unreachable time-exceeded echo-reply
  action:: deny
}
"""

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

EXPIRED_TERM = """
term is_expired {
    expiration:: 2001-01-01
    action:: accept
}
"""

EXPIRING_TERM = """
term is_expiring {
  expiration:: %s
  action:: accept
}
"""

GOOD_TERM_2 = """
term good-term-2 {
  source-address:: SOURCE_NETWORK
  destination-address:: DESTINATION_NETWORK
  action:: accept
}
"""

GOOD_TERM_3 = """
term good-term-3 {
  source-address:: SOURCE_NETWORK
  destination-address:: DESTINATION_NETWORK
  action:: %s
}
"""

GOOD_TERM_4 = """
term good-term-4 {
  action:: accept
  comment:: "comment first line"
  comment:: "comment second line"
  owner:: owner@enterprise.com
}
"""

GOOD_TERM_5 = """
term good-term-5 {
  protocol:: ah esp
  action:: accept
}
"""

GOOD_TERM_6 = """
term good-term-6 {
  protocol:: ah
  action:: accept
}
"""

GOOD_TERM_7 = """
term good-term-7 {
  destination-port:: SMTP
  protocol:: tcp
  action:: accept
}
"""

GOOD_TERM_8 = """
term good-term-8 {
  source-port:: SMTP
  protocol:: tcp
  action:: accept
}
"""

GOOD_TERM_9 = """
term good-term-9 {
  protocol:: icmp
  icmp-type:: echo-request echo-reply
  action:: accept
}
"""

GOOD_TERM_10 = """
term good-term-10 {
  verbatim:: nftables "mary had a little lamb"
  verbatim:: cisco "mary had second lamb"
  verbatim:: juniper "mary had third lamb"
}
"""

GOOD_TERM_11 = """
term good-term-11 {
  source-address:: SOURCE_NETWORK
  source-exclude:: SOURCE_EXCLUDE_NETWORK
  destination-address:: DESTINATION_NETWORK
  destination-exclude:: DESTINATION_EXCLUDE_NETWORK
  action:: accept
}
"""

GOOD_TERM_12 = """
term good-term-12 {
  policer:: batman
  action:: accept
}
"""

GOOD_TERM_13 = """
term good-term-13 {
  protocol:: icmp
  action:: accept
  logging:: true
}
"""

GOOD_TERM_14 = """
term good-term-14 {
  protocol:: icmp
  action:: accept
  log_name:: "my log prefix"
}
"""

GOOD_TERM_15 = """
term good-term-15 {
  protocol:: icmp
  action:: accept
  logging:: true
  log_name:: "my log prefix"
}
"""
GOOD_TERM_16 = """
term good-term-16 {
  protocol:: icmp
  action:: accept
  counter:: string_content_unused
}
"""

GOOD_TERM_17 = """
term good-term-17 {
  action:: accept
  comment:: "%(long_line)s:"
}
""" % {'long_line': 'A' * 128}

SUPPORTED_TOKENS = {
    'action',
    'comment',
    'counter',
    'destination_address',
    'destination_address_exclude',
    'destination_port',
    'expiration',
    'icmp_type',
    'stateless_reply',
    'logging',
    'log_name',
    'name',
    'option',
    'owner',
    'platform',
    'platform_exclude',
    'protocol',
    'source_address',
    'source_address_exclude',
    'source_port',
    'translated',
    'verbatim',
}

SUPPORTED_SUB_TOKENS = {
    'action': {
        'accept',
        'deny',
        'next',
        'reject',
        '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',
    },
}

# 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 NftablesTest(unittest.TestCase):

  def setUp(self):
    super(NftablesTest, self).setUp()
    self.mock_naming = mock.MagicMock()

  def testBadHeader(self):
    cases = [
        'chain_name input 0 inet extraneous_target_option',
        'chain_name input 0 unsupported_af',
        'chain_name input not_an_int_priority',
        'chain_name invalid_hook_name 0',
        'chain_name input'
        'chain_name',
        '',
    ]
    for case in cases:
      logging.info('Testing bad header case %s.', case)
      header = BAD_HEADER % case
      pol = policy.ParsePolicy(header + GOOD_TERM_1, self.mock_naming)
      self.assertRaises(nftables.InvalidTargetOption,
                        nftables.Nftables.__init__,
                        nftables.Nftables.__new__(nftables.Nftables),
                        pol, EXP_INFO)

  def testBadAddressFamily(self):
    cases = [
        'chain_name input 0 mixed',
    ]
    for case in cases:
      logging.info('Testing bad address family case %s.', case)
      header = BAD_HEADER % case
      pol = policy.ParsePolicy(header + GOOD_TERM_1, self.mock_naming)
      self.assertRaises(aclgenerator.UnsupportedAFError,
                        nftables.Nftables.__init__,
                        nftables.Nftables.__new__(nftables.Nftables),
                        pol, EXP_INFO)

  def testGoodHeader(self):
    nftables.Nftables(policy.ParsePolicy(GOOD_HEADER_1 + GOOD_TERM_1,
                                         self.mock_naming), EXP_INFO)
    nft = str(nftables.Nftables(policy.ParsePolicy(GOOD_HEADER_1 + GOOD_TERM_1 +
                                                   GOOD_HEADER_2 + IPV6_TERM_2,
                                                   self.mock_naming),
                                EXP_INFO))
    self.assertIn('flush table ip table_filter', nft)
    self.assertIn(
        'table ip table_filter {\n\tchain chain_name {\n\t\ttype filter '
        'hook input priority 0;\n\t\taccept\n\t}\n}', nft)
    self.assertIn('flush table ip6 table_filter', nft)
    self.assertIn(
        'table ip6 table_filter {\n\tchain chain_name {\n\t\ttype filter '
        'hook input priority 0;\n\t\tdrop\n\t}\n}', nft)

  @mock.patch.object(logging, 'warning')
  def testExpired(self, mock_logging_warn):
    nftables.Nftables(policy.ParsePolicy(GOOD_HEADER_1 + EXPIRED_TERM,
                                         self.mock_naming), EXP_INFO)
    mock_logging_warn.assert_called_once_with('Term %s in policy %s is expired '
                                              'and will not be rendered.',
                                              'is_expired', 'chain_name')

  @mock.patch.object(logging, 'info')
  def testExpiring(self, mock_logging_info):
    exp_date = datetime.date.today() + datetime.timedelta(weeks=EXP_INFO)
    nftables.Nftables(policy.ParsePolicy(GOOD_HEADER_1 + EXPIRING_TERM %
                                         exp_date.strftime('%Y-%m-%d'),
                                         self.mock_naming), EXP_INFO)
    mock_logging_info.assert_called_once_with('Term %s in policy %s '
                                              'expires in less than %d weeks.',
                                              'is_expiring', 'chain_name',
                                              EXP_INFO)

  @mock.patch.object(logging, 'debug')
  def testIcmpv6InetMismatch(self, mock_logging_debug):
    str(nftables.Nftables(policy.ParsePolicy(GOOD_HEADER_1 + IPV6_TERM_1,
                                             self.mock_naming), EXP_INFO))
    mock_logging_debug.assert_called_once_with('Term inet6-icmp will not be '
                                               'rendered, as it has '
                                               'icmpv6 match specified '
                                               'but the ACL is of inet address '
                                               'family.')

  def testSingleSourceDestIp(self):
    source_network = [nacaddr.IPv4('172.16.0.0/24')]
    destination_network = [nacaddr.IPv4('10.0.0.0/24')]
    self.mock_naming.GetNetAddr.side_effect = [source_network,
                                               destination_network]
    nft = str(nftables.Nftables(policy.ParsePolicy(GOOD_HEADER_1 + GOOD_TERM_2,
                                                   self.mock_naming), EXP_INFO))
    self.assertIn('ip saddr 172.16.0.0/24 ip daddr 10.0.0.0/24', nft)

  def testMultipleSourceDestIp(self):
    source_network = [nacaddr.IPv4('172.16.0.0/24'),
                      nacaddr.IPv4('172.16.2.0/24')]
    destination_network = [nacaddr.IPv4('10.0.0.0/24'),
                           nacaddr.IPv4('10.0.2.0/24')]
    self.mock_naming.GetNetAddr.side_effect = [source_network,
                                               destination_network]
    nft = str(nftables.Nftables(policy.ParsePolicy(GOOD_HEADER_1 + GOOD_TERM_2,
                                                   self.mock_naming), EXP_INFO))
    self.assertIn('ip saddr { 172.16.0.0/24, 172.16.2.0/24} ip daddr '
                  '{ 10.0.0.0/24, 10.0.2.0/24}', nft)

  def testSingleProtocol(self):
    nft = str(nftables.Nftables(policy.ParsePolicy(GOOD_HEADER_1 + GOOD_TERM_5,
                                                   self.mock_naming), EXP_INFO))
    self.assertIn('ip protocol { ah, esp}', nft)

  def testMultiProtocol(self):
    nft = str(nftables.Nftables(policy.ParsePolicy(GOOD_HEADER_1 + GOOD_TERM_6,
                                                   self.mock_naming), EXP_INFO))
    self.assertIn('ip protocol ah', nft)

  def testSingleDport(self):
    destination_ports = ['25']
    self.mock_naming.GetServiceByProto.return_value = destination_ports
    nft = str(nftables.Nftables(policy.ParsePolicy(GOOD_HEADER_1 + GOOD_TERM_7,
                                                   self.mock_naming), EXP_INFO))
    self.assertIn('dport 25', nft)

  def testMultiDport(self):
    destination_ports = ['25', '80', '6610', '6611', '6612']
    self.mock_naming.GetServiceByProto.return_value = destination_ports
    nft = str(nftables.Nftables(policy.ParsePolicy(GOOD_HEADER_1 + GOOD_TERM_7,
                                                   self.mock_naming), EXP_INFO))
    self.assertIn('dport { 25, 80, 6610-6612}', nft)

  def testSingleSport(self):
    source_ports = ['25']
    self.mock_naming.GetServiceByProto.return_value = source_ports
    nft = str(nftables.Nftables(policy.ParsePolicy(GOOD_HEADER_1 + GOOD_TERM_8,
                                                   self.mock_naming), EXP_INFO))
    self.assertIn('sport 25', nft)

  def testMultiSport(self):
    source_ports = ['25', '80', '6610', '6611', '6612']
    self.mock_naming.GetServiceByProto.return_value = source_ports
    nft = str(nftables.Nftables(policy.ParsePolicy(GOOD_HEADER_1 + GOOD_TERM_8,
                                                   self.mock_naming), EXP_INFO))
    self.assertIn('sport { 25, 80, 6610-6612}', nft)

  def testIcmpType(self):
    nft = str(nftables.Nftables(policy.ParsePolicy(GOOD_HEADER_1 + GOOD_TERM_9,
                                                   self.mock_naming), EXP_INFO))
    self.assertIn('icmp type { echo-reply, echo-request}', nft)

  def testAction(self):
    cases = {'accept': 'accept',
             'deny': 'drop',
             'reject': 'reject',
             'next': 'continue',
             'reject-with-tcp-rst': 'reject with tcp reset'}
    for case in cases:
      logging.info('Testing action case %s.', case)
      nft = str(nftables.Nftables(policy.ParsePolicy(GOOD_HEADER_1 +
                                                     GOOD_TERM_3 % case,
                                                     self.mock_naming),
                                  EXP_INFO))
      self.assertIn(cases[case], nft)

  def testCommentOwner(self):
    nft = str(nftables.Nftables(policy.ParsePolicy(GOOD_HEADER_1 + GOOD_TERM_4,
                                                   self.mock_naming), EXP_INFO))
    self.assertIn('comment "comment first line comment second line '
                  'Owner: owner@enterprise.com"', nft)

  @mock.patch.object(logging, 'warning')
  def testCommentTruncate(self, mock_logging_warn):
    nft = str(nftables.Nftables(policy.ParsePolicy(GOOD_HEADER_1 + GOOD_TERM_17,
                                                   self.mock_naming), EXP_INFO))
    mock_logging_warn.assert_called_once_with(
        'Term %s in policy is too long (>%d characters) and will be'
        ' truncated', 'good-term-17', nftables.Term.MAX_CHARACTERS)
    # Ensure that the truncate did happen and stripped off the ':'
    self.assertIn('comment "%(long_line)s' % {'long_line': 'A' * 127}, nft)

  def testLogTerm(self):
    nft = str(nftables.Nftables(policy.ParsePolicy(GOOD_HEADER_1 + GOOD_TERM_13,
                                                   self.mock_naming), EXP_INFO))
    self.assertIn(' icmp log accept', nft)

  def testLogNameTerm(self):
    nft = str(nftables.Nftables(policy.ParsePolicy(GOOD_HEADER_1 + GOOD_TERM_14,
                                                   self.mock_naming), EXP_INFO))
    self.assertIn('log prefix "my log prefix: " ', nft)

  def testLogAndLogNameTerm(self):
    nft = str(nftables.Nftables(policy.ParsePolicy(GOOD_HEADER_1 + GOOD_TERM_15,
                                                   self.mock_naming), EXP_INFO))
    self.assertIn('log prefix "my log prefix: " ', nft)

  def testCounterTerm(self):
    nft = str(nftables.Nftables(policy.ParsePolicy(GOOD_HEADER_1 + GOOD_TERM_16,
                                                   self.mock_naming), EXP_INFO))
    self.assertIn(' icmp counter accept', nft)

  def testVerbatimTerm(self):
    nft = str(nftables.Nftables(policy.ParsePolicy(GOOD_HEADER_1 + GOOD_TERM_10,
                                                   self.mock_naming), EXP_INFO))
    self.assertIn('mary had a little lamb', nft)
    # check if another platforms verbatim shows up
    self.assertNotIn('mary had a second lamb', nft)
    self.assertNotIn('mary had a third lamb', nft)

  def testSourceDestExclude(self):
    source_network = [nacaddr.IPv4('192.168.0.0/24')]
    source_exclude_network = [nacaddr.IPv4('192.168.0.0/27')]
    destination_network = [nacaddr.IPv4('10.0.0.0/24')]
    destination_exclude_network = [nacaddr.IPv4('10.0.0.0/27')]
    self.mock_naming.GetNetAddr.side_effect = [source_network,
                                               source_exclude_network,
                                               destination_network,
                                               destination_exclude_network]
    nft = str(nftables.Nftables(policy.ParsePolicy(GOOD_HEADER_1 + GOOD_TERM_11,
                                                   self.mock_naming), EXP_INFO))
    self.assertIn('ip saddr { 192.168.0.32/27, 192.168.0.64/26, '
                  '192.168.0.128/25}', nft)
    self.assertIn('ip daddr { 10.0.0.32/27, 10.0.0.64/26, '
                  '10.0.0.128/25}', nft)

  def testSourceDestExcludeFromAllIps(self):
    source_network = []
    source_exclude_network = [nacaddr.IPv4('192.168.0.0/27')]
    destination_network = []
    destination_exclude_network = [nacaddr.IPv4('10.0.0.0/27')]
    self.mock_naming.GetNetAddr.side_effect = [source_network,
                                               source_exclude_network,
                                               destination_network,
                                               destination_exclude_network]
    nft = str(nftables.Nftables(policy.ParsePolicy(GOOD_HEADER_1 + GOOD_TERM_11,
                                                   self.mock_naming), EXP_INFO))
    self.assertIn('ip saddr { 0.0.0.0/1, 128.0.0.0/2, 192.0.0.0/9, '
                  '192.128.0.0/11, 192.160.0.0/13, 192.168.0.32/27, '
                  '192.168.0.64/26, 192.168.0.128/25, 192.168.1.0/24, '
                  '192.168.2.0/23, 192.168.4.0/22, 192.168.8.0/21, '
                  '192.168.16.0/20, 192.168.32.0/19, 192.168.64.0/18, '
                  '192.168.128.0/17, 192.169.0.0/16, 192.170.0.0/15, '
                  '192.172.0.0/14, 192.176.0.0/12, 192.192.0.0/10, '
                  '193.0.0.0/8, 194.0.0.0/7, 196.0.0.0/6, 200.0.0.0/5, '
                  '208.0.0.0/4, 224.0.0.0/3}', nft)
    self.assertIn('ip daddr { 0.0.0.0/5, 8.0.0.0/7, 10.0.0.32/27, '
                  '10.0.0.64/26, 10.0.0.128/25, 10.0.1.0/24, 10.0.2.0/23, '
                  '10.0.4.0/22, 10.0.8.0/21, 10.0.16.0/20, 10.0.32.0/19, '
                  '10.0.64.0/18, 10.0.128.0/17, 10.1.0.0/16, 10.2.0.0/15, '
                  '10.4.0.0/14, 10.8.0.0/13, 10.16.0.0/12, 10.32.0.0/11, '
                  '10.64.0.0/10, 10.128.0.0/9, 11.0.0.0/8, 12.0.0.0/6, '
                  '16.0.0.0/4, 32.0.0.0/3, 64.0.0.0/2, 128.0.0.0/1}', nft)

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

  def testBuildWarningTokens(self):
    pol1 = nftables.Nftables(policy.ParsePolicy(GOOD_HEADER_1 + GOOD_TERM_12,
                                                self.mock_naming), EXP_INFO)
    st, sst = pol1._BuildTokens()
    self.assertEqual(st, SUPPORTED_TOKENS)
    self.assertEqual(sst, SUPPORTED_SUB_TOKENS)


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