week06

env/lib/python3.12/site-packages/pyshark/packet/layers/base.py

@@ -0,0 +1,71 @@
+import os
+import typing
+import io
+import sys
+
+from pyshark.packet import common
+
+DATA_LAYER_NAME = "DATA"
+
+
+class BaseLayer(common.SlotsPickleable):
+    """An object representing a Packet layer."""
+    __slots__ = ["_layer_name"]
+
+    def __init__(self, layer_name):
+        self._layer_name = layer_name
+
+    def get_field(self, name):
+        raise NotImplementedError()
+
+    @property
+    def field_names(self) -> typing.List[str]:
+        """Gets all XML field names of this layer."""
+        raise NotImplementedError()
+
+    def has_field(self, name):
+        return name in self.field_names
+
+    @property
+    def layer_name(self):
+        return self._layer_name
+
+    def get(self, item, default=None):
+        """Gets a field in the layer, or the default if not found.
+
+        Works the same way as getattr, but returns the given default if not the field was not found"""
+        try:
+            return getattr(self, item)
+        except AttributeError:
+            return default
+
+    def __dir__(self):
+        return dir(type(self)) + self.field_names
+
+    def __getattr__(self, item):
+        val = self.get_field(item)
+        if val is None:
+            raise AttributeError(f"{item} does not exist in Layer")
+        return val
+
+    def pretty_print(self, writer=None):
+        if not writer:
+            writer = sys.stdout
+        if self.layer_name == DATA_LAYER_NAME:
+            writer.write('DATA')
+            return
+
+        text = f'Layer {self.layer_name.upper()}{os.linesep}:'
+        writer.write(common.colored(text, color="yellow", attrs=["bold"]))
+        self._pretty_print_layer_fields(writer)
+
+    def _pretty_print_layer_fields(self, terminal_writer: io.IOBase):
+        raise NotImplementedError()
+
+    def __repr__(self):
+        return f'<{self.layer_name.upper()} Layer>'
+
+    def __str__(self):
+        writer = io.StringIO()
+        self.pretty_print(writer=writer)
+        return writer.getvalue()

env/lib/python3.12/site-packages/pyshark/packet/layers/ek_layer.py

@@ -0,0 +1,185 @@
+import abc
+import os
+import io
+
+import typing
+
+from pyshark.packet.common import colored
+from pyshark import ek_field_mapping
+from pyshark.packet.layers.base import BaseLayer
+
+
+class _EkLayerHelperFuncsMixin(abc.ABC):
+    """For methods shared between the EK layer and sublayers"""
+
+    def get_field_as_list(self, name) -> list:
+        """Helper function to get a certain field always as a list.
+
+        Some fields may appear once or more in the packet. The field will appear as a list if it appears more
+        than once. In order to avoid checking certain fields if they're lists or not, this function will
+        return the field inside a list at all times.
+
+        For example, in a DNS packet there may be one or more responses.
+        A packet with with one response (www.google.com) will return:
+            >>> print(pkt.dns.resp_name)
+            "www.google.com"
+        While a packet with two responses will return:
+            >>> print(pkt.dns.resp_name)
+            ["www.google.com", "www.google2.com"]
+
+        To avoid this changing behaviour, use:
+            >>> print(pkt.dns.get_field_as_list("resp_name"))
+            ["www.google.com"]
+        """
+        field_value = self.get_field(name)
+        if isinstance(field_value, list):
+            return field_value
+        return [field_value]
+
+
+class EkLayer(BaseLayer, _EkLayerHelperFuncsMixin):
+    __slots__ = ["_layer_name", "_fields_dict"]
+
+    def __init__(self, layer_name, layer_dict):
+        super().__init__(layer_name)
+        self._fields_dict = layer_dict
+
+    def get_field(self, name) -> typing.Union["EkMultiField", None, str, int, bool, bytes, list]:
+        name = name.replace(".", "_")
+        if name in self._fields_dict:
+            # For cases like "text"
+            return self._get_field_value(name)
+
+        for prefix in self._get_possible_layer_prefixes():
+            nested_field = self._get_nested_field(prefix, name)
+            if nested_field is not None:
+                return nested_field
+
+        return None
+
+    def has_field(self, name) -> bool:
+        """Checks if the field exists, either a nested field or a regular field"""
+        return name in self.field_names or name in self.all_field_names
+
+    @property
+    def field_names(self):
+        return list({field_name.split("_", 1)[0] for field_name in self.all_field_names})
+
+    @property
+    def all_field_names(self):
+        """Gets all field names, including subfields"""
+        names = set()
+        for field_name in self._fields_dict:
+            for prefix in self._get_possible_layer_prefixes():
+                if field_name.startswith(prefix):
+                    names.add(_remove_ek_prefix(prefix, field_name))
+                    break
+        return list(names)
+
+    def _get_field_value(self, full_field_name):
+        """Gets the field value, optionally casting it using the cached field mapping"""
+        field_value = self._fields_dict[full_field_name]
+        return ek_field_mapping.MAPPING.cast_field_value(self._layer_name, full_field_name, field_value)
+
+    def _get_nested_field(self, prefix, name):
+        """Gets a field that is directly on the layer
+
+        Returns either a multifield or a raw value.
+        """
+        # TODO: Optimize
+        field_ek_name = f"{prefix}_{name}"
+        if field_ek_name in self._fields_dict:
+            if self._field_has_subfields(field_ek_name):
+                return EkMultiField(self, self._fields_dict, name,
+                                    value=self._get_field_value(field_ek_name))
+            return self._get_field_value(field_ek_name)
+
+        for possible_nested_name in self._fields_dict:
+            if possible_nested_name.startswith(f"{field_ek_name}_"):
+                return EkMultiField(self, self._fields_dict, name, value=None)
+
+        return None
+
+    def _field_has_subfields(self, field_ek_name):
+        field_ek_name_with_ext = f"{field_ek_name}_"
+        for field_name in self._fields_dict:
+            if field_name.startswith(field_ek_name_with_ext):
+                return True
+        return False
+
+    def _pretty_print_layer_fields(self, file: io.IOBase):
+        for field_name in self.field_names:
+            field = self.get_field(field_name)
+            self._pretty_print_field(field_name, field, file, indent=1)
+
+    def _pretty_print_field(self, field_name, field, file, indent=0):
+        prefix = "\t" * indent
+        if isinstance(field, EkMultiField):
+            file.write(colored(f"{prefix}{field_name}: ", "green", attrs=["bold"]))
+            if field.value is not None:
+                file.write(str(field.value))
+            file.write(os.linesep)
+            for subfield in field.subfields:
+                self._pretty_print_field(subfield, field.get_field(subfield), file,
+                                         indent=indent + 1)
+        else:
+            file.write(colored(f"{prefix}{field_name}: ", "green", attrs=["bold"]))
+            file.write(f"{field}{os.linesep}")
+
+    def _get_possible_layer_prefixes(self):
+        """Gets the possible prefixes for a field under this layer.
+
+        The order matters, longest must be first
+        """
+        return [f"{self._layer_name}_{self._layer_name}", self._layer_name]
+
+
+class EkMultiField(_EkLayerHelperFuncsMixin):
+    __slots__ = ["_containing_layer", "_full_name", "_all_fields", "value"]
+
+    def __init__(self, containing_layer: EkLayer, all_fields, full_name, value=None):
+        self._containing_layer = containing_layer
+        self._full_name = full_name
+        self._all_fields = all_fields
+        self.value = value
+
+    def get_field(self, field_name):
+        return self._containing_layer.get_field(f"{self._full_name}_{field_name}")
+
+    @property
+    def subfields(self):
+        names = set()
+        for field_name in self._containing_layer.all_field_names:
+            if field_name != self._full_name and field_name.startswith(f"{self._full_name}_"):
+                names.add(field_name[len(self._full_name):].split("_")[1])
+        return list(names)
+
+    @property
+    def field_name(self):
+        return self._full_name.split("_")[-1]
+
+    def __getattr__(self, item):
+        value = self.get_field(item)
+        if value is None:
+            raise AttributeError(f"Subfield {item} not found")
+        return value
+
+    def __repr__(self):
+        value = f": {self.value}" if self.value else ""
+        return f"<EkMultiField {self.field_name}{value}>"
+
+    def __dir__(self) -> typing.Iterable[str]:
+        return dir(type(self)) + self.subfields
+
+
+def _remove_ek_prefix(prefix, value):
+    """Removes prefix given and the underscore after it"""
+    return value[len(prefix) + 1:]
+
+
+def _get_subfields(all_fields, field_ek_name):
+    subfield_names = []
+    for field in all_fields:
+        if field != field_ek_name and field.startswith(field_ek_name):
+            subfield_names.append(_remove_ek_prefix(field_ek_name, field))
+    return subfield_names

env/lib/python3.12/site-packages/pyshark/packet/layers/json_layer.py

@@ -0,0 +1,200 @@
+import os
+import io
+
+from pyshark.packet.common import colored
+from pyshark.packet.fields import LayerField
+from pyshark.packet.fields import LayerFieldsContainer
+from pyshark.packet.layers.base import BaseLayer
+
+
+class JsonLayer(BaseLayer):
+    __slots__ = [
+        "duplicate_layers",
+        "_showname_fields_converted_to_regular",
+        "_full_name",
+        "_is_intermediate",
+        "_wrapped_fields",
+        "value",
+        "_all_fields"
+    ] + BaseLayer.__slots__
+
+    def __init__(self, layer_name, layer_dict, full_name=None, is_intermediate=False):
+        """Creates a JsonLayer. All sublayers and fields are created lazily later."""
+        super().__init__(layer_name)
+        self.duplicate_layers = []
+        self._showname_fields_converted_to_regular = False
+        if not full_name:
+            self._full_name = self._layer_name
+        else:
+            self._full_name = full_name
+        self._is_intermediate = is_intermediate
+        self._wrapped_fields = {}
+        if isinstance(layer_dict, list):
+            self.duplicate_layers = [JsonLayer(layer_name, duplicate_dict,
+                                               full_name=full_name, is_intermediate=is_intermediate)
+                                     for duplicate_dict in layer_dict[1:]]
+            layer_dict = layer_dict[0]
+        if not isinstance(layer_dict, dict):
+            self.value = layer_dict
+            self._all_fields = {}
+            return
+
+        self._all_fields = layer_dict
+
+    def get_field(self, name):
+        """Gets a field by its full or partial name."""
+        # We only make the wrappers here (lazily) to avoid creating a ton of objects needlessly.
+        self._convert_showname_field_names_to_field_names()
+        field = self._wrapped_fields.get(name)
+        if field is None:
+            is_fake = False
+            field = self._get_internal_field_by_name(name)
+            if field is None:
+                # Might be a "fake" field in JSON
+                is_fake = self._is_fake_field(name)
+                if not is_fake:
+                    raise AttributeError(f"No such field {name}")
+            field = self._make_wrapped_field(name, field, is_fake=is_fake)
+            self._wrapped_fields[name] = field
+        return field
+
+    @property
+    def field_names(self):
+        self._convert_showname_field_names_to_field_names()
+        return list(set([self._sanitize_field_name(name) for name in self._all_fields
+                         if name.startswith(self._full_name)] +
+                        [name.rsplit('.', 1)[1] for name in self._all_fields if '.' in name]))
+
+    def has_field(self, dotted_name) -> bool:
+        """Checks whether the layer has the given field name.
+
+        Can get a dotted name, i.e. layer.sublayer.subsublayer.field
+        """
+        parts = dotted_name.split('.')
+        cur_layer = self
+        for part in parts:
+            if part in cur_layer.field_names:
+                cur_layer = cur_layer.get_field(part)
+            else:
+                return False
+        return True
+
+    def _pretty_print_layer_fields(self, file: io.IOBase):
+        for field_line in self._get_all_field_lines():
+            if ':' in field_line:
+                field_name, field_line = field_line.split(':', 1)
+                file.write(colored(field_name + ':', "green", ["bold"]))
+            file.write(colored(field_line, attrs=["bold"]))
+
+    def _get_all_field_lines(self):
+        """Returns all lines that represent the fields of the layer (both their names and values)."""
+        for field in self._get_all_fields_with_alternates():
+            yield from self._get_field_or_layer_repr(field)
+
+    def _get_field_or_layer_repr(self, field):
+        if isinstance(field, JsonLayer):
+            yield "\t" + field.layer_name + ":" + os.linesep
+            for line in field._get_all_field_lines():
+                yield "\t" + line
+        elif isinstance(field, list):
+            for subfield_or_layer in field:
+                yield from self._get_field_or_layer_repr(subfield_or_layer)
+        else:
+            yield f"\t{self._sanitize_field_name(field.name)}: {field.raw_value}{os.linesep}"
+
+    def _sanitize_field_name(self, field_name):
+        return field_name.replace(self._full_name + '.', '')
+
+    def _field_name_from_showname(self, field_name):
+        """Converts a 'showname'-like field key to a regular field name
+
+        Sometimes in the JSON, there are "text" type fields which might look like this:
+        "my_layer":
+            {
+                "my_layer.some_field": 1,
+                "Something Special: it's special": {
+                    "my_layer.special_field": "it's special"
+                }
+            }
+
+        We convert the showname key into the field name. The internals will turn into a fake layer.
+        In this case the field will be accessible by pkt.my_layer.something_special.special_field
+        """
+        showname_key = field_name.split(":", 1)[0]
+        return self._full_name + "." + showname_key.lower().replace(" ", "_")
+
+    def _get_all_fields_with_alternates(self):
+        return [self.get_field(name) for name in self.field_names]
+
+    def _convert_showname_field_names_to_field_names(self):
+        """Converts all fields that don't have a proper name (they have a showname name) to a regular name
+
+        See self._field_name_from_showname docs for more.
+        """
+        if self._showname_fields_converted_to_regular:
+            return
+        for field_name in list(self._all_fields):
+            if ":" in field_name:
+                field_value = self._all_fields.pop(field_name)
+                if isinstance(field_value, dict):
+                    # Save the showname
+                    field_value["showname"] = field_name
+                # Convert the old name to the new name.
+                self._all_fields[
+                    self._field_name_from_showname(field_name)] = field_value
+
+        self._showname_fields_converted_to_regular = True
+
+    def _get_internal_field_by_name(self, name):
+        """Gets the field by name, or None if not found."""
+        field = self._all_fields.get(name, self._all_fields.get(f"{self._full_name}.{name}"))
+        if field is not None:
+            return field
+        for field_name in self._all_fields:
+            # Specific name
+            if field_name.endswith(f'.{name}'):
+                return self._all_fields[field_name]
+
+    def _is_fake_field(self, name):
+        # Some fields include parts that are not reflected in the JSON dictionary
+        # i.e. a possible json is:
+        # {
+        #   foo: {
+        #           foo.bar.baz: {
+        #                   foo.baz: 3
+        #               }
+        # }
+        # So in this case we must create a fake layer for "bar".
+        field_full_name = f"{self._full_name}.{name}."
+        for name, field in self._all_fields.items():
+            if name.startswith(field_full_name):
+                return True
+        return False
+
+    def _make_wrapped_field(self, name, field, is_fake=False, full_name=None):
+        """Creates the field lazily.
+
+        If it's a simple field, wraps it in a container that adds extra features.
+        If it's a nested layer, creates a layer for it.
+        If it's an intermediate layer, copies over the relevant fields and creates a new layer for
+        it.
+        """
+        if not full_name:
+            full_name = f"{self._full_name}.{name}"
+
+        if is_fake:
+            # Populate with all fields that are supposed to be inside of it
+            field = {key: value for key, value in self._all_fields.items()
+                     if key.startswith(full_name)}
+        if isinstance(field, dict):
+            if name.endswith('_tree'):
+                name = name.replace('_tree', '')
+                full_name = f'{self._full_name}.{name}'
+            return JsonLayer(name, field, full_name=full_name, is_intermediate=is_fake)
+        elif isinstance(field, list):
+            # For whatever reason in list-type object it goes back to using the original parent name
+            return [self._make_wrapped_field(name, field_part,
+                                             full_name=self._full_name.split('.')[0])
+                    for field_part in field]
+
+        return LayerFieldsContainer(LayerField(name=name, value=field))

env/lib/python3.12/site-packages/pyshark/packet/layers/xml_layer.py

@@ -0,0 +1,142 @@
+import os
+import typing
+import io
+
+from pyshark.packet.common import colored
+from pyshark.packet.fields import LayerField, LayerFieldsContainer
+from pyshark.packet.layers import base
+
+
+class XmlLayer(base.BaseLayer):
+    __slots__ = [
+        "raw_mode",
+        "_all_fields"
+    ] + base.BaseLayer.__slots__
+
+    def __init__(self, xml_obj=None, raw_mode=False):
+        super().__init__(xml_obj.attrib['name'])
+        self.raw_mode = raw_mode
+
+        self._all_fields = {}
+
+        # We copy over all the fields from the XML object
+        # Note: we don't read lazily from the XML because the lxml objects are very memory-inefficient
+        # so we'd rather not save them.
+        for field in xml_obj.findall('.//field'):
+            attributes = dict(field.attrib)
+            field_obj = LayerField(**attributes)
+            if attributes['name'] in self._all_fields:
+                # Field name already exists, add this field to the container.
+                self._all_fields[attributes['name']].add_field(field_obj)
+            else:
+                self._all_fields[attributes['name']] = LayerFieldsContainer(field_obj)
+
+    def get_field(self, name) -> typing.Union[LayerFieldsContainer, None]:
+        """Gets the XML field object of the given name."""
+        # Quicker in case the exact name was used.
+        field = self._all_fields.get(name)
+        if field is not None:
+            return field
+
+        for field_name, field in self._all_fields.items():
+            if self._sanitize_field_name(name) == self._sanitize_field_name(field_name):
+                return field
+        return None
+
+    def get_field_value(self, name, raw=False) -> typing.Union[LayerFieldsContainer, None]:
+        """Tries getting the value of the given field.
+
+        Tries it in the following order: show (standard nice display), value (raw value),
+        showname (extended nice display).
+
+        :param name: The name of the field
+        :param raw: Only return raw value
+        :return: str of value
+        """
+        field = self.get_field(name)
+        if field is None:
+            return None
+
+        if raw:
+            return field.raw_value
+
+        return field
+
+    @property
+    def field_names(self) -> typing.List[str]:
+        """Gets all XML field names of this layer."""
+        return [self._sanitize_field_name(field_name) for field_name in self._all_fields]
+
+    @property
+    def layer_name(self):
+        if self._layer_name == 'fake-field-wrapper':
+            return base.DATA_LAYER_NAME
+        return super().layer_name
+
+    def __getattr__(self, item):
+        val = self.get_field(item)
+        if val is None:
+            raise AttributeError()
+        if self.raw_mode:
+            return val.raw_value
+        return val
+
+    @property
+    def _field_prefix(self) -> str:
+        """Prefix to field names in the XML."""
+        if self.layer_name == 'geninfo':
+            return ''
+        return self.layer_name + '.'
+
+    def _sanitize_field_name(self, field_name):
+        """Sanitizes an XML field name
+
+        An xml field might have characters which would make it inaccessible as a python attribute).
+        """
+        field_name = field_name.replace(self._field_prefix, '')
+        return field_name.replace('.', '_').replace('-', '_').lower()
+
+    def _pretty_print_layer_fields(self, file: io.IOBase):
+        for field_line in self._get_all_field_lines():
+            if ':' in field_line:
+                field_name, field_line = field_line.split(':', 1)
+                file.write(colored(field_name + ':', "green", attrs=["bold"]))
+            file.write(colored(field_line, attrs=["bold"]))
+
+    def _get_all_fields_with_alternates(self):
+        all_fields = list(self._all_fields.values())
+        all_fields += sum([field.alternate_fields for field in all_fields
+                           if isinstance(field, LayerFieldsContainer)], [])
+        return all_fields
+
+    def _get_all_field_lines(self):
+        """Returns all lines that represent the fields of the layer (both their names and values)."""
+        for field in self._get_all_fields_with_alternates():
+            yield from self._get_field_or_layer_repr(field)
+
+    def _get_field_or_layer_repr(self, field):
+        field_repr = self._get_field_repr(field)
+        if field_repr:
+            yield f"\t{field_repr}{os.linesep}"
+
+    def _get_field_repr(self, field):
+        if field.hide:
+            return
+        if field.showname:
+            return field.showname
+        elif field.show:
+            return field.show
+        elif field.raw_value:
+            return f"{self._sanitize_field_name(field.name)}: {field.raw_value}"
+
+    def get_field_by_showname(self, showname) -> typing.Union[LayerFieldsContainer, None]:
+        """Gets a field by its "showname"
+        This is the name that appears in Wireshark's detailed display i.e. in 'User-Agent: Mozilla...',
+        'User-Agent' is the .showname
+        Returns None if not found.
+        """
+        for field in self._get_all_fields_with_alternates():
+            if field.showname_key == showname:
+                # Return it if "XXX: whatever == XXX"
+                return field
+        return None