1 files changed, 16 insertions, 49 deletions
diff --git a/src/utils/probe.py b/src/utils/probe.py
index ee9a08e..c8218f2 100644
--- a/src/utils/probe.py
+++ b/src/utils/probe.py
@@ -12,6 +12,7 @@ import platform
 import logging
 import sys
 
+from src.utils.winreg import *
 from enum import Enum
 from subprocess import PIPE
 
@@ -45,35 +46,23 @@ def getlinuxversion(osrelease):
 
 def getwindowsversion(winreghives):
     """
-    Tries to obtain windows version information by
-    querying the SOFTWARE registry hive. Registry
-    hives path is a required parameter.
-
-    Runs hivexget(1) to fetch ProductName and
-    ReleaseId. If something fails (hivexget is
-    not installed, or registry is not found) it
-    returns a generic "Microsoft Windows" string.
+    Try to obtain windows version information by querying the SOFTWARE registry
+    hive to fetch ProductName and ReleaseId.
+    Return a generic "Microsoft Windows" string if something fails.
     """
 
-    # XXX: 3.6 friendly
     try:
-        proc_prodname = subprocess.run(['hivexget',
-            f'{winreghives}/SOFTWARE',
-            'microsoft\windows nt\currentversion',
-            'ProductName'], stdout=PIPE)
-        proc_releaseid = subprocess.run(['hivexget',
-            f'{winreghives}/SOFTWARE',
-            'microsoft\windows nt\currentversion',
-            'ReleaseId'], stdout=PIPE)
-
-        prodname = proc_prodname.stdout.decode().replace('\n', '')
-        releaseid = proc_releaseid.stdout.decode().replace('\n', '')
-        bits = ' 64 bits' if windows_is64bit(winreghives) else ''
-
-        if proc_prodname.returncode == 0 and proc_releaseid.returncode == 0:
-            return f'{prodname} {releaseid}{bits}'
-    except FileNotFoundError:  # hivexget command not found
-        pass
+        hivepath = f'{winreghives}/SOFTWARE'
+        hive = hive_handler_open(hivepath, write = False)
+        root_node = hive.root()
+        version_node = get_node_child_from_path(hive, root_node, 'Microsoft/Windows NT/CurrentVersion')
+
+        prodname = get_value_from_node(hive, version_node, 'ProductName')
+        releaseid = get_value_from_node(hive, version_node, 'ReleaseId')
+
+        return f'{prodname} {releaseid}'
+    except (RuntimeError, OgError) as e:
+        logging.error(f'Hivex was not able to operate over {hivepath}. Reported: {e}')
     return 'Microsoft Windows'
 
 
@@ -81,28 +70,6 @@ def interpreter_is64bit():
     return sys.maxsize > 2**32
 
 
-def windows_is64bit(winreghives):
-    """
-    Check for 64 bit Windows by means of retrieving the value of
-    ProgramW6432Dir. This key is set if Windows is running 64 bit.
-
-    If set returns True.
-    If not set or hivexget exits with non-zero, returns False.
-    """
-    try:
-        proc_hivexget = subprocess.run(['hivexget',
-            f'{winreghives}/SOFTWARE',
-            'Microsoft\Windows\CurrentVersion',
-            'ProgramW6432Dir'], stdout=PIPE)
-        stdout = proc_hivexget.stdout.decode().replace('\n', '')
-
-        if proc_hivexget.returncode == 0 and stdout:
-            return True
-    except FileNotFoundError:  # hivexget command not found
-        pass
-    return False
-
-
 def linux_is64bit(mountpoint):
     """
     If /sbin/init is detected, check if compiled for 64-bit machine.
@@ -174,7 +141,7 @@ def os_probe(mountpoint):
 
     Returns a string depending on the OS it detects.
     """
-    winreghives = f'{mountpoint}/Windows/System32/config'
+    winreghives = f'{mountpoint}{WINDOWS_HIVES_PATH}'
     osrelease = f'{mountpoint}/etc/os-release'
 
     if os.path.exists(osrelease):