22
33class Net ::SSH ::CommandStream
44
5- attr_accessor :channel , :thread , :error , :ssh
5+ attr_accessor :channel , :thread , :error , :ssh , :session , :logger
66 attr_accessor :lsock , :rsock , :monitor
77
88 module PeerInfo
@@ -13,7 +13,8 @@ module PeerInfo
1313
1414 def shell_requested ( channel , success )
1515 unless success
16- raise Net ::SSH ::ChannelRequestFailed , 'Shell/exec channel request failed'
16+ error = Net ::SSH ::ChannelRequestFailed . new ( 'Shell/exec channel request failed' )
17+ handle_error ( error : error )
1718 end
1819
1920 self . channel = channel
@@ -40,7 +41,9 @@ def shell_requested(channel, success)
4041 end
4142 end
4243
43- def initialize ( ssh , cmd = nil , pty : false , cleanup : false )
44+ def initialize ( ssh , cmd = nil , pty : false , cleanup : false , session : nil , logger : nil )
45+ self . session = session
46+ self . logger = logger
4447 self . lsock , self . rsock = Rex ::Socket . tcp_socket_pair ( )
4548 self . lsock . extend ( Rex ::IO ::Stream )
4649 self . lsock . extend ( PeerInfo )
@@ -74,31 +77,40 @@ def initialize(ssh, cmd = nil, pty: false, cleanup: false)
7477 end
7578
7679 channel . on_open_failed do |ch , code , desc |
77- raise Net ::SSH ::ChannelOpenFailed . new ( code , 'Session channel open failed' )
80+ error = Net ::SSH ::ChannelOpenFailed . new ( code , 'Session channel open failed' )
81+ handle_error ( error : error )
7882 end
7983
8084 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
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 )
8797 end
8898 end
8999
90- while true
91- rssh . process ( 0.5 ) { true }
100+ begin
101+ Kernel . loop { rssh . process ( 0.5 ) { true } }
102+ rescue ::StandardError => e
103+ handle_error ( error : e )
92104 end
93105
94106 # Shut down the SSH session if requested
95107 if !rcmd . nil? && rcleanup
96108 rssh . close
97109 end
98110 end
99- rescue ::Exception => e
111+ rescue ::StandardError => e
100112 # XXX: This won't be set UNTIL there's a failure from a thread
101- self . error = e
113+ handle_error ( error : e )
102114 ensure
103115 self . monitor . kill if self . monitor
104116 end
@@ -113,7 +125,18 @@ def verify_channel
113125 end
114126 end
115127
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+
116138 def cleanup
139+ self . session . alive = false if self . session
117140 self . monitor . kill
118141 self . lsock . close rescue nil
119142 self . rsock . close rescue nil
0 commit comments