Adds fallback to newly created environment block for known OS variables. (#432)

Also makes it a hard failure for environment variables in paths to be empty.

Fixes #427

Author: zooba

_msbuild.py

@@ -67,7 +67,8 @@ class ResourceFile(CSourceFile):
     PyFile('__init__.py'),
     CPP_SETTINGS,
     ItemDefinition("Link",
-        AdditionalDependencies=Prepend("wintrust.lib;"),
+        AdditionalDependencies=Prepend("wintrust.lib;userenv.lib;"),
+        DelayLoadDLLs='userenv.dll',
     ),
     IncludeFile('*.h'),
     CSourceFile('*.cpp'),
@@ -100,6 +101,7 @@ class ResourceFile(CSourceFile):
     CFunction('read_alias_package'),
     CFunction('broadcast_settings_change'),
     CFunction('get_processor_architecture'),
+    CFunction('get_environment_block'),
     CFunction('verify_trust'),
     source='src/_native',
     RootNamespace='_native',

_msbuild_test.py

@@ -24,7 +24,7 @@
             LanguageStandard="stdcpp20",
         ),
         ItemDefinition("Link",
-            AdditionalDependencies=Prepend("wintrust.lib;"),
+            AdditionalDependencies=Prepend("wintrust.lib;userenv.lib;"),
         ),
         IncludeFile('*.h'),
         CSourceFile('*.cpp'),
@@ -58,6 +58,7 @@
         CFunction('read_alias_package'),
         CFunction('broadcast_settings_change'),
         CFunction('get_processor_architecture'),
+        CFunction('get_environment_block'),
         CFunction('verify_trust'),
         source='src/_native',
     ),

src/_native/misc.cpp

@@ -1,6 +1,7 @@
 #include <Python.h>
 #include <windows.h>
 #include <appmodel.h>
+#include <userenv.h>
 
 #include "helpers.h"
 
@@ -253,4 +254,32 @@ PyObject *get_processor_architecture(PyObject *, PyObject *, PyObject *) {
     return PyUnicode_FromString("-64");
 }
 
+PyObject *get_environment_block(PyObject *, PyObject *, PyObject *) {
+    HANDLE token = NULL;
+    wchar_t *env = NULL;
+
+    if (!OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY | TOKEN_DUPLICATE, &token)
+        || !CreateEnvironmentBlock((LPVOID *)&env, token, FALSE)) {
+        PyErr_SetFromWindowsErr(0);
+        if (token) CloseHandle(token);
+        return NULL;
+    }
+
+    // env is double null terminated, so read each string's length until we find
+    // the zero-length one. Our caller can do the split.
+    wchar_t *p = env;
+    size_t cch = wcslen(p);
+    while (p && cch) {
+        p += cch + 1;
+        cch = wcslen(p);
+    }
+
+    PyObject *r = PyUnicode_FromWideChar(env, p - env);
+
+    DestroyEnvironmentBlock(env);
+    CloseHandle(token);
+
+    return r;
+}
+
 }

src/manage/config.py

@@ -150,18 +150,117 @@ def load_registry_config(key_path, schema):
     return cfg
 
 
+class _EnvironWithFallback(dict):
+    # Key names must be uppercase
+    # These are the keys that are filled in from a default environment if they
+    # are missing in the current process environment. This ensures that a
+    # broken environment doesn't lead to an invalid loaded configuration. It
+    # only applies to configuration files, and does not change the environment.
+    # A launched Python process will probably also fail without some of these,
+    # at which point the caller still has to fix their setup.
+    FALLBACK_KEYS = {
+        "ALLUSERSPROFILE",
+        "APPDATA",
+        "COMMONPROGRAMFILES",
+        "COMMONPROGRAMFILES(X86)",
+        "COMMONPROGRAMW6432",
+        "COMPUTERNAME",
+        "COMSPEC",
+        "HOMEDRIVE",
+        "HOMEPATH",
+        "LOCALAPPDATA",
+        "LOGONSERVER",
+        "NUMBER_OF_PROCESSORS",
+        "OS",
+        "PATH",
+        "PATHEXT",
+        "PROCESSOR_ARCHITECTURE",
+        "PROCESSOR_IDENTIFIER",
+        "PROCESSOR_LEVEL",
+        "PROCESSOR_REVISION",
+        "PROGRAMDATA",
+        "PROGRAMFILES",
+        "PROGRAMFILES(X86)",
+        "PROGRAMW6432",
+        "PUBLIC",
+        "SESSIONNAME",
+        "SYSTEMDRIVE",
+        "SYSTEMROOT",
+        "TEMP",
+        "TMP",
+        "USERDNSDOMAIN",
+        "USERDOMAIN",
+        "USERDOMAIN_ROAMINGPROFILE",
+        "USERNAME",
+        "USERPROFILE",
+        "WINDIR",
+    }
+
+    _default = None
+
+    @classmethod
+    def default(cls):
+        if cls._default is None:
+            cls._default = cls()
+        return cls._default
+
+    def __init__(self, d=os.environ):
+        self._fallback = None
+        for k, v in (d or {}).items():
+            self[k.upper()] = v
+
+    def _get_fallback(self, k):
+        assert k.isupper()
+        if k not in self.FALLBACK_KEYS:
+            return None
+        if self._fallback is None:
+            from _native import get_environment_block
+            self._fallback = {}
+            for var in get_environment_block().split("\0"):
+                k2, sep, v2 = var.partition("=")
+                if k2 and sep and v2:
+                    self._fallback[k2.upper()] = v2
+        return self._fallback.get(k)
+
+    def get(self, key, default=None):
+        k = key.upper()
+        try:
+            return self[k]
+        except KeyError:
+            pass
+
+        v = self._get_fallback(k)
+        if v:
+            LOGGER.warn("Environment variable %%%s%% was not found and so "
+                        "a default was loaded. This usually indicates an "
+                        "issue in your launch conditions and could lead to "
+                        "later failures. Suppress this warning by ensuring "
+                        "system environment variables remain set, or by "
+                        "changing your configuration file to not rely on "
+                        "them.", key)
+
+        return v or default
+
+
+
 def _expand_vars(v, env):
     import re
     def _sub(m):
         v2 = env.get(m.group(1))
         if v2:
             return v2 + (m.group(2) or "")
+        if m.group(2):
+            # Empty values before a path separator are not allowed
+            raise LookupError(m.group(1))
         return ""
     v2 = re.sub(r"%(.*?)%([\\/])?", _sub, v)
     return v2
 
 
-def resolve_config(cfg, source, relative_to, key_so_far="", schema=None, error_unknown=False):
+def resolve_config(cfg, source, relative_to, key_so_far="", schema=None, error_unknown=False, environ=None):
+    if environ is None:
+        environ = _EnvironWithFallback.default()
+
     for k, v in list(cfg.items()):
         if k.startswith("#"):
             continue
@@ -177,18 +276,20 @@ def resolve_config(cfg, source, relative_to, key_so_far="", schema=None, error_u
         if isinstance(subschema, dict):
             if not isinstance(v, dict):
                 raise InvalidConfigurationError(source, key_so_far + k, v)
-            resolve_config(v, source, relative_to, f"{key_so_far}{k}.", subschema)
+            resolve_config(v, source, relative_to, f"{key_so_far}{k}.", subschema, environ=environ)
             continue
 
         kind, merge, *opts = subschema
         from_env = False
         if "env" in opts and isinstance(v, str):
             try:
                 orig_v = v
-                v = _expand_vars(v, os.environ)
+                v = _expand_vars(v, environ)
                 from_env = orig_v != v
             except TypeError:
                 pass
+            except LookupError as ex:
+                raise InvalidConfigurationError(source, key_so_far + k, f'%{ex}%') from ex
             if not v:
                 del cfg[k]
                 continue
@@ -202,7 +303,9 @@ def resolve_config(cfg, source, relative_to, key_so_far="", schema=None, error_u
             if not from_env:
                 v = relative_to / v
             else:
-                v = type(relative_to)(v).absolute()
+                v = type(relative_to)(v)
+                if not v.is_absolute():
+                    v = v.absolute()
         if v and "uri" in opts:
             if hasattr(v, "as_uri"):
                 v = v.as_uri()

src/manage/pathutils.py

@@ -105,6 +105,10 @@ def __ne__(self, other):
             return not _eq(self._p, other._p)
         return not _eq(self._p, str(other))
 
+    def is_absolute(self):
+        drive, root, tail = os.path.splitroot(self._p)
+        return bool(drive and root)
+
     def with_name(self, name):
         return type(self)(os.path.join(self._parent, name))
 

tests/test_config.py

@@ -1,12 +1,17 @@
+import os
 import pytest
 
+from pathlib import Path, PurePath
+
 from manage.config import (
+    _EnvironWithFallback,
     _expand_vars,
     config_append,
     config_bool,
     config_dict_merge,
     config_split,
     config_split_append,
+    resolve_config,
 )
 
 
@@ -88,16 +93,26 @@ def test_expand_vars(value, variables, expected):
     assert _expand_vars(value, variables) == expected
 
 
-def test_expand_vars_removes_missing_variable_and_leading_separator():
-    actual = _expand_vars("%MISSING%\\Python", {})
+def test_expand_vars_fails_on_missing_variable():
+    with pytest.raises(LookupError):
+        _expand_vars("%MISSING%\\Python", {})
 
-    assert actual == "Python"
 
+def test_expand_vars_fails_on_empty_variable():
+    with pytest.raises(LookupError):
+        _expand_vars("%EMPTY%\\Python", {"EMPTY": ""})
 
-def test_expand_vars_removes_empty_variable_and_leading_separator():
-    actual = _expand_vars("%EMPTY%\\Python", {"EMPTY": ""})
 
-    assert actual == "Python"
+def test_environ_with_fallback():
+    env = _EnvironWithFallback({"AllUsersProfile": "test", "APPDATA": ""})
+    # Should be the value we provided, despite case
+    assert env.get("allusersprofile") == "test"
+    # Should be empty, since it was provided as empty
+    assert env.get("AppData", "default") == ""
+    # Should be set, even though it wasn't provided
+    assert env.get("WinDir", "default") not in (None, "", "default")
+    # Should be empty, since it's neither set nor has a fallback
+    assert env.get("PYTHON_MANAGER_DEFAULT", "default") == "default"
 
 
 @pytest.mark.parametrize(
@@ -200,4 +215,40 @@ def test_config_split_append_does_not_mutate_existing_list():
 
     assert actual == ["first", ["second", "third"]]
     assert existing == ["first"]
-    assert actual is not existing
+    assert actual is not existing
+
+
+@pytest.mark.parametrize(
+    ("cfg", "schema", "expect", "env"),
+    [
+        ({}, {"a": (int, None), "b": (int, None)}, {}, None),
+        ({"a": "1"}, {"a": (int, None)}, {"a": 1}, None),
+        ({"a": 1}, {"a": (str, None)}, {"a": "1"}, None),
+        ({"a": "abc"}, {"a": (str, None, "path")}, {"a": PurePath("X:/test/abc")}, None),
+        ({"a": "%X%"}, {"a": (str, None, "path")}, {"a": PurePath("X:/test/%X%")}, {"X": "abc"}),
+        ({"a": "%X%"}, {"a": (str, None, "path", "env")}, {"a": PurePath("X:/env/abc")}, {"X": "X:/env/abc"}),
+        ({"a": "%X%"}, {"a": (str, None, "path")}, {"a": PurePath("X:/test/%X%")}, None),
+    ],
+    ids=[
+        "empty",
+        "to_int",
+        "to_str",
+        "to_path",
+        "to_path_without_env",
+        "to_path_with_env",
+        "to_path_not_in_env",
+    ],
+)
+def test_config_resolve(cfg, schema, expect, env, monkeypatch):
+    resolve_config(cfg, "<test.cfg>", PurePath("X:/test/"), schema=schema, environ=env)
+    assert cfg == expect
+
+
+def test_config_resolve_env_fallback(assert_log):
+    cfg = {"a": "%AppData%/Python"}
+    schema = {"a": (str, None, "path", "env")}
+    expect = {"a": Path(os.environ["AppData"]) / "Python"}
+    env = _EnvironWithFallback({k: v for k, v in os.environ.items() if k.lower() != "appdata"})
+    resolve_config(cfg, "<test.cfg>", Path(), schema=schema, environ=env)
+    assert cfg == expect
+    assert_log(("Environment variable.*not found", ["AppData"]))

tests/test_pathutils.py

@@ -49,3 +49,10 @@ def test_path_relative_to():
 
     assert relative_to(p, r"C:\A\B\C\D") is p
     assert relative_to(p, None) is p
+
+
+def test_path_is_absolute():
+    assert Path(r"C:\A\B\C\python.exe").is_absolute()
+    assert not Path(r"\A\B\C\python.exe").is_absolute()
+    assert not Path(r"C:A\B\C\python.exe").is_absolute()
+    assert not Path(r"A\B\C\python.exe").is_absolute()