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

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

reorganized code
added remaining methods to iterator
added test case for new methods

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