From 4be54584cac276ec1ab5718c21bc2087a5b075c5 Mon Sep 17 00:00:00 2001 From: "Remi GASCOU (Podalirius)" <79218792+p0dalirius@users.noreply.github.com> Date: Tue, 25 Jun 2024 13:16:35 +0200 Subject: [PATCH] Refactored parts of CommandCompleter --- smbclientng/__main__.py | 1 + smbclientng/core/CommandCompleter.py | 75 +++++++++++----------------- 2 files changed, 30 insertions(+), 46 deletions(-) diff --git a/smbclientng/__main__.py b/smbclientng/__main__.py index 67767f2..459372c 100644 --- a/smbclientng/__main__.py +++ b/smbclientng/__main__.py @@ -138,5 +138,6 @@ def main(): if options.debug: print("[debug] Exiting the console.") + if __name__ == "__main__": main() diff --git a/smbclientng/core/CommandCompleter.py b/smbclientng/core/CommandCompleter.py index d70a7f0..64313f2 100644 --- a/smbclientng/core/CommandCompleter.py +++ b/smbclientng/core/CommandCompleter.py @@ -389,9 +389,9 @@ def complete(self, text, state): if '\\' in remainder.strip() or '/' in remainder.strip(): path = remainder.strip().replace(ntpath.sep, '/') path = '/'.join(path.split('/')[:-1]) - + # Get remote directory contents directory_contents = self.smbSession.list_contents(path=path).items() - + # matching_entries = [] for _, entry in directory_contents: if entry.is_directory() and entry.get_longname() not in [".",".."]: @@ -399,15 +399,10 @@ def complete(self, text, state): matching_entries.append(path + '/' + entry.get_longname() + '/') else: matching_entries.append(entry.get_longname() + '/') - - # Add quoting for shlex - matching_entries = [shlex.quote(s) for s in matching_entries] - - self.matches += [ - command + " " + s - for s in matching_entries - if s.lower().startswith(remainder.lower()) - ] + # + for s in matching_entries: + if s.lower().startswith(remainder.lower()) or shlex.quote(s).lower().startswith(remainder.lower()): + self.matches.append(command + " " + shlex.quote(s)) # Autocomplete file if "remote_file" in self.commands[command]["autocomplete"]: @@ -415,26 +410,21 @@ def complete(self, text, state): path = "" if '\\' in remainder.strip() or '/' in remainder.strip(): path = remainder.strip().replace(ntpath.sep, '/') - path = '/'.join(path.split('/')[:-1]) - + path = '/'.join(path.split('/')[:-1]) + # Get remote directory contents directory_contents = self.smbSession.list_contents(path=path).items() - + # matching_entries = [] for _, entry in directory_contents: - if not entry.is_directory() and entry.get_longname() not in [".",".."]: + if (not entry.is_directory()) and entry.get_longname() not in [".",".."]: if len(path) != 0: matching_entries.append(path + '/' + entry.get_longname()) else: matching_entries.append(entry.get_longname()) - - # Add quoting for shlex - matching_entries = [shlex.quote(s) for s in matching_entries] - - self.matches += [ - command + " " + s - for s in matching_entries - if s.lower().startswith(remainder.lower()) - ] + # + for s in matching_entries: + if s.lower().startswith(remainder.lower()) or shlex.quote(s).lower().startswith(remainder.lower()): + self.matches.append(command + " " + shlex.quote(s)) # Autocomplete local_directory if "local_directory" in self.commands[command]["autocomplete"]: @@ -443,11 +433,10 @@ def complete(self, text, state): if os.path.sep in remainder.strip(): path = path.split(os.path.sep)[:-1] path = os.path.sep.join(path) - # Current dir if len(path.strip()) == 0: path = "." - + # directory_contents = os.listdir(path=path + os.path.sep) matching_entries = [] for entry in directory_contents: @@ -455,12 +444,10 @@ def complete(self, text, state): entry_path = path + os.path.sep + entry if os.path.isdir(entry_path): matching_entries.append(entry_path + os.path.sep) - - self.matches += [ - command + " " + s - for s in matching_entries - if s.startswith(remainder) - ] + # + for s in matching_entries: + if s.lower().startswith(remainder.lower()) or shlex.quote(s).lower().startswith(remainder.lower()): + self.matches.append(command + " " + shlex.quote(s)) # Autocomplete local_file if "local_file" in self.commands[command]["autocomplete"]: @@ -469,32 +456,28 @@ def complete(self, text, state): if os.path.sep in remainder.strip(): path = path.split(os.path.sep)[:-1] path = os.path.sep.join(path) - # Current dir if len(path.strip()) == 0: path = "." - - directory_contents = os.listdir(path=path + os.path.sep) + # + directory_contents = os.listdir(path=(path + os.path.sep)) matching_entries = [] for entry in directory_contents: if entry not in [".",".."]: entry_path = path + os.path.sep + entry if not os.path.isdir(entry_path): matching_entries.append(entry_path) - - self.matches += [ - command + " " + s - for s in matching_entries - if s.startswith(remainder) - ] + # + for s in matching_entries: + if s.lower().startswith(remainder.lower()) or shlex.quote(s).lower().startswith(remainder.lower()): + self.matches.append(command + " " + shlex.quote(s)) else: # Generic case for subcommands - self.matches += [ - command + " " + s - for s in self.commands[command]["subcommands"] - if s.startswith(remainder) - ] + for s in self.commands[command]["subcommands"]: + if s.startswith(remainder): + self.matches.append(command + " " + s) + else: # Unknown subcommand, skipping autocomplete pass