1 files changed, 131 insertions, 0 deletions
diff --git a/external/construct/debug.py b/external/construct/debug.py
new file mode 100644
index 0000000..3910cae
--- /dev/null
+++ b/external/construct/debug.py
@@ -0,0 +1,131 @@
+"""
+Debugging utilities for constructs
+"""
+import sys
+import traceback
+import pdb
+import inspect
+from construct.core import Construct, Subconstruct
+from construct.lib import HexString, Container, ListContainer
+
+
+class Probe(Construct):
+    """
+    A probe: dumps the context, stack frames, and stream content to the screen
+    to aid the debugging process.
+
+    .. seealso:: :class:`Debugger`.
+
+    :param name: the display name
+    :param show_stream: whether or not to show stream contents. default is True. the stream must be seekable.
+    :param show_context: whether or not to show the context. default is True.
+    :param show_stack: whether or not to show the upper stack frames. default is True.
+    :param stream_lookahead: the number of bytes to dump when show_stack is set. default is 100.
+    
+    Example::
+    
+        Struct("foo",
+            UBInt8("a"),
+            Probe("between a and b"),
+            UBInt8("b"),
+        )
+    """
+    __slots__ = [
+        "printname", "show_stream", "show_context", "show_stack", 
+        "stream_lookahead"
+    ]
+    counter = 0
+    
+    def __init__(self, name = None, show_stream = True, 
+                 show_context = True, show_stack = True, 
+                 stream_lookahead = 100):
+        Construct.__init__(self, None)
+        if name is None:
+            Probe.counter += 1
+            name = "<unnamed %d>" % (Probe.counter,)
+        self.printname = name
+        self.show_stream = show_stream
+        self.show_context = show_context
+        self.show_stack = show_stack
+        self.stream_lookahead = stream_lookahead
+    def __repr__(self):
+        return "%s(%r)" % (self.__class__.__name__, self.printname)
+    def _parse(self, stream, context):
+        self.printout(stream, context)
+    def _build(self, obj, stream, context):
+        self.printout(stream, context)
+    def _sizeof(self, context):
+        return 0
+    
+    def printout(self, stream, context):
+        obj = Container()
+        if self.show_stream:
+            obj.stream_position = stream.tell()
+            follows = stream.read(self.stream_lookahead)
+            if not follows:
+                obj.following_stream_data = "EOF reached"
+            else:
+                stream.seek(-len(follows), 1)
+                obj.following_stream_data = HexString(follows)
+            print("")
+        
+        if self.show_context:
+            obj.context = context
+        
+        if self.show_stack:
+            obj.stack = ListContainer()
+            frames = [s[0] for s in inspect.stack()][1:-1]
+            frames.reverse()
+            for f in frames:
+                a = Container()
+                a.__update__(f.f_locals)
+                obj.stack.append(a)
+        
+        print("=" * 80)
+        print("Probe %s" % (self.printname,))
+        print(obj)
+        print("=" * 80)
+
+class Debugger(Subconstruct):
+    """
+    A pdb-based debugger. When an exception occurs in the subcon, a debugger
+    will appear and allow you to debug the error (and even fix on-the-fly).
+    
+    :param subcon: the subcon to debug
+    
+    Example::
+    
+        Debugger(
+            Enum(UBInt8("foo"),
+                a = 1,
+                b = 2,
+                c = 3
+            )
+        )
+    """
+    __slots__ = ["retval"]
+    def _parse(self, stream, context):
+        try:
+            return self.subcon._parse(stream, context)
+        except Exception:
+            self.retval = NotImplemented
+            self.handle_exc("(you can set the value of 'self.retval', "
+                "which will be returned)")
+            if self.retval is NotImplemented:
+                raise
+            else:
+                return self.retval
+    def _build(self, obj, stream, context):
+        try:
+            self.subcon._build(obj, stream, context)
+        except Exception:
+            self.handle_exc()
+    def handle_exc(self, msg = None):
+        print("=" * 80)
+        print("Debugging exception of %s:" % (self.subcon,))
+        print("".join(traceback.format_exception(*sys.exc_info())[1:]))
+        if msg:
+            print(msg)
+        pdb.post_mortem(sys.exc_info()[2])
+        print("=" * 80)
+