CoCalc provides the best real-time collaborative environment for Jupyter Notebooks, LaTeX documents, and SageMath, scalable from individual users to large groups and classes!

1
require 'winrm'
2
require 'winrm/wsmv/write_stdin'
3
require 'net/winrm/ctrl_c'
4
require 'net/winrm/receive_response_reader'
5

6
module Net
7
  module MsfWinRM
8
    # WinRM shell to use stdin, rather than sending isolated commands
9
    class StdinShell < WinRM::Shells::Cmd
10
      # We create our own empty finalizers because the built-in one triggers a
11
      # request using the Rex HTTP client, which segfaults; possibly because it
12
      # creates a thread, or something else that is not allowed in a finalizer.
13
      # In this situation (observed only when the user quits MSF with active sessions),
14
      # we'll just let the shell continue.
15
      def remove_finalizer; end
16

17
      def add_finalizer; end
18

19
      def send_command(command, arguments = [])
20
        open unless shell_id
21
        super(command, arguments)
22
      end
23

24
      # Runs a shell command synchronously, and returns the output
25
      def shell_command_synchronous(command, args, timeout)
26
        start_time = Process.clock_gettime(Process::CLOCK_MONOTONIC, :millisecond)
27
        command_id = send_command(command, args)
28
        buffer = []
29
        begin
30
          while (Process.clock_gettime(Process::CLOCK_MONOTONIC, :millisecond) - start_time) < (timeout * 1000)
31
            read_stdout(command_id) do |stdout, stderr|
32
              buffer << stdout if stdout
33
              buffer << stderr if stderr
34
            end
35
          end
36
        rescue EOFError
37
          # Shell terminated of its own accord
38
        ensure
39
          cleanup_command(command_id)
40
        end
41
        buffer.join('')
42
      end
43

44
      # Runs the specified command with optional arguments
45
      # @param block [&block] The optional callback for any realtime output
46
      # @yieldparam [string] standard out response text
47
      # @yieldparam [string] standard error response text
48
      # @yieldreturn [WinRM::Output] The command output
49
      def read_stdout(command_id, &block)
50
        open unless shell_id
51
        begin
52
          response_reader.read_output(command_output_message(shell_id, command_id), &block)
53
        rescue WinRM::WinRMWSManFault => e
54
          # If no output is available before the wsman:OperationTimeout expires,
55
          # the server MUST return a WSManFault with the Code attribute equal to
56
          # 2150858793. When the client receives this fault, it SHOULD issue
57
          # another Receive request.
58
          # http://msdn.microsoft.com/en-us/library/cc251676.aspx
59
          if e.fault_code == '2150858793'
60
            yield nil, nil
61
          else
62
            raise
63
          end
64
        end
65
      end
66

67
      def send_ctrl_c(command_id)
68
        ctrl_c_msg = CtrlC.new(
69
          connection_opts,
70
          shell_uri: shell_uri,
71
          shell_id: shell_id,
72
          command_id: command_id
73
        )
74
        transport.send_request(ctrl_c_msg.build)
75
      end
76

77
      def send_stdin(input, command_id)
78
        open unless shell_id
79

80
        stdin_msg = WinRM::WSMV::WriteStdin.new(
81
          connection_opts,
82
          shell_uri: shell_uri,
83
          shell_id: shell_id,
84
          command_id: command_id,
85
          stdin: input
86
        )
87
        result = transport.send_request(stdin_msg.build)
88
        result
89
      rescue WinRM::WinRMWSManFault => e
90
        raise unless [ERROR_OPERATION_ABORTED, SHELL_NOT_FOUND].include?(e.fault_code)
91
      rescue WinRM::WinRMHTTPTransportError => e
92
        # dont let the cleanup raise so we dont lose any errors from the command
93
        logger.info("[WinRM] #{e.status_code} returned in cleanup with error: #{e.message}")
94
      end
95

96
      def response_reader
97
        @response_reader ||= ReceiveResponseReader.new(transport, logger)
98
      end
99

100
      def open_shell
101
        msg = WinRM::WSMV::CreateShell.new(connection_opts, shell_opts)
102
        resp_doc = transport.send_request(msg.build)
103
        match = REXML::XPath.first(resp_doc, '//rsp:Owner')
104
        self.owner = match.text if match
105
        REXML::XPath.first(resp_doc, "//*[@Name='ShellId']").text
106
      end
107

108
      attr_accessor :owner
109

110
    end
111
  end
112
end
113

114