[`shell`] Better buffering for the output sent to websockets.

This commit is contained in:
Fabio Manganiello 2023-10-17 20:12:48 +02:00
parent a652bd9df8
commit 81e99a0e22
Signed by: blacklight
GPG Key ID: D90FBA7F76362774
1 changed files with 21 additions and 8 deletions

View File

@ -1,7 +1,8 @@
from base64 import b64encode from base64 import b64encode
import json import json
import threading
import subprocess import subprocess
import threading
import time
from typing import Callable, Optional from typing import Callable, Optional
from uuid import uuid1 from uuid import uuid1
@ -54,16 +55,28 @@ class ShellPlugin(Plugin):
shell=True, shell=True,
stdout=subprocess.PIPE, stdout=subprocess.PIPE,
) as proc: ) as proc:
while proc.poll() is None: first_loop = True
if not proc.stdout:
break
while True: while proc.poll() is None and proc.stdout and proc.stdout.readable():
buf = proc.stdout.read(1024) buf = b''
if not buf: bufsize = 1024
while len(buf) < bufsize:
ch = proc.stdout.read(1)
if not ch:
break break
self._send_ws_output(cmd_id, buf) 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: finally:
self._send_ws_output(cmd_id, None) self._send_ws_output(cmd_id, None)