def test_part(self):
file_put_contents('/tmp/test.txt', 'aaa\nbbb\nccc\nddd')
gen = FileGenerator('/tmp/test.txt', parts=2, part=1)
assert gen.get() == "aaa"
assert gen.get() == "bbb"
assert gen.get() is None
def __init__(self, mask, dict_file, parts, part, template):
if not template.count("%m%"):
raise WSException(
"Template '{0}' not contains %m% (mask) marker ".format(
template))
if not template.count("%d%"):
raise WSException(
"Template '{0}' not contains %d% (dict) marker ".format(
template))
if not os.path.exists(dict_file):
raise WSException("Dict '{0}' not exists ".format(dict_file))
self.template = template
self.dict_generator = FileGenerator(dict_file)
self.next_dict_line()
self.mask = mask
self.mask_generator = DictOfMask(mask)
self.lines_count = (self.dict_generator.lines_count * self.template.count("%d%")) * \
(self.mask_generator.all_objects_count * self.template.count("%m%"))
if parts and part:
one_part_count = int(self.lines_count / parts)
self.first_border = one_part_count * (part - 1)
self.second_border = one_part_count * part
while self.current_counter < self.first_border:
self.get()
def test_no_parts(self):
file_put_contents('/tmp/test.txt', 'aaa\nbbb\nccc')
gen = FileGenerator('/tmp/test.txt', parts=0, part=0)
assert gen.get() == "aaa"
assert gen.get() == "bbb"
assert gen.get() == "ccc"
assert gen.get() is None
def load_objects(self, queue):
""" Make generator with objects to check """
generator = FileGenerator(self.options['dict'].value,
int(self.options['parts'].value),
int(self.options['part'].value))
queue.set_generator(generator)
return {
'all': generator.lines_count,
'start': generator.first_border,
'end': generator.second_border
}
def make_queue(self):
"""
Make work queue
:return:
"""
self.queue = FuzzerHeadersJob()
generator = FileGenerator(self.options['urls-file'].value)
self.queue.set_generator(generator)
self.logger.log("Loaded {0} variants.".format(generator.lines_count))
self.counter = WSCounter.factory(generator.lines_count)
def make_queue(self):
"""
Make work queue
:return:
"""
self.queue = FuzzerUrlsJob()
generator = FileGenerator(self.source_temp_file)
self.queue.set_generator(generator)
self.logger.log("Loaded {0} variants.".format(generator.lines_count))
self.counter = WSCounter.factory(generator.lines_count)
class CombineGenerator(object):
""" Class for generate and work with masks """
mask = None
template = None
mask_generator = None
dict_generator = None
first_border = 0
second_border = 0
lines_count = 0
current_counter = 0
current_dict_line = None
def __init__(self, mask, dict_file, parts, part, template):
if not template.count("%m%"):
raise WSException(
"Template '{0}' not contains %m% (mask) marker ".format(
template))
if not template.count("%d%"):
raise WSException(
"Template '{0}' not contains %d% (dict) marker ".format(
template))
if not os.path.exists(dict_file):
raise WSException("Dict '{0}' not exists ".format(dict_file))
self.template = template
self.dict_generator = FileGenerator(dict_file)
self.next_dict_line()
self.mask = mask
self.mask_generator = DictOfMask(mask)
self.lines_count = (self.dict_generator.lines_count * self.template.count("%d%")) * \
(self.mask_generator.all_objects_count * self.template.count("%m%"))
if parts and part:
one_part_count = int(self.lines_count / parts)
self.first_border = one_part_count * (part - 1)
self.second_border = one_part_count * part
while self.current_counter < self.first_border:
self.get()
def _get(self):
""" Get next combine item """
mask = self.get_next_mask()
if mask is None:
self.next_dict_line()
if self.current_dict_line is None:
return None
self.mask_generator = DictOfMask(self.mask)
mask = self.get_next_mask()
return self.template.replace("%m%",
mask).replace("%d%",
self.current_dict_line)
def get(self):
""" Get next combine item and up counters """
if self.second_border and self.current_counter >= self.second_border:
return None
to_return = self._get()
self.current_counter += 1
return to_return
def get_next_mask(self):
""" Get next mask string from mask generator """
return self.mask_generator.get()
def next_dict_line(self):
""" Set next string from dict generator """
self.current_dict_line = self.dict_generator.get()
def get_dict_current_line(self):
""" Get next line from dict and get next if need """
if not len(self.current_dict_line):
self.next_dict_line()
return self.current_dict_line