resources/traffic_scripts/ipfix_sessions.py

   1 #!/usr/bin/env python
   2
   3 # Copyright (c) 2016 Cisco and/or its affiliates.
   4 # Licensed under the Apache License, Version 2.0 (the "License");
   5 # you may not use this file except in compliance with the License.
   6 # You may obtain a copy of the License at:
   7 #
   8 #     http://www.apache.org/licenses/LICENSE-2.0
   9 #
  10 # Unless required by applicable law or agreed to in writing, software
  11 # distributed under the License is distributed on an "AS IS" BASIS,
  12 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13 # See the License for the specific language governing permissions and
  14 # limitations under the License.
  15
  16 """Traffic script - IPFIX listener."""
  17
  18 import sys
  19 from ipaddress import IPv4Address, IPv6Address, AddressValueError
  20
  21 from scapy.layers.inet import IP, TCP, UDP
  22 from scapy.layers.inet6 import IPv6
  23 from scapy.layers.l2 import Ether
  24
  25 from resources.libraries.python.IPFIXUtil import IPFIXHandler, IPFIXData
  26 from resources.libraries.python.PacketVerifier import RxQueue, TxQueue, auto_pad
  27 from resources.libraries.python.TrafficScriptArg import TrafficScriptArg
  28
  29
  30 def valid_ipv4(ip):
  31     """Check if IP address has the correct IPv4 address format.
  32
  33     :param ip: IP address.
  34     :type ip: str
  35     :return: True in case of correct IPv4 address format,
  36     otherwise return false.
  37     :rtype: bool
  38     """
  39     try:
  40         IPv4Address(unicode(ip))
  41         return True
  42     except (AttributeError, AddressValueError):
  43         return False
  44
  45
  46 def valid_ipv6(ip):
  47     """Check if IP address has the correct IPv6 address format.
  48
  49     :param ip: IP address.
  50     :type ip: str
  51     :return: True in case of correct IPv6 address format,
  52     otherwise return false.
  53     :rtype: bool
  54     """
  55     try:
  56         IPv6Address(unicode(ip))
  57         return True
  58     except (AttributeError, AddressValueError):
  59         return False
  60
  61
  62 def verify_data(data, count, src_ip, dst_ip, protocol):
  63     """Compare data in IPFIX flow report against parameters used to send test
  64      packets.
  65
  66      :param data: Dictionary of fields in IPFIX flow report.
  67      :param count: Number of packets expected.
  68      :param src_ip: Expected source IP address.
  69      :param dst_ip: Expected destination IP address.
  70      :param protocol: Expected protocol, TCP or UDP.
  71      :type data: dict
  72      :type count: int
  73      :type src_ip: str
  74      :type dst_ip: str
  75      :type protocol: scapy.layers
  76      """
  77
  78     # verify packet count
  79     if data["packetTotalCount"] != count:
  80         raise RuntimeError(
  81             "IPFIX reported wrong packet count. Count was {0},"
  82             " but should be {1}".format(data["packetTotalCount"], count))
  83     # verify IP addresses
  84     keys = data.keys()
  85     e = "{0} mismatch. Packets used {1}, but were classified as {2}."
  86     if valid_ipv4(src_ip) and valid_ipv4(dst_ip):
  87         if "IPv4_src" in keys:
  88             if data["IPv4_src"] != src_ip:
  89                 raise RuntimeError(
  90                     e.format("Source IP", src_ip, data["IPv4_src"]))
  91         if "IPv4_dst" in keys:
  92             if data["IPv4_dst"] != dst_ip:
  93                 raise RuntimeError(
  94                     e.format("Destination IP", dst_ip, data["IPv4_dst"]))
  95     else:
  96         if "IPv6_src" in keys:
  97             if data["IPv6_src"] != src_ip:
  98                 raise RuntimeError(
  99                     e.format("Source IP", src_ip, data["IPv6_src"]))
 100         if "IPv6_dst" in keys:
 101             if data["IPv6_dst"] != dst_ip:
 102                 raise RuntimeError(
 103                     e.format("Source IP", src_ip, data["IPv6_dst"]))
 104     # verify protocol ID
 105     if "Protocol_ID" in keys:
 106         if protocol == TCP and int(data["Protocol_ID"]) != 6:
 107             raise RuntimeError(
 108                 "TCP Packets were classified as not TCP.")
 109         if protocol == UDP and int(data["Protocol_ID"]) != 17:
 110             raise RuntimeError(
 111                 "UDP Packets were classified as not UDP.")
 112     # return port number
 113     for item in ("src_port", "tcp_src_port", "udp_src_port",
 114                  "dst_port", "tcp_dst_port", "udp_dst_port"):
 115         if item in keys:
 116             return int(data[item])
 117         else:
 118             raise RuntimeError("Data contains no port information.")
 119
 120
 121 def main():
 122     """Send packets to VPP, then listen for IPFIX flow report. Verify that
 123     the correct packet count was reported."""
 124     args = TrafficScriptArg(
 125         ['src_mac', 'dst_mac', 'src_ip', 'dst_ip', 'protocol', 'port', 'count',
 126          'sessions']
 127     )
 128
 129     dst_mac = args.get_arg('dst_mac')
 130     src_mac = args.get_arg('src_mac')
 131     src_ip = args.get_arg('src_ip')
 132     dst_ip = args.get_arg('dst_ip')
 133     tx_if = args.get_arg('tx_if')
 134
 135     protocol = args.get_arg('protocol')
 136     source_port = int(args.get_arg('port'))
 137     destination_port = int(args.get_arg('port'))
 138     count = int(args.get_arg('count'))
 139     sessions = int(args.get_arg('sessions'))
 140
 141     txq = TxQueue(tx_if)
 142     rxq = RxQueue(tx_if)
 143
 144     # generate simple packet based on arguments
 145     ip_version = None
 146     if valid_ipv4(src_ip) and valid_ipv4(dst_ip):
 147         ip_version = IP
 148     elif valid_ipv6(src_ip) and valid_ipv6(dst_ip):
 149         ip_version = IPv6
 150     else:
 151         ValueError("Invalid IP version!")
 152
 153     if protocol.upper() == 'TCP':
 154         protocol = TCP
 155     elif protocol.upper() == 'UDP':
 156         protocol = UDP
 157     else:
 158         raise ValueError("Invalid type of protocol!")
 159
 160     packets = []
 161     for x in range(sessions):
 162         pkt = (Ether(src=src_mac, dst=dst_mac) /
 163                ip_version(src=src_ip, dst=dst_ip) /
 164                protocol(sport=x, dport=x))
 165         pkt = auto_pad(pkt)
 166         packets.append(pkt)
 167
 168     # do not print details for sent packets
 169     verbose = False
 170     print("Sending more than one packet. Details will be filtered for "
 171           "all packets sent.")
 172
 173     ignore = []
 174     for x in range(sessions):
 175         for _ in range(count):
 176             txq.send(packets[x], verbose=verbose)
 177             ignore.append(packets[x])
 178
 179     # allow scapy to recognize IPFIX headers and templates
 180     ipfix = IPFIXHandler()
 181
 182     # clear receive buffer
 183     while True:
 184         pkt = rxq.recv(1, ignore=packets, verbose=verbose)
 185         if pkt is None:
 186             break
 187
 188     data = None
 189     ports = [x for x in range(sessions)]
 190
 191     # get IPFIX template and data
 192     while True:
 193         pkt = rxq.recv(5)
 194         if pkt is None:
 195             raise RuntimeError("RX timeout")
 196         if pkt.haslayer("IPFIXHeader"):
 197             if pkt.haslayer("IPFIXTemplate"):
 198                 # create or update template for IPFIX data packets
 199                 ipfix.update_template(pkt)
 200             elif pkt.haslayer("IPFIXData"):
 201                 for x in range(sessions):
 202                     try:
 203                         data = pkt.getlayer(IPFIXData, x+1).fields
 204                     except AttributeError:
 205                         raise RuntimeError("Could not find data layer "
 206                                            "#{0}".format(x+1))
 207                     port = verify_data(data, count, src_ip, dst_ip, protocol)
 208                     if port in ports:
 209                         ports.remove(port)
 210                     else:
 211                         raise RuntimeError("Unexpected or duplicate port {0} "
 212                                            "in flow report.".format(port))
 213                 print("All {0} sessions verified "
 214                       "with packet count {1}.".format(sessions, count))
 215                 sys.exit(0)
 216             else:
 217                 raise RuntimeError("Unable to parse IPFIX template "
 218                                    "or data set.")
 219         else:
 220             raise RuntimeError("Received non-IPFIX packet or IPFIX header was"
 221                                "not recognized.")
 222
 223 if __name__ == "__main__":
 224
 225     main()