# Copyright 2007 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 nacaddr.py module."""

from absl.testing import absltest

from capirca.lib import nacaddr


class NacaddrUnitTest(absltest.TestCase):
  """Unit Test for nacaddr.py.

     nacaddr class extends ipaddr by adding .text fields to allow
     comments for each of the IPv4 and IPv6 classes.
  """

  def setUp(self):
    super().setUp()
    self.addr1 = nacaddr.IPv4(u'10.0.0.0/8', 'The 10 block')
    self.addr2 = nacaddr.IPv6('DEAD:BEEF:BABE:FACE:DEAF:FEED:C0DE:F001/64',
                              'An IPv6 Address', strict=False)

  def testCollapsing(self):
    ip1 = nacaddr.IPv4(u'1.1.0.0/24', 'foo')
    ip2 = nacaddr.IPv4(u'1.1.1.0/24', 'foo')
    ip3 = nacaddr.IPv4(u'1.1.2.0/24', 'baz')
    ip4 = nacaddr.IPv4(u'1.1.3.0/24')
    ip5 = nacaddr.IPv4(u'1.1.4.0/24')

    # stored in no particular order b/c we want CollapseAddr to call [].sort
    # and we want that sort to call nacaddr.IP.__cmp__() on our array members
    ip6 = nacaddr.IPv4(u'1.1.0.0/22')

    # check that addreses are subsumed properlly.
    collapsed = nacaddr.CollapseAddrList([ip1, ip2, ip3, ip4, ip5, ip6])
    self.assertEqual(len(collapsed), 2)
    # test that the comments are collapsed properlly, and that comments aren't
    # added to addresses that have no comments.
    self.assertListEqual([collapsed[0].text, collapsed[1].text],
                         ['foo, baz', ''])
    self.assertListEqual(collapsed, [nacaddr.IPv4(u'1.1.0.0/22'),
                                     nacaddr.IPv4(u'1.1.4.0/24')])

    # test that two addresses are supernet'ed properlly
    collapsed = nacaddr.CollapseAddrList([ip1, ip2])
    self.assertEqual(len(collapsed), 1)
    self.assertEqual(collapsed[0].text, 'foo')
    self.assertListEqual(collapsed, [nacaddr.IPv4(u'1.1.0.0/23')])

    ip_same1 = ip_same2 = nacaddr.IPv4(u'1.1.1.1/32')
    self.assertListEqual(nacaddr.CollapseAddrList([ip_same1, ip_same2]),
                         [ip_same1])
    ip1 = nacaddr.IPv6(u'::2001:1/100', strict=False)
    ip2 = nacaddr.IPv6(u'::2002:1/120', strict=False)
    ip3 = nacaddr.IPv6(u'::2001:1/96', strict=False)
    # test that ipv6 addresses are subsumed properlly.
    collapsed = nacaddr.CollapseAddrList([ip1, ip2, ip3])
    self.assertListEqual(collapsed, [ip3])

  def testNacaddrV4Comment(self):
    self.assertEqual(self.addr1.text, 'The 10 block')

  def testNacaddrV6Comment(self):
    self.assertEqual(self.addr2.text, 'An IPv6 Address')

  def testSupernetting(self):
    self.assertEqual(self.addr1.Supernet().text, 'The 10 block')
    self.assertEqual(self.addr2.Supernet().text, 'An IPv6 Address')
    self.assertEqual(self.addr1.Supernet().prefixlen, 7)
    self.assertEqual(self.addr2.Supernet().prefixlen, 63)

    token_ip = nacaddr.IP('1.1.1.0/24', token='FOO_TOKEN')
    self.assertEqual(token_ip.Supernet().token, 'FOO_TOKEN')
    self.assertEqual(nacaddr.IP('0.0.0.0/0').Supernet(),
                     nacaddr.IP('0.0.0.0/0'))
    self.assertEqual(nacaddr.IP('::0/0').Supernet(), nacaddr.IP('::0/0'))

    self.assertRaises(nacaddr.PrefixlenDiffInvalidError,
                      nacaddr.IP('1.1.1.0/24').Supernet, 25)
    self.assertRaises(nacaddr.PrefixlenDiffInvalidError,
                      nacaddr.IP('::1/64', strict=False).Supernet, 65)

  def testAddressListExclusion(self):
    a1 = nacaddr.IPv4('1.1.1.0/24')
    a2 = nacaddr.IPv4('10.0.0.0/24')
    b1 = nacaddr.IPv4('1.1.1.1/32')
    b2 = nacaddr.IPv4('10.0.0.25/32')
    b3 = nacaddr.IPv4('192.168.0.0/16')

    expected = [nacaddr.IPv4('1.1.1.0/32'), nacaddr.IPv4('1.1.1.2/31'),
                nacaddr.IPv4('1.1.1.4/30'), nacaddr.IPv4('1.1.1.8/29'),
                nacaddr.IPv4('1.1.1.16/28'), nacaddr.IPv4('1.1.1.32/27'),
                nacaddr.IPv4('1.1.1.64/26'), nacaddr.IPv4('1.1.1.128/25'),
                nacaddr.IPv4('10.0.0.0/28'), nacaddr.IPv4('10.0.0.16/29'),
                nacaddr.IPv4('10.0.0.24/32'), nacaddr.IPv4('10.0.0.26/31'),
                nacaddr.IPv4('10.0.0.28/30'), nacaddr.IPv4('10.0.0.32/27'),
                nacaddr.IPv4('10.0.0.64/26'), nacaddr.IPv4('10.0.0.128/25')]
    self.assertListEqual(nacaddr.AddressListExclude([a1, a2], [b1, b2, b3]),
                         expected)
    # [1,2,3] + [4,5,6] = [1,2,3,4,5,6].  this is basically the same test as
    # above but i think it's a little more readable
    expected_two = list(a1.address_exclude(b1))
    expected_two.extend(a2.address_exclude(b2))
    self.assertListEqual(nacaddr.AddressListExclude([a1, a2], [b1, b2, b3]),
                         sorted(expected_two))

  def testComplexAddressListExcludesion(self):
    # this is a big fugly test. there was a bug in AddressListExclude
    # which manifested itself when more than one member of the excludes
    # list was a part of the same superset token.
    #
    # for example, it used to be like so:
    #  excludes = ['1.1.1.1/32', '1.1.1.2/32']
    #  superset = ['1.1.1.0/30']
    #
    # '1.1.1.0/30'.AddressExclude('1.1.1.1/32') ->
    #    ['1.1.1.0/32', '1.1.1.2/32', '1.1.1.3/32']
    # '1.1.1.0/30'.AddressExclude('1.1.1.2/32') ->
    #    ['1.1.1.0/32', '1.1.1.1/32', '1.1.1.3/32']
    #
    # yet combining those two results gives you
    #   ['1.1.1.0/32', '1.1.1.1/32', '1.1.1.2/32' '1.1.1.3/32'], or
    #   '1.1.1.0/30', which clearly isn't right.

    excludes = [nacaddr.IPv4('10.0.0.0/23'), nacaddr.IPv4('10.1.0.0/16')]
    superset = [nacaddr.IPv4('0.0.0.0/0')]

    expected = [nacaddr.IPv4('0.0.0.0/5'), nacaddr.IPv4('8.0.0.0/7'),
                nacaddr.IPv4('10.0.2.0/23'), nacaddr.IPv4('10.0.4.0/22'),
                nacaddr.IPv4('10.0.8.0/21'), nacaddr.IPv4('10.0.16.0/20'),
                nacaddr.IPv4('10.0.32.0/19'), nacaddr.IPv4('10.0.64.0/18'),
                nacaddr.IPv4('10.0.128.0/17'), nacaddr.IPv4('10.2.0.0/15'),
                nacaddr.IPv4('10.4.0.0/14'), nacaddr.IPv4('10.8.0.0/13'),
                nacaddr.IPv4('10.16.0.0/12'), nacaddr.IPv4('10.32.0.0/11'),
                nacaddr.IPv4('10.64.0.0/10'), nacaddr.IPv4('10.128.0.0/9'),
                nacaddr.IPv4('11.0.0.0/8'), nacaddr.IPv4('12.0.0.0/6'),
                nacaddr.IPv4('16.0.0.0/4'), nacaddr.IPv4('32.0.0.0/3'),
                nacaddr.IPv4('64.0.0.0/2'), nacaddr.IPv4('128.0.0.0/1')]

    self.assertListEqual(nacaddr.AddressListExclude(superset, excludes),
                         expected)

  def testAddressListExcludeCaseOne(self):
    # Small block eliminated by large block, and an extra block that stays.
    # For both IP versions.
    superset = [nacaddr.IPv4('200.0.0.0/24'), nacaddr.IPv4('10.1.0.0/24'),
                nacaddr.IPv6('200::/56'), nacaddr.IPv6('10:1::/56')]
    excludes = [nacaddr.IPv6('10::/16'), nacaddr.IPv4('10.0.0.0/8')]
    expected = [nacaddr.IPv4('200.0.0.0/24'), nacaddr.IPv6('200::/56')]
    self.assertListEqual(nacaddr.AddressListExclude(superset, excludes),
                         expected)

  def testAddressListExcludeCaseTwo(self):
    # Two blocks out of the middle of a large block.
    superset = [nacaddr.IPv4('200.0.0.0/24'), nacaddr.IPv4('10.0.0.0/8'),
                nacaddr.IPv6('200::/56'), nacaddr.IPv6('10::/16')]
    excludes = [nacaddr.IPv6('10:8000::/18'), nacaddr.IPv6('10:4000::/18'),
                nacaddr.IPv4('10.128.0.0/10'), nacaddr.IPv4('10.64.0.0/10')]
    expected = [nacaddr.IPv4('10.0.0.0/10'), nacaddr.IPv4('10.192.0.0/10'),
                nacaddr.IPv4('200.0.0.0/24'),
                nacaddr.IPv6('10::/18'), nacaddr.IPv6('10:c000::/18'),
                nacaddr.IPv6('200::/56')]
    self.assertListEqual(nacaddr.AddressListExclude(superset, excludes),
                         expected)

  def testAddressListExcludeCaseThree(self):
    # Two blocks off both ends of a large block.
    superset = [nacaddr.IPv4('200.0.0.0/24'), nacaddr.IPv4('10.0.0.0/8'),
                nacaddr.IPv6('200::/56'), nacaddr.IPv6('10::/16')]
    excludes = [nacaddr.IPv6('10::/18'), nacaddr.IPv6('10:c000::/18'),
                nacaddr.IPv4('10.0.0.0/10'), nacaddr.IPv4('10.192.0.0/10')]
    expected = [nacaddr.IPv4('10.64.0.0/10'), nacaddr.IPv4('10.128.0.0/10'),
                nacaddr.IPv4('200.0.0.0/24'),
                nacaddr.IPv6('10:4000::/18'), nacaddr.IPv6('10:8000::/18'),
                nacaddr.IPv6('200::/56')]
    self.assertListEqual(nacaddr.AddressListExclude(superset, excludes),
                         expected)

  def testAddressListExcludeCaseFour(self):
    # IPv6 does not affect IPv4
    superset = [nacaddr.IPv4('0.0.0.0/0')]
    excludes = [nacaddr.IPv6('::/0')]
    expected = [nacaddr.IPv4('0.0.0.0/0')]
    self.assertListEqual(nacaddr.AddressListExclude(superset, excludes),
                         expected)

  def testAddressListExcludeCaseFive(self):
    # IPv6 does not affect IPv4
    superset = [nacaddr.IPv6('::/0')]
    excludes = [nacaddr.IPv4('0.0.0.0/0')]
    expected = [nacaddr.IPv6('::/0')]
    self.assertListEqual(nacaddr.AddressListExclude(superset, excludes),
                         expected)

  def testAddressListExcludeCaseSix(self):
    # IPv6 does not affect IPv4
    superset = [nacaddr.IPv6('0::ffff:0.0.0.0/96')]
    excludes = [nacaddr.IPv4('0.0.0.0/0')]
    expected = [nacaddr.IPv6('0::ffff:0.0.0.0/96')]
    self.assertListEqual(nacaddr.AddressListExclude(superset, excludes),
                         expected)

  def testCollapseAddrListPreserveTokens(self):
    addr_list = [nacaddr.IPv4('10.0.1.7/32', token='BIZ'),
                 nacaddr.IPv4('192.168.1.10/32', token='ALSOUNDERSUPER'),
                 nacaddr.IPv4('10.0.0.6/32', token='FOO'),
                 nacaddr.IPv4('10.0.0.9/32', token='BAR'),
                 nacaddr.IPv4('10.0.0.8/32', token='FOO'),
                 nacaddr.IPv4('10.0.0.7/32', token='BAR'),
                 nacaddr.IPv4('192.168.1.1/24', token='SUPER', strict=False),
                 nacaddr.IPv4('10.0.1.6/32', token='BIZ'),
                 nacaddr.IPv4('192.168.1.7/31',
                              token='UNDERSUPER',
                              strict=False)
                 ]
    expected = [nacaddr.IPv4('10.0.0.7/32', token='BAR'),
                nacaddr.IPv4('10.0.0.9/32', token='BAR'),
                nacaddr.IPv4('10.0.1.6/31', token='BIZ'),
                nacaddr.IPv4('10.0.0.6/32', token='FOO'),
                nacaddr.IPv4('10.0.0.8/32', token='FOO'),
                nacaddr.IPv4('192.168.1.1/24', token='SUPER', strict=False)]
    collapsed = nacaddr.CollapseAddrListPreserveTokens(addr_list)
    self.assertListEqual(collapsed, expected)

  def testIsSupernet(self):
    addrs1 = [nacaddr.IPv4('10.0.1.7/32'), nacaddr.IPv4('10.0.1.2/32')]
    addrs2 = [nacaddr.IPv4('10.0.1.0/24')]
    addrs3 = [nacaddr.IPv4('10.0.1.7/32'), nacaddr.IPv4('10.1.1.2/32')]
    addrs4 = [nacaddr.IPv4('192.168.1.1/32', nacaddr.IPv4('172.0.0.1/32'))]
    addrs5 = [nacaddr.IPv4('192.168.1.1/24', strict=False),
              nacaddr.IPv4('10.0.1.0/24')]
    self.assertTrue(nacaddr.IsSuperNet(addrs2, addrs1))
    self.assertFalse(nacaddr.IsSuperNet(addrs2, addrs3))
    self.assertFalse(nacaddr.IsSuperNet(addrs2, addrs4))
    self.assertFalse(nacaddr.IsSuperNet(addrs2, addrs5))
    self.assertTrue(nacaddr.IsSuperNet(addrs5, addrs2))

  def testSafeCollapsing(self):
    test_data = [([nacaddr.IPv4('10.0.0.0/8'),
                   nacaddr.IPv4('10.0.0.0/10')],
                  [nacaddr.IPv4('10.0.0.0/9')],
                  [nacaddr.IPv4('10.0.0.0/8'),
                   nacaddr.IPv4('10.0.0.0/10')]),

                 ([nacaddr.IPv4('192.168.0.0/27'),
                   nacaddr.IPv4('192.168.0.0/24')],
                  [nacaddr.IPv4('192.168.1.0/24')],
                  [nacaddr.IPv4('192.168.0.0/24')]),

                 ([nacaddr.IPv6('10::/56'),
                   nacaddr.IPv6('10::/128')],
                  [nacaddr.IPv6('10::/64')],
                  [nacaddr.IPv6('10::/56'),
                   nacaddr.IPv6('10::/128')]),

                 ([nacaddr.IPv6('10::/128'),
                   nacaddr.IPv6('10::/56')],
                  [nacaddr.IPv6('8::/64')],
                  [nacaddr.IPv6('10::/56')])
                 ]
    for addresses, complement_addresses, result in test_data:
      self.assertEqual(nacaddr.CollapseAddrList(addresses,
                                                complement_addresses), result)


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