Land rapid7#8512, Enable adaptive download with variable block sizes

Author: bwatters-r7

lib/rex/post/meterpreter/extensions/stdapi/fs/file.rb

@@ -7,6 +7,7 @@
 require 'rex/post/meterpreter/extensions/stdapi/fs/io'
 require 'rex/post/meterpreter/extensions/stdapi/fs/file_stat'
 require 'fileutils'
+require 'filesize'
 
 module Rex
 module Post
@@ -25,6 +26,8 @@ class File < Rex::Post::Meterpreter::Extensions::Stdapi::Fs::IO
 
   include Rex::Post::File
 
+  MIN_BLOCK_SIZE = 1024
+
   class << self
     attr_accessor :client
   end
@@ -312,7 +315,7 @@ def File.download(dest, src_files, opts = nil, &stat)
         dest += timestamp
       end
 
-      stat.call('downloading', src, dest) if (stat)
+      stat.call('Downloading', src, dest) if (stat)
       result = download_file(dest, src, opts, &stat)
       stat.call(result, src, dest) if (stat)
     }
@@ -325,8 +328,11 @@ def File.download_file(dest_file, src_file, opts = nil, &stat)
     continue=false
     tries=false
     tries_no=0
+    stat ||= lambda { }
+
     if opts
       continue = true if opts["continue"]
+      adaptive = true if opts['adaptive']
       tries = true if opts["tries"]
       tries_no = opts["tries_no"]
     end
@@ -346,17 +352,19 @@ def File.download_file(dest_file, src_file, opts = nil, &stat)
     dir = ::File.dirname(dest_file)
     ::FileUtils.mkdir_p(dir) if dir and not ::File.directory?(dir)
 
+    src_size = Filesize.new(src_stat.size).pretty
+
     if continue
       # continue downloading the file - skip downloaded part in the source
       dst_fd = ::File.new(dest_file, "ab")
       begin
         dst_fd.seek(0, ::IO::SEEK_END)
         in_pos = dst_fd.pos
         src_fd.seek(in_pos)
-        stat.call('continuing from ', in_pos, src_file) if (stat)
+        stat.call("Continuing from #{Filesize.new(in_pos).pretty} of #{src_size}", src_file, dest_file)
       rescue
         # if we can't seek, download again
-        stat.call('error continuing - downloading from scratch', src_file, dest_file) if (stat)
+        stat.call('Error continuing - downloading from scratch', src_file, dest_file)
         dst_fd.close
         dst_fd = ::File.new(dest_file, "wb")
       end
@@ -365,41 +373,59 @@ def File.download_file(dest_file, src_file, opts = nil, &stat)
     end
 
     # Keep transferring until EOF is reached...
+    block_size = opts['block_size'] || 1024 * 1024
     begin
       if tries
         # resume when timeouts encountered
         seek_back = false
+        adjust_block = false
         tries_cnt = 0
         begin # while
           begin # exception
             if seek_back
               in_pos = dst_fd.pos
               src_fd.seek(in_pos)
               seek_back = false
-              stat.call('resuming at ', in_pos, src_file) if (stat)
+              stat.call("Resuming at #{Filesize.new(in_pos).pretty} of #{src_size}", src_file, dest_file)
             else
               # succesfully read and wrote - reset the counter
               tries_cnt = 0
             end
-            data = src_fd.read
+            adjust_block = true
+            data = src_fd.read(block_size)
+            adjust_block = false
           rescue Rex::TimeoutError
             # timeout encountered - either seek back and retry or quit
             if (tries && (tries_no == 0 || tries_cnt < tries_no))
               tries_cnt += 1
               seek_back = true
-              stat.call('error downloading - retry #', tries_cnt, src_file) if (stat)
+              # try a smaller block size for the next round
+              if adaptive && adjust_block
+                block_size = [block_size >> 1, MIN_BLOCK_SIZE].max
+                adjust_block = false
+                msg = "Error downloading, block size set to #{block_size} - retry # #{tries_cnt}"
+                stat.call(msg, src_file, dest_file)
+              else
+                stat.call("Error downloading - retry # #{tries_cnt}", src_file, dest_file)
+              end
               retry
             else
-              stat.call('error downloading - giving up', src_file, dest_file) if (stat)
+              stat.call('Error downloading - giving up', src_file, dest_file)
               raise
             end
           end
           dst_fd.write(data) if (data != nil)
+          percent = dst_fd.pos.to_f / src_stat.size.to_f * 100.0
+          msg = "Downloaded #{Filesize.new(dst_fd.pos).pretty} of #{src_size} (#{percent.round(2)}%)"
+          stat.call(msg, src_file, dest_file)
         end while (data != nil)
       else
         # do the simple copying quiting on the first error
-        while ((data = src_fd.read) != nil)
+        while ((data = src_fd.read(block_size)) != nil)
           dst_fd.write(data)
+          percent = dst_fd.pos.to_f / src_stat.size.to_f * 100.0
+          msg = "Downloaded #{Filesize.new(dst_fd.pos).pretty} of #{src_size} (#{percent.round(2)}%)"
+          stat.call(msg, src_file, dest_file)
         end
       end
     rescue EOFError

lib/rex/post/meterpreter/ui/console/command_dispatcher/stdapi/fs.rb

@@ -28,6 +28,8 @@ class Console::CommandDispatcher::Stdapi::Fs
   @@download_opts = Rex::Parser::Arguments.new(
     "-h" => [ false, "Help banner." ],
     "-c" => [ false, "Resume getting a partially-downloaded file." ],
+    "-a" => [ false, "Enable adaptive download buffer size." ],
+    "-b" => [ true,  "Set the initial block size for the download." ],
     "-l" => [ true,  "Set the limit of retries (0 unlimits)." ],
     "-r" => [ false, "Download recursively." ],
     "-t" => [ false, "Timestamp downloaded files." ])
@@ -382,6 +384,10 @@ def cmd_download(*args)
 
     @@download_opts.parse(args) { |opt, idx, val|
       case opt
+      when "-a"
+        opts['adaptive'] = true
+      when "-b"
+        opts['block_size'] = val.to_i
       when "-r"
         recursive = true
         opts['recursive'] = true