1 files changed, 46 insertions, 2 deletions
diff --git a/src/utils/probe.py b/src/utils/probe.py
index fa0ca83..796bf55 100644
--- a/src/utils/probe.py
+++ b/src/utils/probe.py
@@ -8,15 +8,20 @@
 
 import os
 import subprocess
+import platform
 
 from subprocess import PIPE
 
+from src.utils.fs import find_mountpoint
+
 def getlinuxversion(osrelease):
     """
     Parses a os-release file to fetch 'PRETTY_NAME' key.
     If file or key are not found, then returns generic
     'Linux' string.
     """
+    mountpoint = find_mountpoint(osrelease)
+
     with open(osrelease, 'r') as f:
         for line in f:
             if line.find('=') == -1:
@@ -25,7 +30,8 @@ def getlinuxversion(osrelease):
             if key == 'PRETTY_NAME':
                 value = value.replace('\n', '')
                 value = value.strip('"')
-                return value
+                bits = ' 64 bits' if linux_is64bit(mountpoint) else ''
+                return f'{value}{bits}'
         return 'Linux'
 
 
@@ -54,14 +60,52 @@ def getwindowsversion(winreghives):
 
         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}'
+            return f'{prodname} {releaseid}{bits}'
     except FileNotFoundError:  # hivexget command not found
         pass
     return 'Microsoft Windows'
 
 
+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.
+
+    If init executable is not found, check for /lib64.
+    If /lib64 is found returns True, otherwise False.
+    """
+    init_path = f'{mountpoint}/sbin/init'
+    lib64_path = f'{mountpoint}/lib64'
+    if os.path.exists(init_path):
+        bits, linkage = platform.architecture(executable=init_path)
+        return True if bits == '64bit' else False
+    return os.path.exists(lib64_path)
+
+
 def cache_probe():
     """
     Runs 'blkid -L CACHE' and returns stripped stdout