diff --git a/base-config.yaml b/base-config.yaml index a91e6184cf4c86cf48cb072a1ff7093c930e8f89..a664942cbd422afba07b39cc24f7f60f4e3f7539 100644 --- a/base-config.yaml +++ b/base-config.yaml @@ -1,5 +1,5 @@ # The message prefix to treat as exec commands -prefix: !exec +prefix: '!exec' # Whether or not to enable "userbot" mode, where commands that the bot's user # sends are handled and responded to with edits instead of replies. # This is intended to be used when you run the plugin on your own account. @@ -8,5 +8,49 @@ userbot: false # sandboxing in maubot or this plugin, keep this list small. whitelist: - '@user:example.com' -# Number of seconds to wait between output update edits. -output_interval: 5 + +output: + # Number of seconds to wait between output update edits. + interval: 5 + # Arguments for the Jinja2 template initialization. + template_args: + lstrip_blocks: true + trim_blocks: yes + # Plaintext output template. + plaintext: | + Input ({{ input }}): + {{ code }} + {% if output %} + + Output: + {{ output }} + {% endif %} + {% if return_value %} + + Return: + {{ return_value }} + {% endif %} + {% if duration %} + + Took {{ duration | round(3) }} seconds + {% else %} + + Running... + {% endif %} + # HTML output template. + html: | + <h4>Input</h4> + <pre><code class="language-{{ language }}">{{ code }}</code></pre> + {% if output %} + <h4>Output</h4> + <pre>{{ output }}</pre> + {% endif %} + {% if return_value %} + <h4>Return</h4> + <pre>{{ return_value }}</pre> + {% endif %} + {% if duration %} + <h4>Took {{ duration | round(3) }} seconds</h4> + {% else %} + <h4>Running...</h4> + {% endif %} diff --git a/exec/bot.py b/exec/bot.py index eb10c64c20be986a05cc61a60dc09e1eeb4e9b41..1cce8d701999ed87d753e5125a22b2ec8bc20296 100644 --- a/exec/bot.py +++ b/exec/bot.py @@ -16,8 +16,11 @@ from typing import Type, Set, Optional, Any from io import StringIO from time import time +from html import escape -from mautrix.types import EventType, UserID +from jinja2 import Template + +from mautrix.types import EventType, UserID, TextMessageEventContent, MessageType, Format from mautrix.util.config import BaseProxyConfig, ConfigUpdateHelper from mautrix.util.formatter import MatrixParser, EntityString, SimpleEntity, EntityType from maubot import Plugin, MessageEvent @@ -35,7 +38,10 @@ class Config(BaseProxyConfig): helper.copy("prefix") helper.copy("userbot") helper.copy("whitelist") - helper.copy("output_interval") + helper.copy("output.interval") + helper.copy("output.template_args") + helper.copy("output.plaintext") + helper.copy("output.html") class ExecBot(Plugin): @@ -43,6 +49,8 @@ class ExecBot(Plugin): userbot: bool prefix: str output_interval: int + plaintext_template: Template + html_template: Template @classmethod def get_config_class(cls) -> Type[BaseProxyConfig]: @@ -56,7 +64,24 @@ class ExecBot(Plugin): self.whitelist = set(self.config["whitelist"]) self.userbot = self.config["userbot"] self.prefix = self.config["prefix"] - self.output_interval = self.config["output_interval"] + self.output_interval = self.config["output.interval"] + template_args = self.config["output.template_args"] + self.plaintext_template = Template(self.config["output.plaintext"], **template_args) + self.html_template = Template(self.config["output.html"], **template_args) + + def format_status(self, code: str, language: str, output: str = "", output_html: str = "", + return_value: Any = None, duration: Optional[float] = None, + msgtype: MessageType = MessageType.NOTICE) -> TextMessageEventContent: + return_value = repr(return_value) if return_value else '' + content = TextMessageEventContent( + msgtype=msgtype, format=Format.HTML, + body=self.plaintext_template.render( + code=code, language=language, output=output, return_value=return_value, + duration=duration), + formatted_body=self.html_template.render( + code=escape(code), language=language, output=output_html, + return_value=escape(return_value), duration=duration)) + return content @event.on(EventType.ROOM_MESSAGE) async def exec(self, evt: MessageEvent) -> None: @@ -76,7 +101,7 @@ class ExecBot(Plugin): if entity.type != EntityType.PREFORMATTED: continue current_lang = entity.extra_info["language"].lower() - value = command.text[entity.offset:entity.offset+entity.length] + value = command.text[entity.offset:entity.offset + entity.length] if not code: code = value lang = current_lang @@ -89,16 +114,42 @@ class ExecBot(Plugin): await evt.respond("Only python is currently supported") return + if self.userbot: + msgtype = MessageType.TEXT + content = self.format_status(code, lang, msgtype=msgtype) + await evt.edit(content) + output_event_id = evt.event_id + else: + msgtype = MessageType.NOTICE + content = self.format_status(code, lang, msgtype=msgtype) + output_event_id = await evt.respond(content) + runner = PythonRunner() - stdout = StringIO() - stderr = StringIO() + output = StringIO() + output_html = StringIO() return_value: Any = None start_time = time() + prev_output = start_time async for out_type, data in runner.run(code, stdin): if out_type == OutputType.STDOUT: - stdout.write(data) + output.write(data) + output_html.write(escape(data)) elif out_type == OutputType.STDERR: - stderr.write(data) + output.write(data) + output_html.write(f'<font color="red" data-mx-color="red">{escape(data)}</font>') elif out_type == OutputType.RETURN: return_value = data + continue + + cur_time = time() + if prev_output + self.output_interval < cur_time: + content = self.format_status(code, lang, output.getvalue(), output_html.getvalue(), + msgtype=msgtype) + content.set_edit(output_event_id) + await self.client.send_message(evt.room_id, content) + prev_output = cur_time duration = time() - start_time + content = self.format_status(code, lang, output.getvalue(), output_html.getvalue(), + return_value, duration, msgtype=msgtype) + content.set_edit(output_event_id) + await self.client.send_message(evt.room_id, content) diff --git a/exec/runner/python.py b/exec/runner/python.py index 65435bda65848979864b5e71add3301f2d97eb9b..7c4e2383177b62a2780f4ab9aaf4caba2ee1ae7d 100644 --- a/exec/runner/python.py +++ b/exec/runner/python.py @@ -37,6 +37,7 @@ class AsyncTextOutput: self.read_task = None self.queue = asyncio.Queue(loop=self.loop) self.closed = False + self.writers = {} def __aiter__(self) -> 'AsyncTextOutput': return self