commit c5701fe8c670344e019fa225051f2cd4ea50a525 Author: Jake Howard Date: Thu May 18 23:55:17 2023 +0100 Init something which kinda works diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..7abfd0e --- /dev/null +++ b/.gitignore @@ -0,0 +1,180 @@ +# Created by https://www.toptal.com/developers/gitignore/api/python +# Edit at https://www.toptal.com/developers/gitignore?templates=python + +### Python ### +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ +cover/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +.pybuilder/ +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +# For a library or package, you might want to ignore these files since the code is +# intended to run in multiple environments; otherwise, check them in: +# .python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# poetry +# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control +#poetry.lock + +# pdm +# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. +#pdm.lock +# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it +# in version control. +# https://pdm.fming.dev/#use-with-ide +.pdm.toml + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ + +# PyCharm +# JetBrains specific template is maintained in a separate JetBrains.gitignore that can +# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore +# and can be added to the global gitignore or merged into this file. For a more nuclear +# option (not recommended) you can uncomment the following to ignore the entire idea folder. +#.idea/ + +### Python Patch ### +# Poetry local configuration file - https://python-poetry.org/docs/configuration/#local-configuration +poetry.toml + +# ruff +.ruff_cache/ + +# LSP config files +pyrightconfig.json + +# End of https://www.toptal.com/developers/gitignore/api/python + +runusb.sock +global-metadata.json +log.txt diff --git a/main.py b/main.py new file mode 100644 index 0000000..57ca9be --- /dev/null +++ b/main.py @@ -0,0 +1,6 @@ +import time +from itertools import count + +for i in count(): + print("Still running", i) + time.sleep(1) diff --git a/metadata.json b/metadata.json new file mode 100644 index 0000000..e63d37b --- /dev/null +++ b/metadata.json @@ -0,0 +1,3 @@ +{ + "foo": "bar" +} diff --git a/runusb.py b/runusb.py new file mode 100644 index 0000000..8bba5a4 --- /dev/null +++ b/runusb.py @@ -0,0 +1,169 @@ +import subprocess +import threading +from pathlib import Path +import logging +import os +import time +import atexit +import json +import zmq + +GLOBAL_METADATA = Path.cwd() / "global-metadata.json" +ZMQ_SOCKET = Path.cwd() / "runusb.sock" + +def get_zmq_publisher(): + context = zmq.Context() + socket = context.socket(zmq.PUB) + socket.bind("ipc://" + str(ZMQ_SOCKET)) + return socket + +def get_zmq_subscriber(topic: str): + context = zmq.Context() + socket = context.socket(zmq.SUB) + socket.setsockopt_string(zmq.SUBSCRIBE, topic) + socket.connect("ipc://" + str(ZMQ_SOCKET)) + return socket + +class UserCodeSupervisor(threading.Thread): + def __init__(self, process: subprocess.Popen, mountpoint: Path): + super().__init__() + self.process = process + self.logfile = mountpoint / "log.txt" + self.zmq_publisher = get_zmq_publisher() + + def run(self): + with self.logfile.open(mode="w") as log_file: + while self.process.returncode is None: + stdout_data = self.process.stdout.readline() + print(stdout_data, end="") + log_file.write(stdout_data) + log_file.flush() + + self.zmq_publisher.send_multipart([b"log", stdout_data.encode()]) + + self.process.poll() + + + +class RunUSBRegistry: + process: subprocess.Popen | None = None + usercode_dir: Path | None = None + + def __init__(self): + self.lock = threading.Lock() + self.metadata = {} + self.update_metadata({}) + + def start_user_code(self, mountpoint: Path): + self.usercode_dir = mountpoint + + entrypoint = mountpoint / "main.py" + self.process = subprocess.Popen( + ["python3", entrypoint], + env={ + **os.environ, + "PYTHONUNBUFFERED": "1", + "METADATA_FILE": str(GLOBAL_METADATA) + }, + cwd=mountpoint, + stdin=subprocess.DEVNULL, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + universal_newlines=True, + ) + + UserCodeSupervisor(self.process, mountpoint).start() + + def restart_user_code(self): + mountpoint = self.usercode_dir + + self.stop_user_code() + self.start_user_code(mountpoint) + + def stop_user_code(self): + self.process.terminate() + + try: + self.process.wait(timeout=5) + except subprocess.TimeoutExpired: + pass + + if self.process.poll() is not None: + self.process.kill() + + # Remove the reference + self.process = None + self.usercode_dir = None + + def handle_metadata_file(self, metadata_file: Path): + with metadata_file.open(mode="r") as f: + self.metadata = json.load(f) + + # Notify with a blank update, as everything is new + self.update_metadata({}) + + def update_metadata(self, new_metadata: dict): + self.metadata.update(new_metadata) + + with GLOBAL_METADATA.open("w") as f: + json.dump(self.metadata, f) + + def __enter__(self): + self.lock.acquire() + + def __exit__(self, *args): + self.lock.release() + + def close(self): + self.stop_user_code() + + +class BusHandler(threading.Thread): + def __init__(self, registry: RunUSBRegistry): + super().__init__() + self.registry = registry + self.zmq_subscriber = get_zmq_subscriber("") + + def run(self): + while True: + topic, message = self.zmq_subscriber.recv_multipart() + + topic = topic.decode() + message = message.decode() + + # if topic == "log": + # continue + + if topic == "restart": + with registry: + registry.restart_user_code() + + else: + print("Ignoring message with topic", topic, message.strip()) + + + +def watch_drive_activity(registry: RunUSBRegistry): + with registry: + registry.start_user_code(Path.cwd()) + +def watch_metadata(registry: RunUSBRegistry): + with registry: + registry.handle_metadata_file(Path.cwd() / "metadata.json") + + +def main(): + logging.basicConfig(level=logging.DEBUG) + + registry = RunUSBRegistry() + atexit.register(registry.close) + + BusHandler(registry).start() + + watch_drive_activity(registry) + watch_metadata(registry) + + # Threads are running now + +if __name__ == '__main__': + main() diff --git a/runusbctl.py b/runusbctl.py new file mode 100644 index 0000000..e66a192 --- /dev/null +++ b/runusbctl.py @@ -0,0 +1,13 @@ +import zmq +from pathlib import Path +import sys +import time + +context = zmq.Context() + +socket = context.socket(zmq.PUB) +socket.connect("ipc://" + str(Path.cwd() / "runusb.sock")) + +# time.sleep(3600) + +socket.send_multipart([sys.argv[1].encode(), b"foo"]) diff --git a/zmq_watch.py b/zmq_watch.py new file mode 100644 index 0000000..a2a2f24 --- /dev/null +++ b/zmq_watch.py @@ -0,0 +1,11 @@ +import zmq +from pathlib import Path + +context = zmq.Context() + +socket = context.socket(zmq.SUB) +socket.setsockopt_string(zmq.SUBSCRIBE, "log") +socket.connect("ipc://" + str(Path.cwd() / "runusb.sock")) + +while True: + print(socket.recv_multipart())