Feature/arg uniqueness

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|default_value>
                 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 ?

arsenal/modules/config.py

@@ -24,4 +24,4 @@
 
 savevarfile = join(HOMEPATH, ".arsenal.json")
 
-PREFIX_GLOBALVAR_NAME = "arsenal_prefix_cmd"
+PREFIX_GLOBALVAR_NAME = "arsenal_prefix_cmd"

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,24 +502,24 @@ 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):
         """
         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,23 +718,26 @@ 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:
                     pass
             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