Added Environment API to winproxy + windows.system.environ : a unicode environ even on py2

tests/test_system.py

@@ -1,3 +1,4 @@
+import ctypes
 import pytest
 import windows
 
@@ -64,4 +65,48 @@ class TestSystemWithCheckGarbageAndHandleLeak(object):
         assert windows.current_process.pid in [p.pid for p in procs]
 
     def test_system_modules(self):
-        return windows.system.modules
+        return windows.system.modules
+
+
+# Test environement dict
+# On py3 this will just test os.environ, but at least we can expect some consistence
+
+UNICODE_STRING_1 = u"\u4e2d\u56fd\u94f6\u884c\u7f51\u94f6\u52a9\u624b" # some chinese
+UNICODE_RU_STRING = u"\u0441\u0443\u043a\u0430\u0020\u0431\u043b\u044f\u0442\u044c" # CYKA BLYAT in Cyrillic
+UNICODE_UNICORD = u'\U0001f984' # Encoded on 4 char in utf-16
+UNICODE_LOIC_ESCAPE = u'lo\xefc' # Loic with trema
+
+def check_env_variable_exist(name):
+    buf = ctypes.create_unicode_buffer(0x1000)
+    try:
+        windows.winproxy.GetEnvironmentVariableW(name, buf, 0x1000)
+    except WindowsError as e:
+        if e.winerror == gdef.ERROR_ENVVAR_NOT_FOUND:
+            return False
+        raise
+    return True
+
+def test_unicode_environ_dict():
+    unicode_environ = windows.system.environ
+
+    unicode_environ["lower"] = "lower"
+    unicode_environ[UNICODE_LOIC_ESCAPE] = UNICODE_UNICORD
+
+    assert "LOWER" in unicode_environ
+    assert "LOwer" in unicode_environ # Case does not count on __contains__
+
+    assert UNICODE_LOIC_ESCAPE in unicode_environ
+    assert unicode_environ[UNICODE_LOIC_ESCAPE] == UNICODE_UNICORD
+
+    assert check_env_variable_exist("lower")
+    del unicode_environ["Lower"]
+    assert not check_env_variable_exist("lower")
+
+    assert check_env_variable_exist(UNICODE_LOIC_ESCAPE)
+    del unicode_environ[UNICODE_LOIC_ESCAPE]
+    assert not check_env_variable_exist(UNICODE_LOIC_ESCAPE)
+
+    assert not check_env_variable_exist(UNICODE_STRING_1)
+    unicode_environ[UNICODE_STRING_1] = UNICODE_RU_STRING
+    assert check_env_variable_exist(UNICODE_STRING_1)
+

windows/generated_def/meta.py

@@ -14236,6 +14236,8 @@ functions = set(['AccessCheck',
      'FindWindowA',
      'FindWindowW',
      'FreeConsole',
+     'FreeEnvironmentStringsA',
+     'FreeEnvironmentStringsW',
      'FreeLibrary',
      'FreeResource',
      'FreeSid',
@@ -14262,6 +14264,10 @@ functions = set(['AccessCheck',
      'GetDesktopWindow',
      'GetDriveTypeA',
      'GetDriveTypeW',
+     'GetEnvironmentStringsA',
+     'GetEnvironmentStringsW',
+     'GetEnvironmentVariableA',
+     'GetEnvironmentVariableW',
      'GetEventLogInformation',
      'GetExitCodeProcess',
      'GetExitCodeThread',
@@ -14626,6 +14632,9 @@ functions = set(['AccessCheck',
      'SetAclInformation',
      'SetClipboardData',
      'SetConsoleCtrlHandler',
+     'SetEnvironmentStringsW',
+     'SetEnvironmentVariableA',
+     'SetEnvironmentVariableW',
      'SetNamedPipeHandleState',
      'SetNamedSecurityInfoA',
      'SetNamedSecurityInfoW',

windows/generated_def/winfuncs.py

@@ -840,6 +840,51 @@ CryptProtectMemoryParams = ((1, 'pDataIn'), (1, 'cbDataIn'), (1, 'dwFlags'))
 CryptUnprotectMemoryPrototype = WINFUNCTYPE(BOOL, LPVOID, DWORD, DWORD)
 CryptUnprotectMemoryParams = ((1, 'pDataIn'), (1, 'cbDataIn'), (1, 'dwFlags'))
 
+#def GetEnvironmentVariableA(lpName, lpBuffer, nSize):
+#    return GetEnvironmentVariableA.ctypes_function(lpName, lpBuffer, nSize)
+GetEnvironmentVariableAPrototype = WINFUNCTYPE(DWORD, LPCSTR, LPSTR, DWORD)
+GetEnvironmentVariableAParams = ((1, 'lpName'), (1, 'lpBuffer'), (1, 'nSize'))
+
+#def GetEnvironmentVariableW(lpName, lpBuffer, nSize):
+#    return GetEnvironmentVariableW.ctypes_function(lpName, lpBuffer, nSize)
+GetEnvironmentVariableWPrototype = WINFUNCTYPE(DWORD, LPCWSTR, LPWSTR, DWORD)
+GetEnvironmentVariableWParams = ((1, 'lpName'), (1, 'lpBuffer'), (1, 'nSize'))
+
+#def SetEnvironmentVariableA(lpName, lpValue):
+#    return SetEnvironmentVariableA.ctypes_function(lpName, lpValue)
+SetEnvironmentVariableAPrototype = WINFUNCTYPE(BOOL, LPCSTR, LPCSTR)
+SetEnvironmentVariableAParams = ((1, 'lpName'), (1, 'lpValue'))
+
+#def SetEnvironmentVariableW(lpName, lpValue):
+#    return SetEnvironmentVariableW.ctypes_function(lpName, lpValue)
+SetEnvironmentVariableWPrototype = WINFUNCTYPE(BOOL, LPCWSTR, LPCWSTR)
+SetEnvironmentVariableWParams = ((1, 'lpName'), (1, 'lpValue'))
+
+#def GetEnvironmentStringsA():
+#    return GetEnvironmentStringsA.ctypes_function()
+GetEnvironmentStringsAPrototype = WINFUNCTYPE(PVOID)
+GetEnvironmentStringsAParams = ()
+
+#def GetEnvironmentStringsW():
+#    return GetEnvironmentStringsW.ctypes_function()
+GetEnvironmentStringsWPrototype = WINFUNCTYPE(PVOID)
+GetEnvironmentStringsWParams = ()
+
+#def SetEnvironmentStringsW(NewEnvironment):
+#    return SetEnvironmentStringsW.ctypes_function(NewEnvironment)
+SetEnvironmentStringsWPrototype = WINFUNCTYPE(BOOL, LPWCH)
+SetEnvironmentStringsWParams = ((1, 'NewEnvironment'),)
+
+#def FreeEnvironmentStringsA(penv):
+#    return FreeEnvironmentStringsA.ctypes_function(penv)
+FreeEnvironmentStringsAPrototype = WINFUNCTYPE(BOOL, PVOID)
+FreeEnvironmentStringsAParams = ((1, 'penv'),)
+
+#def FreeEnvironmentStringsW(penv):
+#    return FreeEnvironmentStringsW.ctypes_function(penv)
+FreeEnvironmentStringsWPrototype = WINFUNCTYPE(BOOL, PVOID)
+FreeEnvironmentStringsWParams = ((1, 'penv'),)
+
 #def EnumerateTraceGuidsEx(TraceQueryInfoClass, InBuffer, InBufferSize, OutBuffer, OutBufferSize, ReturnLength):
 #    return EnumerateTraceGuidsEx.ctypes_function(TraceQueryInfoClass, InBuffer, InBufferSize, OutBuffer, OutBufferSize, ReturnLength)
 EnumerateTraceGuidsExPrototype = WINFUNCTYPE(ULONG, TRACE_QUERY_INFO_CLASS, PVOID, ULONG, PVOID, ULONG, PULONG)

windows/winobject/system.py

@@ -29,6 +29,37 @@ from windows.winobject import bits
 
 from windows.dbgprint import dbgprint
 
+
+class UnicodeEnvironment(dict):
+    def __getitem__(self, key):
+        key = key.upper()
+        return super(UnicodeEnvironment, self).__getitem__(key)
+
+    def __delitem__(self, key):
+        key = key.upper()
+        windows.winproxy.SetEnvironmentVariableW(key, None)
+        return super(UnicodeEnvironment, self).__delitem__(key)
+
+    def __contains__(self, key):
+        key = key.upper()
+        return super(UnicodeEnvironment, self).__contains__(key)
+
+    def __setitem__(self, key, value):
+        if value is None:
+            # SetEnvironmentVariableW
+            # If this parameter is NULL, the variable is deleted from the current process's environment.
+            del self[key]
+            return
+        key = key.upper()
+        windows.winproxy.SetEnvironmentVariableW(key, value)
+        return super(UnicodeEnvironment, self).__setitem__(key, value)
+
+    def update(self, *args, **kwargs):
+        for k, v in dict(*args, **kwargs).items():
+            self[k.upper()] = v
+
+    # Setitem -> SetEnvironmement variable
+
 class System(object):
     """The state of the current ``Windows`` system ``Python`` is running on"""
 
@@ -97,6 +128,35 @@ class System(object):
             return 64
         return 32
 
+    @utils.fixedpropety
+    def environ(self):
+        """A unicode version of os.environ
+        Same as os.environ on py3
+        Custom dict built on GetEnvironmentStringsW() on py2
+
+        :type: :class:`dict` -- {unicode: unicode}
+		"""
+        if windows.pycompat.is_py3:
+            return os.environ # Py3 environ is already unicode
+        # Create a unicode wrapper environment
+        return self._create_unicode_environ_dict()
+
+    def _create_unicode_environ_dict(self):
+        rawdictenv = {}
+        curenv = envptr = windows.winproxy.GetEnvironmentStringsW()
+        while True:
+            envstr = ctypes.c_wchar_p(curenv).value
+            if not envstr:
+                break
+            try:
+                key, value = envstr.split("=", 1)
+            except ValueError as e: # No =
+                pass
+            rawdictenv[key.upper()] = value
+            curenv += (len(envstr) + 1) * 2
+        windows.winproxy.FreeEnvironmentStringsW(envptr)
+        return UnicodeEnvironment(rawdictenv)
+
     @utils.fixedpropety
     def wmi(self):
         r"""An object to perform wmi requests to various namespaces

windows/winproxy/apis/kernel32.py

@@ -932,4 +932,45 @@ def EnumResourceNamesA(hModule, lpType, lpEnumFunc, lParam):
 
 @Kernel32Proxy()
 def EnumResourceNamesW(hModule, lpType, lpEnumFunc, lParam):
-    return EnumResourceNamesW.ctypes_function(hModule, lpType, lpEnumFunc, lParam)
+    return EnumResourceNamesW.ctypes_function(hModule, lpType, lpEnumFunc, lParam)
+
+# Environment
+@Kernel32Proxy()
+def GetEnvironmentVariableA(lpName, lpBuffer, nSize):
+    if nSize is None:
+        nSize = ctypes.sizeof(lpBuffer)
+    return GetEnvironmentVariableA.ctypes_function(lpName, lpBuffer, nSize)
+
+@Kernel32Proxy()
+def GetEnvironmentVariableW(lpName, lpBuffer, nSize):
+    if nSize is None:
+        nSize = ctypes.sizeof(lpBuffer)
+    return GetEnvironmentVariableW.ctypes_function(lpName, lpBuffer, nSize)
+
+@Kernel32Proxy()
+def SetEnvironmentVariableA(lpName, lpValue):
+    return SetEnvironmentVariableA.ctypes_function(lpName, lpValue)
+
+@Kernel32Proxy()
+def SetEnvironmentVariableW(lpName, lpValue):
+    return SetEnvironmentVariableW.ctypes_function(lpName, lpValue)
+
+@Kernel32Proxy()
+def GetEnvironmentStringsA():
+    return GetEnvironmentStringsA.ctypes_function()
+
+@Kernel32Proxy()
+def GetEnvironmentStringsW():
+    return GetEnvironmentStringsW.ctypes_function()
+
+@Kernel32Proxy()
+def SetEnvironmentStringsW(NewEnvironment):
+    return SetEnvironmentStringsW.ctypes_function(NewEnvironment)
+
+@Kernel32Proxy()
+def FreeEnvironmentStringsA(penv):
+    return FreeEnvironmentStringsA.ctypes_function(penv)
+
+@Kernel32Proxy()
+def FreeEnvironmentStringsW(penv):
+    return FreeEnvironmentStringsW.ctypes_function(penv)