source: trunk/python/pyregfi/__init__.py @ 217

Last change on this file since 217 was 217, checked in by tim, 13 years ago

added another test
made is_root check more correct

File size: 16.4 KB
Line 
1#!/usr/bin/env python
2
3## @package pyregfi
4# Python interface to the regfi library.
5#
6
7import sys
8import weakref
9from pyregfi.structures import *
10
11import ctypes
12import ctypes.util
13from ctypes import c_char,c_char_p,c_int,c_uint16,c_uint32,c_bool,POINTER
14
15regfi = ctypes.CDLL(ctypes.util.find_library('regfi'), use_errno=True)
16
17
18regfi.regfi_alloc.argtypes = [c_int, REGFI_ENCODING]
19regfi.regfi_alloc.restype = POINTER(REGFI_FILE)
20
21regfi.regfi_alloc_cb.argtypes = [POINTER(REGFI_RAW_FILE), REGFI_ENCODING]
22regfi.regfi_alloc_cb.restype = POINTER(REGFI_FILE)
23
24regfi.regfi_free.argtypes = [POINTER(REGFI_FILE)]
25regfi.regfi_free.restype = None
26
27regfi.regfi_log_get_str.argtypes = []
28regfi.regfi_log_get_str.restype = c_char_p
29
30regfi.regfi_log_set_mask.argtypes = [c_uint16]
31regfi.regfi_log_set_mask.restype = c_bool
32
33regfi.regfi_get_rootkey.argtypes = [POINTER(REGFI_FILE)]
34regfi.regfi_get_rootkey.restype = POINTER(REGFI_NK)
35
36regfi.regfi_free_record.argtypes = [c_void_p]
37regfi.regfi_free_record.restype = None
38
39regfi.regfi_fetch_num_subkeys.argtypes = [POINTER(REGFI_NK)]
40regfi.regfi_fetch_num_subkeys.restype = c_uint32
41
42regfi.regfi_fetch_num_values.argtypes = [POINTER(REGFI_NK)]
43regfi.regfi_fetch_num_values.restype = c_uint32
44
45regfi.regfi_fetch_classname.argtypes = [POINTER(REGFI_FILE), POINTER(REGFI_NK)]
46regfi.regfi_fetch_classname.restype = POINTER(REGFI_CLASSNAME)
47
48regfi.regfi_fetch_sk.argtypes = [POINTER(REGFI_FILE), POINTER(REGFI_NK)]
49regfi.regfi_fetch_sk.restype = POINTER(REGFI_SK)
50
51regfi.regfi_fetch_data.argtypes = [POINTER(REGFI_FILE), POINTER(REGFI_VK)]
52regfi.regfi_fetch_data.restype = POINTER(REGFI_DATA)
53
54regfi.regfi_find_subkey.argtypes = [POINTER(REGFI_FILE), POINTER(REGFI_NK),
55                                    c_char_p, POINTER(c_uint32)]
56regfi.regfi_find_subkey.restype = c_bool
57
58regfi.regfi_find_value.argtypes = [POINTER(REGFI_FILE), POINTER(REGFI_NK),
59                                    c_char_p, POINTER(c_uint32)]
60regfi.regfi_find_value.restype = c_bool
61
62regfi.regfi_get_subkey.argtypes = [POINTER(REGFI_FILE), POINTER(REGFI_NK),
63                                   c_uint32]
64regfi.regfi_get_subkey.restype = POINTER(REGFI_NK)
65
66regfi.regfi_get_value.argtypes = [POINTER(REGFI_FILE), POINTER(REGFI_NK),
67                                   c_uint32]
68regfi.regfi_get_value.restype = POINTER(REGFI_VK)
69
70regfi.regfi_get_parentkey.argtypes = [POINTER(REGFI_FILE), POINTER(REGFI_NK)]
71regfi.regfi_get_parentkey.restype = POINTER(REGFI_NK)
72
73regfi.regfi_iterator_new.argtypes = [POINTER(REGFI_FILE), REGFI_ENCODING]
74regfi.regfi_iterator_new.restype = POINTER(REGFI_ITERATOR)
75
76regfi.regfi_iterator_free.argtypes = [POINTER(REGFI_ITERATOR)]
77regfi.regfi_iterator_free.restype = None
78
79regfi.regfi_iterator_down.argtypes = [POINTER(REGFI_ITERATOR)]
80regfi.regfi_iterator_down.restype = c_bool
81
82regfi.regfi_iterator_up.argtypes = [POINTER(REGFI_ITERATOR)]
83regfi.regfi_iterator_up.restype = c_bool
84
85regfi.regfi_iterator_to_root.argtypes = [POINTER(REGFI_ITERATOR)]
86regfi.regfi_iterator_to_root.restype = c_bool
87
88regfi.regfi_iterator_walk_path.argtypes = [POINTER(REGFI_ITERATOR)]
89regfi.regfi_iterator_walk_path.restype = c_bool
90
91regfi.regfi_iterator_cur_key.argtypes = [POINTER(REGFI_ITERATOR)]
92regfi.regfi_iterator_cur_key.restype = POINTER(REGFI_NK)
93
94regfi.regfi_iterator_first_subkey.argtypes = [POINTER(REGFI_ITERATOR)]
95regfi.regfi_iterator_first_subkey.restype = c_bool
96
97regfi.regfi_iterator_cur_subkey.argtypes = [POINTER(REGFI_ITERATOR)]
98regfi.regfi_iterator_cur_subkey.restype = POINTER(REGFI_NK)
99
100regfi.regfi_iterator_next_subkey.argtypes = [POINTER(REGFI_ITERATOR)]
101regfi.regfi_iterator_next_subkey.restype = c_bool
102
103regfi.regfi_iterator_find_subkey.argtypes = [POINTER(REGFI_ITERATOR), c_char_p]
104regfi.regfi_iterator_find_subkey.restype = c_bool
105
106regfi.regfi_iterator_first_value.argtypes = [POINTER(REGFI_ITERATOR)]
107regfi.regfi_iterator_first_value.restype = c_bool
108
109regfi.regfi_iterator_cur_value.argtypes = [POINTER(REGFI_ITERATOR)]
110regfi.regfi_iterator_cur_value.restype = POINTER(REGFI_VK)
111
112regfi.regfi_iterator_next_value.argtypes = [POINTER(REGFI_ITERATOR)]
113regfi.regfi_iterator_next_value.restype = c_bool
114
115regfi.regfi_iterator_find_value.argtypes = [POINTER(REGFI_ITERATOR), c_char_p]
116regfi.regfi_iterator_find_value.restype = c_bool
117
118
119regfi.regfi_init.argtypes = []
120regfi.regfi_init.restype = None
121regfi.regfi_init()
122
123
124## Retrieves messages produced by regfi during parsing and interpretation
125#
126def GetLogMessages():
127    msgs = regfi.regfi_log_get_str()
128    if msgs == None:
129        return ''
130    return msgs.decode('utf-8')
131
132
133def _buffer2bytearray(char_pointer, length):
134    if length == 0 or char_pointer == None:
135        return None
136   
137    ret_val = bytearray(length)
138    for i in range(0,length):
139        ret_val[i] = char_pointer[i][0]
140
141    return ret_val
142
143
144def _strlist2charss(str_list):
145    ret_val = []
146    for s in str_list:
147        ret_val.append(s.encode('utf-8', 'replace'))
148
149    ret_val = (c_char_p*(len(str_list)+1))(*ret_val)
150    # Terminate the char** with a NULL pointer
151    ret_val[-1] = 0
152
153    return ret_val
154
155
156def _charss2strlist(chars_pointer):
157    ret_val = []
158    i = 0
159    s = chars_pointer[i]
160    while s != None:
161        ret_val.append(s.decode('utf-8', 'replace'))
162        i += 1
163        s = chars_pointer[i]
164
165    return ret_val
166
167
168## Abstract class which Handles memory management and proxies attribute
169#  access to base structures 
170class _StructureWrapper(object):
171    _hive = None
172    _base = None
173
174    def __init__(self, hive, base):
175        if not hive:
176            raise Exception("Could not create _StructureWrapper,"
177                            + " hive is NULL.  Current log:\n"
178                            + GetLogMessages())
179        if not base:
180            raise Exception("Could not create _StructureWrapper,"
181                            + " base is NULL.  Current log:\n"
182                            + GetLogMessages())
183        self._hive = hive
184        self._base = base
185
186    def __del__(self):
187        regfi.regfi_free_record(self._base)
188
189    def __getattr__(self, name):
190        return getattr(self._base.contents, name)
191
192    def __eq__(self, other):
193        return (type(self) == type(other)) and (self.offset == other.offset)
194
195    def __ne__(self, other):
196        return (not self.__eq__(other))
197
198class Key(_StructureWrapper):
199    pass
200
201class Value(_StructureWrapper):
202    pass
203
204## Registry value data
205class Data(_StructureWrapper):
206    pass
207
208## Registry security record/permissions
209class Security(_StructureWrapper):
210    pass
211
212
213class _GenericList(object):
214    _hive = None
215    _key = None
216    _length = None
217    _current = None
218
219    # implementation-specific functions for _SubkeyList and _ValueList
220    _fetch_num = None
221    _find_element = None
222    _get_element = None
223    _constructor = None
224
225    def __init__(self, key):
226        self._hive = key._hive
227
228        # Normally it's good to avoid cyclic references like this
229        # (key.list.key...) but in this case it makes ctypes memory
230        # management easier to reference the Key instead of the base
231        # structure.  We use a weak reference in order to allow for garbage
232        # collection, since value/subkey lists should not be usable if their
233        # parent Key is freed anyway.
234
235        # XXX: check for NULL here, throw an exception if so.
236        self._key = weakref.proxy(key)
237        self._length = self._fetch_num(key._base)
238
239   
240    def __len__(self):
241        return self._length
242
243    def __getitem__(self, name):
244        index = c_uint32()
245        if isinstance(name, str):
246            name = name.encode('utf-8')
247
248        if name != None:
249            name = create_string_buffer(bytes(name))
250
251        if self._find_element(self._hive.file, self._key._base, name, byref(index)):
252            return self._constructor(self._hive, 
253                                     self._get_element(self._hive.file,
254                                                       self._key._base,
255                                                       index))
256        raise KeyError('')
257
258    def get(self, name, default):
259        try:
260            return self[name]
261        except KeyError:
262            return default
263   
264    def __iter__(self):
265        self._current = 0
266        return self
267   
268    def __next__(self):
269        if self._current >= self._length:
270            raise StopIteration('')
271
272        elem = self._get_element(self._hive.file, self._key._base,
273                                 c_uint32(self._current))
274        self._current += 1
275        return self._constructor(self._hive, elem)
276   
277    # For Python 2.x
278    next = __next__
279
280
281class _SubkeyList(_GenericList):
282    _fetch_num = regfi.regfi_fetch_num_subkeys
283    _find_element = regfi.regfi_find_subkey
284    _get_element = regfi.regfi_get_subkey
285
286
287class _ValueList(_GenericList):
288    _fetch_num = regfi.regfi_fetch_num_values
289    _find_element = regfi.regfi_find_value
290    _get_element = regfi.regfi_get_value
291
292
293## Registry key
294class Key(_StructureWrapper):
295    values = None
296    subkeys = None
297
298    def __init__(self, hive, base):
299        super(Key, self).__init__(hive, base)
300        self.values = _ValueList(self)
301        self.subkeys = _SubkeyList(self)
302
303    def __getattr__(self, name):
304        ret_val = super(Key, self).__getattr__(name)
305       
306        if name == "name":
307            if ret_val == None:
308                ret_val = self.name_raw
309            else:
310                ret_val = ret_val.decode('utf-8', 'replace')
311               
312        elif name == "name_raw":
313            length = super(Key, self).__getattr__('name_length')
314            ret_val = _buffer2bytearray(ret_val, length)
315       
316        return ret_val
317
318
319    def fetch_security(self):
320        return Security(self._hive,
321                        regfi.regfi_fetch_sk(self._hive.file, self._base))
322
323    def get_parent(self):
324        parent_base = regfi.regfi_get_parentkey(self._hive.file, self._base)
325        if parent_base:
326            return Key(self._hive, parent_base)
327
328        return None
329
330    def is_root(self):
331        return (self._hive.get_root() == self)
332
333
334## Registry value (metadata)
335#
336# These represent registry values (@ref REGFI_VK records) and provide
337# access to their associated data.
338#
339class Value(_StructureWrapper):
340    def __getattr__(self, name):
341        ret_val = None
342        if name == "data":
343            data_p = regfi.regfi_fetch_data(self._hive.file, self._base)
344            try:
345                data_struct = data_p.contents
346            except Exception:
347                return None
348
349            if data_struct.interpreted_size == 0:
350                ret_val = None
351            elif data_struct.type in (REG_SZ, REG_EXPAND_SZ, REG_LINK):
352                # Unicode strings
353                ret_val = data_struct.interpreted.string.decode('utf-8', 'replace')
354            elif data_struct.type in (REG_DWORD, REG_DWORD_BE):
355                # 32 bit integers
356                ret_val = data_struct.interpreted.dword
357            elif data_struct.type == REG_QWORD:
358                # 64 bit integers
359                ret_val = data_struct.interpreted.qword
360            elif data_struct.type == REG_MULTI_SZ:
361                ret_val = _charss2strlist(data_struct.interpreted.multiple_string)
362            elif data_struct.type in (REG_NONE, REG_RESOURCE_LIST,
363                                      REG_FULL_RESOURCE_DESCRIPTOR,
364                                      REG_RESOURCE_REQUIREMENTS_LIST,
365                                      REG_BINARY):
366                ret_val = _buffer2bytearray(data_struct.interpreted.none,
367                                            data_struct.interpreted_size)
368
369            regfi.regfi_free_record(data_p)
370           
371        elif name == "data_raw":
372            # XXX: should we load the data without interpretation instead?
373            data_p = regfi.regfi_fetch_data(self._hive.file, self._base)
374            try:
375                data_struct = data_p.contents
376            except Exception:
377                return None
378
379            ret_val = _buffer2bytearray(data_struct.raw,
380                                        data_struct.size)
381            regfi.regfi_free_record(data_p)
382           
383        else:
384            ret_val = super(Value, self).__getattr__(name)
385            if name == "name":
386                if ret_val == None:
387                    ret_val = self.name_raw
388                else:
389                    ret_val = ret_val.decode('utf-8', 'replace')
390
391            elif name == "name_raw":
392                length = super(Value, self).__getattr__('name_length')
393                ret_val = _buffer2bytearray(ret_val, length)
394
395        return ret_val
396
397
398# Avoids chicken/egg class definitions.
399# Also makes for convenient code reuse in these lists' parent classes.
400_SubkeyList._constructor = Key
401_ValueList._constructor = Value
402
403
404
405## Represents a single registry hive (file)
406#
407class Hive():
408    file = None
409    raw_file = None
410   
411    def __init__(self, fh):
412        # The fileno method may not exist, or it may throw an exception
413        # when called if the file isn't backed with a descriptor.
414        try:
415            if hasattr(fh, 'fileno'):
416                self.file = regfi.regfi_alloc(fh.fileno(), REGFI_ENCODING_UTF8)
417                return
418        except:
419            pass
420       
421        self.raw_file = structures.REGFI_RAW_FILE()
422        self.raw_file.fh = fh
423        self.raw_file.seek = seek_cb_type(self.raw_file.cb_seek)
424        self.raw_file.read = read_cb_type(self.raw_file.cb_read)
425        self.file = regfi.regfi_alloc_cb(self.raw_file, REGFI_ENCODING_UTF8)
426
427    def __getattr__(self, name):
428        return getattr(self.file.contents, name)
429   
430    def __del__(self):
431        regfi.regfi_free(self.file)
432        if self.raw_file != None:
433            self.raw_file = None
434
435    def __iter__(self):
436        return HiveIterator(self)
437
438    def get_root(self):
439        return Key(self, regfi.regfi_get_rootkey(self.file))
440
441
442    ## Creates a @ref HiveIterator initialized at the specified path in
443    #  the hive.
444    #
445    # Raises an Exception if the path could not be found/traversed.
446    def subtree(self, path):
447        hi = HiveIterator(self)
448        hi.descend(path)
449        return hi
450
451
452## A special purpose iterator for registry hives
453#
454# Iterating over an object of this type causes all keys in a specific
455# hive subtree to be returned in a depth-first manner. These iterators
456# are typically created using the @ref Hive.subtree() function on a @ref Hive
457# object.
458#
459# HiveIterators can also be used to manually traverse up and down a
460# registry hive as they retain information about the current position in
461# the hive, along with which iteration state for subkeys and values for
462# every parent key.  See the @ref up and @ref down methods for more
463# information.
464class HiveIterator():
465    hive = None
466    iter = None
467    iteration_root = None
468
469    def __init__(self, hive):
470        self.iter = regfi.regfi_iterator_new(hive.file, REGFI_ENCODING_UTF8)
471        if self.iter == None:
472            raise Exception("Could not create iterator.  Current log:\n"
473                            + GetLogMessages())
474        self._hive = hive
475       
476    def __getattr__(self, name):
477        return getattr(self.file.contents, name)
478
479    def __del__(self):   
480        regfi.regfi_iterator_free(self.iter)
481
482    def __iter__(self):
483        self.iteration_root = None
484        return self
485
486    def __next__(self):
487        if self.iteration_root == None:
488            self.iteration_root = self.current_key()
489        elif not regfi.regfi_iterator_down(self.iter):
490            up_ret = regfi.regfi_iterator_up(self.iter)
491            while (up_ret and
492                   not regfi.regfi_iterator_next_subkey(self.iter)):
493                if self.iteration_root == self.current_key():
494                    self.iteration_root = None
495                    raise StopIteration('')
496                up_ret = regfi.regfi_iterator_up(self.iter)
497
498            if not up_ret:
499                raise StopIteration('')
500           
501            # XXX: Use non-generic exception
502            if not regfi.regfi_iterator_down(self.iter):
503                raise Exception('Error traversing iterator downward.'+
504                                ' Current log:\n'+ GetLogMessages())
505
506        regfi.regfi_iterator_first_subkey(self.iter)
507        return self.current_key()
508
509    # For Python 2.x
510    next = __next__
511
512    def down(self):
513        pass
514
515    def up(self):
516        pass
517
518    def descend(self, path):
519        cpath = _strlist2charss(path)
520
521        # XXX: Use non-generic exception
522        if not regfi.regfi_iterator_walk_path(self.iter, cpath):
523            raise Exception('Could not locate path.\n'+GetLogMessages())
524
525    def current_key(self):
526        return Key(self._hive, regfi.regfi_iterator_cur_key(self.iter))
527
528    #XXX Add subkey/value search accessor functions (?)
Note: See TracBrowser for help on using the repository browser.