1 files changed, 118 insertions, 0 deletions
diff --git a/src/utils/inventory.py b/src/utils/inventory.py
new file mode 100644
index 0000000..7f02c31
--- /dev/null
+++ b/src/utils/inventory.py
@@ -0,0 +1,118 @@
+import platform
+import re
+import os
+
+from collections import namedtuple
+
+import hivex
+
+from src.utils.probe import getwindowsversion, getlinuxversion
+
+
+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):
+    """
+    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']
+    for ch in valid_childs:
+        name = h.value_string(h.node_get_value(ch, 'DisplayName'))
+        value = h.node_get_value(ch, 'DisplayVersion')
+        version = h.value_string(value)
+        pkg = Package(name, version)
+        pkg_set.add(pkg)
+
+
+def _fill_package_set_1(h, 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)
+
+
+def _fill_package_set_2(h, pkg_set):
+    """
+    Looks for entries in registry path
+    /Wow6432Node/Microsoft/Windows/CurrentVersion/Uninstall
+    64 bit Windows only.
+
+    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)
+
+
+def _get_package_set_windows(hivepath):
+    packages = set()
+    h = hivex.Hivex(hivepath)
+    _fill_package_set_1(h, packages)
+    _fill_package_set_2(h, packages)
+    return packages
+
+
+def _get_package_set_dpkg(dpkg_status_path):
+    regex_pkg = '(?:^Package: )(?P<name>.*)\n(Essential:.*\n)?(?:Status: install ok installed)'
+    regex_ver = '(?:^Version: )(?P<version>.*)'
+    packages = set()
+    with open(dpkg_status_path, 'r') as f:
+        # Split by empty line
+        for par in re.split('^\n+', f.read(), flags=re.MULTILINE):
+            # Search for package with "Status: install ok installed"
+            result = re.search(regex_pkg, par)
+            if result is None:
+                continue
+            else:
+                pkg_name = result.group('name')
+            # If we hit a properly installed package, search for its version
+            result = re.search(regex_ver, par, flags=re.MULTILINE)
+            if result is None:
+                continue
+            else:
+                pkg_version = result.group('version')
+            pkg = Package(pkg_name, pkg_version)
+            packages.add(pkg)
+    return packages
+
+
+def get_package_set(mountpoint):
+    dpkg_status_path = f'{mountpoint}{DPKG_STATUS_PATH}'
+    winreghives = f'{mountpoint}{WINDOWS_HIVES_PATH}'
+    osrelease = f'{mountpoint}{OSRELEASE_PATH}'
+    softwarehive = f'{mountpoint}{WINDOWS_HIVES_SOFTWARE}'
+    if os.path.exists(softwarehive):
+        pkgset = _get_package_set_windows(softwarehive)
+        osname = getwindowsversion(winreghives)
+    elif os.path.exists(dpkg_status_path):
+        pkgset = _get_package_set_dpkg(dpkg_status_path)
+        osname = getlinuxversion(osrelease)
+    else:
+        raise ValueError(f'Cannot fetch software inventory at {mountpoint}')
+    # Legacy software inventory first element is the OS name
+    return [osname] + list(pkgset)