module Puppet::Util::Windows::Process

# File lib/puppet/util/windows/process.rb, line 183
def elevated_security?
  # default / pre-Vista
  elevated = false
  handle = nil

  begin
    handle = get_current_process
    open_process_token(handle, TOKEN_QUERY) do |token_handle|
      get_token_information(token_handle, :TokenElevation) do |token_info|
        token_elevation = parse_token_information_as_token_elevation(token_info)
        # TokenIsElevated member of the TOKEN_ELEVATION struct
        elevated = token_elevation[:TokenIsElevated] != 0
      end
    end

    elevated
  rescue Puppet::Util::Windows::Error => e
    raise e if e.code != ERROR_NO_SUCH_PRIVILEGE
  ensure
    FFI::WIN32.CloseHandle(handle) if handle
  end
end

# File lib/puppet/util/windows/process.rb, line 11
def execute(command, arguments, stdin, stdout, stderr)
  Process.create( :command_line => command, :startup_info => {:stdin => stdin, :stdout => stdout, :stderr => stderr}, :close_handles => false )
end

# File lib/puppet/util/windows/process.rb, line 39
def get_current_process
  # this pseudo-handle does not require closing per MSDN docs
  GetCurrentProcess()
end

# File lib/puppet/util/windows/process.rb, line 101
def get_token_information(token_handle, token_information, &block)
  # to determine buffer size
  FFI::MemoryPointer.new(:dword, 1) do |return_length_ptr|
    result = GetTokenInformation(token_handle, token_information, nil, 0, return_length_ptr)
    return_length = return_length_ptr.read_dword

    if return_length <= 0
      raise Puppet::Util::Windows::Error.new(
        "GetTokenInformation(#{token_handle}, #{token_information}, nil, 0, #{return_length_ptr})")
    end

    # re-call API with properly sized buffer for all results
    FFI::MemoryPointer.new(return_length) do |token_information_buf|
      result = GetTokenInformation(token_handle, token_information,
        token_information_buf, return_length, return_length_ptr)

      if result == FFI::WIN32_FALSE
        raise Puppet::Util::Windows::Error.new(
          "GetTokenInformation(#{token_handle}, #{token_information}, #{token_information_buf}, " +
            "#{return_length}, #{return_length_ptr})")
      end

      yield token_information_buf
    end
  end

  # GetTokenInformation buffer has been cleaned up by this point, nothing to return
  nil
end

# File lib/puppet/util/windows/process.rb, line 80
def lookup_privilege_value(name, system_name = '', &block)
  FFI::MemoryPointer.new(LUID.size) do |luid_ptr|
    result = LookupPrivilegeValueW(
      wide_string(system_name),
      wide_string(name.to_s),
      luid_ptr
      )

    if result == FFI::WIN32_FALSE
      raise Puppet::Util::Windows::Error.new(
        "LookupPrivilegeValue(#{system_name}, #{name}, #{luid_ptr})")
    end

    yield LUID.new(luid_ptr)
  end

  # the underlying MemoryPointer for LUID is cleaned up by this point
  nil
end

# File lib/puppet/util/windows/process.rb, line 45
def open_process_token(handle, desired_access, &block)
  token_handle = nil
  begin
    FFI::MemoryPointer.new(:handle, 1) do |token_handle_ptr|
      result = OpenProcessToken(handle, desired_access, token_handle_ptr)
      if result == FFI::WIN32_FALSE
        raise Puppet::Util::Windows::Error.new(
          "OpenProcessToken(#{handle}, #{desired_access.to_s(8)}, #{token_handle_ptr})")
      end

      yield token_handle = token_handle_ptr.read_handle
    end

    token_handle
  ensure
    FFI::WIN32.CloseHandle(token_handle) if token_handle
  end

  # token_handle has had CloseHandle called against it, so nothing to return
  nil
end

# File lib/puppet/util/windows/process.rb, line 148
def parse_token_information_as_token_elevation(token_information_buf)
  TOKEN_ELEVATION.new(token_information_buf)
end

# File lib/puppet/util/windows/process.rb, line 132
def parse_token_information_as_token_privileges(token_information_buf)
  raw_privileges = TOKEN_PRIVILEGES.new(token_information_buf)
  privileges = { :count => raw_privileges[:PrivilegeCount], :privileges => [] }

  offset = token_information_buf + TOKEN_PRIVILEGES.offset_of(:Privileges)
  privilege_ptr = FFI::Pointer.new(LUID_AND_ATTRIBUTES, offset)

  # extract each instance of LUID_AND_ATTRIBUTES
  0.upto(privileges[:count] - 1) do |i|
    privileges[:privileges] <<  LUID_AND_ATTRIBUTES.new(privilege_ptr[i])
  end

  privileges
end

# File lib/puppet/util/windows/process.rb, line 155
def process_privilege_symlink?
  privilege_symlink = false
  handle = get_current_process
  open_process_token(handle, TOKEN_ALL_ACCESS) do |token_handle|
    lookup_privilege_value('SeCreateSymbolicLinkPrivilege') do |luid|
      get_token_information(token_handle, :TokenPrivileges) do |token_info|
        token_privileges = parse_token_information_as_token_privileges(token_info)
        privilege_symlink = token_privileges[:privileges].any? { |p| p[:Luid].values == luid.values }
      end
    end
  end

  privilege_symlink
rescue Puppet::Util::Windows::Error => e
  if e.code == ERROR_NO_SUCH_PRIVILEGE
    false # pre-Vista
  else
    raise e
  end
end

# File lib/puppet/util/windows/process.rb, line 16
def wait_process(handle)
  while WaitForSingleObject(handle, 0) == WAIT_TIMEOUT
    sleep(1)
  end

  exit_status = -1
  FFI::MemoryPointer.new(:dword, 1) do |exit_status_ptr|
    if GetExitCodeProcess(handle, exit_status_ptr) == FFI::WIN32_FALSE
      raise Puppet::Util::Windows::Error.new("Failed to get child process exit code")
    end
    exit_status = exit_status_ptr.read_dword

    # $CHILD_STATUS is not set when calling win32/process Process.create
    # and since it's read-only, we can't set it. But we can execute a
    # a shell that simply returns the desired exit status, which has the
    # desired effect.
    %x{#{ENV['COMSPEC']} /c exit #{exit_status}}
  end

  exit_status
end

# File lib/puppet/util/windows/process.rb, line 69
def with_process_token(access, &block)
  handle = get_current_process
  open_process_token(handle, access) do |token_handle|
    yield token_handle
  end

  # all handles have been closed, so nothing to safely return
  nil
end