# Copyright 2011 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.
#

"""Unit tests for cgrep.

   Order doesn't matter for the purposes of these tests, so many
   actual and expected results are sorted/sets to prevent issues relating to the
   order in which items are returned.
"""

import argparse
from absl.testing import absltest

from capirca.lib import nacaddr
from capirca.lib import naming
from tools import cgrep


_NETWORK = """
#
# Sample naming defintions for network objects
#
RFC1918 = 10.0.0.0/8      # non-public
          172.16.0.0/12   # non-public
          192.168.0.0/16  # non-public

INTERNAL = RFC1918

LOOPBACK = 127.0.0.0/8  # loopback
           ::1/128       # ipv6 loopback

RFC_3330 = 169.254.0.0/16  # special use IPv4 addresses - netdeploy

RFC_6598 = 100.64.0.0/10   # Shared Address Space

LINKLOCAL = FE80::/10  # IPv6 link-local

SITELOCAL = FEC0::/10    # Ipv6 Site-local

MULTICAST = 224.0.0.0/4  # IP multicast
            FF00::/8     # IPv6 multicast

CLASS-E   = 240.0.0.0/4

RESERVED = 0.0.0.0/8           # reserved
           RFC1918
           LOOPBACK
           RFC_3330
           RFC_6598
           MULTICAST
           CLASS-E
           0000::/8            # reserved by IETF
           0100::/8            # reserved by IETF
           0200::/7            # reserved by IETF
           0400::/6            # reserved by IETF
           0800::/5            # reserved by IETF
           1000::/4            # reserved by IETF
           4000::/3            # reserved by IETF
           6000::/3            # reserved by IETF
           8000::/3            # reserved by IETF
           A000::/3            # reserved by IETF
           C000::/3            # reserved by IETF
           E000::/4            # reserved by IETF
           F000::/5            # reserved by IETF
           F800::/6            # reserved by IETF
           FC00::/7            # unique local unicast
           FE00::/9            # reserved by IETF
           LINKLOCAL           # link local unicast
           SITELOCAL           # IPv6 site-local

ANY = 0.0.0.0/0

# http://www.team-cymru.org/Services/Bogons/bogon-bn-agg.txt
# 22-Apr-2011
BOGON = 0.0.0.0/8
        192.0.0.0/24
        192.0.2.0/24
        198.18.0.0/15
        198.51.100.0/24
        203.0.113.0/24
        MULTICAST
        CLASS-E
        3FFE::/16      # 6bone
        5F00::/8       # 6bone
        2001:DB8::/32  # IPv6 documentation prefix

GOOGLE_PUBLIC_DNS_ANYCAST = 8.8.4.4/32               # IPv4 Anycast
                            8.8.8.8/32               # IPv4 Anycast
                            2001:4860:4860::8844/128 # IPv6 Anycast
                            2001:4860:4860::8888/128 # IPv6 Anycast
GOOGLE_DNS = GOOGLE_PUBLIC_DNS_ANYCAST


# The following are sample entires intended for us in the included
# sample policy file.  These should be removed.

WEB_SERVERS = 200.1.1.1/32  # Example web server 1
              200.1.1.2/32  # Example web server 2

MAIL_SERVERS = 200.1.1.4/32 # Example mail server 1
               200.1.1.5/32 # Example mail server 2

PUBLIC_NAT = 200.1.1.3/32   # Example company NAT address

NTP_SERVERS = 10.0.0.1/32   # Example NTP server
              10.0.0.2/32   # Example NTP server

TACACS_SERVERS = 10.1.0.1/32  # Example tacacs server
                 10.1.0.2/32  # Example tacacs server

INTERNAL_SERVER = 1.0.0.1/32

PUBLIC_SERVER = 100.0.0.1/32

INTERNAL_SERVERS = INTERNAL_SERVER

PUBLIC_SERVERS = PUBLIC_SERVER

SERVERS = INTERNAL_SERVERS
          PUBLIC_SERVERS

"""

_SERVICE = """
#
# Sample naming service definitions
#
WHOIS = 43/udp
SSH = 22/tcp
TELNET = 23/tcp
SMTP = 25/tcp
MAIL_SERVICES = SMTP
                ESMTP
                SMTP_SSL
                POP_SSL
TIME = 37/tcp 37/udp
TACACS = 49/tcp
DNS = 53/tcp 53/udp
BOOTPS = 67/udp   # BOOTP server
BOOTPC = 68/udp   # BOOTP client
DHCP = BOOTPS
       BOOTPC
TFTP = 69/tcp 69/udp
HTTP = 80/tcp
WEB_SERVICES = HTTP HTTPS
POP3 = 110/tcp
RPC = 111/udp
IDENT = 113/tcp 113/udp
NNTP = 119/tcp
NTP = 123/tcp 123/udp
MS_RPC_EPMAP = 135/udp 135/tcp
MS_137 = 137/udp
MS_138 = 138/udp
MS_139 = 139/tcp
IMAP = 143/tcp
SNMP = 161/udp
SNMP_TRAP = 162/udp
BGP = 179/tcp
IMAP3 = 220/tcp
LDAP = 389/tcp
LDAP_SERVICE = LDAP
               LDAPS
HTTPS = 443/tcp
MS_445 = 445/tcp
SMTP_SSL = 465/tcp
IKE = 500/udp
SYSLOG = 514/udp
RTSP = 554/tcp
ESMTP = 587/tcp
LDAPS = 636/tcp
IMAPS = 993/tcp
POP_SSL = 995/tcp
HIGH_PORTS = 1024-65535/tcp 1024-65535/udp
MSSQL = 1433/tcp
MSSQL_MONITOR = 1434/tcp
RADIUS = 1812/tcp 1812/udp
HSRP = 1985/udp
NFSD = 2049/tcp 2049/udp
NETFLOW = 2056/udp
SQUID_PROXY = 3128/tcp
MYSQL = 3306/tcp
RDP = 3389/tcp
IPSEC = 4500/udp
POSTGRESQL = 5432/tcp
TRACEROUTE = 33434-33534/udp
"""


class Namespace:

  def __init__(self, **kwargs):
    for arg in kwargs:
      setattr(self, arg, kwargs[arg])


class CgrepTest(absltest.TestCase):

  def setUp(self):
    super().setUp()
    self.db = naming.Naming(None)
    self.db.ParseServiceList(_SERVICE.split('\n'))
    self.db.ParseNetworkList(_NETWORK.split('\n'))

  #
  # test ip->token resolution (-i)
  #
  # 1.1.1.1 should only be in 'ANY'
  def test_one_ip(self):
    expected_results = [('ANY', ['0.0.0.0/0'])]
    ip = '1.1.1.1'
    results = cgrep.get_ip_parents(ip, self.db)
    self.assertEqual(results, expected_results)

  # 2001:db8::1 should only be in 'BOGON'
  def test_one_ipv6(self):
    expected_results = [('BOGON', ['2001:db8::/32'])]
    ip = '2001:db8::1'
    results = cgrep.get_ip_parents(ip, self.db)
    self.assertEqual(results, expected_results)

  # 1.1.1.1 should not be in CLASS-E
  def test_one_ip_fail(self):
    expected_results = [('CLASS-E', ['240.0.0.0/4'])]
    ip = nacaddr.IP('1.1.1.1/32')
    results = cgrep.get_ip_parents(ip, self.db)
    self.assertNotEqual(results, expected_results)

  # 2001:db8::1 should not be in LINKLOCAL
  def test_one_ipv6_fail(self):
    expected_results = [('LINKLOCAL', ['FE80::/10'])]
    ip = '2001:db8::1'
    results = cgrep.get_ip_parents(ip, self.db)
    self.assertNotEqual(results, expected_results)

  # 8.8.8.8 is in GOOGLE_PUBLIC_DNS_ANYCAST which is inside GOOGLE_DNS
  def test_one_ip_nested(self):
    expected_results = sorted((('GOOGLE_DNS', ['8.8.8.8/32']),
                               ('GOOGLE_DNS -> GOOGLE_PUBLIC_DNS_ANYCAST',
                                ['8.8.8.8/32']),
                               ('ANY', ['0.0.0.0/0'])))
    ip = '8.8.8.8'
    results = sorted(cgrep.get_ip_parents(ip, self.db))
    self.assertEqual(results, expected_results)

  # 2001:4860:4860::8844/128 is in GOOGLE_PUBLIC_DNS_ANYCAST which is
  # inside GOOGLE_DNS
  def test_one_ipv6_nested(self):
    expected_results = sorted((('GOOGLE_DNS', ['2001:4860:4860::8844/128']),
                               ('GOOGLE_DNS -> GOOGLE_PUBLIC_DNS_ANYCAST',
                                ['2001:4860:4860::8844/128'])))
    ip = '2001:4860:4860::8844/128'
    results = sorted(cgrep.get_ip_parents(ip, self.db))
    self.assertEqual(results, expected_results)

  # 1.0.0.1 is inside INTERNAL_SERVER, which is inside INTERNAL_SERVERS, which
  # is inside SERVERS
  def test_one_ip_multi_nested(self):
    expected_results = sorted((('INTERNAL_SERVERS -> INTERNAL_SERVER',
                                ['1.0.0.1/32']),
                               ('SERVERS -> INTERNAL_SERVER', ['1.0.0.1/32']),
                               ('SERVERS -> INTERNAL_SERVERS', ['1.0.0.1/32']),
                               ('SERVERS', ['1.0.0.1/32']),
                               ('ANY', ['0.0.0.0/0'])))
    ip = '1.0.0.1'
    results = sorted(cgrep.get_ip_parents(ip, self.db))
    self.assertEqual(results, expected_results)

  #
  # test 'ip in token' (-i -t)
  #
  # 8.8.8.8 is inside GOOGLE_DNS
  def test_ip_in_token(self):
    expected_results = r'8.8.8.8 is in GOOGLE_DNS'
    options = Namespace()
    options.ip = ('8.8.8.8',)
    options.token = ('GOOGLE_DNS')
    results = cgrep.compare_ip_token(options, self.db)
    self.assertEqual(results, expected_results)

  # 2001:4860:4860::8844 is inside GOOGLE_DNS
  def test_ipv6_in_token(self):
    expected_results = r'2001:4860:4860::8844 is in GOOGLE_DNS'
    options = Namespace()
    options.ip = ('2001:4860:4860::8844',)
    options.token = ('GOOGLE_DNS')
    results = cgrep.compare_ip_token(options, self.db)
    self.assertEqual(results, expected_results)

  # 69.171.239.12 is not in GOOGLE_DNS
  def test_ip_in_token_fail(self):
    expected_results = r'69.171.239.12 is _not_ in GOOGLE_DNS'
    options = Namespace()
    options.ip = ('69.171.239.12',)
    options.token = ('GOOGLE_DNS')
    results = cgrep.compare_ip_token(options, self.db)
    self.assertEqual(results, expected_results)

  # 2a03:2880:fffe:c:face:b00c:0:35 is not in GOOGLE_DNS
  def test_ipv6_in_token_fail(self):
    expected_results = r'2a03:2880:fffe:c:face:b00c:0:35 is _not_ in GOOGLE_DNS'
    options = Namespace()
    options.ip = ('2a03:2880:fffe:c:face:b00c:0:35',)
    options.token = ('GOOGLE_DNS')
    results = cgrep.compare_ip_token(options, self.db)
    self.assertEqual(results, expected_results)

  #
  # test network token compare (-c)
  #
  # these two tokens are identical and should contain the same nets
  def test_compare_same_token(self):
    expected_results = (
        (
            r'PUBLIC_NAT',
            r'PUBLIC_NAT',
            [
                nacaddr.IPv4('200.1.1.3/32')
            ],
        ),
        [
            r'200.1.1.3/32'
        ]
    )
    options = Namespace()
    options.cmp = ('PUBLIC_NAT', 'PUBLIC_NAT')
    results = cgrep.compare_tokens(options, self.db)
    self.assertEqual(results, expected_results)

  #
  # test network token encapsulations
  #
  def test_ip_contained(self):
    expected_results = True
    results = cgrep.check_encapsulated('network', 'RFC1918', 'RESERVED',
                                       self.db)
    self.assertEqual(results, expected_results)

  def test_ip_not_contained(self):
    expected_results = False
    results = cgrep.check_encapsulated('network', 'RESERVED', 'RFC1918',
                                       self.db)
    self.assertEqual(results, expected_results)

  def test_ipv6_contained(self):
    expected_results = True
    results = cgrep.check_encapsulated('network', 'LINKLOCAL', 'RESERVED',
                                       self.db)
    self.assertEqual(results, expected_results)

  def test_ipv6_not_contained(self):
    expected_results = False
    results = cgrep.check_encapsulated('network', 'RESERVED', 'LINKLOCAL',
                                       self.db)
    self.assertEqual(results, expected_results)

  #
  # test ip->object comparisons (-g)
  #
  # 8.8.8.8 is not present in object RESERVED and
  # 127.0.0.1 is not present in object GOOGLE_DNS and
  # the two IPs both exist in 'ANY'
  def test_group_diff(self):
    expected_results = sorted((
        ['ANY'],
        ['GOOGLE_DNS', 'GOOGLE_DNS -> GOOGLE_PUBLIC_DNS_ANYCAST'],
        ['RESERVED', 'RESERVED -> LOOPBACK']
    ))
    options = Namespace()
    options.gmp = ['8.8.8.8', '127.0.0.1']
    results = sorted(cgrep.group_diff(options, self.db))
    self.assertCountEqual(results, expected_results)

  # test to make sure two IPs share the same groups
  def test_group_diff_identical(self):
    expected_results = sorted((
        ['ANY',
         'INTERNAL',
         'INTERNAL -> RFC1918',
         'RESERVED',
         'RESERVED -> RFC1918'],
        [],
        []
    ))
    options = Namespace()
    options.gmp = ['172.16.0.1', '192.168.0.1']
    results = sorted(cgrep.group_diff(options, self.db))
    self.assertCountEqual(results, expected_results)

  #
  # test token->ip(s) resolution (-o)
  #
  # resolve GOOGLE_DNS to the 4 given IPs
  def test_token_to_ips(self):
    expected_results = [
        (
            r'GOOGLE_DNS',
            [
                nacaddr.IPv4('8.8.4.4/32'),
                nacaddr.IPv4('8.8.8.8/32'),
                nacaddr.IPv6('2001:4860:4860::8844/128'),
                nacaddr.IPv6('2001:4860:4860::8888/128')
            ]
        )
    ]
    options = Namespace()
    options.obj = ('GOOGLE_DNS',)

    results = cgrep.get_nets(options.obj, self.db)
    self.assertEqual(results[0][0], expected_results[0][0])
    self.assertEqual(set(results[0][1]), set(expected_results[0][1]))

  # GOOGLE_DNS does not resole to the given IP
  def test_token_to_ip_fail(self):
    expected_results = [
        (
            r'GOOGLE_DNS',
            [
                nacaddr.IPv4('69.171.239.12/32'),
                nacaddr.IPv6('2a03:2880:fffe:c:face:b00c:0:35/128')
            ]
        )
    ]
    options = Namespace()
    options.obj = ('GOOGLE_DNS',)
    results = cgrep.get_nets(options.obj, self.db)
    # the network object name should match, but not the IPs contained within
    self.assertEqual(results[0][0], expected_results[0][0])
    self.assertNotEqual(set(results[0][1]), set(expected_results[0][1]))

  #
  # test service->port resolution (-s)
  #
  # "SSH" is just '22/tcp'
  def test_svc_to_port(self):
    expected_results = [
        (
            r'SSH',
            [
                '22/tcp'
            ]
        )
    ]
    options = Namespace()
    options.service = ('SSH',)
    results = cgrep.get_ports(options.service, self.db)
    self.assertEqual(results, expected_results)

  # "SSH" does not contain '23/tcp'
  def test_svc_to_port_fail(self):
    expected_results = [
        (
            r'SSH',
            [
                '23/tcp'
            ]
        )
    ]
    options = Namespace()
    options.svc = ('SSH',)
    results = cgrep.get_ports(options.svc, self.db)
    self.assertNotEqual(results, expected_results)

  #
  # test port->service object resolution (-p)
  #
  # '22/tcp' belongs to SSH
  def test_get_port_parents(self):
    expected_results = (r'22', r'tcp', ['SSH'])
    options = Namespace()
    options.port = ('22', 'tcp')
    results = cgrep.get_services(options, self.db)
    self.assertEqual(results, expected_results)

  # 22/tcp does not belong to TELNET
  def test_get_port_parents_fail(self):
    expected_results = (r'22', r'tcp', ['TELNET'])
    options = Namespace()
    options.port = ('22', 'tcp')
    results = cgrep.get_services(options, self.db)
    self.assertNotEqual(results, expected_results)

  # 33434/tcp should only be in HIGH_PORTS (not also TRACEROUTE)
  def test_get_port_parents_range_tcp(self):
    expected_results = (r'33434', r'tcp', ['HIGH_PORTS'])
    options = Namespace()
    options.port = ('33434', 'tcp')
    results = cgrep.get_services(options, self.db)
    self.assertEqual(results, expected_results)

  # 33434/udp should be in HIGH_PORTS and TRACEROUTE
  def test_get_port_parents_range_udp(self):
    expected_results = (r'33434', r'udp', ['HIGH_PORTS', 'TRACEROUTE'])
    options = Namespace()
    options.port = ('33434', 'udp')
    results = cgrep.get_services(options, self.db)
    self.assertEqual(results, expected_results)

  #
  # test IP validity
  #
  def test_invalid_ip(self):
    self.assertRaises(argparse.ArgumentTypeError,
                      cgrep.is_valid_ip, '10.0.0.256')

  def test_invalid_ipv6(self):
    self.assertRaises(argparse.ArgumentTypeError,
                      cgrep.is_valid_ip, '2001:db8::z')

  def test_valid_ips(self):
    arg = '8.8.8.8'
    results = cgrep.is_valid_ip(arg)
    self.assertEqual(results, arg)

  def test_valid_ips_v6(self):
    arg = '2001:4860:4860::8844'
    results = cgrep.is_valid_ip(arg)
    self.assertEqual(results, arg)


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