naksyn-PythonMemoryModule/pythonmemorymodule/windows/winobject/registry.py

import sys
import ctypes
import itertools
import struct
from collections import namedtuple, defaultdict

import windows
from windows.dbgprint import dbgprint
import windows.generated_def as gdef
from windows import winproxy
from windows.pycompat import basestring, int_types, is_py3

WENCODING = "utf-16-le"

# So _winreg does not handle unicode stuff in Py2 :(
# Need to rewrite everything to get it working with unicode

class WinRegistryKey(gdef.HKEY):
    _close_function = staticmethod(winproxy.RegCloseKey)

    def __del__(self):
        if sys is None or sys.path is None: # Late shutdown (not sur winproxy is still up)
            return
        if self: # Not NULL handle ?
            dbgprint(u"Closing registry key handle {0:#x}".format(self.value), 'REGISTRY')
            self._close_function(self)




class ExpectWindowsError(object):
    def __init__(self, errornumber):
        self.errornumber = errornumber

    def __enter__(self):
        pass

    def __exit__(self, etype, e, tb):
        return (etype in (winproxy.WinproxyError, WindowsError) and e.winerror == self.errornumber)

# Translation reg-buffer <-> python methodes
def Reg2Py_QWORD(buffer, size):
    return buffer.cast(gdef.PULONG64)[0]

def Py2Reg_QWORD(obj):
    return struct.pack("<Q", obj)

def Reg2Py_DWORD(buffer, size):
    # Check size ?
    return buffer.cast(gdef.LPDWORD)[0]

def Py2Reg_DWORD(obj):
    return struct.pack("<I", obj)


def Reg2Py_DWORD_BIG_ENDIAN(buffer, size):
    # Check size ?
    return (buffer[0] << 24) + (buffer[1] << 16) + (buffer[2] << 8) + buffer[3]

def Py2Reg_DWORD_BIG_ENDIAN(obj):
    return struct.pack(">I", obj)


def Reg2Py_BINARY(buffer, size):
    return bytes(bytearray(buffer[:size]))

def Py2Reg_BINARY(obj):
    # latin-1 encoding if py3 & type is str ?
    return obj


def Reg2Py_SZ(buffer, size):
    # Buffer is UTF16. buffer is extended-buffer
    if size == 0:
        return u""
    if buffer[size - 1] == 0 and buffer[size - 2] == 0:
        # NULL TERMINATED: EASY
        return buffer.as_wstring()
    # Not null terminated: keep last byte
    assert not size % 2
    return (gdef.WCHAR * (size // 2)).from_buffer(buffer)[:]

def Py2Reg_SZ(obj):
    return obj.encode(WENCODING)

def Reg2Py_Multi_SZ(buffer, size):
    if not size:
        return []
    # Simple path
    if is_py3:
        rawstr = bytes(buffer)
    else:
        rawstr = "".join([chr(c) for c in buffer[:size]])
    try:
        unistr = rawstr.decode(WENCODING)
        return unistr.rstrip(u"\x00").split(u"\x00")
    except UnicodeDecodeError as e:
        pass
    # Complexe-path
    # This is not some valide UTF-16
    # Try our best to extract some stuff from raw
    return rawstr.rstrip(b"\x00").split(b"\x00")


def Py2Reg_Multi_SZ(obj):
    # Work on encoded values (to prevent str/unicode errors)
    uni_list = [s.encode(WENCODING) for s in obj]
    # Separate by UTF-16 NULL BYTE (2 \x00)
    uni_str = b"\x00\x00".join(uni_list)
    # Add UTF-16 NULL byte for final string + final UTF-16 \x00 (4 \x00)
    return uni_str + b"\x00\x00\x00\x00"

DECODE_METHOD = 0
ENCODE_METHOD = 1


KNOWN_ENCODE_DECODE_METHODS = {
    gdef.REG_SZ: (Reg2Py_SZ, Py2Reg_SZ),
    gdef.REG_EXPAND_SZ: (Reg2Py_SZ, Py2Reg_SZ),
    gdef.REG_MULTI_SZ: (Reg2Py_Multi_SZ, Py2Reg_Multi_SZ),
    gdef.REG_DWORD: (Reg2Py_DWORD, Py2Reg_DWORD),
    gdef.REG_DWORD_BIG_ENDIAN: (Reg2Py_DWORD_BIG_ENDIAN, Py2Reg_DWORD_BIG_ENDIAN),
    gdef.REG_QWORD: (Reg2Py_QWORD, Py2Reg_QWORD),
    # Binary formats
    gdef.REG_LINK: (Reg2Py_BINARY, Py2Reg_BINARY), # TESTING
    gdef.REG_BINARY: (Reg2Py_BINARY, Py2Reg_BINARY),
    gdef.REG_NONE: (Reg2Py_BINARY, Py2Reg_BINARY),
}

# All unknown format are seens as binary data
UNKNOWM_FORMAT = (Reg2Py_BINARY, Py2Reg_BINARY)
ENCODE_DECODE_METHODS = defaultdict(lambda: UNKNOWM_FORMAT, KNOWN_ENCODE_DECODE_METHODS)

def decode_registry_buffer(type, buffer, size):
    try:
        return ENCODE_DECODE_METHODS[type][DECODE_METHOD](buffer, size)
    except UnicodeDecodeError as e:
        # Best effort if any decoding error happen
        return "".join(chr(c) for c in buffer[:size])


KeyValue = namedtuple("KeyValue", ["name", "value", "type"])
"""A registry value (name, value, type)"""


class PyHKey(object):
    """A windows registry key"""
    def __init__(self, surkey, name, sam=gdef.KEY_READ):
        self.surkey = surkey
        self.name = name
        self.fullname = self.surkey.fullname + "\\" + self.name if self.name else self.surkey.name
        self.sam = sam
        self._phkey = None
        #self.phkey

    def __repr__(self):
        return '<PyHKey "{0}">'.format(self.fullname)

    def _open_key(self, handle, name, sam):
        result = WinRegistryKey()
        winproxy.RegOpenKeyExW(handle, name, 0, sam, result) # TODO: options REG_OPTION_OPEN_LINK
        dbgprint(u"Opening registry key <{0}> (handle={1:#x})".format(name, result.value), "REGISTRY")
        return result

    def _create_key(self, parent, name, sam):
        result = WinRegistryKey()
        flags = 0
        winproxy.RegCreateKeyExW(parent, name, 0, None, flags, sam, None, result, None)
        dbgprint(u"Creating registry key <{0}> (handle={1:#x})".format(name, result.value), "REGISTRY")
        return result

    @property
    def phkey(self):
        if self._phkey is not None:
            return self._phkey
        try:
            self._phkey = self._open_key(self.surkey.phkey, self.name, self.sam)
        except WindowsError as e:
            raise WindowsError(e.winerror, "Could not open registry key <{0}> ({1})".format(self.fullname, e.strerror))
        return self._phkey

    @property
    def exists(self):
        # May have been deleted in between
        # So <self._phkey> tells use nothing
        if self._phkey: # Not None + pointer not NULL
            try:
                self.get_key_size_info()
            except WindowsError as e:
                if e.winerror == gdef.ERROR_KEY_DELETED:
                    return False
                raise
            return True
        try:
            tmpphkey = self._open_key(self.surkey.phkey, self.name, gdef.KEY_READ)
        except WindowsError as e:
            return False
        # tmpphkey will be garbage collected and auto-closed
        return True

    @property
    def subkeys(self):
        """The subkeys of the registry key

        :type: [:class:`PyHKey`] - A list of keys"""
        res = []
        with ExpectWindowsError(259):
            default_name_size = 256 + 1
            name_size = gdef.DWORD(default_name_size)
            name_buffer = ctypes.create_unicode_buffer(name_size.value)
            for i in itertools.count():
                name_size.value = default_name_size
                winproxy.RegEnumKeyExW(self.phkey, i, name_buffer, name_size, None, None, None, None)
                res.append(name_buffer[:name_size.value]) # Will allow key name with \x00 inside
        return [PyHKey(self, n) for n in  res]

    def get_key_size_info(self):
        max_name_len = gdef.DWORD()
        max_value_len = gdef.DWORD()
        winproxy.RegQueryInfoKeyW(self.phkey, None, None, None, None, None, None, None, max_name_len, max_value_len, None, None)
        return (max_name_len.value, max_value_len.value)

    @property
    def values(self):
        """The values of the registry key

        :type: [:class:`KeyValue`] - A list of values"""
        res = []
        # Get max info keys

        max_name_size, max_data_size = self.get_key_size_info()
        # Null terminators
        max_name_size += 1
        max_data_size += 2
        with ExpectWindowsError(259):
            for i in itertools.count():
                value_type = gdef.DWORD()
                namesize = gdef.DWORD(max_name_size)
                keyname = ctypes.create_unicode_buffer(namesize.value)
                datasize = gdef.DWORD(max_data_size)
                databuffer = windows.utils.BUFFER(gdef.BYTE, nbelt=datasize.value)()
                # A value can have been added in-between.
                # So recheck the size given by get_key_size_info :)
                # But check 10 times max as RegEnumValueW may bug (seen) and always return  ERROR_MORE_DATA even with enought size
                for _ in range(10):
                    try:
                        winproxy.RegEnumValueW(self.phkey, i, keyname, namesize, None, value_type, databuffer, datasize)
                        break
                    except WindowsError as e:
                        if e.winerror != gdef.ERROR_MORE_DATA:
                            raise
                        # I found some strange Windows where even with a big enought buffer:
                        # - the data was filled
                        # - ERROR_MORE_DATA was returned
                        ## To prevent such bug to trigger and infinite loop, two things
                        # - If the retuned namesize <= the passed keysize and keyname is not empty -> return the data`
                        # - Max 10 test to prevent Infinite loop
                        if ((namesize.value <= max_name_size)  and (datasize.value <= max_data_size) and
                            (keyname[:namesize.value].count("\x00") < namesize.value)): # Not just 0 Zero ?
                            break

                        # Update the sizes / buffers & try again :)
                        max_name_size, max_data_size = self.get_key_size_info()
                        max_name_size = max(max_name_size + 1, namesize.value + 1) # namesize.value may be > to max_name_size apparently (guessed)
                        max_data_size = max(max_data_size + 2, datasize.value + 2) # datasize.value may be > to max_data_size apparently (seen)
                        namesize = gdef.DWORD(max_name_size)
                        keyname = ctypes.create_unicode_buffer(namesize.value)
                        datasize = gdef.DWORD(max_data_size)
                        databuffer = windows.utils.BUFFER(gdef.BYTE, nbelt=datasize.value)()
                else:
                    # Probably a windows bug that prevent us from retrieving the data
                    # Raise something (thus preventing getting the other values..) ? ignore it ?
                    raise ValueError("Could not extract registry key values, problably a Windows/hook bug")
                vobj = decode_registry_buffer(value_type.value, databuffer, datasize.value)
                res.append(KeyValue(keyname.value, vobj, value_type.value))
        return res


    @property
    def info(self):
        # Need other stuff ?
        nb_key = gdef.DWORD()
        nb_values = gdef.DWORD()
        last_modif = gdef.FILETIME()
        winproxy.RegQueryInfoKeyW(self.phkey, None, None, None, nb_key, None, None, nb_values, None, None, None, last_modif)
        return nb_key.value, nb_values.value, int(last_modif)

    @property
    def last_write(self):
        return self.info[2]

    def get(self, value_name):
        """Retrieves the value ``value_name``

        :rtype: :class:`KeyValue`
        """

        type = gdef.DWORD(0)
        size = gdef.DWORD(0x100)
        while True:
            buffer = windows.utils.BUFFER(gdef.BYTE, nbelt=size.value)()
            try:
                winproxy.RegQueryValueExW(self.phkey, value_name, None, type, buffer, size)
                break
            except WindowsError as e:
                if e.winerror != gdef.ERROR_MORE_DATA:
                    raise
                size.value *= 2
                buffer = windows.utils.BUFFER(gdef.BYTE, nbelt=size.value)()
                continue
        vobj = decode_registry_buffer(type.value, buffer, size.value)
        return KeyValue(value_name, vobj, type.value)

    def _guess_value_type(self, value):
        if isinstance(value, basestring):
            return gdef.REG_SZ
        elif isinstance(value, int_types):
            return gdef.REG_DWORD
        # elif isinstance(value, (list, tuple)):
            # if all(isinstance(v, basestring) in value):
                # return _winreg.REG_MULTI_SZ
        raise ValueError("Cannot guest registry type of value to set <{0}>".format(value))


    def set(self, name, value, type=None):
        """Set the value for ``name`` to ``value``. if ``type`` is None try to guess items"""
        if type is None:
            type = self._guess_value_type(value)


        buffer = ENCODE_DECODE_METHODS[type][ENCODE_METHOD](value)
        if isinstance(buffer, bytes): # Should not be unicode at this point
            buffer = windows.utils.BUFFER(gdef.BYTE).from_buffer_copy(buffer)
        return winproxy.RegSetValueExW(self.phkey, name, 0, type, buffer, len(buffer))

    def delete_value(self, name):
        """Delete the value with ``name``"""
        return winproxy.RegDeleteValueW(self.phkey, name)


    def open_subkey(self, name, sam=None):
        """Open the subkey ``name``

        :rtype: :class:`PyHKey`
        """
        if sam is None:
            sam = self.sam
        return PyHKey(self, name, sam)

    def reopen(self, sam):
        """Reopen the registry key with a new ``sam``

        :rtype: :class:`PyHKey`
        """
        return PyHKey(self.surkey, self.name, sam)

    def create(self):
        """Create the registry key"""
        try:
            self._phkey = self._create_key(self.surkey.phkey, self.name, self.sam)
        except WindowsError as e:
            raise WindowsError(e.winerror, "Could not create registry key <{0}> ({1})".format(self.fullname, e.strerror))
        return self

    def delete(self):
        """Delete the registry key"""
        # Allow a 'recursive' param to empty before delete ?
        try:
            windows.winproxy.RegDeleteKeyExW(self.surkey.phkey, self.name, self.sam, 0)
        except WindowsError as e:
            raise WindowsError(e.winerror, "Could not delete registry key <{0}> ({1})".format(self.fullname, e.strerror))
        return None

    def empty(self):
        windows.winproxy.RegDeleteTreeW(self.phkey, None)



    def __setitem__(self, name, value):
        rtype = None
        if not (isinstance(value, basestring) or isinstance(value, int_types)):
            value, rtype = value
        return self.set(name, value, rtype)

    __getitem__ = get

    __delitem__ = delete_value

    __call__ = open_subkey


class DummyPHKEY(object):
    def __init__(self, phkey, name):
        self.phkey = phkey
        self.name = name


HKEY_LOCAL_MACHINE = PyHKey(DummyPHKEY(gdef.HKEY_LOCAL_MACHINE, "HKEY_LOCAL_MACHINE"), "", gdef.KEY_READ)
HKEY_CLASSES_ROOT = PyHKey(DummyPHKEY(gdef.HKEY_CLASSES_ROOT, "HKEY_CLASSES_ROOT"), "", gdef.KEY_READ )
HKEY_CURRENT_USER = PyHKey(DummyPHKEY(gdef.HKEY_CURRENT_USER, "HKEY_CURRENT_USER"), "", gdef.KEY_READ)
HKEY_DYN_DATA = PyHKey(DummyPHKEY(gdef.HKEY_DYN_DATA, "HKEY_DYN_DATA"), "", gdef.KEY_READ)
HKEY_PERFORMANCE_DATA = PyHKey(DummyPHKEY(gdef.HKEY_PERFORMANCE_DATA, "HKEY_PERFORMANCE_DATA"), "", gdef.KEY_READ)
HKEY_USERS = PyHKey(DummyPHKEY(gdef.HKEY_USERS, "HKEY_USERS"), "", gdef.KEY_READ )


class Registry(object):
    """The ``Windows`` registry"""

    registry_base_keys = {
        "HKEY_LOCAL_MACHINE" : HKEY_LOCAL_MACHINE,
        "HKEY_CLASSES_ROOT" : HKEY_CLASSES_ROOT,
        "HKEY_CURRENT_USER" : HKEY_CURRENT_USER,
        "HKEY_DYN_DATA" : HKEY_DYN_DATA,
        "HKEY_PERFORMANCE_DATA": HKEY_PERFORMANCE_DATA,
        "HKEY_USERS" : HKEY_USERS
    }

    def __init__(self, sam=gdef.KEY_READ):
        self.sam = sam

    @classmethod
    def reopen(cls, sam):
        """Return a new :class:`Registry` using ``sam`` as the new default

        :rtype: :class:`Registry`
        """
        return cls(sam)

    def __call__(self, name, sam=None):
        """Get a registry key::

            registry(r"HKEY_LOCAL_MACHINE\\Software")
            registry("HKEY_LOCAL_MACHINE")("Software")

        :rtype: :class:`PyHKey`
        """
        if sam is None:
            sam = self.sam

        if name in self.registry_base_keys:
            key = self.registry_base_keys[name]
            if sam != key.sam:
                key = key.reopen(sam)
            return key
        if "\\" not in name:
            raise ValueError("Unknow registry base key <{0}>".format(name))
        base_name, subkey = name.split("\\", 1)
        if base_name not in self.registry_base_keys:
            raise ValueError("Unknow registry base key <{0}>".format(base_name))
        return self.registry_base_keys[base_name](subkey,  sam)