week06

env/lib/python3.12/site-packages/pyshark/ek_field_mapping.py

@@ -0,0 +1,91 @@
+import binascii
+import json
+
+from pyshark import cache
+from pyshark.tshark import tshark
+
+
+_MAPPING_CACHE_NAME = "ek_field_mapping.json"
+
+
+class FieldNotFound(Exception):
+    pass
+
+
+class ProtocolMappingNotInitialized(Exception):
+    pass
+
+
+class _EkFieldMapping:
+
+    def __init__(self):
+        self._protocol_to_mapping = {}
+
+    def load_mapping(self, tshark_version, tshark_path=None):
+        if self._protocol_to_mapping:
+            return
+
+        mapping_cache_file = cache.get_cache_dir(tshark_version).joinpath(_MAPPING_CACHE_NAME)
+        if mapping_cache_file.exists():
+            self._protocol_to_mapping = json.load(mapping_cache_file.open())
+        else:
+            self._protocol_to_mapping = tshark.get_ek_field_mapping(tshark_path=tshark_path)
+            mapping_cache_file.open("w").write(json.dumps(self._protocol_to_mapping))
+
+    def cast_field_value(self, protocol, field_name, field_value):
+        """Casts the field value to its proper type according to the mapping"""
+        if isinstance(field_value, list):
+            return [self.cast_field_value(protocol, field_name, item) for item in field_value]
+        if not isinstance(field_value, str):
+            return field_value
+        field_type = self.get_field_type(protocol, field_name)
+        if field_type == str:
+            return field_value
+        if field_type == int and field_value.startswith("0x"):
+            return int(field_value, 16)
+        if field_type == bytes:
+            try:
+                return binascii.unhexlify(field_value.replace(":", ""))
+            except binascii.Error:
+                return field_value
+
+        try:
+            return field_type(field_value)
+        except ValueError:
+            return field_value
+
+    def get_field_type(self, protocol, field_name):
+        """Gets the Python type for the given field (only for EK fields).
+
+        If we are unfamiliar with the type, str will be returned.
+        """
+        if not self._protocol_to_mapping:
+            raise ProtocolMappingNotInitialized("Protocol mapping not initialized. Call load_mapping() first")
+        if protocol not in self._protocol_to_mapping:
+            raise FieldNotFound(f"Type mapping for protocol {protocol} not found")
+
+        fields = self._protocol_to_mapping[protocol]["properties"]
+        if field_name not in fields:
+            return str
+        return self._get_python_type_for_field_type(fields[field_name]["type"])
+
+    def clear(self):
+        self._protocol_to_mapping.clear()
+
+    @classmethod
+    def _get_python_type_for_field_type(cls, field_type):
+        if field_type in ("integer", "long", "short"):
+            return int
+        if field_type == "float":
+            return float
+        if field_type == "date":
+            # We don't use datetime.datetime because these can be timedeltas as well.
+            # Better let the user decide.
+            return float
+        if field_type == "byte":
+            return bytes
+        # Other known types are IP. Retain as str
+        return str
+
+
+MAPPING = _EkFieldMapping()