forked from fuzzitdev/pythonfuzz
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathdictionary.py
More file actions
64 lines (53 loc) · 1.84 KB
/
Copy pathdictionary.py
File metadata and controls
64 lines (53 loc) · 1.84 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
"""
Basic reader for libfuzzer/AFL style dictionaries.
See documentation at:
https://llvm.org/docs/LibFuzzer.html#dictionaries
https://github.com/google/AFL/blob/master/dictionaries/README.dictionaries
For our use, we only support reading the content of the dictionary values.
"""
import codecs
import random
import re
import os
class Dictionary:
line_re = re.compile('"(.+)"$')
def __init__(self):
self._dict = list()
def load(self, dict_path):
if os.path.isfile(dict_path):
self.load_file(dict_path)
else:
self.load_directory(dict_path)
def load_directory(self, dict_path):
"""
Read a directory of files, which are loaded raw.
"""
for bin_file in os.listdir(dict_path):
filename = os.path.join(dict_path, bin_file)
if os.path.isfile(filename):
with open(filename, 'rb') as fh:
self._dict.append(fh.read())
def load_file(self, dict_path):
"""
Read a dictionary file containing tokens.
"""
# Token names are discarded, as per the AFL documentation
if not dict_path or not os.path.exists(dict_path):
return
_dict = set()
with open(dict_path) as f:
for line in f:
line = line.lstrip()
if not line or line.startswith('#'):
continue
word = self.line_re.search(line)
if word:
# Decode any escaped characters, giving us a bytes object
value = word.group(1)
(value, _) = codecs.escape_decode(value)
_dict.add(value)
self._dict = list(_dict)
def get_word(self):
if not self._dict:
return None
return random.choice(self._dict)