cgrep.py

# Copyright 2011 Google Inc.
#
# 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.
#
#
# Simple util to grep through network and service definitions.
# Examples:
#   To find out which tokens contain "10.4.3.1" use
#   $ cgrep.py -i 10.4.3.1
#
#   To find out if token 'FOO' includes ip "1.2.3.4" use
#   $ cgrep.py -t FOO -i 1.2.3.4
#
#   To find the difference and union of tokens 'FOO' and 'BAR' use
#   $ cgrep.py -c FOO BAR
#
#   To find the difference of network tokens to which 2 IPs belong use
#   $ cgrep.py -g 1.1.1.1 2.2.2.2
#
#   To find which IPs are in the 'FOO' network token use
#   $ cgrep.py -o FOO
#
#   To find which port & protocol pairs are in a service token 'FOO' use
#   $ cgrep.py -s FOO
#
#   To find which service tokens contain port '22' and protocol 'tcp' use
#   $ cgrep.py -p 22 tcp
#
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import argparse
import pprint
from lib import nacaddr
from lib import naming


def is_valid_ip(arg):
  """ Validates a value to be an IP or not.

  Args:
    arg as str: potential IP address

  Returns:
    arg as IP object (if arg is an IP)

  Raises:
    Error (if arg is not an IP)
  """
  try:
    nacaddr.IP(arg)
  except:
    raise argparse.ArgumentTypeError('%s is an invalid ip address' % arg)
  return arg


def cli_options():
  """ Builds the argparse options for cgrep.

    Returns:
      parser as argparse.ArgumentParser: the arguments, ready to be parsed.
  """

  parser = argparse.ArgumentParser(
      description='c[apirca]grep',
      formatter_class=argparse.RawTextHelpFormatter
  )

  parser.add_argument('-d', '--def', dest='defs',
                      help='Network Definitions directory location. \n',
                      default='./def')

  # -i and -t can be used together, but not with any other option.
  ip_group = parser.add_argument_group()
  # take 1 or more IPs
  ip_group.add_argument('-i', '--ip', dest='ip', nargs='+', type=is_valid_ip,
                        help='Return list of definitions containing the '
                        'IP(s).\nMultiple IPs permitted.')

  ip_group.add_argument('-t', '--token', dest='token',
                        help=('See if an IP is contained within the given '
                              'token.\nMust be used in conjunction with '
                              '-i/--ip [addr].'))

  exclusive_group = parser.add_mutually_exclusive_group()
  # the rest of the arguments are mutually exclusive with each other,
  # and -i / -t
  exclusive_group.add_argument('-c', '--cmp', dest='cmp', nargs=2,
                               metavar=('OBJ', 'OBJ'),
                               help=('Compare the two given network '
                                     'definition tokens'))

  exclusive_group.add_argument('-g', '--gmp', dest='gmp', nargs=2,
                               type=is_valid_ip, metavar=('IP', 'IP'),
                               help=('Diff the network objects to'
                                     ' which the given IP(s) belong'))

  exclusive_group.add_argument('-o', '--obj', dest='obj', nargs='+',
                               help=('Return list of IP(s) contained within '
                                     'the given token(s)'))

  exclusive_group.add_argument('-s', '--svc', dest='svc', nargs='+',
                               help=('Return list of port(s) contained '
                                     'within given token(s)'))

  exclusive_group.add_argument('-p', '--port', dest='port', nargs=2,
                               metavar=('PORT', 'PROTO'),
                               help=('Returns a list of tokens containing '
                                     'the given port and protocol'))

  return parser


def main(parser):
  """ Determines the code path based on the arguments passed.

  Args:
    parser as argparse.ArgumentParser: the argument parser, but not parsed yet.
  """
  options = parser.parse_args()
  db = naming.Naming(options.defs)
  p = pprint.PrettyPrinter(indent=1, depth=4, width=1).pprint

  # if -i and any other option:
  if options.ip and any([options.gmp, options.cmp, options.obj, options.svc,
                         options.port]):
    print('You can only use -i with -t or by itself')

  # if -i and -t
  elif options.token and options.ip:
    try:
      get_nets([options.token], db)
    except naming.UndefinedAddressError as e:
      print("Network group '%s' is not defined!" % (e.message.split()[0]))
    else:
      results = compare_ip_token(options, db)
      print(results)

  # if -t, but not -i; invalid!
  elif options.token and not options.ip:
    print('You must specify an IP Address with -i [addr]')

  # if -i
  elif options.ip:
    for ip in options.ip:
      groups = get_ip_parents(ip, db)
      print('Results for IP: %s' % ip)
      # iterate and print the tokens we found.
      for name, networks in groups:
        # print the group name [0], and the networks it was in [1]
        print('%s  %s' % (name, networks))

  elif options.gmp:
    common, diff1, diff2 = group_diff(options, db)
    print_diff(options.gmp[0], common, diff1, diff2)
    print('')
    print_diff(options.gmp[1], common, diff2, diff1)

  # if -c
  elif options.cmp:
    meta, results = compare_tokens(options, db)
    print('Union of %s and %s:\n %s\n' % meta)
    print('Diff of %s and %s:' % meta[:-1])
    for i in results:
      print(' ' + i)
    print('')
    first_obj, sec_obj = options.cmp
    if check_encapsulated('network', first_obj, sec_obj, db):
      print('%s fully encapsulates %s' % (sec_obj, first_obj))
    else:
      print('%s does _not_ fully encapsulate %s' % (sec_obj, first_obj))
    # check the other way around.
    if check_encapsulated('network', sec_obj, first_obj, db):
      print('%s fully encapsulates %s' % (first_obj, sec_obj))
    else:
      print('%s does _not_ fully encapsulate %s' % (first_obj, sec_obj))

  # if -o
  elif options.obj:
    for obj in options.obj:
      try:
        token, ips = get_nets([obj], db)[0]
      except naming.UndefinedAddressError:
        print('%s is an invalid object' % obj)
      else:
        print(token + ':')
        # convert list of ip objects to strings and sort them
        ips.sort(key=lambda x: int(x.ip))
        p([str(x) for x in ips])

  # if -s
  elif options.svc:
    try:
      results = get_ports(options.svc, db)
    except naming.UndefinedServiceError:
      print('%s contains an invalid service object' % str(options.svc))
    else:
      for result in get_ports(options.svc, db):
        svc, port = result
        print(svc + ':')
        p(port)

  # if -p
  elif options.port:
    port, protocol, result = get_services(options, db)
    print('%s/%s:' % (port, protocol))
    p(result)

  # if nothing is passed
  elif not any((options.cmp, options.ip, options.token, options.obj,
                options.svc, options.port)):
    parser.print_help()
  print('')


def check_encapsulated(obj_type, first_obj, second_obj, db):
  """ Checks if a network/service object is entirely contained within another.

  Args:
    obj_type as str: "network" or "service"
    first_obj as str: The name of the first network/service object
    second_obj as str: The name of the secondnetwork/service object
    db as naming.Naming(): The network and service definitions

  Returns:
    Error or bool:
      ValueError if an invalid object type is passed
      True if the first_obj is entirely within second_obj, otherwise False
  """
  if obj_type == 'network':
    # the indexing is to get the list of networks out of the tuple[1] and
    # list[0] returned by get_nets
    first = get_nets([first_obj], db)[0][1]
    second = get_nets([second_obj], db)[0][1]

  elif obj_type == 'service':
    first = get_ports([first_obj], db)[0][1]
    second = get_ports([second_obj], db)[0][1]
  else:
    raise ValueError("check_encapsulated() currently only supports "
                     "'network' and 'service' for the obj_type parameter")
  # iterates over each object in the first group, and then each obj in the
  # second group, making sure each one in the first is contained
  # somewhere in the second.
  for obj in first:
    for sec_obj in second:
      if obj in sec_obj:
        break
    # if we got through every object in the second group, and didn't have
    # a match, then the first group is not entirely contained.
    else:
      return False
  # if we got here, then the group was fully contained.
  return True


def print_diff(ip, common, diff1, diff2):
  """ Print out the common, added, and removed network objects between 2 IPs

  Args:
    ip as str: the IP being compared against
    common as list: the network objects shared between the two IPs
                    ('ip' and the other passed into options.cmp)
    diff1 as list: the network objects present in 'ip' but not in the other IP
                   passed into options.cmp
    diff2 as list: the network objects not present in 'ip' but are present in
                   the other IP passed into options.cmp
  """
  print('IP: %s' % ip)
  if common:
    common = ['  {0}'.format(elem) for elem in common]
    print('\n'.join(common))
  if diff1:
    diff = ['+ {0}'.format(elem) for elem in diff1]
    print('\n'.join(diff))
  if diff2:
    diff = ['- {0}'.format(elem) for elem in diff2]
    print('\n'.join(diff))


def group_diff(options, db):
  """ Diffs two different group objects.

  Args:
    options as ArgumentParser.parse_args(): the options sent to the script
    db as naming.Naming(): network and service definitions

  Returns:
    tuple of three lists: the common lines, the differences from 1 to 2,
                          and the differences from 2 to 1
  """
  nested_rvals = []
  for ip in options.gmp:
    nested_rvals.append(get_ip_parents(ip, db))
  # get just the list of groups, stripping out the networks.
  group1 = [x[0] for x in nested_rvals[0]]
  group2 = [x[0] for x in nested_rvals[1]]
  common = list(set(group1) & set(group2))
  diff1 = list(set(group1) - set(group2))
  diff2 = list(set(group2) - set(group1))
  return common, diff1, diff2


def get_ip_parents(ip, db):
  """ Gets a list of all network objects that include an IP.

  Args:
    ip as str: the IP we're looking for the parents of
    db as naming.Naming(): network and service definitions

  Returns:
    results as list: a list of all groups that include the IP, in the format:
                     [("Group", ["networks", "matched"]), (etc)]
  """
  results = []
  rval = db.GetIpParents(ip)
  for v in rval:
    nested = db.GetNetParents(v)
    prefix_and_nets = get_nets_and_highest_prefix(ip, v, db)
    if nested:
      for n in nested:
        results.append(('%s -> %s' % (n, v), prefix_and_nets))
    else:
      results.append((v, prefix_and_nets))
  # sort the results by prefix length descending
  results = sorted(results, key=lambda x: x[1][0], reverse=True)
  # strip out the no longer needed prefix lengths before handing off
  for index, group in enumerate(results):
    results[index] = (group[0], group[1][1])
  return results


def get_nets_and_highest_prefix(ip, net_group, db):
  """ Find the highest prefix length out of all the networks in net_group,
    but only if they contain the passed IP.

  Args:
    ip as str: the IP address contained in net_group
    net_group as str: the name of the network object we'll be looking through
    db as naming.Naming(): network and service definitions

  Returns:
    highest_prefix_length, networks as tuple
      highest_prefix_length as int: the longest prefix length found,
      networks as list: network objects
  """
  highest_prefix_length = 0
  networks = []
  ip = nacaddr.IP(ip)
  # loop through all the networks in the net_group
  for net in get_nets([net_group], db)[0][1]:
    # find the highest prefix length for the networks that contain the IP
    if ip in net:
      networks.append(str(net))
      if net.prefixlen > highest_prefix_length:
        highest_prefix_length = net.prefixlen
  return highest_prefix_length, networks


def get_nets(objects, db):
  """ Gets a list of all networks that are inside of a network object.

  Args:
    objects as list: network objects
    db as naming.Naming(): network and service definitions

  Returns:
    results as list: all networks inside a network object
  """
  results = []
  for obj in objects:
    net = db.GetNet(obj)
    results.append((obj, net))
  return results


def compare_tokens(options, db):
  """ Compares to network objects against eachother.

  Args:
    options as ArgumentParser.parse_args(): the options sent to the script
    db as naming.Naming(): network and service definitions

  Returns:
    meta, results as tuple:
      ((first object, second object, union of those two),
       diff of those two network objects)
  """
  t1, t2 = options.cmp
  d1 = db.GetNet(t1)
  d2 = db.GetNet(t2)
  union = list(set(d1 + d2))
  meta = (t1, t2, union)
  results = []
  for el in set(d1 + d2):
    el = nacaddr.IP(el)
    if el in d1 and el in d2:
      results.append(str(el))
    elif el in d1:
      results.append(str(el))
    elif el in d2:
      results.append(str(el))
  return meta, results


def compare_ip_token(options, db):
  """ looks to see if a network IP is contained in a network object

  Args:
    options as ArgumentParser.parse_args(): the options sent to the script
    db as naming.Naming(): network and service definitions

  Returns:
    results as str: end-user string stating the results
  """
  token = options.token
  results = []
  for ip in options.ip:
    rval = db.GetIpParents(ip)
    if token in rval:
      results = '%s is in %s' % (ip, token)
    else:
      results = '%s is _not_ in %s' % (ip, token)
  return results


def get_ports(svc_group, db):
  """ Gets the ports and protocols defined in a service group

  Args:
    svc_group as list: a list of strings for each service group
    db as naming.Naming(): network and service definitions

  Returns:
    results as list: a list of tuples for each service defined, in the format:
                     (service name, "<port>/<protocol>")
  """
  results = []
  for svc in svc_group:
    port = db.GetService(svc)
    results.append((svc, port))
  return results


def get_services(options, db):
  """ Finds any services with that include a specific port/protocol pair

  Args:
    options as ArgumentParser.parse_args(): the options sent to the script
    db as naming.Naming(): network and service definitions

  Returns:
    port, protocol, results as tuple in the format:
    (port, protocol, list of the services containing this pair)
  """
  results = []
  port, protocol = options.port
  # swap values if they were passed in wrong order
  if port.isalpha() and protocol.isdigit():
    port, protocol = protocol, port
  results = db.GetPortParents(port, protocol)
  return port, protocol, results


if __name__ == '__main__':
  main(cli_options())