diff --git a/arsenal/modules/command.py b/arsenal/modules/command.py index b60aad7..52c1f43 100644 --- a/arsenal/modules/command.py +++ b/arsenal/modules/command.py @@ -7,12 +7,13 @@ class Command: cmdline = "" description = "" args = [] # [(name, value)] - nb_args = 0 + nb_place_holder = 0 nb_lines_cmd = 1 nb_lines_desc = 0 def __init__(self, cheat, gvars): self.cmdline = cheat.command + self.nb_place_holder = 0 self.cmd_tags = cheat.command_tags self.description = '' @@ -22,8 +23,7 @@ def __init__(self, cheat, gvars): self.description += '\n-----\n' self.description += cheat.description - self.get_args(cheat, gvars) - self.nb_args = len(self.args) + self.compute_args(cheat, gvars) # careful this is not the lines number in GUI self.nb_lines_cmd = len(cheat.command.split('\n')) # careful this is not the lines number in GUI @@ -39,30 +39,54 @@ def get_description_cut_by_size(self, size): result.extend(textwrap.wrap(line, size)) return result - def get_args(self, cheat, gvars): + def compute_args(self, cheat, gvars): """ Process cmdline from the cheatsheet to get args names """ - self.args = [] - # Use a list of tuples here instead of dict in case - # the cmd has multiple args with the same name.. - for arg_name in re.findall(r'<([^ <>]+)>', cheat.command): + self.args = {} + for position, arg_name in enumerate(re.findall(r"<([^ <>]+)>", cheat.command)): if "|" in arg_name: # Format name, var = arg_name.split("|")[:2] - self.args.append([name, var]) + self._add_arg(name, var, position) # Variable has been added to cheat variables before, remove it cheat.command = cheat.command.replace(arg_name, name) self.cmdline = cheat.command - elif arg_name in gvars: - self.args.append([arg_name, gvars[arg_name]]) + else: + self._add_arg(arg_name, "", position) + # compute values + for arg_name in self.args: + if arg_name in gvars: + self.args[arg_name]["value"] = gvars[arg_name] elif arg_name in cheat.variables: - self.args.append([arg_name, cheat.variables[arg_name]]) + self.args[arg_name]["value"] = cheat.variables[arg_name] else: - self.args.append([arg_name, ""]) + continue + + def _add_arg(self, name=None, value="", position=0): + if name in self.args: + self.args[name]["value"] = value + self.args[name]["positions"].append(position) + else: + v = {} + v["value"] = value + v["positions"] = [position] + self.args[name] = v + self.nb_place_holder += 1 + + def get_arg(self, position): + for k, v in self.args.items(): + if position in v["positions"]: + return k, v["value"] + return f"{position}", f"|{self.nb_place_holder}|" + + def set_arg_value(self, position, value): + for k, v in self.args.items(): + if position in v["positions"]: + self.args[k]["value"] = value def get_command_parts(self): - if self.nb_args != 0: - regex = ''.join('<' + arg[0] + '>|' for arg in self.args)[:-1] + if self.nb_place_holder != 0: + regex = "|".join("<" + arg + ">" for arg in self.args) cmdparts = re.split(regex, self.cmdline) else: cmdparts = [self.cmdline] @@ -74,20 +98,21 @@ def build(self): -> if some args values are still empty do nothing -> else build the final command string by adding args values """ - if self.nb_args == 0 : + if self.nb_place_holder == 0 : return True - argsval = [a[1] for a in self.args] + argsval = [a["value"] for a in self.args.values()] if "" not in argsval: # split cmdline at each arg position - regex = ''.join('<' + arg[0] + '>|' for arg in self.args)[:-1] + regex = "|".join("<" + arg + ">" for arg in self.args) cmdparts = re.split(regex, self.cmdline) # concat command parts and arguments values to build the command self.cmdline = "" - for i in range(len(cmdparts) + len(self.args)): + for i in range(len(cmdparts) + self.nb_place_holder): if i % 2 == 0: self.cmdline += cmdparts[i // 2] else: - self.cmdline += argsval[(i - 1) // 2] + _, value = self.get_arg((i - 1) // 2) + self.cmdline += value curses.endwin() # build ok ? diff --git a/arsenal/modules/config.py b/arsenal/modules/config.py index 94dfcf8..b40f90e 100644 --- a/arsenal/modules/config.py +++ b/arsenal/modules/config.py @@ -24,4 +24,4 @@ savevarfile = join(HOMEPATH, ".arsenal.json") -PREFIX_GLOBALVAR_NAME = "arsenal_prefix_cmd" \ No newline at end of file +PREFIX_GLOBALVAR_NAME = "arsenal_prefix_cmd" diff --git a/arsenal/modules/gui.py b/arsenal/modules/gui.py index e4fe63b..4a9e265 100644 --- a/arsenal/modules/gui.py +++ b/arsenal/modules/gui.py @@ -384,6 +384,8 @@ def run(self, stdscr): class ArgslistMenu: current_arg = 0 + current_arg_name = "" + current_arg_value = "" max_preview_size = 0 prev_lastline_len = 0 @@ -397,6 +399,7 @@ class ArgslistMenu: def __init__(self, prev): self.previous_menu = prev + self.current_arg_name, self.current_arg_value = Gui.cmd.get_arg(self.current_arg) def get_nb_preview_new_lines(self): """ @@ -422,18 +425,16 @@ def get_nb_preview_new_lines(self): firstline = False # extract len of args in the current line - i = 0 - for arg_name, arg_val in Gui.cmd.args: + for i, (arg_name, arg_val) in enumerate(Gui.cmd.args.items()): if i == next_arg and nb_args_todo > 0: - if arg_val != "": + if arg_val["value"] != "": # use value len if not empty - nbchar += len(arg_val) + nbchar += len(arg_val["value"]) else: # else use name len + 2 for '<' and '>' nbchar += (len(arg_name) + 2) next_arg += 1 nb_args_todo -= 1 - i += 1 # len of the cmd body for p in parts: @@ -452,10 +453,11 @@ def next_arg(self): self.x_init = None self.y_init = None # change selected arg - if self.current_arg < Gui.cmd.nb_args - 1: + if self.current_arg < len(Gui.cmd.args) - 1: self.current_arg += 1 else: self.current_arg = 0 + self.current_arg_name, self.current_arg_value = Gui.cmd.get_arg(self.current_arg) def previous_arg(self): """ @@ -469,7 +471,8 @@ def previous_arg(self): if self.current_arg > 0: self.current_arg -= 1 else: - self.current_arg = Gui.cmd.nb_args - 1 + self.current_arg = len(Gui.cmd.args) - 1 + self.current_arg_name, self.current_arg_value = Gui.cmd.get_arg(self.current_arg) def draw_preview_part(self, win, text, color): """ @@ -499,12 +502,12 @@ def draw_selected_arg(self, y_pos): """ y, x = self.AB_TOP + y_pos + self.current_arg, self.AB_SIDE + 1 ncols, nlines = self.width - 2 * (self.AB_SIDE + 1), 1 - arg = Gui.cmd.args[self.current_arg] - max_size = self.max_preview_size - 4 - len(arg[0]) + arg_name, arg_value = Gui.cmd.get_arg(self.current_arg) + max_size = self.max_preview_size - 4 - len(arg_name) selectedargline = curses.newwin(nlines, ncols, y, x) selectedargline.addstr(" > ", curses.color_pair(Gui.BASIC_COLOR)) - selectedargline.addstr(arg[0], curses.color_pair(Gui.ARG_NAME_COLOR)) - selectedargline.addstr(" = " + Gui.draw_string(arg[1], max_size), curses.color_pair(Gui.BASIC_COLOR)) + selectedargline.addstr(arg_name, curses.color_pair(Gui.ARG_NAME_COLOR)) + selectedargline.addstr(" = " + Gui.draw_string(arg_value, max_size), curses.color_pair(Gui.BASIC_COLOR)) selectedargline.refresh() def draw_args_list(self, y_pos): @@ -512,11 +515,11 @@ def draw_args_list(self, y_pos): Draw the asked arguments list in the argument menu """ y, x = self.AB_TOP + y_pos, self.AB_SIDE + 1 - ncols, nlines = self.width - 2 * (self.AB_SIDE + 1), Gui.cmd.nb_args + 1 + ncols, nlines = self.width - 2 * (self.AB_SIDE + 1), len(Gui.cmd.args) + 1 argwin = curses.newwin(nlines, ncols, y, x) - for arg in Gui.cmd.args: + for arg_name, arg_data in Gui.cmd.args.items(): max_size = self.max_preview_size + 4 - argline = Gui.draw_string(" {} = {}".format(*arg), max_size) + "\n" + argline = Gui.draw_string(f" {arg_name} = {arg_data["value"]}", max_size) + "\n" argwin.addstr(argline, curses.color_pair(Gui.BASIC_COLOR)) argwin.refresh() @@ -544,24 +547,18 @@ def draw_cmd_preview(self, argprev, p_x, p_y=1): # draw command argprev.addstr(p_y, p_x, "$ ", curses.color_pair(Gui.BASIC_COLOR)) - # draw preview cmdline - for i in range(len(cmdparts) + Gui.cmd.nb_args): + for i in range(len(cmdparts) + Gui.cmd.nb_place_holder): if i % 2 == 0: # draw cmd parts in white self.draw_preview_part(argprev, cmdparts[i // 2], curses.color_pair(Gui.BASIC_COLOR)) else: - # get argument value - if Gui.cmd.args[(i - 1) // 2][1] == "": - # if arg empty use its name - arg = '<' + Gui.cmd.args[(i - 1) // 2][0] + '>' - else: - # else its value - arg = Gui.cmd.args[(i - 1) // 2][1] + arg_name, arg_value = Gui.cmd.get_arg((i - 1) // 2) + arg = "<" + arg_name + ">" if arg_value == "" else arg_value # draw argument - if (i - 1) // 2 == self.current_arg: - # if arg is selected print in blue + if self.current_arg_name == arg_name: + # if arg is selected print in blue COL1_COLOR self.draw_preview_part(argprev, arg, curses.color_pair(Gui.ARG_NAME_COLOR)) else: # else in white @@ -586,10 +583,6 @@ def draw(self, stdscr): # draw argslist menu popup self.prev_lastline_len = 0 nbpreviewnewlines = self.get_nb_preview_new_lines() - # if Gui.cmd.nb_args != 0: - # nbpreviewnewlines = self.get_nb_preview_new_lines() - # else: - # nbpreviewnewlines = 0 # -------------- border # cmd @@ -609,7 +602,7 @@ def draw(self, stdscr): border_height = 1 cmd_height = 1 + nbpreviewnewlines - args_height = (2 + Gui.cmd.nb_args) if (Gui.cmd.nb_args > 0) else 0 + args_height = (2 + len(Gui.cmd.args)) if (len(Gui.cmd.args) > 0) else 0 desc_height = (len(description_lines) + 1 + 1) if (len(description_lines) > 0) else 0 cmd_pos = 1 @@ -633,14 +626,15 @@ def draw(self, stdscr): self.draw_desc_preview(argprev, padding_text_border, desc_pos, description_lines) if len(Gui.cmd.args) > 0: + arg_name, arg_value = Gui.cmd.get_arg(self.current_arg) self.draw_args_list(args_pos) self.draw_selected_arg(args_pos) # init cursor position (if first draw) if self.x_init is None or self.y_init is None or self.xcursor is None: self.y_init, self.x_init = curses.getsyx() # prefill compatibility - self.x_init -= len(Gui.cmd.args[self.current_arg][1]) - self.xcursor = self.x_init + len(Gui.cmd.args[self.current_arg][1]) + self.x_init -= len(arg_value) + self.xcursor = self.x_init + len(arg_value) # set cursor position curses.setsyx(self.y_init, self.xcursor) curses.doupdate() @@ -649,18 +643,17 @@ def draw(self, stdscr): pass def check_move_cursor(self, n): - if Gui.cmd.nb_args == 0: + if len(Gui.cmd.args) == 0: return False - return self.x_init <= (self.xcursor + n) < self.x_init + len(Gui.cmd.args[self.current_arg][1]) + 1 + return self.x_init <= (self.xcursor + n) < self.x_init + len(self.current_arg_value) + 1 def autocomplete_arg(self): """ Autocomplete the current argument """ # current argument value - argument = Gui.cmd.args[self.current_arg][1] # look for all files that match the argument in the working directory - matches = glob.glob('{}*'.format(argument)) + matches = glob.glob(f"{self.current_arg_value}*") if not matches: return False @@ -678,7 +671,7 @@ def autocomplete_arg(self): autocompleted_argument = autocompleted_argument + sep # autocomplete the argument - Gui.cmd.args[self.current_arg][1] = autocompleted_argument + Gui.cmd.set_arg_value(self.current_arg, autocompleted_argument) # update cursor position self.xcursor = self.x_init + len(autocompleted_argument) @@ -712,7 +705,7 @@ def run(self, stdscr): elif c == 9: if Gui.cmd.args: # autocomplete the current argument - if Gui.cmd.args[self.current_arg][1]: + if self.current_arg_value: self.autocomplete_arg() # go to the next argument else: @@ -725,7 +718,8 @@ def run(self, stdscr): files += glob.glob(fuzz_dir, recursive=True) fzf = FzfPrompt().prompt(files) # autocomplete the argument - Gui.cmd.args[self.current_arg][1] = fzf[0] + self.current_arg_value = fzf[0] + Gui.cmd.set_arg_value(self.current_arg, self.current_arg_value) # update cursor position self.xcursor = self.x_init + len(fzf[0]) except ImportError: @@ -733,15 +727,17 @@ def run(self, stdscr): elif c == curses.KEY_BACKSPACE or c == 127 or c == 8: if self.check_move_cursor(-1): i = self.xcursor - self.x_init - 1 - Gui.cmd.args[self.current_arg][1] = Gui.cmd.args[self.current_arg][1][:i] + \ - Gui.cmd.args[self.current_arg][1][i + 1:] + self.current_arg_value = self.current_arg_value[:i] + \ + self.current_arg_value[i + 1:] + Gui.cmd.set_arg_value(self.current_arg, self.current_arg_value) self.xcursor -= 1 elif c == curses.KEY_DC or c == 127: # DELETE key if self.check_move_cursor(1): i = self.xcursor - self.x_init - 1 - Gui.cmd.args[self.current_arg][1] = Gui.cmd.args[self.current_arg][1][:i + 1] + \ - Gui.cmd.args[self.current_arg][1][i + 2:] + self.current_arg_value = self.current_arg_value[:i + 1] + \ + self.current_arg_value[i + 2:] + Gui.cmd.set_arg_value(self.current_arg, self.current_arg_value) elif c == curses.KEY_LEFT: # Move cursor LEFT if self.check_move_cursor(-1): self.xcursor -= 1 @@ -753,11 +749,12 @@ def run(self, stdscr): self.xcursor = self.x_init elif c == curses.KEY_END: # Move cursor to the END - self.xcursor = self.x_init + len(Gui.cmd.args[self.current_arg][1]) - elif 20 <= c < 127 and Gui.cmd.nb_args > 0: + self.xcursor = self.x_init + len(self.current_arg_value) + elif 20 <= c < 127 and len(Gui.cmd.args) > 0: i = self.xcursor - self.x_init - Gui.cmd.args[self.current_arg][1] = Gui.cmd.args[self.current_arg][1][:i] + chr(c) + \ - Gui.cmd.args[self.current_arg][1][i:] + self.current_arg_value = self.current_arg_value[:i] + chr(c) + \ + self.current_arg_value[i:] + Gui.cmd.set_arg_value(self.current_arg, self.current_arg_value) self.xcursor += 1