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

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

added a new testcase to the smoketest for data attributes
fixed data naming issues in pyregfi
File size: 16.6 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        # This is quicker than retrieving the root key for comparison and
332        # is more trustworthy than trusting the key's flags.
333        return ((self._hive.root_cell+REGFI_REGF_SIZE) == self.offset)
334
335
336## Registry value (metadata)
337#
338# These represent registry values (@ref REGFI_VK records) and provide
339# access to their associated data.
340#
341class Value(_StructureWrapper):
342    def __getattr__(self, name):
343        ret_val = None
344        if name == "data":
345            data_p = regfi.regfi_fetch_data(self._hive.file, self._base)
346            try:
347                data_struct = data_p.contents
348            except Exception:
349                return None
350
351            if data_struct.interpreted_size == 0:
352                ret_val = None
353            elif data_struct.type in (REG_SZ, REG_EXPAND_SZ, REG_LINK):
354                # Unicode strings
355                ret_val = data_struct.interpreted.string.decode('utf-8', 'replace')
356            elif data_struct.type in (REG_DWORD, REG_DWORD_BE):
357                # 32 bit integers
358                ret_val = data_struct.interpreted.dword
359            elif data_struct.type == REG_QWORD:
360                # 64 bit integers
361                ret_val = data_struct.interpreted.qword
362            elif data_struct.type == REG_MULTI_SZ:
363                ret_val = _charss2strlist(data_struct.interpreted.multiple_string)
364            elif data_struct.type in (REG_NONE, REG_RESOURCE_LIST,
365                                      REG_FULL_RESOURCE_DESCRIPTOR,
366                                      REG_RESOURCE_REQUIREMENTS_LIST,
367                                      REG_BINARY):
368                ret_val = _buffer2bytearray(data_struct.interpreted.none,
369                                            data_struct.interpreted_size)
370
371            regfi.regfi_free_record(data_p)
372           
373        elif name == "data_raw":
374            # XXX: should we load the data without interpretation instead?
375            data_p = regfi.regfi_fetch_data(self._hive.file, self._base)
376            try:
377                data_struct = data_p.contents
378            except Exception:
379                return None
380
381            ret_val = _buffer2bytearray(data_struct.raw,
382                                        data_struct.size)
383            regfi.regfi_free_record(data_p)
384           
385        else:
386            ret_val = super(Value, self).__getattr__(name)
387            if name == "name":
388                if ret_val == None:
389                    ret_val = self.name_raw
390                else:
391                    ret_val = ret_val.decode('utf-8', 'replace')
392
393            elif name == "name_raw":
394                length = super(Value, self).__getattr__('name_length')
395                ret_val = _buffer2bytearray(ret_val, length)
396
397        return ret_val
398
399
400# Avoids chicken/egg class definitions.
401# Also makes for convenient code reuse in these lists' parent classes.
402_SubkeyList._constructor = Key
403_ValueList._constructor = Value
404
405
406
407## Represents a single registry hive (file)
408#
409class Hive():
410    file = None
411    raw_file = None
412   
413    def __init__(self, fh):
414        # The fileno method may not exist, or it may throw an exception
415        # when called if the file isn't backed with a descriptor.
416        try:
417            if hasattr(fh, 'fileno'):
418                self.file = regfi.regfi_alloc(fh.fileno(), REGFI_ENCODING_UTF8)
419                return
420        except:
421            pass
422       
423        self.raw_file = structures.REGFI_RAW_FILE()
424        self.raw_file.fh = fh
425        self.raw_file.seek = seek_cb_type(self.raw_file.cb_seek)
426        self.raw_file.read = read_cb_type(self.raw_file.cb_read)
427        self.file = regfi.regfi_alloc_cb(self.raw_file, REGFI_ENCODING_UTF8)
428
429    def __getattr__(self, name):
430        return getattr(self.file.contents, name)
431   
432    def __del__(self):
433        regfi.regfi_free(self.file)
434        if self.raw_file != None:
435            self.raw_file = None
436
437    def __iter__(self):
438        return HiveIterator(self)
439
440    def get_root(self):
441        return Key(self, regfi.regfi_get_rootkey(self.file))
442
443
444    ## Creates a @ref HiveIterator initialized at the specified path in
445    #  the hive.
446    #
447    # Raises an Exception if the path could not be found/traversed.
448    def subtree(self, path):
449        hi = HiveIterator(self)
450        hi.descend(path)
451        return hi
452
453
454## A special purpose iterator for registry hives
455#
456# Iterating over an object of this type causes all keys in a specific
457# hive subtree to be returned in a depth-first manner. These iterators
458# are typically created using the @ref Hive.subtree() function on a @ref Hive
459# object.
460#
461# HiveIterators can also be used to manually traverse up and down a
462# registry hive as they retain information about the current position in
463# the hive, along with which iteration state for subkeys and values for
464# every parent key.  See the @ref up and @ref down methods for more
465# information.
466class HiveIterator():
467    hive = None
468    iter = None
469    iteration_root = None
470
471    def __init__(self, hive):
472        self.iter = regfi.regfi_iterator_new(hive.file, REGFI_ENCODING_UTF8)
473        if self.iter == None:
474            raise Exception("Could not create iterator.  Current log:\n"
475                            + GetLogMessages())
476        self._hive = hive
477       
478    def __getattr__(self, name):
479        return getattr(self.file.contents, name)
480
481    def __del__(self):   
482        regfi.regfi_iterator_free(self.iter)
483
484    def __iter__(self):
485        self.iteration_root = None
486        return self
487
488    def __next__(self):
489        if self.iteration_root == None:
490            self.iteration_root = self.current_key()
491        elif not regfi.regfi_iterator_down(self.iter):
492            up_ret = regfi.regfi_iterator_up(self.iter)
493            while (up_ret and
494                   not regfi.regfi_iterator_next_subkey(self.iter)):
495                if self.iteration_root == self.current_key():
496                    self.iteration_root = None
497                    raise StopIteration('')
498                up_ret = regfi.regfi_iterator_up(self.iter)
499
500            if not up_ret:
501                raise StopIteration('')
502           
503            # XXX: Use non-generic exception
504            if not regfi.regfi_iterator_down(self.iter):
505                raise Exception('Error traversing iterator downward.'+
506                                ' Current log:\n'+ GetLogMessages())
507
508        regfi.regfi_iterator_first_subkey(self.iter)
509        return self.current_key()
510
511    # For Python 2.x
512    next = __next__
513
514    def down(self):
515        pass
516
517    def up(self):
518        pass
519
520    def descend(self, path):
521        cpath = _strlist2charss(path)
522
523        # XXX: Use non-generic exception
524        if not regfi.regfi_iterator_walk_path(self.iter, cpath):
525            raise Exception('Could not locate path.\n'+GetLogMessages())
526
527    def current_key(self):
528        return Key(self._hive, regfi.regfi_iterator_cur_key(self.iter))
529
530    #XXX Add subkey/value search accessor functions (?)
Note: See TracBrowser for help on using the repository browser.