Fixed bug when using target_list

This commit is contained in:
Joshua Ogunyinka
2019-08-16 04:06:28 +01:00
parent ab34aca3b2
commit d416e5a689
3 changed files with 15 additions and 12 deletions

View File

@@ -31,7 +31,7 @@ def main():
output.print_banner() output.print_banner()
pool = Pool(arguments.threads, build_queue(arguments, output), arguments.timeout, output, arguments.sober) pool = Pool(arguments.threads, build_queue(arguments, output), arguments.timeout, output, arguments.sober)
# pool.run() pool.run()
if __name__ == "__main__": if __name__ == "__main__":

View File

@@ -123,20 +123,20 @@ class InputHelper(object):
def _replace_variable_with_commands(commands, variable, replacements): def _replace_variable_with_commands(commands, variable, replacements):
data = set() data = set()
for command in commands: for command in commands:
is_task_block = isinstance(command, TaskBlock) is_task = not isinstance(command, TaskBlock)
for replacement in replacements: for replacement in replacements:
if not is_task_block and command.name().find(variable) != -1: if is_task and command.name().find(variable) != -1:
new_task = command.clone() new_task = command.clone()
new_task.replace(variable, replacement) new_task.replace(variable, replacement)
data.add(new_task) data.add(new_task)
elif not is_task_block and command not in data: elif is_task and command not in data:
data.add(command)
elif not is_task:
tasks = [task for task in command.get_tasks()]
command.clear_tasks()
for r in InputHelper._replace_variable_with_commands(tasks, variable, replacements):
command.add_task(r)
data.add(command) data.add(command)
elif is_task_block:
new_task_block = TaskBlock(command.name())
result = InputHelper._replace_variable_with_commands(command.get_tasks(), variable, replacements)
for r in result:
new_task_block.add_task(r)
data.add(new_task_block)
return set(data) return set(data)
@staticmethod @staticmethod
@@ -175,8 +175,8 @@ class InputHelper(object):
ranges.add(arguments.target) ranges.add(arguments.target)
else: else:
target_file = arguments.target_list target_file = arguments.target_list
if not sys.stdin.isatty(): # if not sys.stdin.isatty():
target_file = sys.stdin # target_file = sys.stdin
ranges.update([target.strip() for target in target_file if target.strip()]) ranges.update([target.strip() for target in target_file if target.strip()])
# process exclusions first # process exclusions first

View File

@@ -70,6 +70,9 @@ class TaskBlock(Task):
def add_task(self, task): def add_task(self, task):
self.tasks.append(task) self.tasks.append(task)
def clear_tasks(self):
self.tasks.clear()
def __len__(self): def __len__(self):
return len(self.tasks) return len(self.tasks)