diff --git a/eval/fuzzing/edited/generate-edited-pcaps.py b/eval/fuzzing/edited/generate-edited-pcaps.py
index 4258592076d50c3835d2cce9904b9010e55ea5c2..4bed52bdc4d1743935d126707150c1e1247adca1 100755
--- a/eval/fuzzing/edited/generate-edited-pcaps.py
+++ b/eval/fuzzing/edited/generate-edited-pcaps.py
@@ -4,6 +4,7 @@
Generate edited PCAPs for firewall evaluation.
"""
+# Import libraries
import os
from pathlib import Path
import glob
@@ -11,7 +12,7 @@ import shutil
import argparse
import json
import logging
-import subprocess
+import pcap_fuzzer # Custom PCAP fuzzing library
def strictly_positive_int(value: any) -> int:
@@ -43,7 +44,6 @@ if __name__ == "__main__":
parent_dir = script_path.parents[1]
base_dir = script_path.parents[3]
devices_dir = os.path.join(base_dir, "devices")
- pcap_tweaker_path = os.path.join(base_dir, "src", "pcap_tweaker", "src", "pcap_tweaker.py")
### LOGGING CONFIGURATION ###
logging.basicConfig(level=logging.INFO)
@@ -89,8 +89,7 @@ if __name__ == "__main__":
pcap_edited_path = os.path.join(pcap_edited_dir, pcap_edited_basename)
# Generate edited PCAP
- cmd = f"python3 {pcap_tweaker_path} -o {pcap_edited_path} -r 5 {pcap_path}"
- subprocess.run(cmd.split())
+ pcap_fuzzer.fuzz_pcaps(pcaps=pcap_path, output=pcap_edited_path, random_range=5)
logging.info(f"Generated edited PCAP {pcap_edited_basename}.")
# Move files to correct directories
diff --git a/eval/fuzzing/edited/link-interactions.py b/eval/fuzzing/edited/link-interactions.py
index 4b5d8d49946ad53b42b1f8fc10d98722dd32f29c..78c4d1b71978d4fcfcfe1b78825fad2463ed2a98 100755
--- a/eval/fuzzing/edited/link-interactions.py
+++ b/eval/fuzzing/edited/link-interactions.py
@@ -5,13 +5,15 @@ Link recorded packet verdicts with corresponding device interactions.
Indicate, for each packet, if it should have been accepted or dropped.
"""
+## Import libraries
import os
from pathlib import Path
import json
import csv
import yaml
import logging
-from yaml_loaders.IncludeLoader import IncludeLoader
+# Import custom PyYAML loader
+from pyyaml_loaders import IncludeLoader
# Verdict values
ACCEPT = "ACCEPT"
diff --git a/requirements.txt b/requirements.txt
index a0d9f3c3415d106876a37ccb49e683d92de1891d..ff155d53425a518287ac8fce5bd820d836e8f666 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,2 +1,6 @@
+# Libs
PyYAML
-Jinja2
\ No newline at end of file
+Jinja2
+# Custom
+pyyaml-loaders
+pcap-fuzzer
diff --git a/src/pcap_tweaker/README.md b/src/pcap_tweaker/README.md
deleted file mode 100644
index ec46c0b22ff52a8d04e1ac66a0091f10aed288bd..0000000000000000000000000000000000000000
--- a/src/pcap_tweaker/README.md
+++ /dev/null
@@ -1,70 +0,0 @@
-# pcap-tweaker
-This program randomly edits packets from a PCAP file,
-one field per edited packet.
-
-The edited field will be chosen at random,
-starting from the highest layer, and going down until it finds a supported protocol layer.
-
-Example: a DNS packet will have one of its DNS fields edited,
-and not one of the UDP or IP fields.
-
-
-## Dependencies
-
-* [Scapy](https://scapy.net/)
- * `pip install scapy`
-
-Install all with:
-```bash
-pip install -r requirements.txt
-```
-
-## Usage
-
-```bash
-python3 pcap_tweaker.py [-h] [-o OUTPUT] [-r RANDOM_RANGE] [-n PACKET_NUMBER] [-d] pcap [pcap ...]
-```
-
-The program produces new PCAP file with the same name as the input files,
-but with the suffix `.edit`.
-The output files will be placed in a directory called `edited`,
-in the same directory as the input files.
-It will be created if it doesn't exist.
-
-The program also produces CSV log files,
-indicating which fields were edited for each packet.
-The log files will be placed in a directory called `logs`,
-in the same directory as the input files.
-It will be created if it doesn't exist.
-
-### Positional arguments
-
-* `pcap`: PCAP file(s) to edit
-
-### Optional arguments
-
-* `-h`, `--help`: show help message and exit
-* `-o`, `--output`: output PCAP (and CSV) file path. Used only if a single input file is specified. Default: `edited/<input_pcap>.edit.pcap`
-* `-r`, `--random-range`: upper bound for the random range, which will select for each packet if it will be edited or not. In practice, each packet will be edited with a probability of `1/r`. Must be a strictly positive integer. Default: `1` (edit all packets).
-* `-n`, `--packet-number`: index of the packet to edit, starting from 1. Can be specified multiple times. If this is used, only the specified packets will be edited, and no random editing will be performed.
-* `-d`, `--dry-run`: don't write the output PCAP file (but still write the CSV log file)
-
-
-## Supported protocols
-
-* Datalink Layer (2)
- * ARP
-* Network Layer (3)
- * IPv4
- * IPv6
-* Transport Layer (4)
- * TCP
- * UDP
- * ICMP
- * IGMP(v2 and v3)
-* Application Layer (7)
- * HTTP
- * DNS
- * DHCP
- * SSDP
- * CoAP
diff --git a/src/pcap_tweaker/requirements.txt b/src/pcap_tweaker/requirements.txt
deleted file mode 100644
index 30564abd03adca5c21ce2616d72f151ceb8db129..0000000000000000000000000000000000000000
--- a/src/pcap_tweaker/requirements.txt
+++ /dev/null
@@ -1 +0,0 @@
-scapy
diff --git a/src/pcap_tweaker/src/.gitignore b/src/pcap_tweaker/src/.gitignore
deleted file mode 100644
index bee8a64b79a99590d5303307144172cfe824fbf7..0000000000000000000000000000000000000000
--- a/src/pcap_tweaker/src/.gitignore
+++ /dev/null
@@ -1 +0,0 @@
-__pycache__
diff --git a/src/pcap_tweaker/src/packet/ARP.py b/src/pcap_tweaker/src/packet/ARP.py
deleted file mode 100644
index d3a95541c302caa910d397ba35dacc42cb509362..0000000000000000000000000000000000000000
--- a/src/pcap_tweaker/src/packet/ARP.py
+++ /dev/null
@@ -1,16 +0,0 @@
-import scapy.all as scapy
-from packet.Packet import Packet
-
-class ARP(Packet):
-
- # Class variables
- name = "ARP"
-
- # Modifiable fields
- fields = {
- "op": "int[1,2]",
- "hwsrc": "mac",
- "hwdst": "mac",
- "psrc": "ipv4",
- "pdst": "ipv4",
- }
diff --git a/src/pcap_tweaker/src/packet/BOOTP.py b/src/pcap_tweaker/src/packet/BOOTP.py
deleted file mode 100644
index c796c1b6fcc50cdd4519477a386114793cd878a2..0000000000000000000000000000000000000000
--- a/src/pcap_tweaker/src/packet/BOOTP.py
+++ /dev/null
@@ -1,99 +0,0 @@
-import logging
-from typing import Tuple
-import random
-import scapy.all as scapy
-from scapy.layers import dhcp
-from packet.Packet import Packet
-
-class BOOTP(Packet):
- """
- Class for DHCP packets.
- """
-
- # Class variables
- name = "BOOTP"
-
- # Modifiable fields
- fields = [
- "chaddr",
- "message-type"
- ]
-
-
- def __init__(self, packet: scapy.Packet, id: int = 0, last_layer_index: int = -1) -> None:
- """
- BOOTP/DHCP packet constructor.
-
- :param packet: Scapy Packet to be edited.
- :param id: Packet integer identifier.
- :param last_layer_index: [Optional] Index of the last layer of the packet.
- If not specified, it will be calculated.
- """
- super().__init__(packet, id, last_layer_index)
- self.dhcp_options = packet.getlayer("DHCP options")
-
-
- def get_dhcp_option(self, option_name) -> Tuple[str, any]:
- """
- Retrieve a DHCP option from the packet.
-
- :param option_name: Name of the DHCP option to retrieve.
- :return: DHCP option, as a tuple (name, value).
- """
- dhcp_options = self.dhcp_options.getfieldval("options")
- for option in dhcp_options:
- if option[0] == option_name:
- return option
-
-
- def set_dhcp_option(self, option_name, option_value) -> None:
- """
- Set a DHCP option in the packet.
-
- :param option_name: Name of the DHCP option to set.
- :param option_value: Value of the DHCP option to set.
- """
- dhcp_options = self.dhcp_options.getfieldval("options")
- for i in range(len(dhcp_options)):
- if dhcp_options[i][0] == option_name:
- dhcp_options[i] = option_name, option_value
- break
- self.dhcp_options.setfieldval("options", dhcp_options)
-
-
- def tweak(self) -> dict:
- """
- Randomly edit a BOOTP/DHCP field, among the following:
- - chaddr (client hardware address)
- - message-type (DHCP message type)
-
- :return: Dictionary containing tweak information.
- """
- # Store old hash value
- old_hash = self.get_hash()
- # Get field which will be modified
- field = random.choice(self.fields)
-
- # Initialize old and new values
- old_value = None
- new_value = None
-
- if field == "chaddr":
- old_value = self.layer.getfieldval("chaddr") # Store old value of field
- new_value = Packet.bytes_edit_char(old_value[:6]) + old_value[6:] # Randomly change one byte in the MAC address
- self.layer.setfieldval("chaddr", new_value) # Set new value for field
-
- elif field == "message-type":
- old_value = self.get_dhcp_option(field)[1] # Store old value of field
- # Modify field value until it is different from old value
- new_value = old_value
- while new_value == old_value:
- # Message type is an integer between 1 and 8
- new_value = random.randint(1, 8)
- self.set_dhcp_option(field, new_value) # Set new value for field
-
- # Update checksums
- self.update_fields()
-
- # Return value: dictionary containing tweak information
- return self.get_dict_log(field, old_value, new_value, old_hash)
diff --git a/src/pcap_tweaker/src/packet/CoAP.py b/src/pcap_tweaker/src/packet/CoAP.py
deleted file mode 100644
index 1855d93bd4f39ca967d70cb7ec41daac125488ba..0000000000000000000000000000000000000000
--- a/src/pcap_tweaker/src/packet/CoAP.py
+++ /dev/null
@@ -1,109 +0,0 @@
-import logging
-import random
-import scapy.all as scapy
-from scapy.contrib import coap
-from packet.Packet import Packet
-
-class CoAP(Packet):
-
- # Class variables
- name = "CoAP"
-
- # Modifiable fields
- fields = {
- "type": "int[0,3]",
- "code": "int[1,4]",
- }
- fields = [
- "type",
- "code",
- "uri"
- ]
-
-
- @staticmethod
- def new_int_value(old_value: int, start: int, end: int) -> int:
- """
- Generate a new random integer value between start and end, different from old_value.
-
- :param old_value: Old value of the integer.
- :param start: Start of the range.
- :param end: End of the range.
- :return: New random integer value.
- :raises ValueError: If start is greater than end.
- """
- # Invalid parameters handling
- if start > end:
- raise ValueError("Start value must be smaller than end value.")
-
- # Generate new random int value
- new_value = old_value
- while new_value == old_value:
- new_value = random.randint(start, end)
- return new_value
-
-
- @staticmethod
- def edit_uri(options: list) -> dict:
- """
- Randomly edit one character in each part of the URI of a CoAP packet.
-
- :param options: List of CoAP options.
- :return: Edited list of CoAP options.
- """
- result = {
- "new_options": [],
- "old_uri": b"",
- "new_uri": b""
- }
- for i in range(len(options)):
- if options[i][0] == "Uri-Path" or options[i][0] == "Uri-Query":
- new_value = Packet.bytes_edit_char(options[i][1])
- result["new_options"].append((options[i][0], new_value))
- prefix = b"/?" if options[i][0] == "Uri-Query" else b"/"
- result["old_uri"] += prefix + options[i][1]
- result["new_uri"] += prefix + new_value
- else:
- result["new_options"].append(options[i])
- return result
-
-
- def tweak(self) -> dict:
- """
- Randomly edit one field of the CoAP packet, among the following:
- - type
- - code
- - uri
-
- :return: Dictionary containing tweak information.
- """
- # Store old hash value
- old_hash = self.get_hash()
- # Get field which will be modified
- field = random.choice(self.fields)
-
- # Initialize old and new values
- old_value = None
- new_value = None
-
- # Chosen field is an integer
- if field == "type" or field == "code":
- old_value = self.layer.getfieldval(field)
- if field == "type":
- new_value = CoAP.new_int_value(old_value, 0, 3)
- elif field == "code":
- new_value = CoAP.new_int_value(old_value, 1, 4)
- self.layer.setfieldval(field, new_value)
-
- # Chosen field is the URI
- elif field == "uri":
- result = CoAP.edit_uri(self.layer.getfieldval("options"))
- old_value = result["old_uri"]
- new_value = result["new_uri"]
- self.layer.setfieldval("options", result["new_options"])
-
- # Update checksums
- self.update_fields()
-
- # Return value: dictionary containing tweak information
- return self.get_dict_log(field, old_value, new_value, old_hash)
diff --git a/src/pcap_tweaker/src/packet/DNS.py b/src/pcap_tweaker/src/packet/DNS.py
deleted file mode 100644
index cfb79b95d96df8c88394cc6e21596c1d1f137241..0000000000000000000000000000000000000000
--- a/src/pcap_tweaker/src/packet/DNS.py
+++ /dev/null
@@ -1,128 +0,0 @@
-import random
-import scapy.all as scapy
-from scapy.layers import dns
-from packet.Packet import Packet
-
-class DNS(Packet):
-
- # Class variables
- name = "DNS"
- qtypes = [
- 1, # A
- 2, # NS
- 3, # MD
- 4, # MF
- 5, # CNAME
- 6, # SOA
- 7, # MB
- 8, # MG
- 9, # MR
- 10, # NULL
- 11, # WKS
- 12, # PTR
- 13, # HINFO
- 14, # MINFO
- 15, # MX
- 16, # TXT
- 28, # AAAA
- 41, # OPT
- 255 # ANY
- ]
-
- # Modifiable fields
- fields = [
- "qr",
- "qtype",
- "qname"
- ]
-
-
- @staticmethod
- def iter_question_records(question_records: dns.DNSQRField) -> iter:
- """
- Iterate over question records.
-
- :param question_records: List of question records.
- :return: Iterator over question records.
- """
- layer_idx = 0
- question_record = question_records.getlayer(layer_idx)
- while question_record is not None:
- yield question_record
- layer_idx += 1
- question_record = question_records.getlayer(layer_idx)
-
-
- def get_field(self) -> str:
- """
- Randomly pick a DNS field to be modified.
-
- :return: Field name.
- """
- return random.choice(self.fields)
-
-
- def tweak(self) -> dict:
- """
- Randomly edit one DNS field, among the following:
- - QR flag
- - Query type
- - Query name
-
- :return: Dictionary containing tweak information.
- """
- # Store old hash value
- old_hash = self.get_hash()
- # Get field which will be modified
- field = self.get_field()
-
- # Get auxiliary fields
- qdcount = self.layer.getfieldval("qdcount")
- question_records = self.layer.getfieldval("qd") if qdcount > 0 else None
-
- # Initialize old and new values
- old_value = None
- new_value = None
-
- # Field is QR flag
- if field == "qr":
- # Flip QR flag value
- old_value = self.layer.getfieldval("qr")
- new_value = int(not old_value)
- self.layer.setfieldval("qr", new_value)
-
- # Field is query type
- elif field == "qtype" and question_records is not None:
- old_value = question_records.getfieldval("qtype")
- # Randomly pick new query type
- new_value = old_value
- while new_value == old_value:
- new_value = random.choice(self.qtypes)
- question_records.setfieldval("qtype", new_value)
-
- # Field is query name
- elif field == "qname" and question_records is not None:
- old_value = ""
- new_value = ""
- for question_record in DNS.iter_question_records(question_records):
- if old_value != "":
- old_value += " + "
- old_value_single = question_record.getfieldval("qname")
- old_value += old_value_single.decode("utf-8")
- suffix = old_value_single[-1]
- old_value_trimmed = old_value_single[:-1]
- # Randomly change one character in query name
- new_value_trimmed = old_value_trimmed
- while new_value_trimmed == old_value_trimmed:
- new_value_trimmed = Packet.bytes_edit_char(old_value_trimmed)
- new_value_single = new_value_trimmed + bytes(chr(suffix), "utf-8")
- if new_value != "":
- new_value += " + "
- new_value += new_value_single.decode("utf-8")
- question_record.setfieldval("qname", new_value_single)
-
- # Update checksums
- self.update_fields()
-
- # Return value: dictionary containing tweak information
- return self.get_dict_log(field, old_value, new_value, old_hash)
diff --git a/src/pcap_tweaker/src/packet/HTTP_Request.py b/src/pcap_tweaker/src/packet/HTTP_Request.py
deleted file mode 100644
index 83ca1518b880aaaf0395fca7fb7f91489565190c..0000000000000000000000000000000000000000
--- a/src/pcap_tweaker/src/packet/HTTP_Request.py
+++ /dev/null
@@ -1,23 +0,0 @@
-import scapy.all as scapy
-from scapy.layers import http
-from packet.Packet import Packet
-
-class HTTP_Request(Packet):
-
- # Class variables
- name = "HTTP Request"
-
- # Modifiable fields
- fields = {
- "Method": [
- b"GET",
- b"POST",
- b"PUT",
- b"DELETE",
- b"HEAD",
- b"OPTIONS",
- b"TRACE",
- b"CONNECT"
- ],
- "Path": "bytes"
- }
diff --git a/src/pcap_tweaker/src/packet/ICMP.py b/src/pcap_tweaker/src/packet/ICMP.py
deleted file mode 100644
index 4cdb3c45c27ba7b5deccfa95bcfbf790db154bb1..0000000000000000000000000000000000000000
--- a/src/pcap_tweaker/src/packet/ICMP.py
+++ /dev/null
@@ -1,19 +0,0 @@
-import scapy.all as scapy
-from packet.Packet import Packet
-
-class ICMP(Packet):
-
- # Class variables
- name = "ICMP"
-
- # Modifiable fields
- fields = {
- "type": [
- 0, # Echo Reply
- 3, # Destination Unreachable
- 5, # Redirect
- 8, # Echo Request
- 13, # Timestamp
- 14 # Timestamp Reply
- ]
- }
diff --git a/src/pcap_tweaker/src/packet/IGMP.py b/src/pcap_tweaker/src/packet/IGMP.py
deleted file mode 100644
index b413787d68452909abc851b088ce9ae3617755ea..0000000000000000000000000000000000000000
--- a/src/pcap_tweaker/src/packet/IGMP.py
+++ /dev/null
@@ -1,22 +0,0 @@
-import scapy.all as scapy
-from scapy.contrib import igmp
-from packet.Packet import Packet
-
-class IGMP(Packet):
- """
- IGMP Version 2 packet.
- """
-
- # Class variables
- name = "IGMP"
-
- # Modifiable fields
- fields = {
- "type": [
- 0x11, # Membership Query
- 0x12, # Version 1 Membership Report
- 0x16, # Version 2 Membership Report
- 0x17 # Leave Group
- ],
- "gaddr": "ipv4"
- }
diff --git a/src/pcap_tweaker/src/packet/IGMPv3mr.py b/src/pcap_tweaker/src/packet/IGMPv3mr.py
deleted file mode 100644
index 40d3b819312b9f647db5226e1490fc11f9a4cc12..0000000000000000000000000000000000000000
--- a/src/pcap_tweaker/src/packet/IGMPv3mr.py
+++ /dev/null
@@ -1,43 +0,0 @@
-import logging
-import scapy.all as scapy
-from scapy.contrib import igmpv3
-from packet.Packet import Packet
-
-class IGMPv3mr(Packet):
- """
- IGMP Version 3 Membership Report packet.
- """
-
- # Class variables
- name = "IGMPv3mr"
-
-
- def tweak(self) -> dict:
- """
- Tweak the IGMPv3 Membership Report packet,
- by randomizing all group addresses.
-
- :return: Dictionary containing tweak information.
- """
- # Store old hash value
- old_hash = self.get_hash()
- # Set random IP address for all group records
- old_value = ""
- new_value = ""
- groups = self.packet.getfieldval("records")
- i = 0
- for group in groups:
- if i != 0:
- old_value += "-"
- new_value += "-"
- old_value += group.getfieldval("maddr")
- new_address = Packet.random_ip_address(version=4)
- new_value += new_address
- group.setfieldval("maddr", new_address)
- i += 1
-
- # Update checksums
- self.update_fields()
-
- # Return value: dictionary containing tweak information
- return self.get_dict_log("maddr", old_value, new_value, old_hash)
diff --git a/src/pcap_tweaker/src/packet/IPv4.py b/src/pcap_tweaker/src/packet/IPv4.py
deleted file mode 100644
index 47be7673468a479feb26b1ae893d5e3c73be9a9e..0000000000000000000000000000000000000000
--- a/src/pcap_tweaker/src/packet/IPv4.py
+++ /dev/null
@@ -1,13 +0,0 @@
-import scapy.all as scapy
-from packet.Packet import Packet
-
-class IPv4(Packet):
-
- # Class variables
- name = "IPv4"
-
- # Modifiable fields
- fields = {
- "src": "ipv4",
- "dst": "ipv4"
- }
diff --git a/src/pcap_tweaker/src/packet/IPv6.py b/src/pcap_tweaker/src/packet/IPv6.py
deleted file mode 100644
index acc1aa174558365dfaa9af59ac5c4586d9c826bb..0000000000000000000000000000000000000000
--- a/src/pcap_tweaker/src/packet/IPv6.py
+++ /dev/null
@@ -1,13 +0,0 @@
-import scapy.all as scapy
-from packet.Packet import Packet
-
-class IPv6(Packet):
-
- # Class variables
- name = "IPv6"
-
- # Modifiable fields
- fields = {
- "src": "ipv6",
- "dst": "ipv6"
- }
diff --git a/src/pcap_tweaker/src/packet/Packet.py b/src/pcap_tweaker/src/packet/Packet.py
deleted file mode 100644
index 43f9064545710355f46d9297ec01080a63dd8a23..0000000000000000000000000000000000000000
--- a/src/pcap_tweaker/src/packet/Packet.py
+++ /dev/null
@@ -1,372 +0,0 @@
-from __future__ import annotations
-import importlib
-import logging
-import string
-import re
-import random
-from ipaddress import IPv4Address, IPv6Address
-import scapy.all as scapy
-import hashlib
-
-
-class Packet:
- """
- Wrapper around the Scapy `Packet` class.
- """
-
- ##### CLASS VARIABLES #####
-
- # List of all alphanumerical characters
- ALPHANUM_CHARS = list(string.ascii_letters + string.digits)
- ALPHANUM_BYTES = list(bytes(string.ascii_letters + string.digits, "utf-8"))
- # Minimun payload length (in bytes)
- MIN_PAYLOAD_LENGTH = 46
-
- # Protocol name correspondences
- protocols = {
- "DHCP": "BOOTP"
- }
-
- # Modifiable fields, will be overridden by child classes
- fields = {}
-
-
-
- ##### STATIC METHODS #####
-
-
- @staticmethod
- def string_edit_char(s: str) -> str:
- """
- Randomly change one character in a string.
-
- :param s: String to be edited.
- :return: Edited string.
- """
- char = random.choice(Packet.ALPHANUM_CHARS)
- new_value = list(s)
- new_value[random.randint(0, len(new_value) - 1)] = char
- return "".join(new_value)
-
-
- @staticmethod
- def bytes_edit_char(s: bytes) -> bytes:
- """
- Randomly change one character in a byte array.
-
- :param s: Byte array to be edited.
- :return: Edited byte array.
- """
- byte = random.choice(Packet.ALPHANUM_BYTES)
- new_value = list(s)
- new_value[random.randint(0, len(new_value) - 1)] = byte
- return bytes(new_value)
-
-
- @staticmethod
- def random_mac_address() -> str:
- """
- Generate a random MAC address.
-
- :return: Random MAC address.
- """
- return ":".join(["%02x" % random.randint(0, 255) for _ in range(6)])
-
-
- @staticmethod
- def random_ip_address(version: int = 4) -> str:
- """
- Generate a random IP address.
-
- :param version: IP version (4 or 6).
- :return: Random IP address.
- :raises ValueError: If IP version is not 4 or 6.
- """
- if version == 4:
- return str(IPv4Address(random.randint(0, IPv4Address._ALL_ONES)))
- elif version == 6:
- return str(IPv6Address(random.randint(0, IPv6Address._ALL_ONES)))
- else:
- raise ValueError("Invalid IP version (should be 4 or 6).")
-
-
- @staticmethod
- def get_last_layer_index(packet: scapy.Packet) -> int:
- """
- Get the index of the last layer of a Scapy packet.
-
- :param packet: Scapy Packet.
- :return: index of the last packet layer.
- """
- i = 0
- layer = packet.getlayer(i)
- while layer is not None:
- i += 1
- layer = packet.getlayer(i)
- return i - 1
-
-
- @staticmethod
- def rebuild_packet(packet: scapy.Packet) -> scapy.Packet:
- """
- Rebuild a Scapy packet from its bytes representation,
- but keep its old timestamp.
-
- :param packet: Scapy packet
- :return: Rebuilt Scapy packet, with old timestamp
- """
- timestamp = packet.time
- new_packet = packet.__class__(bytes(packet))
- new_packet.time = timestamp
- return new_packet
-
-
- @classmethod
- def init_packet(c, packet: scapy.Packet, id: int = 0, last_layer_index: int = -1) -> Packet:
- """
- Factory method to create a packet of a given protocol.
-
- :param packet: Scapy Packet to be edited.
- :param id: [Optional] Packet integer identifier. Default is 0.
- :param last_layer_index: [Optional] Index of the last layer of the packet.
- If not specified, it will be calculated.
- :return: Packet of given protocol,
- or generic Packet if protocol is not supported.
- """
- # Try creating specific packet if possible
- if last_layer_index == -1:
- last_layer_index = Packet.get_last_layer_index(packet)
- for i in range(last_layer_index, -1, -1):
- layer = packet.getlayer(i)
- try:
- protocol = layer.name.replace(" ", "_")
- if protocol == "IP" and packet.getfieldval("version") == 4:
- # IPv4 packet
- protocol = "IPv4"
- elif protocol == "IP" and packet.getfieldval("version") == 6:
- # IPv6 packet
- protocol = "IPv6"
- elif protocol == "DNS" and packet.getfieldval("sport") == 5353 and packet.getfieldval("sport") == 5353:
- # mDNS packet
- protocol = "mDNS"
- else:
- protocol = Packet.protocols.get(protocol, protocol)
- module = importlib.import_module(f"packet.{protocol}")
- cls = getattr(module, protocol)
- return cls(packet, id, i)
- except ModuleNotFoundError:
- # Layer protocol not supported
- continue
- # No supported protocol found, raise ValueError
- raise ValueError(f"No supported protocol found for packet: {packet.summary()}")
-
-
-
- ##### INSTANCE METHODS #####
-
-
- def __init__(self, packet: scapy.Packet, id: int = 0, last_layer_index: int = -1) -> None:
- """
- Generic packet constructor.
-
- :param packet: Scapy Packet to be edited.
- :param id: Packet integer identifier.
- :param last_layer_index: [Optional] Index of the last layer of the packet.
- If not specified, it will be calculated.
- """
- self.id = id
- self.packet = packet
- self.layer_index = last_layer_index if last_layer_index != -1 else Packet.get_last_layer_index(packet)
- self.layer = packet.getlayer(self.name)
- if self.layer is None:
- self.layer = packet.getlayer(self.layer_index)
-
-
- def get_packet(self) -> scapy.Packet:
- """
- Get Scapy packet.
-
- :return: Scapy Packet.
- """
- return self.packet
-
-
- def get_length(self) -> int:
- """
- Get packet length.
-
- :return: Packet length.
- """
- return len(self.packet)
-
-
- def get_length_from_layer(self, layer: int | str) -> int:
- """
- Get packet length, starting from a given layer.
-
- :param layer: Layer index or name.
- :return: Packet length starting from the given layer.
- """
- return len(self.packet.getlayer(layer))
-
-
- def get_layer_index(self) -> int:
- """
- Get packet layer index.
-
- :return: Packet layer index.
- """
- return self.layer_index
-
-
- def get_hash(self) -> str:
- """
- Get packet payload SHA256 hash.
- The payload is first padded with null bytes to reach the minimum Ethernet payload length of 46 bytes.
-
- :return: Packet payload SHA256 hash.
- """
- pad_bytes_to_add = Packet.MIN_PAYLOAD_LENGTH - len(self.packet.payload)
- payload = bytes(self.packet.payload) + bytes(pad_bytes_to_add) if pad_bytes_to_add > 0 else bytes(self.packet.payload)
- return hashlib.sha256(payload).hexdigest()
-
-
- def rebuild(self) -> None:
- """
- Rebuild packet, but keep old timestamp.
- """
- timestamp = self.packet.time
- self.packet = self.packet.__class__(bytes(self.packet))
- self.packet.time = timestamp
-
-
- def update_fields(self) -> None:
- """
- Update checksum and length fields on all relevant layers,
- and rebuild packet.
- """
- # Loop on all packet layers
- i = 0
- while True:
- layer = self.packet.getlayer(i)
- if layer is None:
- break
-
- # Delete checksum field
- if hasattr(layer, "chksum") and layer.getfieldval("chksum") is not None:
- layer.delfieldval("chksum")
-
- # Delete length field
- if hasattr(layer, "len") and layer.getfieldval("len") is not None:
- layer.delfieldval("len")
-
- i += 1
-
- # Rebuild packet, to update deleted fields
- self.rebuild()
-
-
- def get_dict_log(self, field: str, old_value: str, new_value: str, old_hash: str) -> dict:
- """
- Log packet field modification,
- and return a dictionary containing tweak information.
-
- :param field: Field name.
- :param old_value: Old field value.
- :param new_value: New field value.
- :param old_hash: Old packet hash (before tweak).
- :return: Dictionary containing tweak information.
- """
- timestamp = self.packet.time
- logging.info(f"Packet {self.id}, timestamp {timestamp}: {self.name}.{field} = {old_value} -> {new_value}")
- d = {
- "id": self.id,
- "timestamp": timestamp,
- "protocol": self.name,
- "field": field,
- "old_value": old_value,
- "new_value": new_value,
- "old_hash": old_hash,
- "new_hash": self.get_hash()
- }
- return d
-
-
- def tweak(self) -> dict:
- """
- Randomly edit one packet field.
-
- :return: Dictionary containing tweak information,
- or None if no tweak was performed.
- """
- # Store old hash value
- old_hash = self.get_hash()
- # Get field which will be modified
- field, value_type = random.choice(list(self.fields.items()))
- # Store old value of field
- old_value = self.layer.getfieldval(field)
-
- # Modify field value until it is different from old value
- new_value = old_value
- while new_value == old_value:
-
- if isinstance(value_type, list):
- # Field value is a list
- # Choose randomly a value from the list
- values = value_type
- new_value = old_value
- # Randomly pick new value
- new_value = random.choice(values)
-
- elif "int" in value_type:
- # Field value is an integer
- # Generate a random integer between given range
- if value_type == "int":
- # No range given, default is 0-65535
- new_value = random.randint(0, 65535)
- else:
- # Range given
- pattern = re.compile(r"int\[\s*(?P<start>\d+),\s*(?P<end>\d+)\s*\]")
- match = pattern.match(value_type)
- start = int(match.group("start"))
- end = int(match.group("end"))
- new_value = random.randint(start, end)
-
- elif value_type == "str":
- # Field value is a string
- # Randomly change one character
- new_value = Packet.string_edit_char(old_value)
-
- elif value_type == "bytes":
- # Field value is a byte array
- # Randomly change one byte
- new_value = Packet.bytes_edit_char(old_value)
-
- elif value_type == "port":
- # Field value is an port number
- # Generate a random port number between 1024 and 65535
- new_value = random.randint(1024, 65535)
-
- elif value_type == "ipv4":
- # Field value is an IPv4 address
- # Generate a random IPv4 address
- new_value = Packet.random_ip_address(version=4)
-
- elif value_type == "ipv6":
- # Field value is an IPv6 address
- # Generate a random IPv6 address
- new_value = Packet.random_ip_address(version=6)
-
- elif value_type == "mac":
- # Field value is a MAC address
- # Generate a random MAC address
- new_value = Packet.random_mac_address()
-
- # Set new value for field
- self.layer.setfieldval(field, new_value)
-
- # Update checksums
- self.update_fields()
-
- # Return value: dictionary containing tweak information
- return self.get_dict_log(field, old_value, new_value, old_hash)
diff --git a/src/pcap_tweaker/src/packet/TCP.py b/src/pcap_tweaker/src/packet/TCP.py
deleted file mode 100644
index 590b7295ac7fcebe139ec01133c5b1c994408c94..0000000000000000000000000000000000000000
--- a/src/pcap_tweaker/src/packet/TCP.py
+++ /dev/null
@@ -1,15 +0,0 @@
-import scapy.all as scapy
-from packet.Transport import Transport
-
-class TCP(Transport):
-
- # Class variables
- name = "TCP"
-
- # Well-known ports
- ports = [
- 80, # HTTP
- 443, # HTTPS
- 8080, # HTTP alternate
- 9999 # TP-Link
- ]
diff --git a/src/pcap_tweaker/src/packet/Transport.py b/src/pcap_tweaker/src/packet/Transport.py
deleted file mode 100644
index da73a5e99fbefd2696458b6b0c23c98c47324f3e..0000000000000000000000000000000000000000
--- a/src/pcap_tweaker/src/packet/Transport.py
+++ /dev/null
@@ -1,58 +0,0 @@
-import random
-import scapy.all as scapy
-from scapy.layers import http
-from packet.Packet import Packet
-
-class Transport(Packet):
- """
- Transport layer (layer 4) packet, i.e. TCP or UDP.
- """
-
- # Modifiable fields
- fields = {
- "sport": "port",
- "dport": "port"
- }
-
- # Well-known ports, will be overridden by child classes
- ports = []
-
-
- def tweak(self) -> dict:
- """
- If one of the ports is a well-known port,
- randomly edit destination or source port,
- in this respective order of priority.
-
- :return: Dictionary containing tweak information,
- or None if no tweak was performed.
- """
- # Store old hash value
- old_hash = self.get_hash()
- # Check if destination port is a well-known port
- if self.layer.getfieldval("dport") in self.ports:
- field = "dport"
- # Check if source port is a well-known port
- elif self.layer.getfieldval("sport") in self.ports:
- field = "sport"
- else:
- # No well-known port, do not tweak
- return None
-
- # Store old value of field
- old_value = self.layer.getfieldval(field)
-
- # Modify field value until it is different from old value
- new_value = old_value
- while new_value == old_value:
- # Generate a random port number between 1024 and 65535
- new_value = random.randint(1024, 65535)
-
- # Set new value for field
- self.layer.setfieldval(field, new_value)
-
- # Update checksums, if needed
- self.update_fields()
-
- # Return value: dictionary containing tweak information
- return self.get_dict_log(field, old_value, new_value, old_hash)
diff --git a/src/pcap_tweaker/src/packet/UDP.py b/src/pcap_tweaker/src/packet/UDP.py
deleted file mode 100644
index 5c6f444c1cfedd2250c04602449c13a540dbf129..0000000000000000000000000000000000000000
--- a/src/pcap_tweaker/src/packet/UDP.py
+++ /dev/null
@@ -1,22 +0,0 @@
-import scapy.all as scapy
-from packet.Transport import Transport
-
-class UDP(Transport):
-
- # Class variables
- name = "UDP"
-
- # Well-known ports
- ports = [
- 53, # DNS
- 5353, # mDNS
- 67, # DHCP client
- 68, # DHCP server
- 123, # NTP
- 1900, # SSDP
- 3478, # STUN
- 5683, # CoAP
- 9999, # TP-Link
- 20002, # TP-Link
- 54321 # Xiaomi
- ]
diff --git a/src/pcap_tweaker/src/packet/mDNS.py b/src/pcap_tweaker/src/packet/mDNS.py
deleted file mode 100644
index 1c566b47f20e82f77d62f214bbd117335445ad5a..0000000000000000000000000000000000000000
--- a/src/pcap_tweaker/src/packet/mDNS.py
+++ /dev/null
@@ -1,44 +0,0 @@
-import random
-import scapy.all as scapy
-from scapy.layers import dns
-from packet.DNS import DNS
-
-class mDNS(DNS):
-
- # Class variables
- name = "mDNS"
-
- # Modifiable fields
- fields = {
- "query": [
- "qr",
- "qtype",
- "qname"
- ],
- "response": [
- "qr"
- ]
- }
-
-
- def __init__(self, packet: scapy.Packet, id: int = 0, last_layer_index: int = -1) -> None:
- """
- mDNS packet constructor.
-
- :param packet: Scapy packet to be edited.
- :param id: Packet integer identifier.
- :param last_layer_index: [Optional] Index of the last layer of the packet.
- If not specified, it will be calculated.
- """
- super().__init__(packet, id, last_layer_index)
- qr = self.layer.getfieldval("qr")
- self.qr_str = "query" if qr == 0 else "response"
-
-
- def get_field(self) -> str:
- """
- Randomly pick a DNS field to be modified.
-
- :return: Field name.
- """
- return random.choice(self.fields[self.qr_str])
diff --git a/src/pcap_tweaker/src/pcap_tweaker.py b/src/pcap_tweaker/src/pcap_tweaker.py
deleted file mode 100644
index 521d9f0812691f9b68fc4402c54aaaacf8f8575e..0000000000000000000000000000000000000000
--- a/src/pcap_tweaker/src/pcap_tweaker.py
+++ /dev/null
@@ -1,171 +0,0 @@
-"""
-Randomly edit packet fields in a PCAP file.
-"""
-
-import os
-import argparse
-import random
-import logging
-import csv
-import scapy.all as scapy
-from scapy.layers import dhcp, dns, http
-from scapy.contrib import coap, igmp, igmpv3
-from packet.Packet import Packet
-
-
-def strictly_positive_int(value: any) -> int:
- """
- Custom argparse type for a strictly positive integer value.
-
- :param value: argument value to check
- :return: argument as integer if it is strictly positive
- :raises argparse.ArgumentTypeError: if argument does not represent a strictly positive integer
- """
- try:
- ivalue = int(value)
- except ValueError:
- raise argparse.ArgumentTypeError(f"{value} does not represent an integer.")
- else:
- if ivalue < 1:
- raise argparse.ArgumentTypeError(f"{value} does not represent a strictly positive integer.")
- return ivalue
-
-
-def must_edit_packet(i: int, packet_numbers: list, random_range: int) -> bool:
- """
- Check if a packet must be edited.
-
- :param i: packet number (starting from 1)
- :param packet_numbers: list of packet numbers to edit
- :param random_range: upper bound for random range (not included)
- :return: True if packet must be edited, False otherwise
- """
- is_specified = packet_numbers is not None and i in packet_numbers
- is_random = packet_numbers is None and random.randrange(0, random_range) == 0
- return is_specified or is_random
-
-
-def tweak_pcaps(pcaps: list, output: str, random_range: int = 1, packet_numbers: list = None, dry_run: bool = False) -> None:
- """
- Main functionality of the program:
- (Randomly) edit packet fields in a (list of) PCAP file(s).
-
- :param pcaps: list of input PCAP files
- :param output: output PCAP file path. Used only if a single input file is specified.
- :param random_range: upper bound for random range (not included)
- :param packet_numbers: list of packet numbers to edit (starting from 1)
- :param dry_run: if True, do not write output PCAP file
- """
-
- # Loop on given input PCAP files
- for input_pcap in pcaps:
- # PCAP file directory
- input_dir = os.path.dirname(input_pcap)
-
- # Read input PCAP file
- packets = scapy.rdpcap(input_pcap)
- new_packets = []
- logging.info(f"Read input PCAP file: {input_pcap}")
-
- # Open log CSV file
- csv_log = ""
- if output is not None and len(pcaps) == 1:
- csv_log = output.replace(".pcap", ".csv")
- else:
- csv_dir = os.path.join(input_dir, "csv")
- os.makedirs(csv_dir, exist_ok=True)
- csv_log = os.path.basename(input_pcap).replace(".pcap", ".edit.csv")
- csv_log = os.path.join(csv_dir, csv_log)
- with open(csv_log, "w") as csv_file:
- field_names = ["id", "timestamp", "protocol", "field", "old_value", "new_value", "old_hash", "new_hash"]
- writer = csv.DictWriter(csv_file, fieldnames=field_names)
- writer.writeheader()
-
- i = 1
- for packet in packets:
-
- if must_edit_packet(i, packet_numbers, random_range):
- # Edit packet, if possible
- last_layer_index = Packet.get_last_layer_index(packet)
- while True:
- try:
- my_packet = Packet.init_packet(packet, i, last_layer_index)
- except ValueError:
- # No supported protocol found in packet, skip it
- new_packets.append(Packet.rebuild_packet(packet))
- break
- else:
- d = my_packet.tweak()
- if d is None:
- # Packet was not edited, try editing one layer lower
- last_layer_index = my_packet.get_layer_index() - 1
- else:
- # Packet was edited
- new_packets.append(my_packet.get_packet())
- writer.writerow(d)
- break
- else:
- # Packet won't be edited
- new_packets.append(Packet.rebuild_packet(packet))
-
- i += 1
-
- # Write output PCAP file
- output_pcap = ""
- if output is not None and len(pcaps) == 1:
- output_pcap = output
- else:
- output_dir = os.path.join(os.path.dirname(input_pcap), "edited")
- os.makedirs(output_dir, exist_ok=True)
- output_pcap = os.path.basename(input_pcap).replace(".pcap", ".edit.pcap")
- output_pcap = os.path.join(output_dir, output_pcap)
- if dry_run:
- logging.info(f"Dry run: did not write output PCAP file: {output_pcap}")
- else:
- scapy.wrpcap(output_pcap, new_packets)
- logging.info(f"Wrote output PCAP file: {output_pcap}")
-
-
-if __name__ == "__main__":
-
- # This script's name
- script_name = os.path.basename(__file__)
-
- ### LOGGING CONFIGURATION ###
- logging.basicConfig(level=logging.INFO)
- logging.info(f"Starting {script_name}")
-
-
- ### ARGUMENT PARSING ###
- parser = argparse.ArgumentParser(
- prog=script_name,
- description="Randomly edit packet fields in a PCAP file."
- )
- # Positional arguments: input PCAP file(s)
- parser.add_argument("input_pcaps", metavar="pcap", type=str, nargs="+", help="Input PCAP file(s).")
- # Optional flag: -o / --output
- parser.add_argument("-o", "--output", type=str, help="Output PCAP (and CSV) file path. Used only if a single input file is specified. Default: edited/<input_pcap>.edit.pcap")
- # Optional flag: -r / --random-range
- parser.add_argument("-r", "--random-range", type=strictly_positive_int, default=1,
- help="Upper bound for random range (not included). Must be a strictly positive integer. Default: 1 (edit each packet).")
- # Optional flag: -n / --packet-number
- parser.add_argument("-n", "--packet-number", type=int, action="append",
- help="Index of the packet to edit, starting form 1. Can be specifed multiple times.")
- # Optional flag: -d / --dry-run
- parser.add_argument("-d", "--dry-run", action="store_true",
- help="Dry run: do not write output PCAP file.")
- # Parse arguments
- args = parser.parse_args()
- # Verify arguments
- if args.output is not None and len(args.input_pcaps) > 1:
- logging.warning("Multiple input PCAP files specified, ignoring output PCAP file name.")
-
-
- ### MAIN PROGRAM ###
- tweak_pcaps(
- pcaps=args.input_pcaps,
- output=args.output,
- random_range=args.random_range,
- packet_numbers=args.packet_number,
- dry_run=args.dry_run
- )
diff --git a/src/translator/expand.py b/src/translator/expand.py
index efd8563a4a4e856d873a86d303b417d59b659e29..2a1ecbbf20516748869111998eaa43d3a9cb0370 100644
--- a/src/translator/expand.py
+++ b/src/translator/expand.py
@@ -1,7 +1,9 @@
+## Import libraries
import os
import argparse
import yaml
-from yaml_loaders.IncludeLoader import IncludeLoader
+# Import custom PyYAML loader
+from pyyaml_loaders import IncludeLoader
##### MAIN #####
diff --git a/src/translator/translator.py b/src/translator/translator.py
index 9483df073cd6dee170463ddde4add8162693da2f..ac5499898adffba39db43f61f9e01a19a321c638 100644
--- a/src/translator/translator.py
+++ b/src/translator/translator.py
@@ -22,7 +22,7 @@ sys.path.insert(0, os.path.join(script_dir, "protocols"))
from LogType import LogType
from Policy import Policy
from NFQueue import NFQueue
-from yaml_loaders.IncludeLoader import IncludeLoader
+from pyyaml_loaders import IncludeLoader
##### Custom Argparse types #####
diff --git a/src/translator/yaml_loaders/IgnoreLoader.py b/src/translator/yaml_loaders/IgnoreLoader.py
deleted file mode 100644
index f0a8852a80649f66961c69b87ad0e091dcc5c754..0000000000000000000000000000000000000000
--- a/src/translator/yaml_loaders/IgnoreLoader.py
+++ /dev/null
@@ -1,30 +0,0 @@
-"""
-PyYAML loader which ignores tags.
-Adapted from https://stackoverflow.com/questions/33048540/pyyaml-safe-load-how-to-ignore-local-tags.
-"""
-
-import yaml
-
-
-class IgnoreLoader(yaml.SafeLoader):
- """
- Custom PyYAML loader, which ignores tags.
- """
- def __init__(self, stream) -> None:
- # Use parent constructor
- super().__init__(stream)
-
-
-def construct_ignore(loader: IgnoreLoader, tag_suffix: str, node: yaml.Node) -> None:
- """
- PyYAML constructor which ignores tags.
-
- :param loader: PyYAML IgnoreLoader
- :param tag_suffix: YAML tag suffix
- :param node: YAML node, i.e. the value occurring after the tag
- """
- return None
-
-
-# Add custom constructor
-yaml.add_multi_constructor("!", construct_ignore, IgnoreLoader)
diff --git a/src/translator/yaml_loaders/IncludeLoader.py b/src/translator/yaml_loaders/IncludeLoader.py
deleted file mode 100644
index 29a7613cb1e590947b3e221eecd5ca743dbf1b39..0000000000000000000000000000000000000000
--- a/src/translator/yaml_loaders/IncludeLoader.py
+++ /dev/null
@@ -1,148 +0,0 @@
-"""
-PyYAML loader which supports inclusion of external members.
-Adapted from https://gist.github.com/joshbode/569627ced3076931b02f.
-"""
-
-import sys
-import os
-import yaml
-import collections.abc
-
-# Import IgnoreLoader
-sys.path.append(os.path.abspath(os.path.dirname(__file__)))
-from IgnoreLoader import IgnoreLoader
-
-
-class IncludeLoader(yaml.SafeLoader):
- """
- Custom PyYAML loader, which supports inclusion of members defined in other YAML files.
- """
- def __init__(self, stream) -> None:
- # Use parent constructor
- super().__init__(stream)
-
-
-def update_dict_aux(d: dict, key: str, parent_key: str, current_parent_key: str, old_val: str, new_val: str) -> None:
- """
- Helper recursive function for `update_dict`.
-
- :param d: dictionary to update
- :param key: key to update the value of
- :param parent_key: parent key of `key`
- :param current_parent_key: current parent key
- :param old_val: value to replace
- :param new_val: value to replace with
- """
- for k, v in d.items():
- if isinstance(v, collections.abc.Mapping):
- # Value is a dictionary itself, recursion time
- update_dict_aux(d.get(k, {}), key, parent_key, k, old_val, new_val)
- else:
- # Value is a scalar
- if k == key and current_parent_key == parent_key and v == old_val:
- d[k] = new_val
-
-
-def update_dict(d: dict, key: str, parent_key: str, old_val: str, new_val: str) -> None:
- """
- Recursively update all occurrences of value `old_val`,
- which are nested under key `key` and parent key `parent_key`,
- with `new_val` in dictionary `d`.
-
- :param d: dictionary to update
- :param key: key to update the value of
- :param parent_key: parent key of `key`
- :param old_val: value to replace
- :param new_val: value to replace with
- """
- update_dict_aux(d, key, parent_key, "", old_val, new_val)
-
-
-def replace_self_addrs(d: dict, mac: str = "", ipv4: str = "", ipv6: str = "") -> None:
- """
- Replace all occurrences of "self" with the given addresses.
-
- :param d: dictionary to update
- :param mac (optional): MAC address to replace "self" with
- :param ipv4 (optional): IPv4 address to replace "self" with
- :param ipv6 (optional): IPv6 address to replace "self" with
- """
- if mac:
- update_dict(d, "sha", "arp", "self", mac)
- update_dict(d, "tha", "arp", "self", mac)
- if ipv4:
- update_dict(d, "src", "ipv4", "self", ipv4)
- update_dict(d, "dst", "ipv4", "self", ipv4)
- if ipv6:
- update_dict(d, "src", "ipv6", "self", ipv6)
- update_dict(d, "dst", "ipv6", "self", ipv6)
-
-
-def construct_include(loader: IncludeLoader, node: yaml.Node) -> dict:
- """
- Include member defined in another YAML file.
-
- :param loader: PyYAML IncludeLoader
- :param node: YAML node, i.e. the value occurring after the tag
- :return: included pattern (from this or another YAML profile)
- """
- scalar = loader.construct_scalar(node)
-
- # Split profile and values
- split1 = scalar.split(" ")
- profile = split1[0]
- values = split1[1:]
-
- # Parse values into dictionary
- values_dict = {}
- for value in values:
- split_value = value.split(":")
- if len(split_value) == 2:
- values_dict[split_value[0]] = split_value[1]
-
- # Split path and pattern from profile
- split2 = profile.split('#')
- path = os.path.abspath(loader.stream.name) # Default path, the current profile
- if len(split2) == 1:
- members = split2[0]
- elif len(split2) == 2:
- if split2[0] != "self":
- path = os.path.join(os.path.dirname(path), split2[0])
- members = split2[1]
-
- # Load member to include
- addrs = {}
- data = {}
- with open(path, 'r') as f:
- data = yaml.load(f, IgnoreLoader)
-
- # Populate addrs
- addrs["mac"] = data["device-info"].get("mac", "")
- addrs["ipv4"] = data["device-info"].get("ipv4", "")
- addrs["ipv6"] = data["device-info"].get("ipv6", "")
-
- for member in members.split('.'):
- data = data[member]
-
- # Populate values
- data_top = data
- for key, value in values_dict.items():
- data = data_top
- split_key = key.split('.')
- i = 0
- for sub_key in split_key:
- if i == len(split_key) - 1:
- data[sub_key] = value
- else:
- data = data[sub_key]
- i += 1
-
- # Replace "self" with actual addresses
- if isinstance(data_top, collections.abc.Mapping):
- replace_self_addrs(data_top, addrs["mac"], addrs["ipv4"], addrs["ipv6"])
-
- return data_top
-
-
-# Add custom constructor
-yaml.add_constructor("!include", construct_include, IncludeLoader)