#
# author: Dominik Richter
# author: Christoph Hartmann

require_relative "../plugins"
require_relative "../errors"
require "mixlib/shellout" unless defined?(Mixlib::ShellOut)
require "ostruct" unless defined?(OpenStruct)

module Train::Transports
  class Local < Train.plugin(1)
    name "local"

    class PipeError < Train::TransportError; end

    def connection(_ = nil)
      @connection ||= Connection.new(@options)
    end

    class Connection < BaseConnection
      def initialize(options)
        super(options)

        @runner = if options[:command_runner]
                    force_runner(options[:command_runner])
                  else
                    select_runner(options)
                  end
      end

      def login_command
        nil # none, open your shell
      end

      def close
        @runner.close
      end

      def uri
        "local://"
      end

      def upload(locals, remote)
        FileUtils.mkdir_p(remote)

        Array(locals).each do |local|
          FileUtils.cp_r(local, remote)
        end
      end

      def download(remotes, local)
        upload(remotes, local)
      end

      private

      def select_runner(options)
        if os.windows?
          # Force a 64 bit poweshell if needed
          if RUBY_PLATFORM == "i386-mingw32" && os.arch == "x86_64"
            powershell_cmd = "#{ENV["SystemRoot"]}\\sysnative\\WindowsPowerShell\\v1.0\\powershell.exe"
          else
            powershell_cmd = "powershell"
          end

          # Attempt to use a named pipe but fallback to ShellOut if that fails
          begin
            WindowsPipeRunner.new(powershell_cmd)
          rescue PipeError
            WindowsShellRunner.new(powershell_cmd)
          end
        else
          GenericRunner.new(self, options)
        end
      end

      def force_runner(command_runner)
        case command_runner
        when :generic
          GenericRunner.new(self, options)
        when :windows_pipe
          WindowsPipeRunner.new
        when :windows_shell
          WindowsShellRunner.new
        else
          raise "Runner type `#{command_runner}` not supported"
        end
      end

      def run_command_via_connection(cmd, opts, &_data_handler)
        # Use the runner if it is available
        return @runner.run_command(cmd, opts) if defined?(@runner)

        # If we don't have a runner, such as at the beginning of setting up the
        # transport and performing the first few steps of OS detection, fall
        # back to shelling out.
        res = Mixlib::ShellOut.new(cmd)
        res.run_command
        Local::CommandResult.new(res.stdout, res.stderr, res.exitstatus)
      rescue Errno::ENOENT => _
        CommandResult.new("", "", 1)
      end

      def file_via_connection(path)
        if os.windows?
          Train::File::Local::Windows.new(self, path)
        else
          Train::File::Local::Unix.new(self, path)
        end
      end

      class GenericRunner
        include_options Train::Extras::CommandWrapper

        def initialize(connection, options)
          @cmd_wrapper = Local::CommandWrapper.load(connection, options)
        end

        def run_command(cmd, opts = {})
          if defined?(@cmd_wrapper) && !@cmd_wrapper.nil?
            cmd = @cmd_wrapper.run(cmd)
          end

          res = Mixlib::ShellOut.new(cmd)
          res.timeout = opts[:timeout]
          begin
            res.run_command
          rescue Mixlib::ShellOut::CommandTimeout
            raise Train::CommandTimeoutReached
          end
          Local::CommandResult.new(res.stdout, res.stderr, res.exitstatus)
        end

        def close
          # nothing to do at the moment
        end
      end

      class WindowsShellRunner
        require "json" unless defined?(JSON)
        require "base64" unless defined?(Base64)

        def initialize(powershell_cmd = "powershell")
          @powershell_cmd = powershell_cmd
        end

        def run_command(script, opts)
          # Prevent progress stream from leaking into stderr
          script = "$ProgressPreference='SilentlyContinue';" + script

          # Encode script so PowerShell can use it
          script = script.encode("UTF-16LE", "UTF-8")
          base64_script = Base64.strict_encode64(script)

          cmd = "#{@powershell_cmd} -NoProfile -EncodedCommand #{base64_script}"

          res = Mixlib::ShellOut.new(cmd)
          res.timeout = opts[:timeout]
          begin
            res.run_command
          rescue Mixlib::ShellOut::CommandTimeout
            raise Train::CommandTimeoutReached
          end
          Local::CommandResult.new(res.stdout, res.stderr, res.exitstatus)
        end

        def close
          # nothing to do at the moment
        end
      end

      class WindowsPipeRunner
        require "json" unless defined?(JSON)
        require "base64" unless defined?(Base64)
        require "securerandom" unless defined?(SecureRandom)

        def initialize(powershell_cmd = "powershell")
          @powershell_cmd = powershell_cmd
          @server_pid = nil
          @pipe = acquire_pipe
          raise PipeError if @pipe.nil?
        end

        # @param  cmd The command to execute
        # @return Local::CommandResult with stdout, stderr and exitstatus
        #         Note that exitstatus ($?) in PowerShell is boolean, but we use a numeric exit code.
        #         A command that succeeds without setting an exit code will have exitstatus 0
        #         A command that exits with an exit code will have that value as exitstatus
        #         A command that fails (e.g. throws exception) before setting an exit code will have exitstatus 1
        def run_command(cmd, _opts)
          script = "$ProgressPreference='SilentlyContinue';" + cmd
          encoded_script = Base64.strict_encode64(script)
          # TODO: no way to safely implement timeouts here.
          begin
            @pipe.puts(encoded_script)
            @pipe.flush
          rescue Errno::EPIPE
            # Retry once if the pipe went away
            begin
              # Maybe the pipe went away, but the server didn't? Reset it, to get a clean start.
              close
            rescue Errno::EIO
              # Ignore - server already went away
            end
            @pipe = acquire_pipe
            raise PipeError if @pipe.nil?

            @pipe.puts(encoded_script)
            @pipe.flush
          end
          res = OpenStruct.new(JSON.parse(Base64.decode64(@pipe.readline)))
          Local::CommandResult.new(res.stdout, res.stderr, res.exitstatus)
        end

        def close
          Process.kill("KILL", @server_pid) unless @server_pid.nil?
          @server_pid = nil
        end

        private

        def acquire_pipe
          require "win32/process"
          pipe_name = "inspec_#{SecureRandom.hex}"

          @server_pid = start_pipe_server(pipe_name)

          # Ensure process is killed when the Train process exits
          at_exit { close rescue Errno::EIO }

          pipe = nil
          ownership_verified = false

          # PowerShell needs time to create pipe.
          100.times do
            unless ownership_verified
              owner, current_user, is_owner = pipe_owned_by_current_user?(pipe_name)
              if owner.nil?
                sleep 0.1
                next
              end

              unless is_owner
                raise PipeError, "Unauthorized user '#{current_user}' tried to connect to pipe '#{pipe_name}'. Pipe is owned by '#{owner}'."
              end

              ownership_verified = true
            end

            pipe = open("//./pipe/#{pipe_name}", "r+")
            break
          rescue PipeError
            raise
          rescue
            sleep 0.1
          end

          pipe
        end

        def start_pipe_server(pipe_name)
          require "win32/process"

          script = <<-EOF
            $ErrorActionPreference = 'Stop'
            $ProgressPreference = 'SilentlyContinue'
            $user = [System.Security.Principal.WindowsIdentity]::GetCurrent().Name
            $pipeSecurity = New-Object System.IO.Pipes.PipeSecurity
            $rule = New-Object System.IO.Pipes.PipeAccessRule($user, "FullControl", "Allow")
            $pipeSecurity.AddAccessRule($rule)
            $pipeServer = New-Object System.IO.Pipes.NamedPipeServerStream('#{pipe_name}', [System.IO.Pipes.PipeDirection]::InOut, 1, [System.IO.Pipes.PipeTransmissionMode]::Byte, [System.IO.Pipes.PipeOptions]::None, 4096, 4096, $pipeSecurity)

            try {
              $pipeServer.WaitForConnection()
            } catch {
              exit 1
            }

            $pipeReader = New-Object System.IO.StreamReader($pipeServer)
            $pipeWriter = New-Object System.IO.StreamWriter($pipeServer)

            # Create loop to receive and process user commands/scripts
            $clientConnected = $true
            while($clientConnected) {
              $input = $pipeReader.ReadLine()
              if ($input -eq $null) {
                $clientConnected = $false
                break
              }
              $command = [System.Text.Encoding]::UTF8.GetString([System.Convert]::FromBase64String($input))

              # Execute user command/script and convert result to JSON
              $scriptBlock = $ExecutionContext.InvokeCommand.NewScriptBlock($command)
              try {
                $stdout = & $scriptBlock | Out-String
                $exit_code = $LastExitCode
                if ($exit_code -eq $null)
                {
                  $exit_code = 0
                }
                $result = @{ 'stdout' = $stdout ; 'stderr' = ''; 'exitstatus' = $exit_code }
              } catch {
                $stderr = $_ | Out-String
                $exit_code = $LastExitCode
                $result = @{ 'stdout' = ''; 'stderr' = $stderr; 'exitstatus' = $exit_code }
              }
              $resultJSON = $result | ConvertTo-JSON

              # Encode JSON in Base64 and write to pipe
              $encodedResult = [System.Convert]::ToBase64String([System.Text.Encoding]::UTF8.GetBytes($resultJSON))
              try {
                $pipeWriter.WriteLine($encodedResult)
                $pipeWriter.Flush()
              } catch [System.IO.IOException] {
                # Pipe was closed by client, exit gracefully
                $clientConnected = $false
                break
              }
            }
          EOF

          utf8_script = script.encode("UTF-16LE", "UTF-8")
          base64_script = Base64.strict_encode64(utf8_script)
          cmd = "#{@powershell_cmd} -NoProfile -ExecutionPolicy bypass -NonInteractive -EncodedCommand #{base64_script}"
          Process.create(command_line: cmd).process_id
        end

        def current_windows_user
          user = `powershell -Command "[System.Security.Principal.WindowsIdentity]::GetCurrent().Name"`.strip
          if user.nil? || user.empty?
            user = `whoami`.strip
          end
          if user.nil? || user.empty?
            raise "Unable to determine current Windows user"
          end

          user
        end

        # Verify pipe ownership before connecting
        def pipe_owned_by_current_user?(pipe_name)
          exists = `powershell -Command "Test-Path \\\\.\\pipe\\#{pipe_name}"`.strip.downcase == "true"
          current_user = current_windows_user
          return [nil, current_user, false] unless exists

          owner = `powershell -Command "(Get-Acl \\\\.\\pipe\\#{pipe_name}).Owner" 2>&1`.strip
          is_owner = !owner.nil? && !current_user.nil? && owner.casecmp(current_user) == 0
          [owner, current_user, is_owner]
        end
      end
    end
  end
end
