133 lines
4.2 KiB
Python
133 lines
4.2 KiB
Python
from base64 import b64encode
|
|
import json
|
|
import subprocess
|
|
import threading
|
|
import time
|
|
from typing import Callable, Optional
|
|
from uuid import uuid1
|
|
|
|
from platypush.common.cmd_stream import redis_topic
|
|
from platypush.plugins import Plugin, action
|
|
from platypush.utils import get_redis
|
|
|
|
|
|
class ShellPlugin(Plugin):
|
|
"""
|
|
Plugin to run shell commands.
|
|
"""
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
|
|
def _exec(self, func: Callable, cmd: str, ignore_errors: bool = False):
|
|
try:
|
|
return func(cmd)
|
|
except subprocess.CalledProcessError as e:
|
|
if ignore_errors:
|
|
self.logger.warning(
|
|
'Command %s failed with error: %s', cmd, e.output.decode()
|
|
)
|
|
return None
|
|
|
|
raise RuntimeError(e.output.decode()) from e
|
|
|
|
def _exec_simple(self, cmd: str):
|
|
return subprocess.check_output(
|
|
cmd, stderr=subprocess.STDOUT, shell=True
|
|
).decode()
|
|
|
|
@staticmethod
|
|
def _send_ws_output(cmd_id: str, buf: Optional[bytes]):
|
|
get_redis().publish(
|
|
redis_topic,
|
|
json.dumps(
|
|
{
|
|
'id': cmd_id,
|
|
'output': (b64encode(buf).decode() if buf is not None else None),
|
|
}
|
|
),
|
|
)
|
|
|
|
def _exec_ws_thread(self, cmd: str, cmd_id: str):
|
|
try:
|
|
with subprocess.Popen(
|
|
cmd,
|
|
shell=True,
|
|
stdout=subprocess.PIPE,
|
|
) as proc:
|
|
first_loop = True
|
|
|
|
while proc.poll() is None and proc.stdout and proc.stdout.readable():
|
|
buf = b''
|
|
bufsize = 1024
|
|
|
|
while len(buf) < bufsize:
|
|
ch = proc.stdout.read(1)
|
|
if not ch:
|
|
break
|
|
|
|
buf += ch
|
|
if ch in (b'\r', b'\n'):
|
|
break
|
|
|
|
if first_loop:
|
|
first_loop = False
|
|
# An artificial delay to let the client connect to the websocket,
|
|
# otherwise the first few lines of output might be lost.
|
|
time.sleep(2)
|
|
|
|
self._send_ws_output(cmd_id, buf)
|
|
finally:
|
|
self._send_ws_output(cmd_id, None)
|
|
|
|
def _exec_ws(self, cmd: str):
|
|
cmd_id = str(uuid1())
|
|
threading.Thread(target=self._exec_ws_thread, args=(cmd, cmd_id)).start()
|
|
return {'ws_path': f'/ws/shell?id={cmd_id}'}
|
|
|
|
@action
|
|
def exec(
|
|
self,
|
|
cmd: str,
|
|
background: bool = False,
|
|
ws: bool = False,
|
|
ignore_errors: bool = False,
|
|
):
|
|
"""
|
|
Run a command.
|
|
|
|
:param cmd: Command to execute
|
|
:param background: If set to True, execute the process in the
|
|
background, otherwise wait for the process termination and return
|
|
its output (default: False).
|
|
:param ignore_errors: If set, then any errors in the command execution
|
|
will be ignored. Otherwise a RuntimeError will be thrown (default
|
|
value: False)
|
|
:param ws: If set to True then the output of the command will be
|
|
sent asynchronously over a websocket channel (default: False).
|
|
In this case, the method will return a response in the format:
|
|
|
|
.. code-block:: json
|
|
|
|
{
|
|
"ws_path": "/ws/shell?id=<cmd_id>"
|
|
}
|
|
|
|
Where ``ws_path`` is the websocket path where the output of the
|
|
command will be sent. The websocket will be closed when the command
|
|
terminates.
|
|
|
|
:returns: A response object where the ``output`` field will contain the
|
|
command output as a string, and the ``errors`` field will contain
|
|
whatever was sent to stderr.
|
|
"""
|
|
if background:
|
|
subprocess.Popen(cmd, shell=True)
|
|
return None
|
|
|
|
func = self._exec_ws if ws else self._exec_simple
|
|
return self._exec(func, cmd, ignore_errors=ignore_errors)
|
|
|
|
|
|
# vim:sw=4:ts=4:et:
|