From 0eb78f21bd39bcd02287595665afd526d8898368 Mon Sep 17 00:00:00 2001 From: clement rouault Date: Fri, 26 May 2023 09:25:19 +0200 Subject: [PATCH] Added Environment API to winproxy + windows.system.environ : a unicode environ even on py2 --- tests/test_system.py | 47 +++++++++++++++++++++++- windows/generated_def/meta.py | 9 +++++ windows/generated_def/winfuncs.py | 45 +++++++++++++++++++++++ windows/winobject/system.py | 60 +++++++++++++++++++++++++++++++ windows/winproxy/apis/kernel32.py | 43 +++++++++++++++++++++- 5 files changed, 202 insertions(+), 2 deletions(-) diff --git a/tests/test_system.py b/tests/test_system.py index 0fbef90..e563845 100644 --- a/tests/test_system.py +++ b/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 \ No newline at end of file + 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) + diff --git a/windows/generated_def/meta.py b/windows/generated_def/meta.py index c4fa117..c5e573d 100644 --- a/windows/generated_def/meta.py +++ b/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', diff --git a/windows/generated_def/winfuncs.py b/windows/generated_def/winfuncs.py index a9f4420..bea99e1 100644 --- a/windows/generated_def/winfuncs.py +++ b/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) diff --git a/windows/winobject/system.py b/windows/winobject/system.py index 7476c51..67bfda5 100644 --- a/windows/winobject/system.py +++ b/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 diff --git a/windows/winproxy/apis/kernel32.py b/windows/winproxy/apis/kernel32.py index b763d41..25aa3a0 100644 --- a/windows/winproxy/apis/kernel32.py +++ b/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) \ No newline at end of file + 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) \ No newline at end of file