btree: fix handling of nodes with a single child node

Author: williballenthin

idb/fileformat.py

@@ -284,7 +284,7 @@ def find_index(self, key):
         #  the linear scan below is simpler to read, so we'll use that until it becomes an issue.
         if self.is_leaf():
             for i, entry in enumerate(self.get_entries()):
-                if key == entry.key:
+                if key == bytes(entry.key):
                     return i
         else:
             for i, entry in enumerate(self.get_entries()):
@@ -365,18 +365,18 @@ def _find(self, cursor, page_number, key):
 
         entry = page.get_entry(entry_number)
 
-        if entry.key == key:
+        if bytes(entry.key) == key:
             cursor.entry = entry
             cursor.entry_number = entry_number
             return
         elif page.is_leaf():
             # no matches!
             raise KeyError(key)
         else:
-            if entry_number == 0:
-                next_page_number = page.ppointer
-            elif is_largest:
+            if is_largest:
                 next_page_number = page.get_entry(page.entry_count - 1).page
+            elif entry_number == 0:
+                next_page_number = page.ppointer
             else:
                 next_page_number = page.get_entry(entry_number - 1).page
             self._find(cursor, next_page_number, key)

scripts/explore_btree.py

@@ -43,7 +43,7 @@ def __init__(self, db):
 
     @property
     def prompt(self):
-        return '/' + '/'.join(map(h, self.path)) + " > "
+        return '/'.join(map(h, self.path)) + "/ > "
 
     @property
     def current_page(self):

tests/test_netnode.py

@@ -49,3 +49,10 @@ def test_sups(kernel32_idb, version, bitness, expected):
 def test_alts(kernel32_idb, version, bitness, expected):
     root = idb.netnode.Netnode(kernel32_idb, ROOT_NODEID)
     assert list(root.alts()) == [-8, -6, -5, -4, -3, -2, -1]
+
+
+# the small netnode has a root btree node with a single child.
+# this is a little tricky to handle, so we ensure it works as expected.
+def test_small(small_idb):
+    root = idb.netnode.Netnode(small_idb, ROOT_NODEID)
+    assert list(root.alts()) == [-8, -5, -4, -3, -2, -1]