From 46c047b7b34472e28ca50f0a3c71cf4a6fbc9efb Mon Sep 17 00:00:00 2001
From: Mike Auty <mike.auty@gmail.com>
Date: Tue, 10 Sep 2019 14:36:38 +0100
Subject: [PATCH] Update API slightly to return objects not strings.

---
 .../framework/plugins/windows/registry/printkey.py | 14 +++++---------
 .../plugins/windows/registry/certificates.py       |  7 +------
 2 files changed, 6 insertions(+), 15 deletions(-)

diff --git a/volatility/framework/plugins/windows/registry/printkey.py b/volatility/framework/plugins/windows/registry/printkey.py
index 2fe123831..923c3f0f6 100644
--- a/volatility/framework/plugins/windows/registry/printkey.py
+++ b/volatility/framework/plugins/windows/registry/printkey.py
@@ -4,7 +4,7 @@
 
 import datetime
 import logging
-from typing import Generator, List, Sequence
+from typing import Generator, List, Sequence, Iterable
 
 from volatility.framework import objects, renderers, exceptions, interfaces, constants
 from volatility.framework.configuration import requirements
@@ -37,7 +37,7 @@ class PrintKey(interfaces.plugins.PluginInterface):
 
     @classmethod
     def key_iterator(cls, hive: RegistryHive, node_path: Sequence[objects.StructType] = None,
-                     recurse: bool = False) -> Generator:
+                     recurse: bool = False) -> Iterable[int, bool, datetime.datetime, str, bool, bytes]:
         """Walks through a set of nodes from a given node (last one in
         node_path). Avoids loops by not traversing into nodes already present
         in the node_path.
@@ -85,7 +85,7 @@ class PrintKey(interfaces.plugins.PluginInterface):
                       base_config_path: str,
                       layer_name: str,
                       symbol_table: str,
-                      hive_offsets: List[int] = None):
+                      hive_offsets: List[int] = None) -> Iterable[RegistryHive]:
         """Walks through a registry, hive by hive returning the constructed
         registry layer name.
 
@@ -118,7 +118,7 @@ class PrintKey(interfaces.plugins.PluginInterface):
 
             hive = RegistryHive(context, reg_config_path, name = 'hive' + hex(hive_offset))
             context.layers.add_layer(hive)
-            yield hive.name
+            yield hive
 
     def _printkey_iterator(self,
                            hive: RegistryHive,
@@ -172,17 +172,13 @@ class PrintKey(interfaces.plugins.PluginInterface):
                          key: str = None,
                          recurse: bool = False):
 
-        for hive_name in self.hive_iterator(
+        for hive in self.hive_iterator(
                 self.context,
                 self.config_path,
                 layer_name = layer_name,
                 symbol_table = symbol_table,
                 hive_offsets = hive_offsets):
 
-            hive = self.context.layers[hive_name]
-            if not isinstance(hive, RegistryHive):
-                continue
-
             try:
                 # Walk it
                 if key is not None:
diff --git a/volatility/plugins/windows/registry/certificates.py b/volatility/plugins/windows/registry/certificates.py
index 72115ce06..52230c63b 100644
--- a/volatility/plugins/windows/registry/certificates.py
+++ b/volatility/plugins/windows/registry/certificates.py
@@ -3,7 +3,6 @@ from typing import List, Iterator, Tuple
 
 from volatility.framework import interfaces, renderers
 from volatility.framework.configuration import requirements
-from volatility.framework.layers import registry
 from volatility.framework.symbols.windows.extensions.registry import RegValueTypes
 from volatility.plugins.windows.registry import printkey
 
@@ -33,16 +32,12 @@ class Certificates(interfaces.plugins.PluginInterface):
         return (name, certificate_data)
 
     def _generator(self) -> Iterator[Tuple[int, Tuple[int, str]]]:
-        for hive_name in printkey.PrintKey.hive_iterator(
+        for hive in printkey.PrintKey.hive_iterator(
                 self.context,
                 base_config_path = self.config_path,
                 layer_name = self.config['primary'],
                 symbol_table = self.config['nt_symbols']):
 
-            hive = self.context.layers[hive_name]
-            if not isinstance(hive, registry.RegistryHive):
-                pass
-
             try:
                 # Walk it
                 top_key = "Microsoft\\SystemCertificates"