From adcf45b0ff313ea5fa3ea6f73d03ecced452ac0a Mon Sep 17 00:00:00 2001
From: Spencer McIntyre <zeroSteiner@gmail.com>
Date: Fri, 27 May 2022 16:41:25 -0400
Subject: [PATCH] Fix the arch in #handle_connection too

This fixes an issue with the adated peinject stage which supported both
x86 and x64 via a library that checked its own #arch.
---
 lib/msf/core/payload/windows/pe_inject.rb             | 11 ++++++-----
 modules/payloads/adapters/cmd/windows/powershell.rb   |  5 +++++
 .../payloads/adapters/cmd/windows/powershell/x64.rb   |  5 +++++
 3 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/lib/msf/core/payload/windows/pe_inject.rb b/lib/msf/core/payload/windows/pe_inject.rb
index 9d25e2952c..adab1be52f 100644
--- a/lib/msf/core/payload/windows/pe_inject.rb
+++ b/lib/msf/core/payload/windows/pe_inject.rb
@@ -67,8 +67,9 @@ module Msf
   module Payload::Windows::PEInject
     def initialize(info = {})
       super
+
       register_options([
-        OptInjectablePE.new('PE', [ true, 'The local path to the PE file to upload' ], arch: arch.first)
+        OptInjectablePE.new('PE', [ true, 'The local path to the PE file to upload' ], arch: info.fetch('AdaptedArch', arch.first))
       ], self.class)
     end
 
@@ -83,7 +84,7 @@ module Msf
     # Transmits the reflective PE payload to the remote
     # computer so that it can be loaded into memory.
     #
-    def handle_connection(conn, _opts = {})
+    def handle_connection(conn, opts = {})
       data = ''
       begin
         File.open(pe_path, 'rb') do |f|
@@ -96,7 +97,7 @@ module Msf
       end
 
       print_status('Premapping PE file...')
-      pe_map = create_pe_memory_map(data)
+      pe_map = create_pe_memory_map(data, opts)
       print_status("Mapped PE size #{pe_map[:bytes].length}")
       opts = {}
       opts[:is_dll] = pe_map[:is_dll]
@@ -113,10 +114,10 @@ module Msf
       conn.close
     end
 
-    def create_pe_memory_map(file)
+    def create_pe_memory_map(file, opts = {})
       pe = Rex::PeParsey::Pe.new(Rex::ImageSource::Memory.new(file))
       begin
-        OptInjectablePE.assert_compatible(pe, arch.first)
+        OptInjectablePE.assert_compatible(pe, opts.fetch(:arch, arch.first))
       rescue Msf::ValidationError => e
         print_error("PE validation error: #{e.message}")
         raise
diff --git a/modules/payloads/adapters/cmd/windows/powershell.rb b/modules/payloads/adapters/cmd/windows/powershell.rb
index cfc139a076..4dd84b276a 100644
--- a/modules/payloads/adapters/cmd/windows/powershell.rb
+++ b/modules/payloads/adapters/cmd/windows/powershell.rb
@@ -51,4 +51,9 @@ module MetasploitModule
     conf[:platform] ||= module_info['AdaptedPlatform']
     super
   end
+
+  def handle_connection(conn, opts = {})
+    opts[:arch] ||= module_info['AdaptedArch']
+    super
+  end
 end
diff --git a/modules/payloads/adapters/cmd/windows/powershell/x64.rb b/modules/payloads/adapters/cmd/windows/powershell/x64.rb
index 4f6afda42e..e540c4c3a6 100644
--- a/modules/payloads/adapters/cmd/windows/powershell/x64.rb
+++ b/modules/payloads/adapters/cmd/windows/powershell/x64.rb
@@ -51,4 +51,9 @@ module MetasploitModule
     conf[:platform] ||= module_info['AdaptedPlatform']
     super
   end
+
+  def handle_connection(conn, opts = {})
+    opts[:arch] ||= module_info['AdaptedArch']
+    super
+  end
 end