From c22f0c7c47815534c9d07d529524a383e4a20030 Mon Sep 17 00:00:00 2001
From: "Remi GASCOU (Podalirius)"
 <79218792+p0dalirius@users.noreply.github.com>
Date: Wed, 5 Jun 2024 13:35:24 +0200
Subject: [PATCH] Created 'sizeof' command

---
 smbclientng/core/CommandCompleter.py | 16 +++++-
 smbclientng/core/InteractiveShell.py | 79 ++++++++++++++++++++++++++++
 2 files changed, 94 insertions(+), 1 deletion(-)

diff --git a/smbclientng/core/CommandCompleter.py b/smbclientng/core/CommandCompleter.py
index a08e68a..9502c67 100644
--- a/smbclientng/core/CommandCompleter.py
+++ b/smbclientng/core/CommandCompleter.py
@@ -61,6 +61,13 @@ class CommandCompleter(object):
             ], 
             "subcommands": []
         },
+        "debug": {
+            "description": [
+                "Command for dev debugging.",
+                "Syntax: 'debug'"
+            ], 
+            "subcommands": []
+        },
         "dir": {
             "description": [
                 "List the contents of the current working directory.",
@@ -201,6 +208,13 @@ class CommandCompleter(object):
             ], 
             "subcommands": []
         },
+        "sizeof": {
+            "description": [
+                "Recursively compute the size of a folder.", 
+                "Syntax: 'sizeof [directory|file]'"
+            ], 
+            "subcommands": []
+        },
         "shares": {
             "description": [
                 "Lists the SMB shares served by the remote machine.", 
@@ -292,7 +306,7 @@ def complete(self, text, state):
                                 if s.lower().startswith(remainder.lower())
                             ]
 
-                        elif command in ["bat", "cat", "get", "rm"]:
+                        elif command in ["bat", "cat", "debug", "get", "rm"]:
                             # Choose local files and directories
                             path = ""
                             if '\\' in remainder.strip() or '/' in remainder.strip():
diff --git a/smbclientng/core/InteractiveShell.py b/smbclientng/core/InteractiveShell.py
index 61cf86c..8da0b75 100644
--- a/smbclientng/core/InteractiveShell.py
+++ b/smbclientng/core/InteractiveShell.py
@@ -149,6 +149,10 @@ def process_command(self, command, arguments=[]):
         # Change directory in the current share
         elif command == "cd":
             self.command_cd(arguments, command)
+        
+        # debug
+        elif command == "debug":
+            self.command_debug(arguments, command)
 
         # Get a file
         elif command == "get":
@@ -218,6 +222,10 @@ def process_command(self, command, arguments=[]):
         elif command == "rmdir":
             self.command_rmdir(arguments, command)
 
+        # List shares
+        elif command == "sizeof":
+            self.command_sizeof(arguments, command)
+
         # List shares
         elif command == "shares":
             self.command_shares(arguments, command)
@@ -232,6 +240,12 @@ def process_command(self, command, arguments=[]):
 
     # Commands ================================================================
 
+    def command_debug(self, arguments, command):
+        try:
+            pass
+        except Exception as e:
+            traceback.print_exc()
+
     @command_arguments_required
     @active_smb_connection_needed
     @smb_share_is_set
@@ -611,6 +625,71 @@ def command_rmdir(self, arguments, command):
         else:
             print("[!] Remote directory '%s' does not exist." % path)
 
+    @active_smb_connection_needed
+    @smb_share_is_set
+    def command_sizeof(self, arguments, command):
+        # Command arguments required   : Yes
+        # Active SMB connection needed : Yes
+        # SMB share needed             : Yes
+
+        class RecursiveSizeOfPrint(object):
+            def __init__(self, entry, smbSession, config):
+                self.entry = entry
+                self.config = config
+                self.smbSession = smbSession
+                self.size = 0
+            
+            def update(self, entry, fullpath, depth):
+                self.size += entry.get_filesize()
+                self.print(end='\r')
+            
+            def print(self, end='\n'):
+                #
+                if self.entry.is_directory():
+                    if self.config.no_colors:
+                        path = "%s\\" % self.entry.get_longname()
+                    else:
+                        path = "\x1b[1;96m%s\x1b[0m\\" % self.entry.get_longname()
+                # 
+                else:
+                    if self.config.no_colors:
+                        path = "%s" % self.entry.get_longname()
+                    else:
+                        path = "\x1b[1m%s\x1b[0m" % self.entry.get_longname()
+                print("%10s  %s" % (b_filesize(self.size), path), end=end)
+
+        entries = []
+        if len(arguments) == 0:
+            entries = self.smbSession.list_contents()
+            entries = [entry for name, entry in entries.items() if name not in ['.', '..']]
+        else:
+            entry = self.smbSession.get_entry(path=' '.join(arguments))
+            entries = []
+            if entry is not None:
+                entries = [entry]
+            else:
+                print("[!] Path '%s' does not exist." % ' '.join(arguments))
+
+        total = 0
+        for entry in entries:
+            rsop = RecursiveSizeOfPrint(entry=entry, smbSession=self.smbSession, config=self.config)
+            # Directory
+            if entry.is_directory():
+                self.smbSession.find(
+                    paths=[entry.get_longname()],
+                    callback=rsop.update
+                )
+            # File
+            else:
+                rsop.update(entry=entry, fullpath=entry.get_longname(), depth=0)
+            # Close the print
+            rsop.print()
+            total += rsop.size
+        
+        if len(entries) > 1:
+            print("──────────────────────")
+            print("     total  %s" % b_filesize(total))
+
     @active_smb_connection_needed
     def command_shares(self, arguments, command):
         # Command arguments required   : No