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