Merge pull request #73 from iamOgunyinka/clean

_cleantarget_ feature addition
This commit is contained in:
Michael Skelton
2019-10-09 10:11:03 +10:00
committed by GitHub
2 changed files with 36 additions and 4 deletions

View File

@@ -11,6 +11,7 @@ def build_queue(arguments, output):
task_list = InputHelper.process_commands(arguments)
for task in task_list:
output.terminal(Level.THREAD, task.name(), "Added to Queue")
print('Generated {} commands in total'.format(len(task_list)))
return task_list

View File

@@ -162,21 +162,51 @@ class InputHelper(object):
else:
destination_set.add(ips)
@staticmethod
def _process_clean_targets(commands, dirty_targets):
def add_task(t, item_list, my_command_set):
if t not in my_command_set:
my_command_set.add(t)
item_list.append(t)
variable = '_cleantarget_'
tasks = []
temp = set() # this helps avoid command duplication and re/deconstructing of temporary set
for command in commands:
for dirty_target in dirty_targets:
if command.name().find(variable) != -1:
# replace all https:// or https:// with nothing
dirty_target = dirty_target.replace('http://', '')
dirty_target = dirty_target.replace('https://', '')
# chop off all trailing '/', if any.
while dirty_target.endswith('/'):
dirty_target = dirty_target.strip('/')
# replace all remaining '/' with '-' and that's enough cleanup for the day
clean_target = dirty_target.replace('/', '-')
new_task = command.clone()
new_task.replace(variable, clean_target)
add_task(new_task, tasks, temp)
else:
add_task(command, tasks, temp)
return tasks
@staticmethod
def _replace_variable_with_commands(commands, variable, replacements):
def add_task(t, item_list):
if t not in set(item_list):
def add_task(t, item_list, my_set):
if t not in my_set:
my_set.add(t)
item_list.append(t)
tasks = []
temp_set = set() # to avoid duplicates
for command in commands:
for replacement in replacements:
if command.name().find(variable) != -1:
new_task = command.clone()
new_task.replace(variable, replacement)
add_task(new_task, tasks)
add_task(new_task, tasks, temp_set)
else:
add_task(command, tasks)
add_task(command, tasks, temp_set)
return tasks
@staticmethod
@@ -245,6 +275,7 @@ class InputHelper(object):
commands = InputHelper._replace_variable_with_commands(commands, "_target_", targets)
commands = InputHelper._replace_variable_with_commands(commands, "_host_", targets)
commands = InputHelper._process_clean_targets(commands, targets)
if arguments.port:
commands = InputHelper._replace_variable_with_commands(commands, "_port_", ports)