def exec_remote_cmd(self, cmd, probe_cmd=None):
"""
Execute command(cmd or probe-command) on remote host
:param cmd: Command to execute
:param probe_cmd: Probe command to execute before 'cmd'
:return: None
"""
minion_id = self.get_value('minion', arg_type='query')
LOG.debug(f'[exec_remote_cmd] Minion ID: {minion_id}')
minion = MINIONS.get(minion_id)
args = minion['args']
LOG.debug(f'[exec_remote_cmd] Minion args: {args}')
self.ssh_transport_client = self.create_ssh_client(args)
# Use probe_cmd to detect file's existence
if probe_cmd:
chan = self.ssh_transport_client.get_transport().open_session()
chan.exec_command(probe_cmd)
ext = (chan.recv_exit_status())
if ext:
raise tornado.web.HTTPError(404, "Not found")
transport = self.ssh_transport_client.get_transport()
self.channel = transport.open_channel(kind='session')
self.channel.exec_command(cmd)
def open(self):
self.src_addr = self.get_client_endpoint()
LOG.info('Connected from {}:{}'.format(*self.src_addr))
try:
# Get id from query argument from
minion_id = self.get_value('id')
LOG.debug(f"############ minion id: {minion_id}")
minion = MINIONS.get(minion_id)
if not minion:
self.close(reason='Websocket failed.')
return
minion_obj = minion.get('minion', None)
if minion_obj:
# minions[minion_id]["minion"] = None
self.set_nodelay(True)
minion_obj.set_handler(self)
self.minion_ref = weakref.ref(minion_obj)
self.loop.add_handler(minion_obj.fd, minion_obj, IOLoop.READ)
else:
self.close(reason='Websocket authentication failed.')
except (tornado.web.MissingArgumentError, InvalidValueError) as err:
self.close(reason=str(err))
def _extract_filename(data: bytes) -> str:
LOG.debug(data)
ptn = re.compile(b'filename="(.*)"')
m = ptn.search(data)
if m:
name = m.group(1).decode()
else:
name = "untitled"
# Replace spaces with underscore
return re.sub(r'\s+', '_', name)
def get_args(self):
data = json_decode(self.request.body)
LOG.debug(data)
# Minion login won't pass hostname in form data
hostname = data.get("hostname", "localhost")
username = data["username"]
password = data["password"]
port = int(data["port"])
args = (hostname, port, username, password)
LOG.debug(f"Args for SSH: {args}")
return args
def close(self, reason=None):
if self.closed:
return
self.closed = True
LOG.info(f'Closing minion {self.id}: {reason}')
if self.handler:
self.loop.remove_handler(self.fd)
self.handler.close(reason=reason)
self.chan.close()
self.ssh.close()
LOG.info('Connection to {}:{} lost'.format(*self.dst_addr))
clear_minion(self)
LOG.debug(MINIONS)
async def data_received(self, data):
# A simple multipart/form-data
# b'------WebKitFormBoundarysiqXYmhALsFpsMuh\r\nContent-Disposition: form-data; name="upload";
# filename="hello.txt"\r\nContent-Type: text/plain\r\n\r\n
# hello\r\n\r\nworld\r\n\r\n------WebKitFormBoundarysiqXYmhALsFpsMuh--\r\n'
"""
:param data:
:return: None
"""
if not self.boundary:
self.boundary = self._get_boundary()
LOG.debug(f"multipart/form-data boundary: {self.boundary}")
# Split data with multipart/form-data boundary
sep = f'--{self.boundary}'
chunks = data.split(sep.encode('ISO-8859-1'))
chunks_len = len(chunks)
# DEBUG
# print("=====================================")
# print(f"Stream idx: {self.stream_idx}")
# print(f"CHUNKS length: {len(chunks)}")
# Data is small enough in one stream
if chunks_len == 3:
form_data_info, raw = self._partition_chunk(chunks[1])
self.filename = self._extract_filename(form_data_info)
await run_async_func(self.exec_remote_cmd,
f'cat > /tmp/{self.filename}')
await run_async_func(self._write_chunk, raw)
await run_async_func(self.ssh_transport_client.close)
else:
if self.stream_idx == 0:
form_data_info, raw = self._partition_chunk(chunks[1])
self.filename = self._extract_filename(form_data_info)
await run_async_func(self.exec_remote_cmd,
f'cat > /tmp/{self.filename}')
await run_async_func(self._write_chunk, raw)
else:
# Form data in the middle data stream
if chunks_len == 1:
await run_async_func(self._write_chunk, chunks[0])
else:
# 'chunks_len' == 2, the LAST stream
await run_async_func(self._write_chunk, chunks[0])
await run_async_func(self.ssh_transport_client.close)
self.stream_idx += 1
async def get(self):
chunk_size = 1024 * 1024 * 1 # 1 MiB
remote_file_path = self.get_value("filepath", arg_type="query")
filename = os.path.basename(remote_file_path)
LOG.debug(remote_file_path)
try:
self.exec_remote_cmd(cmd=f'cat {remote_file_path}',
probe_cmd=f'ls {remote_file_path}')
except tornado.web.HTTPError:
self.write(f'Not found: {remote_file_path}')
await self.finish()
return
self.set_header("Content-Type", "application/octet-stream")
self.set_header("Accept-Ranges", "bytes")
self.set_header("Content-Disposition",
f"attachment; filename={filename}")
while True:
chunk = self.channel.recv(chunk_size)
if not chunk:
break
try:
# Write the chunk to response
self.write(chunk)
# Send the chunk to client
await self.flush()
except tornado.iostream.StreamClosedError:
break
finally:
del chunk
await tornado.web.gen.sleep(0.000000001) # 1 nanosecond
self.ssh_transport_client.close()
try:
await self.finish()
except tornado.iostream.StreamClosedError as err:
LOG.error(err)
LOG.debug("Maybe user cancelled download")
LOG.info(f"Download ended: {remote_file_path}")
def on_message(self, message):
LOG.debug(f'{message} from {self.src_addr}')
minion = self.minion_ref()
try:
msg = json.loads(message)
except JSONDecodeError:
return
if not isinstance(msg, dict):
return
resize = msg.get('resize')
if resize and len(resize) == 2:
try:
minion.chan.resize_pty(*resize)
except (TypeError, struct.error, paramiko.SSHException):
pass
data = msg.get('data')
if data and isinstance(data, str):
minion.data_to_dst.append(data)
minion.do_write()
def do_write(self):
LOG.debug('minion {} on write'.format(self.id))
if not self.data_to_dst:
return
data = ''.join(self.data_to_dst)
LOG.debug(f'{data} to {self.dst_addr}')
try:
sent = self.chan.send(data)
except (OSError, IOError) as e:
LOG.error(e)
if errno_from_exception(e) in _ERRNO_CONNRESET:
self.close(reason='chan error on writing')
else:
self.update_handler(IOLoop.WRITE)
else:
self.data_to_dst = []
data = data[sent:]
if data:
self.data_to_dst.append(data)
self.update_handler(IOLoop.WRITE)
else:
self.update_handler(IOLoop.READ)
def do_read(self):
LOG.debug('minion {} on read'.format(self.id))
try:
data = self.chan.recv(self.BUFFER_SIZE)
except (OSError, IOError) as e:
LOG.error(e)
if errno_from_exception(e) in _ERRNO_CONNRESET:
self.close(reason='CHAN ERROR DOING READ ')
else:
LOG.debug(f'{data} from {self.dst_addr}')
if not data:
self.close(reason='BYE ~')
return
LOG.debug(f'{data} to {self.handler.src_addr}')
try:
self.handler.write_message(data, binary=True)
except tornado.websocket.WebSocketClosedError:
self.close(reason='WEBSOCKET CLOSED')
def get(self):
LOG.debug(f"MINIONS: {MINIONS}")
self.render('index.html', mode=conf.mode)