Add upload/download status output

lib/rex/ntpath.rb

@@ -0,0 +1,13 @@
+# frozen_string_literal: true
+
+module Rex
+  module Ntpath
+    def self.as_ntpath(path)
+      Pathname.new(path)
+              .cleanpath
+              .each_filename
+              .drop_while { |file| file == '.' }
+              .join('\\')
+    end
+  end
+end

lib/rex/post/smb/ui/console.rb

@@ -137,7 +137,7 @@ def log_error(msg)
           def format_prompt(val)
             if active_share
               share_name = active_share.share[/[^\\].*$/, 0]
-              cwd = self.cwd.blank? ? '' : "\\#{as_ntpath(self.cwd)}"
+              cwd = self.cwd.blank? ? '' : "\\#{Rex::Ntpath.as_ntpath(self.cwd)}"
               return substitute_colors("%undSMB%clr (#{share_name}#{cwd}) > ", true)
             end
 
@@ -148,14 +148,6 @@ def format_prompt(val)
 
           attr_writer :session, :client, :simple_client # :nodoc: # :nodoc:
           attr_accessor :commands # :nodoc:
-
-          def as_ntpath(path)
-            Pathname.new(path)
-                    .cleanpath
-                    .each_filename
-                    .drop_while { |file| file == '.' }
-                    .join('\\')
-          end
         end
       end
     end

lib/rex/post/smb/ui/console/command_dispatcher/shares.rb

@@ -2,6 +2,7 @@
 
 require 'pathname'
 require 'rex/post/file'
+require 'filesize'
 
 module Rex
   module Post
@@ -196,7 +197,7 @@ def cmd_ls(*args)
               end
             end
 
-            full_path = as_ntpath(Pathname.new(shell.cwd).join(remote_path).to_s)
+            full_path = Rex::Ntpath.as_ntpath(Pathname.new(shell.cwd).join(remote_path).to_s)
 
             files = active_share.list(directory: full_path)
             table = Rex::Text::Table.new(
@@ -297,7 +298,7 @@ def cmd_cd(*args)
 
             path = args[0]
             native_path = Pathname.new(shell.cwd).join(path).to_s
-            new_path = as_ntpath(native_path)
+            new_path = Rex::Ntpath.as_ntpath(native_path)
             begin
               response = active_share.open_directory(directory: new_path)
               directory = RubySMB::SMB2::File.new(name: new_path, tree: active_share, response: response, encrypt: @tree_connect_encrypt_data)
@@ -341,7 +342,7 @@ def cmd_cat(*args)
 
             path = args[0]
 
-            new_path = as_ntpath(Pathname.new(shell.cwd).join(path).to_s)
+            new_path = Rex::Ntpath.as_ntpath(Pathname.new(shell.cwd).join(path).to_s)
 
             begin
               file = simple_client.open(new_path, 'o')
@@ -395,7 +396,7 @@ def cmd_upload(*args)
             end
 
             remote_path = Rex::Post::File.basename(local_path) if remote_path.nil?
-            full_path = as_ntpath(Pathname.new(shell.cwd).join(remote_path).to_s)
+            full_path = Rex::Ntpath.as_ntpath(Pathname.new(shell.cwd).join(remote_path).to_s)
 
             upload_file(full_path, local_path)
 
@@ -443,7 +444,7 @@ def cmd_download(*args)
             end
 
             local_path = Rex::Post::File.basename(remote_path) if local_path.nil?
-            full_path = as_ntpath(Pathname.new(shell.cwd).join(remote_path).to_s)
+            full_path = Rex::Ntpath.as_ntpath(Pathname.new(shell.cwd).join(remote_path).to_s)
 
             download_file(local_path, full_path)
 
@@ -475,7 +476,7 @@ def cmd_delete(*args)
               end
             end
 
-            full_path = as_ntpath(Pathname.new(shell.cwd).join(remote_path).to_s)
+            full_path = Rex::Ntpath.as_ntpath(Pathname.new(shell.cwd).join(remote_path).to_s)
             fd = simple_client.open(full_path, 'o')
             fd.delete
             print_good("Deleted #{full_path}")
@@ -509,7 +510,7 @@ def cmd_mkdir(*args)
               end
             end
 
-            full_path = as_ntpath(Pathname.new(shell.cwd).join(remote_path).to_s)
+            full_path = Rex::Ntpath.as_ntpath(Pathname.new(shell.cwd).join(remote_path).to_s)
 
             response = active_share.open_directory(directory: full_path, disposition: RubySMB::Dispositions::FILE_CREATE)
             directory = RubySMB::SMB2::File.new(name: full_path, tree: active_share, response: response, encrypt: @tree_connect_encrypt_data)
@@ -546,7 +547,7 @@ def cmd_rmdir(*args)
               end
             end
 
-            full_path = as_ntpath(Pathname.new(shell.cwd).join(remote_path).to_s)
+            full_path = Rex::Ntpath.as_ntpath(Pathname.new(shell.cwd).join(remote_path).to_s)
 
             response = active_share.open_directory(directory: full_path, write: true, delete: true)
             directory = RubySMB::SMB2::File.new(name: full_path, tree: active_share, response: response, encrypt: @tree_connect_encrypt_data)
@@ -579,9 +580,14 @@ def upload_file(dest_file, src_file)
             begin
               dest_fd = simple_client.open(dest_file, 'wct', write: true)
               src_fd = ::File.open(src_file, "rb")
+              src_size = src_fd.stat.size
               offset = 0
               while (buf = src_fd.read(buf_size))
                 offset = dest_fd.write(buf, offset)
+                percent = offset / src_size.to_f * 100.0
+                msg = "Uploaded #{Filesize.new(offset).pretty} of " \
+                "#{Filesize.new(src_size).pretty} (#{percent.round(2)}%)"
+                print_status(msg)
               end
             ensure
               src_fd.close unless src_fd.nil?
@@ -593,24 +599,28 @@ def upload_file(dest_file, src_file)
           # @param dest_file [String] The path for the destination file
           # @param src_file [String] The path for the source file
           def download_file(dest_file, src_file)
+            buf_size = 8 * 1024 * 1024
             src_fd = simple_client.open(src_file, 'o')
-
             # Make the destination path if necessary
             dir = ::File.dirname(dest_file)
-            ::FileUtils.mkdir_p(dir) if dir and not ::File.directory?(dir)
-
+            ::FileUtils.mkdir_p(dir) if dir && !::File.directory?(dir)
             dst_fd = ::File.new(dest_file, "wb")
-            offset = 0
 
+            offset = 0
+            src_size = client.open_files[src_fd.file_id].size
             begin
-              while offset < client.open_files[src_fd.file_id].size
-                data = src_fd.read(src_fd.chunk_size, offset)
+              while offset < src_size
+                data = src_fd.read(buf_size, offset)
                 dst_fd.write(data)
-                offset += src_fd.chunk_size
+                offset += data.length
+                percent = offset / src_size.to_f * 100.0
+                msg = "Downloaded #{Filesize.new(offset).pretty} of " \
+                  "#{Filesize.new(src_size).pretty} (#{percent.round(2)}%)"
+                print_status(msg)
               end
             ensure
-              src_fd.close
-              dst_fd.close
+              src_fd.close unless src_fd.nil?
+              dst_fd.close unless dst_fd.nil?
             end
           end
         end

lib/rex/proto/smb/simple_client/open_file.rb

@@ -30,35 +30,36 @@ def close
       end
 
       def read_ruby_smb(length, offset, depth = 0)
+        file_size = client.open_files[client.last_file_id].size
+        file_size_remaining = file_size - offset
         if length.nil?
-          max_size = client.open_files[client.last_file_id].size
-          fptr = offset
+          max_size = file_size_remaining
+        else
+          max_size = [length, file_size_remaining].min
+        end
 
-          chunk = [max_size, chunk_size].min
+        fptr = offset
+        chunk = [max_size, chunk_size].min
 
-          data = client.read(file_id, fptr, chunk).pack('C*')
-          fptr = data.length
+        data = client.read(file_id, fptr, chunk).pack('C*')
+        fptr += data.length
 
-          while data.length < max_size
-            if (max_size - data.length) < chunk
-              chunk = max_size - data.length
-            end
-            data << client.read(file_id, fptr, chunk).pack('C*')
-            fptr = data.length
-          end
-        else
-          begin
-            data = client.read(file_id, offset, length).pack('C*')
-          rescue RubySMB::Error::UnexpectedStatusCode => e
-            if e.message == 'STATUS_PIPE_EMPTY' && depth < 20
-              data = read_ruby_smb(length, offset, depth + 1)
-            else
-              raise e
-            end
+        while data.length < max_size
+          if (max_size - data.length) < chunk
+            chunk = max_size - data.length
           end
+          new_data = client.read(file_id, fptr, chunk).pack('C*')
+          data << new_data
+          fptr += new_data.length
         end
 
         data
+      rescue RubySMB::Error::UnexpectedStatusCode => e
+        if e.message == 'STATUS_PIPE_EMPTY' && depth < 20
+          read_ruby_smb(max_size, offset, depth + 1)
+        else
+          raise e
+        end
       end
 
       def read_rex_smb(length, offset)

spec/lib/rex/ntpath_spec.rb

@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+require 'rspec'
+
+RSpec.describe Rex::Ntpath do
+
+  describe '#as_ntpath' do
+    let(:valid_windows_path) { 'some\\path\\that\\is\\valid' }
+
+    [
+      'some\\path\\that\\is\\valid',
+      'some/path/that/is/valid',
+      'some/./path/that/./is/valid',
+      'some/extra/../path/that/extra/../is/valid',
+      '/some/path/that/is/valid'
+    ].each do |path|
+      context "when the path is #{path}" do
+        it 'formats it as a valid ntpath' do
+          formatted_path = described_class.as_ntpath(path)
+          expect(formatted_path).to eq valid_windows_path
+        end
+      end
+    end
+  end
+end

spec/lib/rex/post/smb/ui/console/command_dispatcher/shares_spec.rb

@@ -27,23 +27,4 @@
   end
 
   subject(:command_dispatcher) { described_class.new(session.console) }
-
-  describe '#as_ntpath' do
-    let(:valid_windows_path) { 'some\\path\\that\\is\\valid' }
-
-    [
-      'some\\path\\that\\is\\valid',
-      'some/path/that/is/valid',
-      'some/./path/that/./is/valid',
-      'some/extra/../path/that/extra/../is/valid',
-      '/some/path/that/is/valid'
-    ].each do |path|
-      context "when the path is #{path}" do
-        it 'formats it as a valid ntpath' do
-          formatted_path = subject.send(:as_ntpath, path)
-          expect(formatted_path).to eq valid_windows_path
-        end
-      end
-    end
-  end
 end