source: test/pyregfi-smoketest.py @ 228

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

added a test case for pyregfi multithreaded use
updated the regfi multithreaded test case
fixed several ctypes interface problems in pyregfi
added locking to pyregfi iterators for thread safety
fixed regfi talloc race conditions with an additional lock

  • Property svn:executable set to *
File size: 7.9 KB
Line 
1#!/usr/bin/env python3
2
3import sys
4import gc
5import io
6import time
7import threading
8import pyregfi
9
10
11
12pyregfi.SetLogMask((pyregfi.LOG_TYPES.INFO, pyregfi.LOG_TYPES.WARN, pyregfi.LOG_TYPES.ERROR))
13
14# Uses the HiveIterator to walk all keys
15# Gathers various (meaningless) statistics to exercise simple attribute access
16# and to hopefully smoke out any bugs that can be identified by changing stats
17def iterTallyNames(hive, fh):
18    key_count = 0
19    key_lens = 0
20    key_rawlens = 0
21    value_count = 0
22    value_lens = 0
23    value_rawlens = 0
24
25    for k in hive:
26        key_count += 1
27        if k.name != None:
28            key_lens += len(k.name)
29        if k.name_raw != None:
30            key_rawlens += len(k.name_raw)
31
32        for v in k.values:
33            value_count += 1
34            if v.name != None:
35                value_lens += len(v.name)
36            if v.name_raw != None:
37                value_rawlens += len(v.name_raw)
38
39    print("  Counts: keys=%d, values=%d" % (key_count, value_count))
40    print("  Total name length: keys=%d, values=%d" % (key_lens, value_lens))
41    print("  Total raw name lengths: keys=%d, values=%d" % (key_rawlens, value_rawlens))
42
43
44# helper function
45def getCurrentPath(key):
46    if key == None:
47        return ''
48    path = []
49    p = key
50    while p != None:
51        path.append(p.name)
52        p = p.get_parent()
53    path.reverse()
54    del path[0]
55
56    return path
57
58# For each key in the hive, this traverses the parent links up to the root,
59# recording the path as it goes, and then uses the subtree/descend method
60# to find the same key again, verifying it is the same.  This test is currently
61# very slow because no key caching is used.
62def iterParentWalk(hive, fh):
63    i = 1
64    for k in hive:
65        path = getCurrentPath(k)
66        try:
67            hive_iter = hive.subtree(path)
68            if hive_iter.current_key() != k:
69                print("WARNING: k != current_key for path '%s'." % path)
70            else:
71                i += 1
72        except Exception as e:
73            print("WARNING: Could not decend to path '%s'.\nError:\n %s\n%s" % (path,e.args,e))
74    print("   Successfully tested paths on %d keys." % i)
75
76
77# Uses the HiveIterator to walk all keys
78# Gathers various (meaningless) statistics about data/data_raw attributes
79def iterTallyData(hive, fh):
80    data_stat = 0.0
81    dataraw_stat = 0.0
82   
83    for k in hive:
84        for v in k.values:
85            d = v.fetch_data()
86            if d == None:
87                data_stat += 0.1
88            elif hasattr(d, "__len__"):
89                data_stat += len(d)
90            else:
91                data_stat += d/2.0**64
92
93            d = v.fetch_raw_data()
94            if d == None:
95                dataraw_stat += 0.1
96            else:
97                dataraw_stat += len(d)
98
99    print("  Data stat: %f" % data_stat)
100    print("  Raw data stat: %f" % dataraw_stat)
101
102
103recurseKey_stat = 0.0
104recurseValue_stat = 0.0
105def checkValues(key):
106    global recurseKey_stat
107    global recurseValue_stat
108    recurseKey_stat += (key.mtime.low^key.mtime.high - key.max_bytes_subkeyname) * key.flags
109    for v in key.values:
110        recurseValue_stat += (v.data_off - v.data_size) / (1.0 + v.flags) + v.data_in_offset
111        value = key.values[v.name]
112        if v != value:
113            print("WARNING: iterator value '%s' does not match dictionary value '%s'." 
114                  % (v.name, value.name))
115
116def recurseTree(cur, operation):
117    for k in cur.subkeys:
118        key = cur.subkeys[k.name]
119        if k != key:
120            print("WARNING: iterator subkey '%s' does not match dictionary subkey '%s'." 
121                  % (k.name, key.name))
122        del key
123        operation(k)
124        recurseTree(k, operation)
125
126# Retrieves all keys by recursion, rather than the iterator, and validates
127# list dictionary access.  Also builds nonsensical statistics as an excuse
128# to access various base structure attributes.
129def recurseKeyTally(hive, fh):
130    checkValues(hive.root)
131    recurseTree(hive.root, checkValues)
132    print("  Key stat: %f" % recurseKey_stat)
133    print("  Value stat: %f" % recurseValue_stat)
134
135
136# Iterates hive gathering stats about security and classname records
137def iterFetchRelated(hive, fh):
138    security_stat = 0.0
139    classname_stat = 0.0
140    modified_stat = 0.0
141
142    for k in hive:
143        cn = k.fetch_classname()
144        if cn == None:
145            classname_stat += 0.000001
146        elif type(cn) == bytearray:
147            classname_stat += len(cn)/2**32
148        else:
149            classname_stat += len(cn)
150
151        modified_stat += k.modified
152       
153    print("  Security stat: %f" % security_stat)
154    print("  Classname stat: %f" % classname_stat)
155    print("  Modified stat: %f" % modified_stat)
156
157
158
159def iterIterWalk(hive, fh):
160    sk_stat = 0.0
161    v_stat = 0.0
162    iter = pyregfi.HiveIterator(hive)
163    for k in iter:
164        path = getCurrentPath(k)
165        try:
166            hive_iter = hive.subtree(path)
167            sk = hive_iter.first_subkey()
168            while sk != None:
169                ssk = hive_iter.find_subkey(sk.name)
170                sk_stat += len(ssk.name)
171                sk = hive_iter.next_subkey()
172
173            v = hive_iter.first_value()
174            while v != None:
175                vv = hive_iter.find_value(v.name)
176                v_stat += len(vv.name)
177                v = hive_iter.next_value()
178
179        except Exception as e:
180            print("WARNING: Could not decend to path '%s'.\nError:\n %s\n%s" % (path,e.args,e))
181    print("   Subkey stat: %f" % sk_stat)
182    print("   Value stat: %f" % v_stat)
183
184
185def iterCallbackIO(hive, fh):
186    fh.seek(0)
187    new_fh = io.BytesIO(fh.read())
188    new_hive = pyregfi.Hive(new_fh)
189    for k in new_hive:
190        pass
191
192
193def threadIterMain(iter):
194    x = 0
195    try:
196        for k in iter:
197            #x += len(k.name) + len(k.subkeys)
198            pass
199    except Exception as e:
200        print("%s dying young: %s" % (threading.current_thread().name, repr(e)))
201        # Exceptions are thrown on iteration because python state gets out of
202        # whack.  That's fine, because we're really just interested in finding
203        # segfaults.  People should not use iterators without locks, but it
204        # should at least not segfault on them.
205        pass
206    print("%s finished" % threading.current_thread().name)
207
208def iterMultithread(hive, fh):
209    num_threads = 10
210    iter = pyregfi.HiveIterator(hive)
211    threads = []
212    for t in range(0,num_threads):
213        threads.append(threading.Thread(target=threadIterMain, args=(iter,)))
214    for t in threads:
215        t.start()
216    for t in threads:
217        t.join()
218   
219
220tests = {
221    "iterTallyNames":iterTallyNames,
222    "iterParentWalk":iterParentWalk,
223    "iterTallyData":iterTallyData,
224    "recurseKeyTally":recurseKeyTally,
225    "iterFetchRelated":iterFetchRelated,
226    "iterIterWalk":iterIterWalk,
227    "iterCallbackIO":iterCallbackIO,
228    "iterMultithread":iterMultithread,
229    }
230
231def usage():
232    sys.stderr.write("USAGE: pyregfi-smoketest.py test1[,test2[,...]] hive1 [hive2 ...]\n")
233    sys.stderr.write("\tAvailable tests:\n")
234    for t in tests.keys():
235        sys.stderr.write("\t\t%s\n" % t)
236
237
238if len(sys.argv) < 3:
239    usage()
240    sys.exit(1)
241
242selected_tests = sys.argv[1].split(',')
243for st in selected_tests:
244    if st not in tests:
245        usage()
246        sys.stderr.write("ERROR: %s not a valid test type" % st)
247        sys.exit(1)
248
249files = []
250for f in sys.argv[2:]:
251    files.append((f, open(f,"rb")))
252
253
254start_time = time.time()
255for hname,fh in files:
256    hive = pyregfi.Hive(fh)
257    for tname in selected_tests:
258        t = tests[tname]
259        teststart = time.time()
260        tstr = "'%s' on '%s'" % (tname,hname)
261        print("##BEGIN %s:" % tstr)
262        t(hive, fh)
263        print("##END %s; runtime=%f; messages:" % (tstr, time.time() - teststart))
264        print(pyregfi.GetLogMessages())
265        print
266        sys.stdout.flush()
267    fh.close()
268
269hive = None
270files = None
271tests = None
272gc.collect()
273print("### Tests Completed, runtime: %f ###" % (time.time() -  start_time))
274#print(gc.garbage)
Note: See TracBrowser for help on using the repository browser.