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
# -*- coding: binary -*-
2

3
class Net::SSH::CommandStream
4

5
  attr_accessor :channel, :thread, :error, :ssh
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
      raise Net::SSH::ChannelRequestFailed, 'Shell/exec channel request failed'
17
    end
18

19
    self.channel = channel
20

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

24
    channel.on_eof do
25
      cleanup
26
    end
27

28
    channel.on_close do
29
      cleanup
30
    end
31

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

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

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

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

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

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

69
        if rcmd.nil?
70
          rch.send_channel_request('shell', &method(:shell_requested))
71
        else
72
          rch.exec(rcmd, &method(:shell_requested))
73
        end
74
      end
75

76
      channel.on_open_failed do |ch, code, desc|
77
        raise Net::SSH::ChannelOpenFailed.new(code, 'Session channel open failed')
78
      end
79

80
      self.monitor = Thread.new do
81
        while(true)
82
          next if not self.rsock.has_read_data?(1.0)
83
          buff = self.rsock.read(16384)
84
          break if not buff
85
          verify_channel
86
          self.channel.send_data(buff) if buff
87
        end
88
      end
89

90
      while true
91
        rssh.process(0.5) { true }
92
      end
93

94
      # Shut down the SSH session if requested
95
      if !rcmd.nil? && rcleanup
96
        rssh.close
97
      end
98
    end
99
  rescue ::Exception => e
100
    # XXX: This won't be set UNTIL there's a failure from a thread
101
    self.error = e
102
  ensure
103
    self.monitor.kill if self.monitor
104
  end
105

106
  #
107
  # Prevent a race condition
108
  #
109
  def verify_channel
110
    while ! self.channel
111
      raise EOFError if ! self.thread.alive?
112
      ::IO.select(nil, nil, nil, 0.10)
113
    end
114
  end
115

116
  def cleanup
117
    self.monitor.kill
118
    self.lsock.close rescue nil
119
    self.rsock.close rescue nil
120
    self.ssh.close rescue nil
121
    self.thread.kill
122
  end
123

124
end
125

126