#!/router/bin/python import external_packages import dpkt import socket import binascii import copy import random import string import struct import re import itertools from abc import ABCMeta, abstractmethod from collections import namedtuple class CTRexPktBuilder(object): """ This class defines the TRex API of building a packet using dpkt package. Using this class the user can also define how TRex will handle the packet by specifying the VM setting. """ def __init__(self, max_pkt_size=dpkt.ethernet.ETH_LEN_MAX): """ Instantiate a CTRexPktBuilder object :parameters: None """ super(CTRexPktBuilder, self).__init__() self._packet = None self._pkt_by_hdr = {} self._pkt_top_layer = None self._max_pkt_size = max_pkt_size self.payload_gen = CTRexPktBuilder.CTRexPayloadGen(self._packet, self._max_pkt_size) self.vm = CTRexPktBuilder.CTRexVM() self.metadata = "" def add_pkt_layer(self, layer_name, pkt_layer): """ This method adds additional header to the already existing packet :parameters: layer_name: str a string representing the name of the layer. Example: "l2", "l4_tcp", etc. pkt_layer : dpkt.Packet obj a dpkt object, generally from higher layer, that will be added on top of existing layer. :raises: + :exc:`ValueError`, in case the desired layer_name already exists. """ assert isinstance(pkt_layer, dpkt.Packet) if layer_name in self._pkt_by_hdr: raise ValueError("Given layer name '{0}' already exists.".format(layer_name)) else: dup_pkt = copy.copy(pkt_layer) # using copy of layer to avoid cyclic packets that may lead to infinite loop if not self._pkt_top_layer: # this is the first header added self._packet = dup_pkt else: self._pkt_top_layer.data = dup_pkt self._pkt_top_layer = dup_pkt self._pkt_by_hdr[layer_name] = dup_pkt return def set_ip_layer_addr(self, layer_name, attr, ip_addr, ip_type="ipv4"): """ This method sets the IP address fields of an IP header (source or destination, for both IPv4 and IPv6) using a human readable addressing representation. :parameters: layer_name: str a string representing the name of the layer. Example: "l3_ip", etc. attr: str a string representation of the sub-field to be set: + "src" for source + "dst" for destination ip_addr: str a string representation of the IP address to be set. Example: "10.0.0.1" for IPv4, or "5001::DB8:1:3333:1:1" for IPv6 ip_type : str a string representation of the IP version to be set: + "ipv4" for IPv4 + "ipv6" for IPv6 Default: **ipv4** :raises: + :exc:`ValueError`, in case the desired layer_name is not an IP layer + :exc:`KeyError`, in case the desired layer_name does not exists. """ try: layer = self._pkt_by_hdr[layer_name.lower()] if not (isinstance(layer, dpkt.ip.IP) or isinstance(layer, dpkt.ip6.IP6)): raise ValueError("The specified layer '{0}' is not of IPv4/IPv6 type.".format(layer_name)) else: decoded_ip = CTRexPktBuilder._decode_ip_addr(ip_addr, ip_type) setattr(layer, attr, decoded_ip) except KeyError: raise KeyError("Specified layer '{0}' doesn't exist on packet.".format(layer_name)) def set_ipv6_layer_addr(self, layer_name, attr, ip_addr): """ This method sets the IPv6 address fields of an IP header (source or destination) :parameters: layer_name: str a string representing the name of the layer. Example: "l3_ip", etc. attr: str a string representation of the sub-field to be set: + "src" for source + "dst" for destination ip_addr: str a string representation of the IP address to be set. Example: "5001::DB8:1:3333:1:1" :raises: + :exc:`ValueError`, in case the desired layer_name is not an IPv6 layer + :exc:`KeyError`, in case the desired layer_name does not exists. """ self.set_ip_layer_addr(layer_name, attr, ip_addr, ip_type="ipv6") def set_eth_layer_addr(self, layer_name, attr, mac_addr): """ This method sets the ethernet address fields of an Ethernet header (source or destination) using a human readable addressing representation. :parameters: layer_name: str a string representing the name of the layer. Example: "l2", etc. attr: str a string representation of the sub-field to be set: + "src" for source + "dst" for destination mac_addr: str a string representation of the MAC address to be set. Example: "00:de:34:ef:2e:f4". :raises: + :exc:`ValueError`, in case the desired layer_name is not an Ethernet layer + :exc:`KeyError`, in case the desired layer_name does not exists. """ try: layer = self._pkt_by_hdr[layer_name.lower()] if not isinstance(layer, dpkt.ethernet.Ethernet): raise ValueError("The specified layer '{0}' is not of Ethernet type.".format(layer_name)) else: decoded_mac = CTRexPktBuilder._decode_mac_addr(mac_addr) setattr(layer, attr, decoded_mac) except KeyError: raise KeyError("Specified layer '{0}' doesn't exist on packet.".format(layer_name)) def set_layer_attr(self, layer_name, attr, val, toggle_bit=False): """ This method enables the user to change a value of a previously defined packet layer. This method isn't to be used to set the data attribute of a packet with payload. Use :func:`packet_builder.CTRexPktBuilder.set_payload` instead. :parameters: layer_name: str a string representing the name of the layer. Example: "l2", "l4_tcp", etc. attr : str a string representing the attribute to be changed on desired layer val : value of attribute. toggle_bit : bool Indicating if trying to set a specific bit of a field, such as "do not fragment" bit of IP layer. Default: **False** :raises: + :exc:`KeyError`, in case of missing layer (the desired layer isn't part of packet) + :exc:`ValueError`, in case invalid attribute to the specified layer. """ try: layer = self._pkt_by_hdr[layer_name.lower()] if attr == 'data' and not isinstance(val, dpkt.Packet): # Don't allow setting 'data' attribute raise ValueError("Set a data attribute with object that is not dpkt.Packet is not allowed using " "set_layer_attr method.\nUse set_payload method instead.") if hasattr(layer, attr): if toggle_bit: setattr(layer, attr, val | getattr(layer, attr, 0)) else: setattr(layer, attr, val) if attr == 'data': # re-evaluate packet from the start, possible broken link between layers self._reevaluate_packet(layer_name.lower()) else: raise ValueError("Given attr name '{0}' doesn't exists on specified layer ({1}).".format(layer_name, attr)) except KeyError: raise KeyError("Specified layer '{0}' doesn't exist on packet.".format(layer_name)) def set_layer_bit_attr(self, layer_name, attr, val): """ This method enables the user to set the value of a field smaller that 1 Byte in size. This method isn't used to set full-sized fields value (>= 1 byte). Use :func:`packet_builder.CTRexPktBuilder.set_layer_attr` instead. :parameters: layer_name: str a string representing the name of the layer. Example: "l2", "l4_tcp", etc. attr : str a string representing the attribute to be set on desired layer val : int value of attribute. This value will be set "ontop" of the existing value using bitwise "OR" operation. .. tip:: It is very useful to use dpkt constants to define the values of these fields. :raises: + :exc:`KeyError`, in case of missing layer (the desired layer isn't part of packet) + :exc:`ValueError`, in case invalid attribute to the specified layer. """ return self.set_layer_attr(layer_name, attr, val, True) def set_pkt_payload(self, payload): """ This method sets a payload to the topmost layer of the generated packet. This method isn't to be used to set another networking layer to the packet. Use :func:`packet_builder.CTRexPktBuilder.set_layer_attr` instead. :parameters: payload: a payload to be added to the packet at the topmost layer. this object cannot be of type dpkt.Packet. :raises: + :exc:`AttributeError`, in case no underlying header to host the payload. """ assert isinstance(payload, str) try: self._pkt_top_layer.data = payload except AttributeError: raise AttributeError("The so far built packet doesn't contain an option for payload attachment.\n" "Make sure to set appropriate underlying header before adding payload") def load_packet(self, packet): """ This method enables the user to change a value of a previously defined packet layer. :parameters: packet: dpkt.Packet obj a dpkt object that represents a packet. :raises: + :exc:`CTRexPktBuilder.IPAddressError`, in case invalid ip type option specified. """ assert isinstance(packet, dpkt.Packet) self._packet = copy.copy(packet) self._pkt_by_hdr.clear() self._pkt_top_layer = self._packet # analyze packet to layers tmp_layer = self._packet while True: if isinstance(tmp_layer, dpkt.Packet): layer_name = self._gen_layer_name(type(tmp_layer).__name__) self._pkt_by_hdr[layer_name] = tmp_layer self._pkt_top_layer = tmp_layer try: # check existence of upper layer tmp_layer = tmp_layer.data except AttributeError: # this is the most upper header self._pkt_by_hdr['pkt_final_payload'] = tmp_layer.data break else: self._pkt_by_hdr['pkt_final_payload'] = tmp_layer break return def load_packet_from_pcap(self, pcap_path): """ This method loads a pcap file into a parsed packet builder object. :parameters: pcap_path: str a path to a pcap file, containing a SINGLE packet. :raises: + :exc:`IOError`, in case provided path doesn't exists. """ with open(pcap_path, 'r') as f: pcap = dpkt.pcap.Reader(f) first_packet = True for _, buf in pcap: # this is an iterator, can't evaluate the number of files in advance if first_packet: self.load_packet(dpkt.ethernet.Ethernet(buf)) else: raise ValueError("Provided pcap file contains more than single packet.") # arrive here ONLY if pcap contained SINGLE packet return def load_from_stream_obj(self, stream_obj): self.load_packet_from_byte_list(stream_obj['packet']['binary']) def load_packet_from_byte_list(self, byte_list): # convert byte array into buffer buf = struct.pack('B'*len(byte_list), *byte_list) # thn, load it based on dpkt parsing self.load_packet(dpkt.ethernet.Ethernet(buf)) def get_packet(self, get_ptr=False): """ This method provides access to the built packet, as an instance or as a pointer to packet itself. :parameters: get_ptr : bool indicate whether to get a reference to packet or a copy. Use only in advanced modes if set to true, metadata for packet is cleared, and any further modification is not guaranteed. default value : False :return: + the current packet built by CTRexPktBuilder object. + None if packet is empty """ if get_ptr: self._pkt_by_hdr = {} self._pkt_top_layer = None return self._packet else: return copy.copy(self._packet) def get_packet_length(self): return len(self._packet) def get_layer(self, layer_name): """ This method provides access to a specific layer of the packet, as a **copy of the layer instance**. :parameters: layer_name : str the name given to desired layer :return: + a copy of the desired layer of the current packet if exists. + None if no such layer """ layer = self._pkt_by_hdr.get(layer_name) return copy.copy(layer) if layer else None # VM access methods def set_vm_ip_range(self, ip_layer_name, ip_field, ip_init, ip_start, ip_end, add_value, operation, is_big_endian=False, val_size=4, ip_type="ipv4", add_checksum_inst=True): if ip_field not in ["src", "dst"]: raise ValueError("set_vm_ip_range only available for source ('src') or destination ('dst') ip addresses") # set differences between IPv4 and IPv6 if ip_type == "ipv4": ip_class = dpkt.ip.IP ip_addr_size = val_size if val_size <= 4 else 4 elif ip_type == "ipv6": ip_class = dpkt.ip6.IP6 ip_addr_size = val_size if val_size <= 8 else 4 else: raise CTRexPktBuilder.IPAddressError() self._verify_layer_prop(ip_layer_name, ip_class) trim_size = ip_addr_size*2 init_val = int(binascii.hexlify(CTRexPktBuilder._decode_ip_addr(ip_init, ip_type))[-trim_size:], 16) start_val = int(binascii.hexlify(CTRexPktBuilder._decode_ip_addr(ip_start, ip_type))[-trim_size:], 16) end_val = int(binascii.hexlify(CTRexPktBuilder._decode_ip_addr(ip_end, ip_type))[-trim_size:], 16) # All validations are done, start adding VM instructions flow_var_name = "{layer}__{field}".format(layer=ip_layer_name, field=ip_field) hdr_offset, field_abs_offset = self._calc_offset(ip_layer_name, ip_field, ip_addr_size) self.vm.add_flow_man_inst(flow_var_name, size=ip_addr_size, operation=operation, init_value=init_val, min_value=start_val, max_value=end_val) self.vm.add_write_flow_inst(flow_var_name, field_abs_offset) self.vm.set_vm_off_inst_field(flow_var_name, "add_value", add_value) self.vm.set_vm_off_inst_field(flow_var_name, "is_big_endian", is_big_endian) if ip_type == "ipv4" and add_checksum_inst: self.vm.add_fix_checksum_inst(self._pkt_by_hdr.get(ip_layer_name), hdr_offset) def set_vm_eth_range(self, eth_layer_name, eth_field, mac_init, mac_start, mac_end, add_value, operation, val_size=4, is_big_endian=False): if eth_field not in ["src", "dst"]: raise ValueError("set_vm_eth_range only available for source ('src') or destination ('dst') eth addresses") self._verify_layer_prop(eth_layer_name, dpkt.ethernet.Ethernet) eth_addr_size = val_size if val_size <= 4 else 4 trim_size = eth_addr_size*2 init_val = int(binascii.hexlify(CTRexPktBuilder._decode_mac_addr(mac_init))[-trim_size:], 16) start_val = int(binascii.hexlify(CTRexPktBuilder._decode_mac_addr(mac_start))[-trim_size:], 16) end_val = int(binascii.hexlify(CTRexPktBuilder._decode_mac_addr(mac_end))[-trim_size:], 16) # All validations are done, start adding VM instructions flow_var_name = "{layer}__{field}".format(layer=eth_layer_name, field=eth_field) hdr_offset, field_abs_offset = self._calc_offset(eth_layer_name, eth_field, eth_addr_size) self.vm.add_flow_man_inst(flow_var_name, size=8, operation=operation, init_value=init_val, min_value=start_val, max_value=end_val) self.vm.add_write_flow_inst(flow_var_name, field_abs_offset) self.vm.set_vm_off_inst_field(flow_var_name, "add_value", add_value) self.vm.set_vm_off_inst_field(flow_var_name, "is_big_endian", is_big_endian) def set_vm_custom_range(self, layer_name, hdr_field, init_val, start_val, end_val, add_val, val_size, operation, is_big_endian=False, range_name="", add_checksum_inst=True): # verify input validity for init/start/end values for val in [init_val, start_val, end_val]: if not isinstance(val, int): raise ValueError("init/start/end values are expected integers, but received type '{0}'". format(type(val))) self._verify_layer_prop(layer_name=layer_name, field_name=hdr_field) if not range_name: range_name = "{layer}__{field}".format(layer=layer_name, field=hdr_field) trim_size = val_size*2 hdr_offset, field_abs_offset = self._calc_offset(layer_name, hdr_field, val_size) self.vm.add_flow_man_inst(range_name, size=val_size, operation=operation, init_value=init_val, min_value=start_val, max_value=end_val) self.vm.add_write_flow_inst(range_name, field_abs_offset) self.vm.set_vm_off_inst_field(range_name, "add_value", add_val) self.vm.set_vm_off_inst_field(range_name, "is_big_endian", is_big_endian) if isinstance(self._pkt_by_hdr.get(layer_name), dpkt.ip.IP) and add_checksum_inst: self.vm.add_fix_checksum_inst(self._pkt_by_hdr.get(layer_name), hdr_offset) def get_vm_data(self): return self.vm.dump() def dump_pkt(self): """ Dumps the packet as a decimal array of bytes (each item x gets value between 0-255) :parameters: None :return: + packet representation as array of bytes :raises: + :exc:`CTRexPktBuilder.EmptyPacketError`, in case packet is empty. """ if self._packet is None: raise CTRexPktBuilder.EmptyPacketError() pkt_in_hex = binascii.hexlify(str(self._packet)) return {"binary": [int(pkt_in_hex[i:i+2], 16) for i in range(0, len(pkt_in_hex), 2)], "meta": self.metadata} # return [pkt_in_hex[i:i+2] for i in range(0, len(pkt_in_hex), 2)] def dump_pkt_to_pcap(self, file_path, ts=None): """ Dumps the packet as a decimal array of bytes (each item x gets value between 0-255) :parameters: file_path : str a path (including filename) to which to write to pcap file to. ts : int a timestamp to attach to the packet when dumped to pcap file. if ts in None, then time.time() is used to set the timestamp. Default: **None** :return: None :raises: + :exc:`CTRexPktBuilder.EmptyPacketError`, in case packet is empty. """ if self._packet is None: raise CTRexPktBuilder.EmptyPacketError() try: with open(file_path, 'wb') as f: pcap_wr = dpkt.pcap.Writer(f) pcap_wr.writepkt(self._packet, ts) return except IOError: raise IOError(2, "The provided path could not be accessed") def get_packet_layers(self, depth_limit=Ellipsis): if self._packet is None: raise CTRexPktBuilder.EmptyPacketError() cur_layer = self._packet layer_types = [] if depth_limit == Ellipsis: iterator = itertools.count(1) else: iterator = xrange(depth_limit) for _ in iterator: # append current layer type if isinstance(cur_layer, dpkt.Packet): layer_types.append(type(cur_layer).__name__) else: # if not dpkt layer, refer as payload layer_types.append("PLD") # advance to next layer if not hasattr(cur_layer, "data"): break else: cur_layer = cur_layer.data return layer_types def export_pkt(self, file_path, link_pcap=False, pcap_name=None, pcap_ts=None): pass # ----- internal methods ----- # def _reevaluate_packet(self, layer_name): cur_layer = self._packet known_layers = set(self._pkt_by_hdr.keys()) found_layers = set() while True: pointing_layer_name = self._find_pointing_layer(known_layers, cur_layer) found_layers.add(pointing_layer_name) if self._pkt_by_hdr[layer_name] is cur_layer: self._pkt_top_layer = cur_layer disconnected_layers = known_layers.difference(found_layers) # remove disconnected layers for layer in disconnected_layers: self._pkt_by_hdr.pop(layer) break else: cur_layer = cur_layer.data def _gen_layer_name(self, layer_class_name): assert isinstance(layer_class_name, str) layer_name = layer_class_name.lower() idx = 1 while True: tmp_name = "{name}_{id}".format(name=layer_name, id=idx) if tmp_name not in self._pkt_by_hdr: return tmp_name else: idx += 1 def _find_pointing_layer(self, known_layers, layer_obj): assert isinstance(known_layers, set) for layer in known_layers: if self._pkt_by_hdr[layer] is layer_obj: return layer def _calc_offset(self, layer_name, hdr_field, hdr_field_size): pkt_header = self._pkt_by_hdr.get(layer_name) hdr_offset = len(self._packet) - len(pkt_header) inner_hdr_offsets = [] for field in pkt_header.__hdr__: if field[0] == hdr_field: field_size = struct.calcsize(field[1]) if field_size == hdr_field_size: break elif field_size < hdr_field_size: raise CTRexPktBuilder.PacketLayerError(layer_name, "The specified field '{0}' size is smaller than given range" " size ('{1}')".format(hdr_field, hdr_field_size)) else: inner_hdr_offsets.append(field_size - hdr_field_size) break else: inner_hdr_offsets.append(struct.calcsize(field[1])) return hdr_offset, hdr_offset + sum(inner_hdr_offsets) def _verify_layer_prop(self, layer_name, layer_type=None, field_name=None): if layer_name not in self._pkt_by_hdr: raise CTRexPktBuilder.PacketLayerError(layer_name) pkt_layer = self._pkt_by_hdr.get(layer_name) if layer_type: # check for layer type if not isinstance(pkt_layer, layer_type): raise CTRexPktBuilder.PacketLayerTypeError(layer_name, type(pkt_layer), layer_type) if field_name and not hasattr(pkt_layer, field_name): # check if field exists on certain header raise CTRexPktBuilder.PacketLayerError(layer_name, "The specified field '{0}' does not exists on " "given packet layer ('{1}')".format(field_name, layer_name)) return @staticmethod def _decode_mac_addr(mac_addr): """ Static method to test for MAC address validity. :parameters: mac_addr : str a string representing an MAC address, separated by ':' or '-'. examples: '00:de:34:ef:2e:f4', '00-de-34-ef-2e-f4 :return: + an hex-string representation of the MAC address. for example, ip 00:de:34:ef:2e:f4 will return '\x00\xdeU\xef.\xf4' :raises: + :exc:`CTRexPktBuilder.MACAddressError`, in case invalid ip type option specified. """ tmp_mac = mac_addr.lower().replace('-', ':') if re.match("[0-9a-f]{2}([-:])[0-9a-f]{2}(\\1[0-9a-f]{2}){4}$", tmp_mac): return binascii.unhexlify(tmp_mac.replace(':', '')) # another option for both Python 2 and 3: # codecs.decode(tmp_mac.replace(':', ''), 'hex') else: raise CTRexPktBuilder.MACAddressError() @staticmethod def _decode_ip_addr(ip_addr, ip_type): """ Static method to test for IPv4/IPv6 address validity. :parameters: ip_addr : str a string representing an IP address (IPv4/IPv6) ip_type : str The type of IP to be checked. Valid types: "ipv4", "ipv6". :return: + an hex-string representation of the ip address. for example, ip 1.2.3.4 will return '\x01\x02\x03\x04' :raises: + :exc:`CTRexPktBuilder.IPAddressError`, in case invalid ip type option specified. """ if ip_type == "ipv4": try: return socket.inet_pton(socket.AF_INET, ip_addr) except AttributeError: # no inet_pton here, sorry # try: return socket.inet_aton(ip_addr) # except socket.error: # return False # return ip_addr.count('.') == 3 except socket.error: # not a valid address raise CTRexPktBuilder.IPAddressError() elif ip_type == "ipv6": try: return socket.inet_pton(socket.AF_INET6, ip_addr) except socket.error: # not a valid address raise CTRexPktBuilder.IPAddressError() else: raise CTRexPktBuilder.IPAddressError() # ------ private classes ------ # class CTRexPayloadGen(object): def __init__(self, packet_ref, max_pkt_size): self._pkt_ref = packet_ref self._max_pkt_size = max_pkt_size def gen_random_str(self): gen_length = self._calc_gen_length() # return a string of size gen_length bytes, to pad the packet to its max_size return ''.join(random.SystemRandom().choice(string.ascii_letters + string.digits) for _ in range(gen_length)) def gen_repeat_ptrn(self, ptrn_to_repeat): gen_length = self._calc_gen_length() if isinstance(ptrn_to_repeat, str): # generate repeated string return (ptrn_to_repeat * (gen_length/len(ptrn_to_repeat) + 1))[:gen_length] elif isinstance(ptrn_to_repeat, int): ptrn = binascii.unhexlify(hex(ptrn_to_repeat)[2:]) return (ptrn * (gen_length/len(ptrn) + 1))[:gen_length] elif isinstance(ptrn_to_repeat, tuple): if not all((isinstance(x, int) and (x < 255) and (x >= 0)) for x in ptrn_to_repeat): raise ValueError("All numbers in tuple must be in range 0 <= number <= 255 ") # generate repeated sequence to_pack = (ptrn_to_repeat * (gen_length/len(ptrn_to_repeat) + 1))[:gen_length] return struct.pack('B'*gen_length, *to_pack) else: raise ValueError("Given ptrn_to_repeat argument type ({0}) is illegal.". format(type(ptrn_to_repeat))) def _calc_gen_length(self): return self._max_pkt_size - len(self._pkt_ref) class CTRexVM(object): """ This class defines the TRex VM which represents how TRex will regenerate packets. The packets will be regenerated based on the built packet containing this class. """ InstStore = namedtuple('InstStore', ['type', 'inst']) def __init__(self): """ Instantiate a CTRexVM object :parameters: None """ super(CTRexPktBuilder.CTRexVM, self).__init__() self.vm_variables = {} self._inst_by_offset = {} # this data structure holds only offset-related instructions, ordered in tuples self._off_inst_by_name = {} self.split_by_var = '' def set_vm_var_field(self, var_name, field_name, val, offset_inst=False): """ Set VM variable field. Only existing variables are allowed to be changed. :parameters: var_name : str a string representing the name of the VM variable to be changed. field_name : str a string representing the field name of the VM variable to be changed. val : a value to be applied to field_name field of the var_name VM variable. :raises: + :exc:`KeyError`, in case invalid var_name has been specified. + :exc:`CTRexPktBuilder.VMVarFieldTypeError`, in case mismatch between `val` and allowed type. + :exc:`CTRexPktBuilder.VMVarValueError`, in case val isn't one of allowed options of field_name. """ if offset_inst: return self._off_inst_by_name[var_name].inst.set_field(field_name, val) else: return self.vm_variables[var_name].set_field(field_name, val) def set_vm_off_inst_field(self, var_name, field_name, val): return self.set_vm_var_field(var_name, field_name, val, True) def add_flow_man_inst(self, name, **kwargs): """ Adds a new flow manipulation object to the VM instance. :parameters: name : str name of the manipulation, must be distinct. Example: 'source_ip_change' **kwargs** : dict optional, set flow_man fields on initialization (key = field_name, val = field_val). Must be used with legit fields, see :func:`CTRexPktBuilder.CTRexVM.CTRexVMVariable.set_field`. :return: None :raises: + :exc:`CTRexPktBuilder.VMVarNameExistsError`, in case of desired flow_man name already taken. + Exceptions from :func:`CTRexPktBuilder.CTRexVM.CTRexVMVariable.set_field` method. Will rise when VM variables were misconfiguration. """ if name not in self.vm_variables: self.vm_variables[name] = self.CTRexVMFlowVariable(name) # try configuring VM instruction attributes for (field, value) in kwargs.items(): self.vm_variables[name].set_field(field, value) else: raise CTRexPktBuilder.VMVarNameExistsError(name) def add_fix_checksum_inst(self, linked_ipv4_obj, offset_to_obj=14, name=None): # check if specified linked_ipv4_obj is indeed an ipv4 object if not (isinstance(linked_ipv4_obj, dpkt.ip.IP)): raise ValueError("The provided layer object is not of IPv4.") if not name: name = "checksum_{off}".format(off=offset_to_obj) # name will override previous checksum inst, OK new_checksum_inst = self.CTRexVMChecksumInst(name, offset_to_obj) # store the checksum inst in the end of the IP header (20 Bytes long) inst = self.InstStore('checksum', new_checksum_inst) self._inst_by_offset[offset_to_obj + 20] = inst self._off_inst_by_name[name] = inst def add_write_flow_inst(self, name, pkt_offset, **kwargs): if name not in self.vm_variables: raise KeyError("Trying to add write_flow_var instruction to a not-exists VM flow variable ('{0}')". format(name)) else: new_write_inst = self.CTRexVMWrtFlowVarInst(name, pkt_offset) # try configuring VM instruction attributes for (field, value) in kwargs.items(): new_write_inst.set_field(field, value) # add the instruction to the date-structure inst = self.InstStore('write', new_write_inst) self._inst_by_offset[pkt_offset] = inst self._off_inst_by_name[name] = inst def load_flow_man(self, flow_obj): """ Loads an outer VM variable (instruction) into current VM. The outer VM variable must contain different name than existing VM variables currently registered on VM. :parameters: flow_obj : CTRexVMVariable a CTRexVMVariable to be loaded into VM variable sets. :return: list holds variables data of VM """ assert isinstance(flow_obj, CTRexPktBuilder.CTRexVM.CTRexVMFlowVariable) if flow_obj.name not in self.vm_variables.keys(): self.vm_variables[flow_obj.name] = flow_obj else: raise CTRexPktBuilder.VMVarNameExistsError(flow_obj.name) def dump(self): """ dumps a VM variables (instructions) and split_by_var into a dict data structure. :parameters: None :return: dict with VM instructions as list and split_by_var as str """ # at first, dump all CTRexVMFlowVariable instructions inst_array = [var.dump() if hasattr(var, 'dump') else var for key, var in self.vm_variables.items()] # then, dump all the CTRexVMWrtFlowVarInst and CTRexVMChecksumInst instructions inst_array += [self._inst_by_offset.get(key).inst.dump() for key in sorted(self._inst_by_offset)] return {'instructions': inst_array, 'split_by_var': self.split_by_var} class CVMAbstractInstruction(object): __metaclass__ = ABCMeta def __init__(self, name): """ Instantiate a CTRexVMVariable object :parameters: name : str a string representing the name of the VM variable. """ super(CTRexPktBuilder.CTRexVM.CVMAbstractInstruction, self).__init__() self.name = name def set_field(self, field_name, val): if not hasattr(self, field_name): raise CTRexPktBuilder.VMFieldNameError(field_name) setattr(self, field_name, val) @abstractmethod def dump(self): pass class CTRexVMFlowVariable(CVMAbstractInstruction): """ This class defines a single VM variable to be used as part of CTRexVar object. """ VALID_SIZE = [1, 2, 4, 8] # size in Bytes VALID_OPERATION = ["inc", "dec", "random"] def __init__(self, name): """ Instantiate a CTRexVMVariable object :parameters: name : str a string representing the name of the VM variable. """ super(CTRexPktBuilder.CTRexVM.CTRexVMFlowVariable, self).__init__(name) # self.name = name self.size = 4 self.big_endian = True self.operation = "inc" # self.split_by_core = False self.init_value = 1 self.min_value = self.init_value self.max_value = self.init_value def set_field(self, field_name, val): """ Set VM variable field. Only existing variables are allowed to be changed. :parameters: field_name : str a string representing the field name of the VM variable to be changed. val : a value to be applied to field_name field of the var_name VM variable. :return: None :raises: + :exc:`CTRexPktBuilder.VMVarNameError`, in case of illegal field name. + :exc:`CTRexPktBuilder.VMVarFieldTypeError`, in case mismatch between `val` and allowed type. + :exc:`CTRexPktBuilder.VMVarValueError`, in case val isn't one of allowed options of field_name. """ if not hasattr(self, field_name): raise CTRexPktBuilder.VMFieldNameError(field_name) elif field_name == "size": if type(val) != int: raise CTRexPktBuilder.VMFieldTypeError("size", int) elif val not in self.VALID_SIZE: raise CTRexPktBuilder.VMFieldValueError("size", self.VALID_SIZE) elif field_name in ["init_value", "min_value", "max_value"]: if type(val) != int: raise CTRexPktBuilder.VMFieldTypeError(field_name, int) elif field_name == "operation": if type(val) != str: raise CTRexPktBuilder.VMFieldTypeError("operation", str) elif val not in self.VALID_OPERATION: raise CTRexPktBuilder.VMFieldValueError("operation", self.VALID_OPERATION) # elif field_name == "split_by_core": # val = bool(val) # update field value on success setattr(self, field_name, val) def dump(self): """ dumps a variable fields in a dictionary data structure. :parameters: None :return: dictionary holds variable data of VM variable """ return {"type": "flow_var", # VM variable dump always refers to manipulate instruction. "name": self.name, "size": self.size, "op": self.operation, # "split_by_core": self.split_by_core, "init_value": str(self.init_value), "min_value": str(self.min_value), "max_value": str(self.max_value)} class CTRexVMChecksumInst(CVMAbstractInstruction): def __init__(self, name, offset): """ Instantiate a CTRexVMChecksumInst object :parameters: name : str a string representing the name of the VM variable. """ super(CTRexPktBuilder.CTRexVM.CTRexVMChecksumInst, self).__init__(name) self.pkt_offset = offset def dump(self): return {"type": "fix_checksum_ipv4", "pkt_offset": int(self.pkt_offset)} class CTRexVMWrtFlowVarInst(CVMAbstractInstruction): def __init__(self, name, pkt_offset): """ Instantiate a CTRexVMWrtFlowVarInst object :parameters: name : str a string representing the name of the VM variable. """ super(CTRexPktBuilder.CTRexVM.CTRexVMWrtFlowVarInst, self).__init__(name) self.pkt_offset = int(pkt_offset) self.add_value = 0 self.is_big_endian = False def set_field(self, field_name, val): if not hasattr(self, field_name): raise CTRexPktBuilder.VMFieldNameError(field_name) elif field_name == 'pkt_offset': raise ValueError("pkt_offset value cannot be changed") cur_attr_type = type(getattr(self, field_name)) if cur_attr_type == type(val): setattr(self, field_name, val) else: CTRexPktBuilder.VMFieldTypeError(field_name, cur_attr_type) def dump(self): return {"type": "write_flow_var", "name": self.name, "pkt_offset": self.pkt_offset, "add_value": int(self.add_value), "is_big_endian": bool(self.is_big_endian) } class CTRexVMChecksumInst(CVMAbstractInstruction): def __init__(self, name, offset): """ Instantiate a CTRexVMChecksumInst object :parameters: name : str a string representing the name of the VM variable. """ super(CTRexPktBuilder.CTRexVM.CTRexVMChecksumInst, self).__init__(name) self.pkt_offset = offset def dump(self): return {"type": "fix_checksum_ipv4", "pkt_offset": int(self.pkt_offset)} class CTRexVMWrtFlowVarInst(CVMAbstractInstruction): def __init__(self, name, pkt_offset): """ Instantiate a CTRexVMWrtFlowVarInst object :parameters: name : str a string representing the name of the VM variable. """ super(CTRexPktBuilder.CTRexVM.CTRexVMWrtFlowVarInst, self).__init__(name) self.pkt_offset = int(pkt_offset) self.add_value = 0 self.is_big_endian = False def set_field(self, field_name, val): if not hasattr(self, field_name): raise CTRexPktBuilder.VMFieldNameError(field_name) elif field_name == 'pkt_offset': raise ValueError("pkt_offset value cannot be changed") cur_attr_type = type(getattr(self, field_name)) if cur_attr_type == type(val): setattr(self, field_name, val) else: CTRexPktBuilder.VMFieldTypeError(field_name, cur_attr_type) def dump(self): return {"type": "write_flow_var", "name": self.name, "pkt_offset": self.pkt_offset, "add_value": int(self.add_value), "is_big_endian": bool(self.is_big_endian) } class CPacketBuildException(Exception): """ This is the general Packet Building error exception class. """ def __init__(self, code, message): self.code = code self.message = message def __str__(self): return self.__repr__() def __repr__(self): return u"[errcode:%r] %r" % (self.code, self.message) class EmptyPacketError(CPacketBuildException): """ This exception is used to indicate an error caused by operation performed on an empty packet. """ def __init__(self, message=''): self._default_message = 'Illegal operation on empty packet.' self.message = message or self._default_message super(CTRexPktBuilder.EmptyPacketError, self).__init__(-10, self.message) class IPAddressError(CPacketBuildException): """ This exception is used to indicate an error on the IP addressing part of the packet. """ def __init__(self, message=''): self._default_message = 'Illegal type or value of IP address has been provided.' self.message = message or self._default_message super(CTRexPktBuilder.IPAddressError, self).__init__(-11, self.message) class MACAddressError(CPacketBuildException): """ This exception is used to indicate an error on the MAC addressing part of the packet. """ def __init__(self, message=''): self._default_message = 'Illegal MAC address has been provided.' self.message = message or self._default_message super(CTRexPktBuilder.MACAddressError, self).__init__(-12, self.message) class PacketLayerError(CPacketBuildException): """ This exception is used to indicate an error caused by operation performed on an non-exists layer of the packet. """ def __init__(self, name, message=''): self._default_message = "The given packet layer name ({0}) does not exists.".format(name) self.message = message or self._default_message super(CTRexPktBuilder.PacketLayerError, self).__init__(-13, self.message) class PacketLayerTypeError(CPacketBuildException): """ This exception is used to indicate an error caused by operation performed on an non-exists layer of the packet. """ def __init__(self, name, layer_type, ok_type, message=''): self._default_message = "The type of packet layer {layer_name} is of type {layer_type}, " \ "and not of the expected {allowed_type}.".format(layer_name=name, layer_type=layer_type, allowed_type=ok_type.__name__) self.message = message or self._default_message super(CTRexPktBuilder.PacketLayerTypeError, self).__init__(-13, self.message) class VMVarNameExistsError(CPacketBuildException): """ This exception is used to indicate a duplicate usage of VM variable. """ def __init__(self, name, message=''): self._default_message = 'The given VM name ({0}) already exists as part of the stream.'.format(name) self.message = message or self._default_message super(CTRexPktBuilder.VMVarNameExistsError, self).__init__(-21, self.message) class VMFieldNameError(CPacketBuildException): """ This exception is used to indicate that an undefined VM var field name has been accessed. """ def __init__(self, name, message=''): self._default_message = "The given VM field name ({0}) is not defined and isn't legal.".format(name) self.message = message or self._default_message super(CTRexPktBuilder.VMFieldNameError, self).__init__(-22, self.message) class VMFieldTypeError(CPacketBuildException): """ This exception is used to indicate an illegal value has type has been given to VM variable field. """ def __init__(self, name, ok_type, message=''): self._default_message = "The desired value of field {field_name} is of type {field_type}, " \ "and not of the allowed {allowed_type}.".format(field_name=name, field_type=type(name).__name__, allowed_type=ok_type.__name__) self.message = message or self._default_message super(CTRexPktBuilder.VMFieldTypeError, self).__init__(-31, self.message) class VMFieldValueError(CPacketBuildException): """ This exception is used to indicate an error an illegal value has been assigned to VM variable field. """ def __init__(self, name, ok_opts, message=''): self._default_message = "The desired value of field {field_name} is illegal.\n" \ "The only allowed options are: {allowed_opts}.".format(field_name=name, allowed_opts=ok_opts) self.message = message or self._default_message super(CTRexPktBuilder.VMFieldValueError, self).__init__(-32, self.message) if __name__ == "__main__": pass