From a4130d310241ae6f0db2f197db5b26e8773afdd4 Mon Sep 17 00:00:00 2001
From: Mikaela Szekely <mikaela.szekely@qyriad.me>
Date: Wed, 30 Sep 2020 12:34:23 -0600
Subject: [PATCH] descriptors: add types.descriptors.partial for parsing
 potentially incomplete data

This re-allows parsing partial descriptors, which had previously been allowed until
7672fb522182dabe6afa71a1ee2653a792507359.
---
 usb_protocol/types/descriptor.py              | 59 +++++++++++++++++++
 .../types/descriptors/partial/__init__.py     |  1 +
 .../types/descriptors/partial/standard.py     | 11 ++++
 usb_protocol/types/descriptors/standard.py    |  7 ++-
 4 files changed, 77 insertions(+), 1 deletion(-)
 create mode 100644 usb_protocol/types/descriptors/partial/__init__.py
 create mode 100644 usb_protocol/types/descriptors/partial/standard.py

diff --git a/usb_protocol/types/descriptor.py b/usb_protocol/types/descriptor.py
index f4519f3..93aeec8 100644
--- a/usb_protocol/types/descriptor.py
+++ b/usb_protocol/types/descriptor.py
@@ -7,6 +7,65 @@ import unittest
 import construct
 
 class DescriptorFormat(construct.Struct):
+    """
+    Creates a Construct structure for a USB descriptor, and a corresponding version that
+    supports parsing incomplete binary as `DescriptorType.Partial`, e.g. `DeviceDescriptor.Partial`.
+    """
+
+    def __init__(self, *subcons, _create_partial=True, **subconskw):
+
+        if _create_partial:
+            self.Partial = self._create_partial(*subcons, **subconskw) # pylint: disable=invalid-name
+
+        super().__init__(*subcons, **subconskw)
+
+
+    @classmethod
+    def _get_subcon_field_type(cls, subcon):
+        """ Gets the actual field type for a Subconstruct behind arbitrary levels of `Renamed`s."""
+
+        # DescriptorFields are usually `<Renamed <Renamed <FormatField>>>`.
+        # The type behind the `Renamed`s is the one we're interested in, so let's recursively examine
+        # the child Subconstruct until we get to it.
+
+        if not isinstance(subcon, construct.Renamed):
+            return subcon
+        else:
+            return cls._get_subcon_field_type(subcon.subcon)
+
+
+    @classmethod
+    def _create_partial(cls, *subcons, **subconskw):
+        """ Creates a version of the descriptor format for parsing incomplete binary data as a descriptor.
+
+        This essentially wraps every field after bLength and bDescriptorType in a `construct.Optional`.
+        """
+
+        def _apply_optional(subcon):
+
+            subcon_type = cls._get_subcon_field_type(subcon)
+
+            #
+            # If it's already Optional then we don't need to apply it again.
+            #
+            if isinstance(subcon_type, construct.Select):
+                # construct uses a weird singleton to define Pass. `construct.core.Pass` would normally be
+                # the type's name, but then they create a singleton of that same name, replacing that name and
+                # making the type technically unnamable and only accessable via `type()`.
+                if isinstance(subcon_type.subcons[1], type(construct.Pass)):
+                    return subcon
+
+            return (subcon.name / construct.Optional(subcon_type)) * subcon.docs
+
+        # First store the Subconstructs we don't want to modify: bLength and bDescriptorType,
+        # as these are never optional.
+        new_subcons = list(subcons[0:2])
+
+        # Then apply Optional to all of the rest of the Subconstructs.
+        new_subcons.extend([_apply_optional(subcon) for subcon in subcons[2:]])
+
+        return DescriptorFormat(*new_subcons, _create_partial=False, **subconskw)
+
 
     @staticmethod
     def _to_detail_dictionary(descriptor, use_pretty_names=True):
diff --git a/usb_protocol/types/descriptors/partial/__init__.py b/usb_protocol/types/descriptors/partial/__init__.py
new file mode 100644
index 0000000..eb66653
--- /dev/null
+++ b/usb_protocol/types/descriptors/partial/__init__.py
@@ -0,0 +1 @@
+from .standard import *
diff --git a/usb_protocol/types/descriptors/partial/standard.py b/usb_protocol/types/descriptors/partial/standard.py
new file mode 100644
index 0000000..aa097ff
--- /dev/null
+++ b/usb_protocol/types/descriptors/partial/standard.py
@@ -0,0 +1,11 @@
+""" Convenience aliases for versions of descriptor structs that support parsing incomplete binary data. """
+
+from .. import standard
+
+DeviceDescriptor          = standard.DeviceDescriptor.Partial
+ConfigurationDescriptor   = standard.ConfigurationDescriptor.Partial
+StringDescriptor          = standard.StringDescriptor.Partial
+StringLanguageDescriptor  = standard.StringLanguageDescriptor.Partial
+InterfaceDescriptor       = standard.InterfaceDescriptor.Partial
+EndpointDescriptor        = standard.EndpointDescriptor.Partial
+DeviceQualifierDescriptor = standard.DeviceQualifierDescriptor.Partial
diff --git a/usb_protocol/types/descriptors/standard.py b/usb_protocol/types/descriptors/standard.py
index 4acbd99..25454dd 100644
--- a/usb_protocol/types/descriptors/standard.py
+++ b/usb_protocol/types/descriptors/standard.py
@@ -1,7 +1,12 @@
 #
 # This file is part of usb-protocol.
 #
-""" Structures describing standard USB descriptors. """
+"""
+Structures describing standard USB descriptors. Versions that support parsing incomplete binary data
+are available as `DescriptorType`.Partial, e.g. `DeviceDescriptor.Partial`, and are collectively available
+in the `usb_protocol.types.descriptors.partial.standard` module (which, like the structs in this module,
+can also be imported without `.standard`).
+"""
 
 import unittest
 from enum import IntEnum