From 0d6bd0d445eca317ebf05ab696833463ed7578cf Mon Sep 17 00:00:00 2001
From: malwarefrank <42877127+malwarefrank@users.noreply.github.com>
Date: Sat, 23 Mar 2024 05:43:53 +0000
Subject: [PATCH] update examples

---
 examples/dndump.py    |  7 +++++++
 examples/dnstrings.py | 22 ++++++++++++----------
 2 files changed, 19 insertions(+), 10 deletions(-)

diff --git a/examples/dndump.py b/examples/dndump.py
index c15af82..4e31493 100644
--- a/examples/dndump.py
+++ b/examples/dndump.py
@@ -243,6 +243,13 @@ def render_pe(ostream: Formatter, dn):
                                         value = v
                                 elif isinstance(v, int):
                                     value = "0x%x" % (v)
+                                elif isinstance(v, dnfile.stream.HeapItemString):
+                                    if v.value is None:
+                                        value = "(invalid){!r}".format(v.value_bytes)
+                                    else:
+                                        value = v.value
+                                elif isinstance(v, dnfile.stream.HeapItemBinary):
+                                    value = v.value_bytes()
                                 else:
                                     value = str(v)
                             rows.append(("%s:" % (field), value))
diff --git a/examples/dnstrings.py b/examples/dnstrings.py
index 3497351..6fe698e 100755
--- a/examples/dnstrings.py
+++ b/examples/dnstrings.py
@@ -23,21 +23,23 @@ def show_strings(fname):
         offset = 1
         # while there is still data in the stream
         while offset < size:
+            # check if we are at padding bytes near end of stream
+            if offset + 4 >= size:
+                if b"\x00" == dn.get_data(us.rva + offset, 1):
+                    break
             # read the raw string bytes, and provide number of bytes read (includes encoded length)
-            ret = us.get_with_size(offset)
-            if ret is None:
+            item = us.get(offset)
+            if item is None:
+                print(f"Bad string: offset=0x{offset:08x}")
                 break
 
-            buf, readlen = ret
-            try:
-                # convert to a UserString object
-                s = dnfile.stream.UserString(buf)
+            if item.value is None:
+                print(f"Bad string: {item.raw_data}")
+            else:
                 # display the decoded string
-                print(s.value)
-            except UnicodeDecodeError:
-                print(f"Bad string: {buf}")
+                print(item.value)
             # continue to next entry
-            offset += readlen
+            offset += item.raw_size
 
 
 # for each filepath provided on command-line