From bf72ceb0ac36467caa59523dfbcea13a7dbdeb95 Mon Sep 17 00:00:00 2001
From: Nathan Rusch <nathan.rusch@scanlinevfx.com>
Date: Thu, 27 Jun 2024 11:56:58 -0700
Subject: [PATCH] Convert Windows registry queries to use `winreg` module

Signed-off-by: Nathan Rusch <nathan.rusch@scanlinevfx.com>
---
 .../shell/_utils/powershell_base.py           | 24 +------
 src/rezplugins/shell/_utils/windows.py        | 72 ++++---------------
 2 files changed, 17 insertions(+), 79 deletions(-)

diff --git a/src/rezplugins/shell/_utils/powershell_base.py b/src/rezplugins/shell/_utils/powershell_base.py
index f059bb248..e7f638868 100644
--- a/src/rezplugins/shell/_utils/powershell_base.py
+++ b/src/rezplugins/shell/_utils/powershell_base.py
@@ -13,7 +13,7 @@
 from rez.utils.platform_ import platform_
 from rez.utils.execution import Popen
 from rez.util import shlex_join
-from .windows import to_windows_path
+from .windows import get_syspaths_from_registry, to_windows_path
 
 
 class PowerShellBase(Shell):
@@ -78,27 +78,7 @@ def get_syspaths(cls):
             cls.syspaths = config.standard_system_paths
             return cls.syspaths
 
-        # TODO: Research if there is an easier way to pull system PATH from
-        # registry in powershell
-        paths = []
-
-        cmds = [
-            [
-                "powershell",
-                '(Get-ItemProperty "HKLM:SYSTEM/CurrentControlSet/Control/Session Manager/Environment").Path',
-            ], [
-                "powershell", "(Get-ItemProperty -Path HKCU:Environment).Path"
-            ]
-        ]
-
-        for cmd in cmds:
-            p = Popen(cmd, stdout=PIPE, stderr=PIPE,
-                      text=True)
-            out_, _ = p.communicate()
-            out_ = out_.strip()
-
-            if p.returncode == 0:
-                paths.extend(out_.split(os.pathsep))
+        paths = get_syspaths_from_registry()
 
         cls.syspaths = [x for x in paths if x]
 
diff --git a/src/rezplugins/shell/_utils/windows.py b/src/rezplugins/shell/_utils/windows.py
index 924c64f3a..813a8ddda 100644
--- a/src/rezplugins/shell/_utils/windows.py
+++ b/src/rezplugins/shell/_utils/windows.py
@@ -4,8 +4,7 @@
 
 import os
 import re
-import subprocess
-from rez.utils.execution import Popen
+import winreg
 
 
 _drive_start_regex = re.compile(r"^([A-Za-z]):\\")
@@ -48,63 +47,22 @@ def to_windows_path(path):
 
 
 def get_syspaths_from_registry():
+    paths = []
 
-    def gen_expected_regex(parts):
-        whitespace = r"[\s]+"
-        return whitespace.join(parts)
-
-    entries = (
-        # local machine
-        dict(
-            cmd=[
-                "REG",
-                "QUERY",
-                "HKLM\\SYSTEM\\CurrentControlSet\\Control\\Session Manager\\Environment",
-                "/v",
-                "PATH"
-            ],
-            expected=gen_expected_regex([
-                "HKEY_LOCAL_MACHINE\\\\SYSTEM\\\\CurrentControlSet\\\\Control\\\\Session Manager\\\\Environment",
-                "PATH",
-                "REG_(EXPAND_)?SZ",
-                "(.*)"
-            ])
-        ),
-        # current user
-        dict(
-            cmd=[
-                "REG",
-                "QUERY",
-                "HKCU\\Environment",
-                "/v",
-                "PATH"
-            ],
-            expected=gen_expected_regex([
-                "HKEY_CURRENT_USER\\\\Environment",
-                "PATH",
-                "REG_(EXPAND_)?SZ",
-                "(.*)"
-            ])
-        )
+    path_query_keys = (
+        (winreg.HKEY_LOCAL_MACHINE, 'SYSTEM\\CurrentControlSet\\Control\\Session Manager\\Environment'),
+        (winreg.HKEY_CURRENT_USER, 'Environment')
     )
 
-    paths = []
-
-    for entry in entries:
-        p = Popen(
-            entry["cmd"],
-            stdout=subprocess.PIPE,
-            stderr=subprocess.PIPE,
-            shell=True,
-            text=True
-        )
-
-        out_, _ = p.communicate()
-        out_ = out_.strip()
-
-        if p.returncode == 0:
-            match = re.match(entry["expected"], out_)
-            if match:
-                paths.extend(match.group(2).split(os.pathsep))
+    for root_key, sub_key in path_query_keys:
+        try:
+            with winreg.OpenKey(root_key, sub_key) as key:
+                reg_value, _ = winreg.QueryValueEx(key, 'Path')
+        except OSError:
+            # Key does not exist
+            pass
+        else:
+            expanded_value = winreg.ExpandEnvironmentStrings(reg_value)
+            paths.extend(expanded_value.split(os.pathsep))
 
     return [x for x in paths if x]