smbclientng.core.CommandCompleter

  1#!/usr/bin/env python3
  2# -*- coding: utf-8 -*-
  3# File name          : CommandCompleter.py
  4# Author             : Podalirius (@podalirius_)
  5# Date created       : 20 may 2024
  6
  7
  8import ntpath
  9import os
 10
 11
 12class CommandCompleter(object):
 13    """
 14    A class to handle command completion for the smbclient-ng shell.
 15
 16    This class provides a command completion feature that suggests possible command names based on the current input.
 17    It uses a dictionary to store commands and their descriptions, which helps in providing hints during the command line
 18    interaction in the smbclient-ng shell.
 19
 20    Attributes:
 21        smbSession (SMBSession): An instance of SMBSession which maintains the current SMB session.
 22        commands (dict): A dictionary containing command names as keys and their descriptions and subcommands as values.
 23
 24    Methods:
 25        __init__(self, smbSession): Initializes the CommandCompleter with an SMBSession.
 26    """
 27
 28    commands = {
 29        "bat": {
 30            "description": [
 31                "Pretty prints the contents of a file.", 
 32                "Syntax: 'bat <file>'"
 33            ], 
 34            "subcommands": []
 35        },
 36        "cat": {
 37            "description": [
 38                "Get the contents of a file.", 
 39                "Syntax: 'cat <file>'"
 40            ], 
 41            "subcommands": []
 42        },
 43        "cd": {
 44            "description": [
 45                "Change the current working directory.", 
 46                "Syntax: 'cd <directory>'"
 47            ], 
 48            "subcommands": []
 49        },
 50        "close": {
 51            "description": [
 52                "Closes the SMB connection to the remote machine.", 
 53                "Syntax: 'close'"
 54            ], 
 55            "subcommands": []
 56        },
 57        "connect": {
 58            "description": [
 59                "Connect to the remote machine (useful if connection timed out).", 
 60                "Syntax: 'connect'"
 61            ], 
 62            "subcommands": []
 63        },
 64        "dir": {
 65            "description": [
 66                "List the contents of the current working directory.",
 67                "Syntax: 'dir'"
 68            ], 
 69            "subcommands": []
 70        },
 71        "exit": {
 72            "description": [
 73                "Exits the smbclient-ng script.",
 74                "Syntax: 'exit'"
 75            ], 
 76            "subcommands": []
 77        },
 78        "get": {
 79            "description": [
 80                "Get a remote file.",
 81                "Syntax: 'get [-r] <directory or file>'"
 82            ], 
 83            "subcommands": []
 84        },
 85        "help": {
 86            "description": [
 87                "Displays this help message.",
 88                "Syntax: 'help'"
 89            ], 
 90            "subcommands": ["format"]
 91        },
 92        "info": {
 93            "description": [
 94                "Get information about the server and or the share.",
 95                "Syntax: 'info [server|share]'"
 96            ], 
 97            "subcommands": ["server", "share"]
 98        },
 99        "lcd": {
100            "description": [
101                "Changes the current local directory.",
102                "Syntax: 'lcd <directory>'"
103            ], 
104            "subcommands": []
105        },
106        "lls": {
107            "description": [
108                "Lists the contents of the current local directory.", 
109                "Syntax: 'lls'"
110            ],
111            "subcommands": []
112        },
113        "lmkdir": {
114            "description": [
115                "Creates a new local directory.", 
116                "Syntax: 'lmkdir <directory>'"
117            ],
118            "subcommands": []
119        },
120        "lpwd": {
121            "description": [
122                "Shows the current local directory.", 
123                "Syntax: 'lpwd'"
124            ],
125            "subcommands": []
126        },
127        "lrm": {
128            "description": [
129                "Removes a local file.", 
130                "Syntax: 'lrm <file>'"
131            ], 
132            "subcommands": []
133        },
134        "lrmdir": {
135            "description": [
136                "Removes a local directory.", 
137                "Syntax: 'lrmdir <directory>'"
138            ], 
139            "subcommands": []
140        },
141        "ls": {
142            "description": [
143                "List the contents of the current remote working directory.", 
144                "Syntax: 'ls'"
145            ], 
146            "subcommands": []
147        },
148        "ltree": {
149            "description": [
150                "Displays a tree view of the local directories.",
151                "Syntax: 'ltree [directory]'"
152            ], 
153            "subcommands": []
154        },
155        "mkdir": {
156            "description": [
157                "Creates a new remote directory.", 
158                "Syntax: 'mkdir <directory>'"
159            ], 
160            "subcommands": []
161        },
162        "module": {
163            "description": [
164                "Loads a specific module for additional functionalities.",
165                "Syntax: 'module <name>'"
166            ], 
167            "subcommands": []
168        },
169        "put": {
170            "description": [
171                "Put a local file or directory in a remote directory.", 
172                "Syntax: 'put [-r] <directory or file>'"
173            ], 
174            "subcommands": []
175        },
176        "reconnect": {
177            "description": [
178                "Reconnect to the remote machine (useful if connection timed out).", 
179                "Syntax: 'reconnect'"
180            ], 
181            "subcommands": []
182        },
183        "reset": {
184            "description": [
185                "Reset the TTY output, useful if it was broken after printing a binary file on stdout.",
186                "Syntax: 'reset'"
187            ], 
188            "subcommands": []
189        },
190        "rmdir": {
191            "description": [
192                "Removes a remote directory.", 
193                "Syntax: 'rmdir <directory>'"
194            ], 
195            "subcommands": []
196        },
197        "rm": {
198            "description": [
199                "Removes a remote file.", 
200                "Syntax: 'rm <file>'"
201            ], 
202            "subcommands": []
203        },
204        "shares": {
205            "description": [
206                "Lists the SMB shares served by the remote machine.", 
207                "Syntax: 'shares'"
208            ], 
209            "subcommands": []
210        },
211        "use": {
212            "description": [
213                "Use a SMB share.", 
214                "Syntax: 'use <sharename>'"
215            ], 
216            "subcommands": []
217        },
218        "tree": {
219            "description": [
220                "Displays a tree view of the remote directories.",
221                "Syntax: 'tree [directory]'"
222            ], 
223            "subcommands": []
224        },
225    }
226
227    def __init__(self, smbSession, config):
228        # Objects
229        self.smbSession = smbSession
230        self.config = config
231        # Pre computing for some commands 
232        self.commands["help"]["subcommands"] = ["format"] + list(self.commands.keys())
233        self.commands["help"]["subcommands"].remove("help")
234
235    def complete(self, text, state):
236        """
237        Function to handle command completion in the LDAP console.
238
239        This function completes the user"s input based on the available options for commands in the LDAP console.
240
241        Args:
242            text (str): The current text input by the user.
243            state (int): The current state of completion.
244
245        Returns:
246            str: The next completion suggestion based on the user"s input state.
247        """
248
249        if state == 0:
250            
251            # No text typed yet, need the list of commands available
252            if len(text) == 0:
253                self.matches = [s for s in self.commands.keys()]
254
255            elif len(text) != 0:
256                # This is for the main command
257                if text.count(" ") == 0:
258                    self.matches = [s for s in self.commands.keys() if s and s.startswith(text)]
259                
260                # This is for subcommands
261                elif text.count(" ") >= 1:
262                    command, remainder = text.split(" ", 1)
263                    if command in self.commands.keys():
264                        if command == "use":
265                            # Choose SMB Share to connect to
266                            self.matches = [
267                                command + " " + s.lower()
268                                for s in self.smbSession.list_shares().keys()
269                                if s.lower().startswith(remainder.lower())
270                            ]
271
272                        elif command in ["cd", "dir", "ls", "mkdir", "rmdir", "tree"]:
273                            # Choose remote directory
274                            path = ""
275                            if '\\' in remainder.strip() or '/' in remainder.strip():
276                                path = remainder.strip().replace('/', ntpath.sep)
277                                path = ntpath.sep.join(path.split(ntpath.sep)[:-1]) 
278
279                            directory_contents = self.smbSession.list_contents(path=path).items()
280
281                            matching_entries = []
282                            for _, entry in directory_contents:
283                                if entry.is_directory() and entry.get_longname() not in [".",".."]:
284                                    if len(path) != 0:
285                                        matching_entries.append(path + ntpath.sep + entry.get_longname() + ntpath.sep)
286                                    else:
287                                        matching_entries.append(entry.get_longname() + ntpath.sep)
288
289                            self.matches = [
290                                command + " " + s 
291                                for s in matching_entries
292                                if s.lower().startswith(remainder.lower())
293                            ]
294
295                        elif command in ["bat", "cat", "get", "rm"]:
296                            # Choose local files and directories
297                            path = ""
298                            if '\\' in remainder.strip() or '/' in remainder.strip():
299                                path = remainder.strip().replace('/', ntpath.sep)
300                                path = ntpath.sep.join(path.split(ntpath.sep)[:-1]) 
301
302                            directory_contents = self.smbSession.list_contents(path=path).items()
303
304                            matching_entries = []
305                            for _, entry in directory_contents:
306                                if entry.get_longname() not in [".",".."]:
307                                    if len(path) != 0:
308                                        if entry.is_directory():
309                                            matching_entries.append(path + ntpath.sep + entry.get_longname() + ntpath.sep)
310                                        else:
311                                            matching_entries.append(path + ntpath.sep + entry.get_longname())
312                                    else:
313                                        if entry.is_directory():
314                                            matching_entries.append(entry.get_longname() + ntpath.sep)
315                                        else:
316                                            matching_entries.append(entry.get_longname())
317
318                            self.matches = [
319                                command + " " + s 
320                                for s in matching_entries
321                                if s.lower().startswith(remainder.lower())
322                            ]
323
324                        elif command in ["lcd", "lls", "put", "lmkdir", "lrm", "lrmdir"]:
325                            # Choose directory
326                            path = ""
327                            if os.path.sep in remainder.strip():
328                                path = path.split(os.path.sep)[:-1]
329                                path = os.path.sep.join(path)
330                            
331                            # Current dir
332                            if len(path.strip()) == 0:
333                                path = "."
334
335                            directory_contents = os.listdir(path=path + os.path.sep)
336                            matching_entries = []
337                            for entry in directory_contents:
338                                if entry not in [".",".."]:
339                                    entry_path = path + os.path.sep + entry
340                                    if os.path.isdir(entry_path):
341                                        matching_entries.append(entry_path + os.path.sep)
342                                    else:
343                                        matching_entries.append(entry_path)
344
345                            self.matches = [
346                                command + " " + s
347                                for s in matching_entries
348                                if s.startswith(remainder)
349                            ]
350                            
351                        else:
352                            # Generic case for subcommands
353                            self.matches = [
354                                command + " " + s
355                                for s in self.commands[command]["subcommands"]
356                                if s.startswith(remainder)
357                            ]
358                    else:
359                        # Unknown subcommand, skipping autocomplete
360                        pass
361                else:
362                    self.matches = []
363            else:
364                self.matches = self.commands.keys()[:]
365
366        try:
367            return self.matches[state] + " "
368        except IndexError:
369            return None
370
371    def print_help(self, command=None):
372        """
373        Prints help information for a specific command or all commands if no command is specified.
374
375        This method displays the help information for the command passed as an argument. If no command is specified,
376        it prints the help information for all available commands. The help information includes the command syntax,
377        description, and any subcommands associated with it. This method is designed to provide users with the necessary
378        guidance on how to use the commands in the smbclient-ng shell.
379
380        Args:
381            command (str, optional): The command to display help information for. If None, help for all commands is displayed.
382
383        Returns:
384            None
385        """
386
387        if command is not None:
388            if command not in list(self.commands.keys())+["format"]:
389                command = None
390        
391        # Print help for a specific command
392        if command is not None:
393            if command == "format":
394                self.print_help_format()
395            else:
396                print("│")
397                if self.config.no_colors:
398                    command_str = command + "─"* (15 - len(command))
399                    if len(self.commands[command]["description"]) == 0:
400                        print("│ ■ %s┤  " % command_str)
401                    elif len(self.commands[command]["description"]) == 1:
402                        print("│ ■ %s┤ %s " % (command_str, self.commands[command]["description"][0]))
403                    else:
404                        print("│ ■ %s┤ %s " % (command_str, self.commands[command]["description"][0]))
405                        for line in self.commands[command]["description"][1:]:
406                            print("│ %s│ %s " % (" "*(15+2), line))
407                else:
408                    command_str = command + " \x1b[90m" + "─"* (15 - len(command)) + "\x1b[0m"
409                    if len(self.commands[command]["description"]) == 0:
410                        print("│ ■ %s\x1b[90m┤\x1b[0m  " % command_str)
411                    elif len(self.commands[command]["description"]) == 1:
412                        print("│ ■ %s\x1b[90m┤\x1b[0m %s " % (command_str, self.commands[command]["description"][0]))
413                    else:
414                        print("│ ■ %s\x1b[90m┤\x1b[0m %s " % (command_str, self.commands[command]["description"][0]))
415                        for line in self.commands[command]["description"][1:]:
416                            print("│ %s\x1b[90m│\x1b[0m %s " % (" "*(15+3), line))
417                print("│")
418        # Generic help
419        else:
420            print("│")
421            commands = sorted(self.commands.keys())
422            for command in commands:
423                if self.config.no_colors:
424                    command_str = command + "─"* (15 - len(command))
425                    if len(self.commands[command]["description"]) == 0:
426                        print("│ ■ %s┤  " % command_str)
427                    elif len(self.commands[command]["description"]) == 1:
428                        print("│ ■ %s┤ %s " % (command_str, self.commands[command]["description"][0]))
429                    else:
430                        print("│ ■ %s┤ %s " % (command_str, self.commands[command]["description"][0]))
431                        for line in self.commands[command]["description"][1:]:
432                            print("│ %s│ %s " % (" "*(15+2), line))
433                else:
434                    command_str = command + " \x1b[90m" + "─"* (15 - len(command)) + "\x1b[0m"
435                    if len(self.commands[command]["description"]) == 0:
436                        print("│ ■ %s\x1b[90m┤\x1b[0m  " % command_str)
437                    elif len(self.commands[command]["description"]) == 1:
438                        print("│ ■ %s\x1b[90m┤\x1b[0m %s " % (command_str, self.commands[command]["description"][0]))
439                    else:
440                        print("│ ■ %s\x1b[90m┤\x1b[0m %s " % (command_str, self.commands[command]["description"][0]))
441                        for line in self.commands[command]["description"][1:]:
442                            print("│ %s\x1b[90m│\x1b[0m %s " % (" "*(15+3), line))
443                print("│")
444
445    def print_help_format(self):
446        """
447        Prints the help information for the 'format' used in remote 'ls' and 'dir' commands.
448
449        This function displays the format of file attributes used in the smbclient-ng shell. It explains the meaning
450        of each character in the file attribute string, such as whether a file is read-only, hidden, or a directory.
451        """
452
453        print("File attributes format:\n")
454        print("\x1b[1mdachnrst\x1b[0m")
455        print("\x1b[90m│││││││└──>\x1b[0m Temporary")
456        print("\x1b[90m││││││└───>\x1b[0m System")
457        print("\x1b[90m│││││└────>\x1b[0m Read-Only")
458        print("\x1b[90m││││└─────>\x1b[0m Normal")
459        print("\x1b[90m│││└──────>\x1b[0m Hidden")
460        print("\x1b[90m││└───────>\x1b[0m Compressed")
461        print("\x1b[90m│└────────>\x1b[0m Archived")
462        print("\x1b[90m└─────────>\x1b[0m Directory")