Implement the RunPyJob for eval/exec

Closes #23

Author: thatbudakguy

mqterm/jobs.py

@@ -1,5 +1,6 @@
 import asyncio
 import logging
+import os
 from binascii import hexlify
 from hashlib import sha256
 from io import BytesIO
@@ -22,6 +23,7 @@ def __init__(self, cmd, args=[], client_id="localhost", **kwargs):
                 f"Wrong number of arguments for {cmd}; expected {self.argc}"
             )
 
+        self.globals = kwargs.get("globals", {})
         self.cmd = cmd
         self.args = args
         self.client_id = client_id
@@ -45,13 +47,13 @@ def ready(self):
         return True
 
     @classmethod
-    def from_cmd(cls, cmd_str, client_id=None):
+    def from_cmd(cls, cmd_str, client_id=None, globals={}):
         """Create a job from a command string, e.g. 'get_file file1.txt'."""
         cmd, *args = cmd_str.split()
         if cmd not in COMMANDS:
             raise ValueError(f"Unknown command: '{cmd}'")
         job_cls = COMMANDS[cmd]
-        return job_cls(cmd, args, client_id)
+        return job_cls(cmd, args, client_id, globals=globals)
 
 
 class SequentialJob(Job):
@@ -282,6 +284,39 @@ async def reboot_callback(op):
         return BytesIO(msg.encode("utf-8"))
 
 
+class RunPyJob(Job):
+    """A job to evaluate Python script on the device."""
+
+    argc = 1
+
+    def output(self):
+        """Eval or exec given Python and return the result."""
+        expr = self.args[0]
+        try:
+            result = self.do_eval(expr)
+        except SyntaxError:  # Not an expression, try exec
+            result = self.do_exec(expr)
+        return BytesIO(result.encode("utf-8"))
+
+    def do_eval(self, expr):
+        """Evaluate a Python expression and return the result."""
+        op = compile(expr, "<string>", "eval")
+        result = eval(op, self.globals, None)
+        return repr(result)
+
+    def do_exec(self, expr):
+        """Execute a Python statement and return the output."""
+        out_buf = BytesIO()
+        old_term = os.dupterm(out_buf)
+        try:
+            op = compile(expr, "<string>", "exec")
+            exec(op, self.globals, None)
+            result = out_buf.getvalue()
+        finally:
+            os.dupterm(old_term)
+        return result
+
+
 # Map commands to associated job names
 COMMANDS = {
     "whoami": WhoAmIJob,
@@ -291,4 +326,5 @@ async def reboot_callback(op):
     "cp": PutFileJob,
     "ota": FirmwareUpdateJob,
     "reboot": RebootJob,
+    "eval": RunPyJob,
 }

mqterm/terminal.py

@@ -44,7 +44,11 @@ class MqttTerminal:
     BUFLEN = PKTLEN * 2  # payload size for MQTT messages
 
     def __init__(
-        self, mqtt_client, topic_prefix=None, logger=logging.getLogger("mqterm")
+        self,
+        mqtt_client,
+        topic_prefix=None,
+        logger=logging.getLogger("mqterm"),
+        globals={},
     ):
         self.mqtt_client = mqtt_client
         self.topic_prefix = topic_prefix
@@ -55,6 +59,7 @@ def __init__(
         self.out_view = memoryview(self.out_buffer)
         self.logger = logger
         self.jobs = {}
+        self.globals = globals
 
     async def connect(self):
         """Start processing messages in the input stream."""
@@ -102,7 +107,7 @@ async def update_job(self, client_id, seq, payload):
             self.logger.debug(f"Updated {job}, seq: {seq}")
         else:
             cmd = payload.decode("utf-8")
-            job = Job.from_cmd(cmd, client_id=client_id)
+            job = Job.from_cmd(cmd, client_id=client_id, globals=self.globals)
             self.jobs[client_id] = job
             self.logger.info(f"Created {job}")
 

tests/test_jobs.py

@@ -14,6 +14,7 @@
     PlatformInfoJob,
     PutFileJob,
     RebootJob,
+    RunPyJob,
     WhoAmIJob,
 )
 
@@ -230,3 +231,23 @@ def test_run_soft(self):
         job = RebootJob("reboot", ["soft"])
         output = job.output().read().decode("utf-8").strip()
         self.assertEqual(output, "Performing soft reboot")
+
+
+class TestRunPyJob(TestCase):
+    def test_eval(self):
+        """RunPyJob should evaluate a Python expression and return result"""
+        job = RunPyJob("eval", ["1 + 2"])
+        output = job.output().read().decode("utf-8").strip()
+        self.assertEqual(output, "3")
+
+    def test_globals(self):
+        """RunPyJob should use provided globals for evaluation"""
+        job = RunPyJob("eval", ["x + y"], globals={"x": 1, "y": 2})
+        output = job.output().read().decode("utf-8").strip()
+        self.assertEqual(output, "3")
+
+    def test_exec(self):
+        """RunPyJob should execute a Python script and return output"""
+        job = RunPyJob("eval", ['print("Hello, World!")'])
+        output = job.output().read().decode("utf-8").strip()
+        self.assertEqual(output, "Hello, World!")