diff -r 9aba5d75ce94 Doc/library/winreg.rst --- a/Doc/library/winreg.rst Thu Jun 12 01:03:35 2014 -0400 +++ b/Doc/library/winreg.rst Fri Jun 13 19:08:25 2014 +0300 @@ -274,6 +274,24 @@ If *key* is a handle returned by :func:`ConnectRegistry`, then the path specified in *file_name* is relative to the remote computer. +.. function:: UnloadKey(key, sub_key) + + Unloads the specified registry key and its subkeys from the registry. + It's the counterpart of :func:`LoadKey`. + + *key* is a handle returned by :func:`ConnectRegistry` or one of the constants + :const:`HKEY_USERS` or :const:`HKEY_LOCAL_MACHINE`. + + *sub_key* is a string that identifies the subkey to load. + + A call to :func:`UnloadKey` fails if the calling process does not have the + :const:`SE_RESTORE_PRIVILEGE` privilege. Note that privileges are different + from permissions -- see the `RegUnLoadKey documentation + `__ for + more details. + + .. versionadded:: 3.5 + .. function:: OpenKey(key, sub_key, reserved=0, access=KEY_READ) OpenKeyEx(key, sub_key, reserved=0, access=KEY_READ) diff -r 9aba5d75ce94 Lib/test/test_winreg.py --- a/Lib/test/test_winreg.py Thu Jun 12 01:03:35 2014 -0400 +++ b/Lib/test/test_winreg.py Fri Jun 13 19:08:25 2014 +0300 @@ -9,8 +9,11 @@ # Do this first so test will be skipped if module doesn't exist support.import_module('winreg', required_on=['win']) +support.import_module('ctypes', required_on=['win']) # Now import everything +import winreg from winreg import * +from test.windows_helper import acquired_privilege, skip_unless_privilege try: REMOTE_NAME = sys.argv[sys.argv.index("--remote")+1] @@ -34,6 +37,8 @@ test_key_name = "SOFTWARE\\" + test_key_base # On OS'es that support reflection we should test with a reflected key test_reflect_key_name = "SOFTWARE\\Classes\\" + test_key_base +ERROR_PRIVILEGE_NOT_HELD = 1314 +ERROR_INVALID_PARAMETER = 87 test_data = [ ("Int Value", 45, REG_DWORD), @@ -354,6 +359,48 @@ finally: DeleteKey(HKEY_CURRENT_USER, test_key_name) + def test_unload_key_no_privilege(self): + # Use this idiom in order to avoid failures in refleak mode + # where the process already has the privileges, thus + # leading to the next test failure. + try: + winreg.UnloadKey(winreg.HKEY_USERS, test_key_base) + except OSError as exc: + found = (exc.winerror in + (ERROR_PRIVILEGE_NOT_HELD, ERROR_INVALID_PARAMETER)) + if not found: + self.fail("Unexpected failure: should have failed with: " + "'A required privilege is not held by the client'") + + @skip_unless_privilege("SeRestorePrivilege") + @skip_unless_privilege("SeBackupPrivilege") + def test_unload_key(self): + self.addCleanup(support.unlink, support.TESTFN) + with acquired_privilege("SeRestorePrivilege"): + # only HKEY_USERS and HKEY_LOCAL_MACHINE are allowed + for key in (HKEY_CURRENT_USER, HKEY_CLASSES_ROOT): + with self.assertRaisesRegex(OSError, + "The parameter is incorrect"): + UnloadKey(key, test_key_base) + # the key does not exist + with self.assertRaisesRegex(OSError, + "The parameter is incorrect"): + UnloadKey(HKEY_USERS, "nanana") + + with acquired_privilege("SeBackupPrivilege"): + SaveKey(HKEY_CURRENT_USER, support.TESTFN) + for key in (HKEY_USERS, HKEY_LOCAL_MACHINE): + LoadKey(key, test_key_base, support.TESTFN) + try: + UnloadKey(key, test_key_base) + finally: + # If test_key_base was mounted in HKEY_USERS, + # the reflected key will appear in HKEY_CURRENT_USER + # as well. The only problem is that we can't use + # UnloadKey to detach it, so we'll just remove it. + if key == HKEY_USERS: + DeleteKey(HKEY_CURRENT_USER, test_key_base) + @unittest.skipUnless(REMOTE_NAME, "Skipping remote registry tests") diff -r 9aba5d75ce94 Lib/test/windows_helper.py --- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/Lib/test/windows_helper.py Fri Jun 13 19:08:25 2014 +0300 @@ -0,0 +1,189 @@ +# Various Windows helpers for tests + +import ctypes +import unittest +from functools import wraps +from ctypes import wintypes +from contextlib import contextmanager +from types import SimpleNamespace + +__all__ = ( + 'skip_unless_privilege', + 'adjust_privilege', + 'acquired_privilege', +) + +class TOKEN_INFORMATION_CLASS: + TokenUser = 1 + TokenGroups = 2 + TokenPrivileges = 3 + +class LUID(ctypes.Structure): + _fields_ = [ + ('low_part', wintypes.DWORD), + ('high_part', wintypes.LONG), + ] + +class LUID_AND_ATTRIBUTES(ctypes.Structure): + _fields_ = [ + ('LUID', LUID), + ('attributes', wintypes.DWORD), + ] + + def enable(self): + self.attributes |= SE_PRIVILEGE_ENABLED + +class TOKEN_PRIVILEGES(ctypes.Structure): + _fields_ = [ + ('count', wintypes.DWORD), + ('privileges', LUID_AND_ATTRIBUTES * 0), + ] + + def get_array(self): + array_type = LUID_AND_ATTRIBUTES * self.count + privileges = ctypes.cast(self.privileges, + ctypes.POINTER(array_type)).contents + return privileges + +SE_PRIVILEGE_ENABLED_BY_DEFAULT = 0x00000001 +SE_PRIVILEGE_ENABLED = 0x00000002 +SE_PRIVILEGE_REMOVED = 0x00000004 +SE_PRIVILEGE_USED_FOR_ACCESS = 0x80000000 +TOKEN_ADJUST_PRIVILEGES = 0x20 +TOKEN_QUERY = 8 + +PTOKEN_PRIVILEGES = ctypes.POINTER(TOKEN_PRIVILEGES) +windll = ctypes.LibraryLoader(ctypes.WinDLL) +LookupPrivilegeName = windll.advapi32.LookupPrivilegeNameW +LookupPrivilegeValueW = windll.advapi32.LookupPrivilegeValueW +LookupPrivilegeName.argtypes = ( + wintypes.LPWSTR, + ctypes.POINTER(LUID), + wintypes.LPWSTR, + ctypes.POINTER(wintypes.DWORD), +) +LookupPrivilegeName.restype = wintypes.BOOL +LookupPrivilegeValueW .argtypes = ( + wintypes.LPWSTR, + wintypes.LPWSTR, + ctypes.POINTER(LUID), +) +LookupPrivilegeValueW .restype = wintypes.BOOL +GetTokenInformation = windll.advapi32.GetTokenInformation +GetTokenInformation.argtypes = [ + wintypes.HANDLE, + ctypes.c_uint, + ctypes.c_void_p, + wintypes.DWORD, + ctypes.POINTER(wintypes.DWORD), +] +GetTokenInformation.restype = wintypes.BOOL +AdjustTokenPrivileges = windll.advapi32.AdjustTokenPrivileges +AdjustTokenPrivileges.restype = wintypes.BOOL +AdjustTokenPrivileges.argtypes = [ + wintypes.HANDLE, + wintypes.BOOL, + PTOKEN_PRIVILEGES, + wintypes.DWORD, + PTOKEN_PRIVILEGES, + ctypes.POINTER(wintypes.DWORD), +] +OpenProcessToken = windll.advapi32.OpenProcessToken +OpenProcessToken.argtypes = ( + wintypes.HANDLE, wintypes.DWORD, + ctypes.POINTER(wintypes.HANDLE)) +OpenProcessToken.restype = wintypes.BOOL +GetCurrentProcess = windll.kernel32.GetCurrentProcess +GetCurrentProcess.restype = wintypes.HANDLE + +def get_current_process_token(): + """ Get a token to the current process. """ + token = wintypes.HANDLE() + res = OpenProcessToken( + GetCurrentProcess(), + TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, token) + if not res > 0: + raise RuntimeError("Couldn't get process token") + return token + +def get_luid(privilege): + luid = LUID() + + # system name, privilege name and the LUID structure + if not LookupPrivilegeValueW(None, privilege, luid) > 0: + raise RuntimeError("Couldn't lookup privilege value") + return luid + +def adjust_privilege(privilege, old_token=None): + """ Enable the given Windows privilege. + + The privileges can be found in + http://msdn.microsoft.com/en-us/library/windows/desktop/bb530716%28v=vs.85%29.aspx. + The functions uses the text representation of the privileges, + e.g. SeBackupPrivilege, SeShutdownPrivilege etc. + Returns a namespace with attributes: + - result: the result of AdjustTokenPrivileges + - success: True if the privilege was acquired + - token_privileges: The token privileges used. They can be reused + to restore the old privileges the current process had. + """ + size = ctypes.sizeof(TOKEN_PRIVILEGES) + size += ctypes.sizeof(LUID_AND_ATTRIBUTES) + buffer = ctypes.create_string_buffer(size) + tp = ctypes.cast(buffer, ctypes.POINTER(TOKEN_PRIVILEGES)).contents + tp.count = 1 + + tp.get_array()[0].enable() + tp.get_array()[0].LUID = get_luid(privilege) + token = get_current_process_token() + + buffer_length = size if old_token else 0 + # current token handle, disable all privileges, new state, + # buffer length, previous state and return length + res = AdjustTokenPrivileges(token, False, tp, + buffer_length, old_token, + ctypes.c_ulong(buffer_length)) + return SimpleNamespace(result=res, token_privileges=tp, + success=not windll.kernel32.GetLastError()) + +def skip_unless_privilege(privilege): + """ Raise an unittest.SkipTest exception if the given + privilege can't be acquired. The privileges supported + are the same privileges supported by `adjust_privilege`. + """ + + def decorator(func): + @wraps(func) + def wrapper(*args, **kw): + result = adjust_privilege(privilege) + if not result.success: + raise unittest.SkipTest( + "Privilege {} wasn't acquired.".format(privilege)) + + # disable the privilege. + adjust_privilege(privilege, old_token=result.token_privileges) + return func(*args, **kw) + return wrapper + return decorator + +@contextmanager +def acquired_privilege(privilege): + """ Context manager for acquiring and releasing + a privilege for the current process using the `with` statement. + If the privilege can't be acquired, a RuntimeError will be + raised. + + :: + with acquired_privilege('SeBackupPrivilege'): + winreg.SaveKey(key, "file") + # no longer has this privilege + """ + result = adjust_privilege(privilege) + if not result.success: + raise RuntimeError("Could not acquire privilege {!r}" + .format(privilege)) + + try: + yield + finally: + adjust_privilege(privilege, old_token=result.token_privileges) diff -r 9aba5d75ce94 PC/winreg.c --- a/PC/winreg.c Thu Jun 12 01:03:35 2014 -0400 +++ b/PC/winreg.c Fri Jun 13 19:08:25 2014 +0300 @@ -50,6 +50,8 @@ "FlushKey() - Writes all the attributes of the specified key to the registry.\n" "LoadKey() - Creates a subkey under HKEY_USER or HKEY_LOCAL_MACHINE and stores\n" " registration information from a specified file into that subkey.\n" +"UnloadKey() - Unloads the specified registry key and its subkeys \n" +" from the registry.\n" "OpenKey() - Opens the specified key.\n" "OpenKeyEx() - Alias of OpenKey().\n" "QueryValue() - Retrieves the value associated with the unnamed value for a\n" @@ -229,6 +231,20 @@ "\n" "The docs imply key must be in the HKEY_USER or HKEY_LOCAL_MACHINE tree"); +PyDoc_STRVAR(UnloadKey_doc, +"UnloadKey(key, sub_key)\n" +"Unloads the specified registry key and its subkeys from the registry.\n" +"\n" +"key is an already open key, or any one of the predefined HKEY_* constants.\n" +"sub_key is a string that identifies the sub_key to unload.\n" +"\n" +"A call to UnloadKey() fails if the calling process does not have the\n" +"SE_RESTORE_PRIVILEGE privilege.\n" +"If key is a handle returned by ConnectRegistry(), then the path specified\n" +"in fileName is relative to the remote computer.\n" +"\n" +"The docs imply key must be in the HKEY_USER or HKEY_LOCAL_MACHINE tree"); + PyDoc_STRVAR(OpenKey_doc, "OpenKey(key, sub_key, reserved=0, access=KEY_READ) -> key\n" "Opens the specified key.\n" @@ -1332,6 +1348,27 @@ } static PyObject * +PyUnloadKey(PyObject *self, PyObject *args) +{ + HKEY hKey; + PyObject *obKey; + wchar_t *subKey; + + long rc; + if (!PyArg_ParseTuple(args, "Ou:UnloadKey", &obKey, &subKey)) + return NULL; + if (!PyHKEY_AsHKEY(obKey, &hKey, FALSE)) + return NULL; + Py_BEGIN_ALLOW_THREADS + rc = RegUnLoadKeyW(hKey, subKey); + Py_END_ALLOW_THREADS + if (rc != ERROR_SUCCESS) + return PyErr_SetFromWindowsErrWithFunction(rc, "RegUnloadKey"); + Py_INCREF(Py_None); + return Py_None; +} + +static PyObject * PyOpenKey(PyObject *self, PyObject *args, PyObject *kwargs) { HKEY hKey; @@ -1733,6 +1770,7 @@ ExpandEnvironmentStrings_doc }, {"FlushKey", PyFlushKey, METH_VARARGS, FlushKey_doc}, {"LoadKey", PyLoadKey, METH_VARARGS, LoadKey_doc}, + {"UnloadKey", PyUnloadKey, METH_VARARGS, UnloadKey_doc}, {"OpenKey", (PyCFunction)PyOpenKey, METH_VARARGS | METH_KEYWORDS, OpenKey_doc}, {"OpenKeyEx", (PyCFunction)PyOpenKey, METH_VARARGS | METH_KEYWORDS,