1
# -*- coding: binary -*-
2

3
class Net::SSH::CommandStream
4

5
  attr_accessor :channel, :thread, :error, :ssh, :session, :logger
6
  attr_accessor :lsock, :rsock, :monitor
7

8
  module PeerInfo
9
    include ::Rex::IO::Stream
10
    attr_accessor :peerinfo
11
    attr_accessor :localinfo
12
  end
13

14
  def shell_requested(channel, success)
15
    unless success
16
      error = Net::SSH::ChannelRequestFailed.new('Shell/exec channel request failed')
17
      handle_error(error: error)
18
    end
19

20
    self.channel = channel
21

22
    channel[:data] = ''
23
    channel[:extended_data] = ''
24

25
    channel.on_eof do
26
      cleanup
27
    end
28

29
    channel.on_close do
30
      cleanup
31
    end
32

33
    channel.on_data do |ch, data|
34
      self.rsock.write(data)
35
      channel[:data] << data
36
    end
37

38
    channel.on_extended_data do |ch, ctype, data|
39
      self.rsock.write(data)
40
      channel[:extended_data] << data
41
    end
42
  end
43

44
  def initialize(ssh, cmd = nil, pty: false, cleanup: false, session: nil, logger: nil)
45
    self.session = session
46
    self.logger = logger
47
    self.lsock, self.rsock = Rex::Socket.tcp_socket_pair()
48
    self.lsock.extend(Rex::IO::Stream)
49
    self.lsock.extend(PeerInfo)
50
    self.rsock.extend(Rex::IO::Stream)
51

52
    self.ssh = ssh
53
    self.thread = Thread.new(ssh, cmd, pty, cleanup) do |rssh, rcmd, rpty, rcleanup|
54
      info = rssh.transport.socket.getpeername_as_array
55
      if Rex::Socket.is_ipv6?(info[1])
56
        self.lsock.peerinfo = "[#{info[1]}]:#{info[2]}"
57
      else
58
        self.lsock.peerinfo = "#{info[1]}:#{info[2]}"
59
      end
60

61
      info = rssh.transport.socket.getsockname
62
      if Rex::Socket.is_ipv6?(info[1])
63
        self.lsock.localinfo = "[#{info[1]}]:#{info[2]}"
64
      else
65
        self.lsock.localinfo = "#{info[1]}:#{info[2]}"
66
      end
67

68
      channel = rssh.open_channel do |rch|
69
        # A PTY will write us to {u,w}tmp and lastlog
70
        rch.request_pty if rpty
71

72
        if rcmd.nil?
73
          rch.send_channel_request('shell', &method(:shell_requested))
74
        else
75
          rch.exec(rcmd, &method(:shell_requested))
76
        end
77
      end
78

79
      channel.on_open_failed do |ch, code, desc|
80
        error = Net::SSH::ChannelOpenFailed.new(code, 'Session channel open failed')
81
        handle_error(error: error)
82
      end
83

84
      self.monitor = Thread.new do
85
        begin
86
          Kernel.loop do
87
            next if not self.rsock.has_read_data?(1.0)
88

89
            buff = self.rsock.read(16384)
90
            break if not buff
91

92
            verify_channel
93
            self.channel.send_data(buff) if buff
94
          end
95
        rescue ::StandardError => e
96
          handle_error(error: e)
97
        end
98
      end
99

100
      begin
101
        Kernel.loop { rssh.process(0.5) { true } }
102
      rescue ::StandardError => e
103
        handle_error(error: e)
104
      end
105

106
      # Shut down the SSH session if requested
107
      if !rcmd.nil? && rcleanup
108
        rssh.close
109
      end
110
    end
111
  rescue ::StandardError => e
112
    # XXX: This won't be set UNTIL there's a failure from a thread
113
    handle_error(error: e)
114
  ensure
115
    self.monitor.kill if self.monitor
116
  end
117

118
  #
119
  # Prevent a race condition
120
  #
121
  def verify_channel
122
    while ! self.channel
123
      raise EOFError if ! self.thread.alive?
124
      ::IO.select(nil, nil, nil, 0.10)
125
    end
126
  end
127

128
  def handle_error(error: nil)
129
    self.error = error if error
130

131
    if self.logger
132
      self.logger.print_error("SSH Command Stream encountered an error: #{self.error} (Server Version: #{self.ssh.transport.server_version.version})")
133
    end
134

135
    cleanup
136
  end
137

138
  def cleanup
139
    self.session.alive = false if self.session
140
    self.monitor.kill
141
    self.lsock.close rescue nil
142
    self.rsock.close rescue nil
143
    self.ssh.close rescue nil
144
    self.thread.kill
145
  end
146

147
end
148

149