1 files changed, 31 insertions, 29 deletions

diff --git a/src/utils/sw_inventory.py b/src/utils/sw_inventory.py
index bcf9f09..ccb70fe 100644
--- a/src/utils/sw_inventory.py
+++ b/src/utils/sw_inventory.py
@@ -16,33 +16,39 @@ from collections import namedtuple
 import hivex
 
 from src.utils.probe import os_probe
+from src.utils.winreg import *
 
 
 Package = namedtuple('Package', ['name', 'version'])
 Package.__str__ = lambda pkg: f'{pkg.name} {pkg.version}'
 
-WINDOWS_HIVES_PATH = '/Windows/System32/config'
-WINDOWS_HIVES_SOFTWARE = f'{WINDOWS_HIVES_PATH}/SOFTWARE'
 DPKG_STATUS_PATH = '/var/lib/dpkg/status'
 OSRELEASE_PATH = '/etc/os-release'
 
 
-def _fill_package_set(h, key, pkg_set):
+def _fill_package_set(hive, key, pkg_set):
     """
     Fill the package set looking for entries at the current registry
     node childs.
 
     Any valid node child must have "DisplayVersion" or "DisplayName" keys.
     """
-    childs = h.node_children(key)
-    valid_childs = [h.node_get_child(key, h.node_name(child))
-                    for child in childs
-                    for value in h.node_values(child) if h.value_key(value) == 'DisplayVersion']
+    childs = hive.node_children(key)
+    valid_childs = []
+    for child in childs:
+        child_name = hive.node_name(child)
+        values = hive.node_values(child)
+
+        for value in values:
+            if hive.value_key(value) == 'DisplayVersion':
+                valid_child = hive.node_get_child(key, child_name)
+                valid_childs.append(valid_child)
+
     for ch in valid_childs:
         try:
-            name = h.value_string(h.node_get_value(ch, 'DisplayName'))
-            value = h.node_get_value(ch, 'DisplayVersion')
-            version = h.value_string(value)
+            name = hive.value_string(hive.node_get_value(ch, 'DisplayName'))
+            value = hive.node_get_value(ch, 'DisplayVersion')
+            version = hive.value_string(value)
             pkg = Package(name, version)
             pkg_set.add(pkg)
         except RuntimeError:
@@ -50,22 +56,19 @@ def _fill_package_set(h, key, pkg_set):
             pass
 
 
-def _fill_package_set_1(h, pkg_set):
+def _fill_package_set_1(hive, pkg_set):
     """
     Looks for entries in registry path
     /Microsoft/Windows/CurrentVersion/Uninstall
 
     Fills the given set with Package instances for each program found.
     """
-    key = h.root()
-    key = h.node_get_child(key, 'Microsoft')
-    key = h.node_get_child(key, 'Windows')
-    key = h.node_get_child(key, 'CurrentVersion')
-    key = h.node_get_child(key, 'Uninstall')
-    _fill_package_set(h, key, pkg_set)
+    root_node = hive.root()
+    key = get_node_child_from_path(hive, root_node, 'Microsoft/Windows/CurrentVersion/Uninstall')
+    _fill_package_set(hive, key, pkg_set)
 
 
-def _fill_package_set_2(h, pkg_set):
+def _fill_package_set_32_bit_compat(hive, pkg_set):
     """
     Looks for entries in registry path
     /Wow6432Node/Microsoft/Windows/CurrentVersion/Uninstall
@@ -73,23 +76,22 @@ def _fill_package_set_2(h, pkg_set):
 
     Fills the given set with Package instances for each program found.
     """
-    key = h.root()
-    key = h.node_get_child(key, 'Wow6432Node')
-    key = h.node_get_child(key, 'Microsoft')
-    key = h.node_get_child(key, 'Windows')
-    key = h.node_get_child(key, 'CurrentVersion')
-    key = h.node_get_child(key, 'Uninstall')
-    _fill_package_set(h, key, pkg_set)
+    root_node = hive.root()
+    key = get_node_child_from_path(hive, root_node, 'Wow6432Node/Windows/CurrentVersion/Uninstall')
+    _fill_package_set(hive, key, pkg_set)
 
 
 def _get_package_set_windows(hivepath):
     packages = set()
     try:
-        h = hivex.Hivex(hivepath)
+        h = hive_handler_open(hivepath, write = False)
         _fill_package_set_1(h, packages)
-        _fill_package_set_2(h, packages)
-    except RuntimeError as e:
+    except (RuntimeError, OgError) as e:
         logging.error(f'Hivex was not able to operate over {hivepath}. Reported: {e}')
+    try:
+        _fill_package_set_32_bit_compat(h, packages)
+    except (RuntimeError, OgError) as e:
+        pass
     return packages
 
 
@@ -119,7 +121,7 @@ def _get_package_set_dpkg(dpkg_status_path):
 
 def get_package_set(mountpoint):
     dpkg_status_path = f'{mountpoint}{DPKG_STATUS_PATH}'
-    softwarehive = f'{mountpoint}{WINDOWS_HIVES_SOFTWARE}'
+    softwarehive = f'{mountpoint}{WINDOWS_HIVE_SOFTWARE}'
     if os.path.exists(softwarehive):
         pkgset = _get_package_set_windows(softwarehive)
     elif os.path.exists(dpkg_status_path):