Merge pull request #224 from MsysTechnologiesllc/Kapil/MSYS-1131_Retrieves_the_environment_variables_for_the_specified_user

Signed-off-by: Tim Smith <tsmith@chef.io>

4 files changed, 117 insertions, 1 deletions

diff --git a/Gemfile b/Gemfile
index 841998c..40d317d 100644
--- a/Gemfile
+++ b/Gemfile
@@ -2,6 +2,8 @@ source "https://rubygems.org"
 
 gemspec name: "mixlib-shellout"
 
+gem "parallel", "< 1.20" # pin until we drop ruby < 2.4
+
 group :docs do
   gem "github-markup"
   gem "redcarpet"
diff --git a/lib/mixlib/shellout/windows/core_ext.rb b/lib/mixlib/shellout/windows/core_ext.rb
index 1c2830b..fdc3d9b 100644
--- a/lib/mixlib/shellout/windows/core_ext.rb
+++ b/lib/mixlib/shellout/windows/core_ext.rb
@@ -18,6 +18,7 @@
 #
 
 require "win32/process"
+require "ffi/win32/extensions"
 
 # Add new constants for Logon
 module Process::Constants
@@ -44,6 +45,8 @@ module Process::Constants
   WIN32_PROFILETYPE_PT_MANDATORY           = 0x04
   WIN32_PROFILETYPE_PT_ROAMING_PREEXISTING = 0x08
 
+  # The environment block list ends with two nulls (\0\0).
+  ENVIRONMENT_BLOCK_ENDS = "\0\0".freeze
 end
 
 # Structs required for data handling
@@ -77,6 +80,12 @@ module Process::Functions
   attach_pfunc :UnloadUserProfile,
     %i{handle handle}, :bool
 
+  attach_pfunc :CreateEnvironmentBlock,
+    %i{pointer ulong bool}, :bool
+
+  attach_pfunc :DestroyEnvironmentBlock,
+    %i{pointer}, :bool
+
   ffi_lib :advapi32
 
   attach_pfunc :LogonUserW,
@@ -169,9 +178,25 @@ module Process
 
       env = nil
 
+      # Retrieve the environment variables for the specified user.
+      if hash["with_logon"]
+        logon, passwd, domain = format_creds_from_hash(hash)
+        logon_type = hash["elevated"] ? LOGON32_LOGON_BATCH : LOGON32_LOGON_INTERACTIVE
+        token = logon_user(logon, domain, passwd, logon_type)
+        logon_ptr = FFI::MemoryPointer.from_string(logon)
+        profile = PROFILEINFO.new.tap do |dat|
+          dat[:dwSize]     = dat.size
+          dat[:dwFlags]    = 1
+          dat[:lpUserName] = logon_ptr
+        end
+
+        load_user_profile(token, profile.pointer)
+        env_list = retrieve_environment_variables(token)
+      end
+
       # The env string should be passed as a string of ';' separated paths.
       if hash["environment"]
-        env = hash["environment"]
+        env = env_list.nil? ? hash["environment"] : merge_env_variables(env_list, hash["environment"])
 
         unless env.respond_to?(:join)
           env = hash["environment"].split(File::PATH_SEPARATOR)
@@ -393,6 +418,33 @@ module Process
       true
     end
 
+    # Retrieves the environment variables for the specified user.
+    #
+    # @param env_pointer [Pointer] The environment block is an array of null-terminated Unicode strings.
+    # @param token [Integer] User token handle.
+    # @return [Boolean] true if successfully retrieves the environment variables for the specified user.
+    #
+    def create_environment_block(env_pointer, token)
+      unless CreateEnvironmentBlock(env_pointer, token, false)
+        raise SystemCallError.new("CreateEnvironmentBlock", FFI.errno)
+      end
+
+      true
+    end
+
+    # Frees environment variables created by the CreateEnvironmentBlock function.
+    #
+    # @param env_pointer [Pointer] The environment block is an array of null-terminated Unicode strings.
+    # @return [Boolean] true if successfully frees environment variables created by the CreateEnvironmentBlock function.
+    #
+    def destroy_environment_block(env_pointer)
+      unless DestroyEnvironmentBlock(env_pointer)
+        raise SystemCallError.new("DestroyEnvironmentBlock", FFI.errno)
+      end
+
+      true
+    end
+
     def create_process_as_user(token, app, cmd, process_security,
       thread_security, inherit, creation_flags, env, cwd, startinfo, procinfo)
 
@@ -527,5 +579,51 @@ module Process
       [ logon, passwd, domain ]
     end
 
+    # Retrieves the environment variables for the specified user.
+    #
+    # @param token [Integer] User token handle.
+    # @return env_list [Array<String>] Environment variables of specified user.
+    #
+    def retrieve_environment_variables(token)
+      env_list = []
+      env_pointer = FFI::MemoryPointer.new(:pointer)
+      create_environment_block(env_pointer, token)
+      str_ptr = env_pointer.read_pointer
+      offset = 0
+      loop do
+        new_str_pointer = str_ptr + offset
+        break if new_str_pointer.read_string(2) == ENVIRONMENT_BLOCK_ENDS
+
+        environment = new_str_pointer.read_wstring
+        env_list << environment
+        offset = offset + environment.length * 2 + 2
+      end
+
+      # To free the buffer when we have finished with the environment block
+      destroy_environment_block(str_ptr)
+      env_list
+    end
+
+    # Merge environment variables of specified user and current environment variables.
+    #
+    # @param fetched_env [Array<String>] environment variables of specified user.
+    # @param current_env [Array<String>] current environment variables.
+    # @return [Array<String>] Merged environment variables.
+    #
+    def merge_env_variables(fetched_env, current_env)
+      env_hash_1 = environment_list_to_hash(fetched_env)
+      env_hash_2 = environment_list_to_hash(current_env)
+      merged_env = env_hash_2.merge(env_hash_1)
+      merged_env.map { |k, v| "#{k}=#{v}" }
+    end
+
+    # Convert an array to a hash.
+    #
+    # @param env_var [Array<String>] Environment variables.
+    # @return [Hash] Converted an array to hash.
+    #
+    def environment_list_to_hash(env_var)
+      Hash[ env_var.map { |pair| pair.split("=", 2) } ]
+    end
   end
 end
diff --git a/mixlib-shellout-universal-mingw32.gemspec b/mixlib-shellout-universal-mingw32.gemspec
index 705366a..f6fd802 100644
--- a/mixlib-shellout-universal-mingw32.gemspec
+++ b/mixlib-shellout-universal-mingw32.gemspec
@@ -4,5 +4,6 @@ gemspec.platform = Gem::Platform.new(%w{universal mingw32})
 
 gemspec.add_dependency "win32-process", "~> 0.9"
 gemspec.add_dependency "wmi-lite", "~> 1.0"
+gemspec.add_dependency "ffi-win32-extensions", "~> 1.0.3"
 
 gemspec
diff --git a/spec/mixlib/shellout_spec.rb b/spec/mixlib/shellout_spec.rb
index 9e79fa8..867644e 100644
--- a/spec/mixlib/shellout_spec.rb
+++ b/spec/mixlib/shellout_spec.rb
@@ -658,6 +658,21 @@ describe Mixlib::ShellOut do
           expect(running_user).to eql("#{ENV["COMPUTERNAME"].downcase}\\#{user}")
         end
 
+        context "when an alternate user is passed" do
+          let(:env_list) { ["ALLUSERSPROFILE=C:\\ProgramData", "TEMP=C:\\Windows\\TEMP", "TMP=C:\\Windows\\TEMP", "USERDOMAIN=WIN-G06ENRTTKF9", "USERNAME=testuser", "USERPROFILE=C:\\Users\\Default", "windir=C:\\Windows"] }
+          let(:current_env) { ["ALLUSERSPROFILE=C:\\ProgramData", "TEMP=C:\\Users\\ADMINI~1\\AppData\\Local\\Temp\\2", "TMP=C:\\Users\\ADMINI~1\\AppData\\Local\\Temp\\2", "USER=Administrator", "USERDOMAIN=WIN-G06ENRTTKF9", "USERDOMAIN_ROAMINGPROFILE=WIN-G06ENRTTKF9", "USERNAME=Administrator", "USERPROFILE=C:\\Users\\Administrator", "windir=C:\\Windows"] }
+          let(:merged_env) { ["ALLUSERSPROFILE=C:\\ProgramData", "TEMP=C:\\Windows\\TEMP", "TMP=C:\\Windows\\TEMP", "USER=Administrator", "USERDOMAIN=WIN-G06ENRTTKF9", "USERDOMAIN_ROAMINGPROFILE=WIN-G06ENRTTKF9", "USERNAME=testuser", "USERPROFILE=C:\\Users\\Default", "windir=C:\\Windows"] }
+          let(:converted) { { "ALLUSERSPROFILE" => "C:\\ProgramData", "TEMP" => "C:\\Windows\\TEMP", "TMP" => "C:\\Windows\\TEMP", "USERDOMAIN" => "WIN-G06ENRTTKF9", "USERNAME" => "testuser", "USERPROFILE" => "C:\\Users\\Default", "windir" => "C:\\Windows" } }
+
+          it "merge environment variables" do
+            expect(Process.merge_env_variables(env_list, current_env)).to eql(merged_env)
+          end
+
+          it "Convert an array to a hash" do
+            expect(Process.environment_list_to_hash(env_list)).to eql(converted)
+          end
+        end
+
         context "when :elevated => true" do
           context "when user and password are passed" do
             let(:options) { { user: user, password: password, elevated: true } }