From 589baf7b5d856409722081b9d448a5d20ecbfa13 Mon Sep 17 00:00:00 2001 From: Kovid Goyal Date: Tue, 17 Feb 2015 22:04:54 +0530 Subject: [PATCH] Implement a replacement for the python _winreg module Much nicer interface and full support for unicode. Uses ctypes. Requires at least Windows Vista. --- src/calibre/utils/winreg/__init__.py | 1 + src/calibre/utils/winreg/lib.py | 240 +++++++++++++++++++++++++++ 2 files changed, 241 insertions(+) create mode 100644 src/calibre/utils/winreg/__init__.py create mode 100644 src/calibre/utils/winreg/lib.py diff --git a/src/calibre/utils/winreg/__init__.py b/src/calibre/utils/winreg/__init__.py new file mode 100644 index 0000000000..8b13789179 --- /dev/null +++ b/src/calibre/utils/winreg/__init__.py @@ -0,0 +1 @@ + diff --git a/src/calibre/utils/winreg/lib.py b/src/calibre/utils/winreg/lib.py new file mode 100644 index 0000000000..c512734abc --- /dev/null +++ b/src/calibre/utils/winreg/lib.py @@ -0,0 +1,240 @@ +#!/usr/bin/env python2 +# vim:fileencoding=utf-8 +from __future__ import (unicode_literals, division, absolute_import, + print_function) + +__license__ = 'GPL v3' +__copyright__ = '2015, Kovid Goyal ' + +import ctypes, ctypes.wintypes as types, _winreg as winreg, struct +import winerror, win32con + +# Binding to C library {{{ +advapi32 = ctypes.windll.advapi32 +HKEY = types.HKEY +PHKEY = ctypes.POINTER(HKEY) +DWORD = types.DWORD +BYTE = types.BYTE +LONG = types.LONG +ULONG = types.ULONG +LPDWORD = ctypes.POINTER(DWORD) +LPBYTE = ctypes.POINTER(BYTE) +LPCWSTR = types.LPCWSTR +LPWSTR = types.LPWSTR +LPCVOID = types.LPCVOID + +HKEY_CURRENT_USER = HKCU = HKEY(ULONG(winreg.HKEY_CURRENT_USER).value) +HKEY_CLASSES_ROOT = HKCR = HKEY(ULONG(winreg.HKEY_CLASSES_ROOT).value) +HKEY_LOCAL_MACHINE = HKLM = HKEY(ULONG(winreg.HKEY_LOCAL_MACHINE).value) +KEY_READ = winreg.KEY_READ +KEY_ALL_ACCESS = winreg.KEY_ALL_ACCESS + +class FILETIME(ctypes.Structure): + _fields_ = [("dwLowDateTime", DWORD), ("dwHighDateTime", DWORD)] + +def default_errcheck(result, func, args): + if result != getattr(winerror, 'ERROR_SUCCESS', 0): # On shutdown winerror becomes None + raise ctypes.WinError(result) + return args + +null = object() +class a(object): + + def __init__(self, name, typ, default=null, in_arg=True): + self.typ = typ + if default is null: + self.spec = ((1 if in_arg else 2), name) + else: + self.spec = ((1 if in_arg else 2), name, default) + +def cwrap(name, restype, *args, **kw): + params = (restype,) + tuple(x.typ for x in args) + paramflags = tuple(x.spec for x in args) + func = ctypes.WINFUNCTYPE(*params)((name, kw.get('lib', advapi32)), paramflags) + func.errcheck = kw.get('errcheck', default_errcheck) + return func + +RegOpenKey = cwrap( + 'RegOpenKeyExW', LONG, a('key', HKEY), a('sub_key', LPCWSTR), a('options', DWORD, 0), a('access', ULONG, KEY_READ), a('result', PHKEY, in_arg=False)) +RegCreateKey = cwrap( + 'RegCreateKeyExW', LONG, a('key', HKEY), a('sub_key', LPCWSTR, ''), a('reserved', DWORD, 0), a('cls', LPWSTR, None), a('options', DWORD, 0), + a('access', ULONG, KEY_ALL_ACCESS), a('security', ctypes.c_void_p, 0), a('result', PHKEY, in_arg=False), a('disposition', LPDWORD, in_arg=False)) +RegCloseKey = cwrap('RegCloseKey', LONG, a('key', HKEY)) + +def enum_value_errcheck(result, func, args): + if result == winerror.ERROR_SUCCESS: + return args + if result == winerror.ERROR_MORE_DATA: + raise ValueError('buffer too small') + if result == winerror.ERROR_NO_MORE_ITEMS: + raise StopIteration() + raise ctypes.WinError(result) +RegEnumValue = cwrap( + 'RegEnumValueW', LONG, a('key', HKEY), a('index', DWORD), a('value_name', LPWSTR), a('value_name_size', LPDWORD), a('reserved', LPDWORD), + a('value_type', LPDWORD), a('data', LPBYTE), a('data_size', LPDWORD), errcheck=enum_value_errcheck) + +def last_error_errcheck(result, func, args): + if result == 0: + raise ctypes.WinError() + return args +ExpandEnvironmentStrings = cwrap( + 'ExpandEnvironmentStringsW', DWORD, a('src', LPCWSTR), a('dest', LPWSTR), a('size', DWORD), errcheck=last_error_errcheck, lib=ctypes.windll.kernel32) + +def expand_environment_strings(src): + buf = ctypes.create_unicode_buffer(32 * 1024) + ExpandEnvironmentStrings(src, buf, len(buf)) + return buf.value + +def convert_to_registry_data(value, has_expansions=False): + if value is None: + return None, winreg.REG_NONE, 0 + if isinstance(value, (type(''), bytes)): + buf = ctypes.create_unicode_buffer(value) + return buf, (winreg.REG_EXPAND_SZ if has_expansions else winreg.REG_SZ), len(buf) * 2 + if isinstance(value, (list, tuple)): + buf = ctypes.create_unicode_buffer('\0'.join(map(type(''), value)) + '\0\0') + return buf, winreg.REG_MULTI_SZ, len(buf) * 2 + if isinstance(value, (int, long)): + try: + raw, dtype = struct.pack(str('L'), value), winreg.REG_DWORD + except struct.error: + raw = struct.pack(str('Q'), value), win32con.REG_QWORD + buf = ctypes.create_string_buffer(raw) + return buf, dtype, len(buf) + raise ValueError('Unknown data type: %r' % value) + +def convert_registry_data(raw, size, dtype): + if dtype == winreg.REG_NONE: + return None + if dtype == winreg.REG_BINARY: + return ctypes.string_at(raw, size) + if dtype in (winreg.REG_SZ, winreg.REG_EXPAND_SZ, winreg.REG_MULTI_SZ): + ans = ctypes.wstring_at(raw, size // 2).rstrip('\0') + if dtype == winreg.REG_MULTI_SZ: + ans = tuple(ans.split('\0')) + elif dtype == winreg.REG_EXPAND_SZ: + ans = expand_environment_strings(ans) + return ans + if dtype == winreg.REG_DWORD: + if size == 0: + return 0 + return ctypes.cast(raw, LPDWORD).contents.value + if dtype == win32con.REG_QWORD: + if size == 0: + return 0 + return ctypes.cast(raw, ctypes.POINTER(types.QWORD)).contents.value + raise ValueError('Unsupported data type: %r' % dtype) + +RegSetKeyValue = cwrap( + 'RegSetKeyValueW', LONG, a('key', HKEY), a('sub_key', LPCWSTR, None), a('name', LPCWSTR, None), + a('dtype', DWORD, winreg.REG_SZ), a('data', LPCVOID, None), a('size', DWORD)) + +RegDeleteTree = cwrap( + 'RegDeleteTreeW', LONG, a('key', HKEY), a('sub_key', LPCWSTR, None)) +# }}} + +class Key(object): + + def __init__(self, create_at=None, open_at=None, root=HKEY_CURRENT_USER, open_mode=KEY_READ): + root = getattr(root, 'hkey', root) + self.was_created = False + if create_at is not None: + self.hkey, self.was_created = RegCreateKey(root, create_at) + elif open_at is not None: + self.hkey = RegOpenKey(root, open_at, 0, open_mode) + else: + self.hkey = HKEY() + + def delete_tree(self, sub_key=None): + ''' Delete this key and all its children. Note that a key is not + actually deleted till the last handle to it is closed. ''' + RegDeleteTree(self.hkey, sub_key) + + def set(self, name=None, value=None, sub_key=None, has_expansions=False): + ''' Set a value for this key (with optional sub-key). If name is None, + the Default value is set. value can be an integer, a string or a list + of strings. If you want to use expansions, set has_expansions=True. ''' + value, dtype, size = convert_to_registry_data(value, has_expansions=has_expansions) + RegSetKeyValue(self.hkey, sub_key, name, dtype, value, size) + + def set_default_value(self, sub_key=None, value=None, has_expansions=False): + self.set(sub_key=sub_key, value=value, has_expansions=has_expansions) + + def sub_key(self, path, allow_create=True, open_mode=KEY_READ): + ' Create (or open) a sub-key at the specified relative path. When opening instead of creating, use open_mode ' + if allow_create: + return Key(create_at=path, root=self.hkey) + return Key(open_at=path, root=self.hkey) + + def itervalues(self, get_data=False, sub_key=None): + '''Iterate over all values in this key (or optionally the specified + sub-key. If get_data is True also return the data for every value, + otherwise, just the name.''' + key = self.hkey + if sub_key is not None: + try: + key = RegOpenKey(key, sub_key) + except WindowsError: + return + try: + name_buf = ctypes.create_unicode_buffer(16385) + lname_buf = DWORD(len(name_buf)) + if get_data: + data_buf = (BYTE * 1024)() + ldata_buf = DWORD(len(data_buf)) + vtype = DWORD(0) + i = 0 + while True: + lname_buf.value = len(name_buf) + if get_data: + ldata_buf.value = len(data_buf) + try: + RegEnumValue( + key, i, name_buf, ctypes.byref(lname_buf), None, ctypes.byref(vtype), data_buf, ctypes.byref(ldata_buf)) + except ValueError: + data_buf = (BYTE * ldata_buf.value)() + continue + data = convert_registry_data(data_buf, ldata_buf.value, vtype.value) + yield name_buf.value[:lname_buf.value], data + else: + RegEnumValue( + key, i, name_buf, ctypes.byref(lname_buf), None, None, None, None) + yield name_buf.value[:lname_buf.value] + + i += 1 + finally: + if sub_key is not None: + RegCloseKey(key) + + def __enter__(self): + return self + + def __exit__(self, *args): + self.close() + + def __nonzero__(self): + return bool(self.hkey) + + def close(self): + if not self.hkey: + return + if RegCloseKey is None or HKEY is None: + return # globals become None during exit + RegCloseKey(self.hkey) + self.hkey = HKEY() + + def __del__(self): + self.close() + +if __name__ == '__main__': + from pprint import pprint + k = Key(open_at=r'Software\RegisteredApplications', root=HKLM) + pprint(tuple(k.itervalues(get_data=True))) + k = Key(r'Software\calibre\winregtest') + k.set('Moose.Cat.1') + k.set('unicode test', 'fällen粗楷体简a\U0001f471') + k.set('none test') + k.set_default_value(r'other\key', '%PATH%', has_expansions=True) + pprint(tuple(k.itervalues(get_data=True))) + k.set_default_value(r'delete\me\please', 'xxx') + k.delete_tree('delete')