#--
# Author:: Daniel DeLeo (<dan@opscode.com>)
# Author:: John Keiser (<jkeiser@opscode.com>)
# Author:: Ho-Sheng Hsiao (<hosh@opscode.com>)
# Copyright:: Copyright (c) 2011, 2012 Opscode, Inc.
# License:: Apache License, Version 2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

require 'win32/process'
require 'mixlib/shellout/windows/core_ext'

module Mixlib
  class ShellOut
    module Windows

      include Process::Functions
      include Process::Constants

      TIME_SLICE = 0.05

      # Option validation that is windows specific
      def validate_options(opts)
        if opts[:user]
          unless opts[:password]
            raise InvalidCommandOption, "You must supply both a username and password when supplying a user in windows"
          end
        end
      end

      #--
      # Missing lots of features from the UNIX version, such as
      # uid, etc.
      def run_command
        #
        # Create pipes to capture stdout and stderr,
        #
        stdout_read, stdout_write = IO.pipe
        stderr_read, stderr_write = IO.pipe
        stdin_read, stdin_write = IO.pipe
        open_streams = [ stdout_read, stderr_read ]

        begin

          #
          # Set cwd, environment, appname, etc.
          #
          app_name, command_line = command_to_run(self.command)
          create_process_args = {
            :app_name => app_name,
            :command_line => command_line,
            :startup_info => {
              :stdout => stdout_write,
              :stderr => stderr_write,
              :stdin => stdin_read
            },
            :environment => inherit_environment.map { |k,v| "#{k}=#{v}" },
            :close_handles => false
          }
          create_process_args[:cwd] = cwd if cwd
          # default to local account database if domain is not specified
          create_process_args[:domain] = domain.nil? ? "." : domain
          create_process_args[:with_logon] = with_logon if with_logon
          create_process_args[:password] = password if password

          #
          # Start the process
          #
          process = Process.create(create_process_args)
          logger.debug(Utils.format_process(process, app_name, command_line, timeout)) if logger
          begin
            # Start pushing data into input
            stdin_write << input if input

            # Close pipe to kick things off
            stdin_write.close

            #
            # Wait for the process to finish, consuming output as we go
            #
            start_wait = Time.now
            while true
              wait_status = WaitForSingleObject(process.process_handle, 0)
              case wait_status
              when WAIT_OBJECT_0
                # Get process exit code
                exit_code = [0].pack('l')
                unless GetExitCodeProcess(process.process_handle, exit_code)
                  raise get_last_error
                end
                @status = ThingThatLooksSortOfLikeAProcessStatus.new
                @status.exitstatus = exit_code.unpack('l').first

                return self
              when WAIT_TIMEOUT
                # Kill the process
                if (Time.now - start_wait) > timeout
                  begin
                    require 'wmi-lite/wmi'
                    wmi = WmiLite::Wmi.new
                    Utils.kill_process_tree(process.process_id, wmi, logger)
                    Process.kill(:KILL, process.process_id)
                  rescue Errno::EIO, SystemCallError
                    logger.warn("Failed to kill timed out process #{process.process_id}") if logger
                  end

                  raise Mixlib::ShellOut::CommandTimeout, [
                    "command timed out:",
                    format_for_exception,
                    Utils.format_process(process, app_name, command_line, timeout)
                  ].join("\n")
                end

                consume_output(open_streams, stdout_read, stderr_read)
              else
                raise "Unknown response from WaitForSingleObject(#{process.process_handle}, #{timeout*1000}): #{wait_status}"
              end

            end

          ensure
            CloseHandle(process.thread_handle) if process.thread_handle
            CloseHandle(process.process_handle) if process.process_handle
          end

        ensure
          #
          # Consume all remaining data from the pipes until they are closed
          #
          stdout_write.close
          stderr_write.close

          while consume_output(open_streams, stdout_read, stderr_read)
          end
        end
      end

      private

      class ThingThatLooksSortOfLikeAProcessStatus
        attr_accessor :exitstatus
        def success?
          exitstatus == 0
        end
      end

      def consume_output(open_streams, stdout_read, stderr_read)
        return false if open_streams.length == 0
        ready = IO.select(open_streams, nil, nil, READ_WAIT_TIME)
        return true if ! ready

        if ready.first.include?(stdout_read)
          begin
            next_chunk = stdout_read.readpartial(READ_SIZE)
            @stdout << next_chunk
            @live_stdout << next_chunk if @live_stdout
          rescue EOFError
            stdout_read.close
            open_streams.delete(stdout_read)
          end
        end

        if ready.first.include?(stderr_read)
          begin
            next_chunk = stderr_read.readpartial(READ_SIZE)
            @stderr << next_chunk
            @live_stderr << next_chunk if @live_stderr
          rescue EOFError
            stderr_read.close
            open_streams.delete(stderr_read)
          end
        end

        return true
      end

      IS_BATCH_FILE = /\.bat"?$|\.cmd"?$/i

      def command_to_run(command)
        return _run_under_cmd(command) if Utils.should_run_under_cmd?(command)

        candidate = candidate_executable_for_command(command)

        # Don't do searching for empty commands.  Let it fail when it runs.
        return [ nil, command ] if candidate.length == 0

        # Check if the exe exists directly.  Otherwise, search PATH.
        exe = Utils.find_executable(candidate)
        exe = Utils.which(unquoted_executable_path(command)) if exe.nil? && exe !~ /[\\\/]/

        # Batch files MUST use cmd; and if we couldn't find the command we're looking for,
        # we assume it must be a cmd builtin.
        if exe.nil? || exe =~ IS_BATCH_FILE
          _run_under_cmd(command)
        else
          _run_directly(command, exe)
        end
      end

      # cmd does not parse multiple quotes well unless the whole thing is wrapped up in quotes.
      # https://github.com/opscode/mixlib-shellout/pull/2#issuecomment-4837859
      # http://ss64.com/nt/syntax-esc.html
      def _run_under_cmd(command)
        [ ENV['COMSPEC'], "cmd /c \"#{command}\"" ]
      end

      def _run_directly(command, exe)
        [ exe, command ]
      end

      def unquoted_executable_path(command)
        command[0,command.index(/\s/) || command.length]
      end

      def candidate_executable_for_command(command)
        if command =~ /^\s*"(.*?)"/
          # If we have quotes, do an exact match
          $1
        else
          # Otherwise check everything up to the first space
          unquoted_executable_path(command).strip
        end
      end

      def inherit_environment
        result = {}
        ENV.each_pair do |k,v|
          result[k] = v
        end

        environment.each_pair do |k,v|
          if v == nil
            result.delete(k)
          else
            result[k] = v
          end
        end
        result
      end

      module Utils
        # api: semi-private
        # If there are special characters parsable by cmd.exe (such as file redirection), then
        # this method should return true.
        #
        # This parser is based on
        # https://github.com/ruby/ruby/blob/9073db5cb1d3173aff62be5b48d00f0fb2890991/win32/win32.c#L1437
        def self.should_run_under_cmd?(command)
          return true if command =~ /^@/

          quote = nil
          env = false
          env_first_char = false

          command.dup.each_char do |c|
            case c
            when "'", '"'
              if (!quote)
                quote = c
              elsif quote == c
                quote = nil
              end
              next
            when '>', '<', '|', '&', "\n"
              return true unless quote
            when '%'
              return true if env
              env = env_first_char = true
              next
            else
              next unless env
              if env_first_char
                env_first_char = false
                env = false and next if c !~ /[A-Za-z_]/
              end
              env = false if c !~ /[A-Za-z1-9_]/
            end
          end
          return false
        end

        def self.pathext
          @pathext ||= ENV['PATHEXT'] ? ENV['PATHEXT'].split(';') + [''] : ['']
        end

        # which() mimicks the Unix which command
        # FIXME: it is not working
        def self.which(cmd)
          ENV['PATH'].split(File::PATH_SEPARATOR).each do |path|
            exe = find_executable("#{path}/#{cmd}")
            return exe if exe
          end
          return nil
        end

        # Windows has a different notion of what "executable" means
        # The OS will search through valid the extensions and look
        # for a binary there.
        def self.find_executable(path)
          return path if executable? path

          pathext.each do |ext|
            exe = "#{path}#{ext}"
            return exe if executable? exe
          end
          return nil
        end

        def self.executable?(path)
          File.executable?(path) && !File.directory?(path)
        end

        # recursively kills all child processes of given pid
        # calls itself querying for children child procs until
        # none remain. Important that a single WmiLite instance
        # is passed in since each creates its own WMI rpc process
        def self.kill_process_tree(pid, wmi, logger)
          wmi.query("select * from Win32_Process where ParentProcessID=#{pid}").each do |instance|
            child_pid = instance.wmi_ole_object.processid
            kill_process_tree(child_pid, wmi, logger)
            begin
              logger.debug([
                "killing child process #{child_pid}::",
                "#{instance.wmi_ole_object.Name} of parent #{pid}"
                ].join) if logger
              kill_process(instance)
            rescue Errno::EIO, SystemCallError
              logger.debug([
                "Failed to kill child process #{child_pid}::",
                "#{instance.wmi_ole_object.Name} of parent #{pid}"
              ].join) if logger
            end
          end
        end

        def self.kill_process(instance)
          Process.kill(:KILL, instance.wmi_ole_object.processid)
        end

        def self.format_process(process, app_name, command_line, timeout)
          msg = []
          msg << "ProcessId: #{process.process_id}"
          msg << "app_name: #{app_name}"
          msg << "command_line: #{command_line}"
          msg << "timeout: #{timeout}"
          msg.join("\n")
        end
      end
    end # class
  end
end
