Implement attachment support

2 files changed, 32 insertions, 8 deletions

diff --git a/snekbox/memfs.py b/snekbox/memfs.py
index 50196ce..520954b 100644
--- a/snekbox/memfs.py
+++ b/snekbox/memfs.py
@@ -12,7 +12,14 @@ from types import TracebackType
 from typing import Type
 from uuid import uuid4
 
+from snekbox.snekio import FileAttachment
+
+# Size of the memory filesystem
 MEMFS_SIZE = "2G"
+# Maximum number of files attachments will be scanned for
+MAX_FILES = 6
+# Maximum size of a file attachment (64 MB)
+MAX_FILE_SIZE = 64 * 1024 * 1024
 
 log = logging.getLogger(__name__)
 
@@ -86,6 +93,14 @@ class MemoryTempDir:
         yield
         self.path.chmod(0o555)
 
+    def attachments(self) -> list[FileAttachment] | None:
+        """Return a list of attachments in the tempdir."""
+        # First look for any file named `output` (any extension)
+        output = next((f for f in self.home.glob("output*") if f.is_file()), None)
+        if output:
+            return [FileAttachment.from_path(output, MAX_FILE_SIZE)]
+        return None
+
     def cleanup(self) -> None:
         """Remove files in temp dir, releases name."""
         if self.path is None:
diff --git a/snekbox/nsjail.py b/snekbox/nsjail.py
index c7103bd..d8b642a 100644
--- a/snekbox/nsjail.py
+++ b/snekbox/nsjail.py
@@ -15,6 +15,8 @@ from snekbox.memfs import MemoryTempDir
 
 __all__ = ("NsJail",)
 
+from snekbox.snekio import FileAttachment
+
 log = logging.getLogger(__name__)
 
 # [level][timestamp][PID]? function_signature:line_no? message
@@ -131,7 +133,7 @@ class NsJail:
 
     def python3(
         self, code: str, *, nsjail_args: Iterable[str] = (), py_args: Iterable[str] = ("",)
-    ) -> CompletedProcess:
+    ) -> tuple[CompletedProcess, list[FileAttachment]]:
         """
         Execute Python 3 code in an isolated environment and return the completed process.
 
@@ -195,16 +197,19 @@ class NsJail:
                     args, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True
                 )
             except ValueError:
-                return CompletedProcess(args, None, "ValueError: embedded null byte", None)
+                return CompletedProcess(args, None, "ValueError: embedded null byte", None), []
 
             try:
                 output = self._consume_stdout(nsjail)
             except UnicodeDecodeError:
-                return CompletedProcess(
-                    args,
-                    None,
-                    "UnicodeDecodeError: invalid Unicode in output pipe",
-                    None,
+                return (
+                    CompletedProcess(
+                        args,
+                        None,
+                        "UnicodeDecodeError: invalid Unicode in output pipe",
+                        None,
+                    ),
+                    [],
                 )
 
             # When you send signal `N` to a subprocess to terminate it using Popen, it
@@ -212,6 +217,10 @@ class NsJail:
             # convert negative exit codes to the `N + 128` form.
             returncode = -nsjail.returncode + 128 if nsjail.returncode < 0 else nsjail.returncode
 
+            # Parse attachments
+            attachments = temp_dir.attachments() or []
+            log.info(f"Found {len(attachments)} attachments.")
+
             log_lines = nsj_log.read().decode("utf-8").splitlines()
             if not log_lines and returncode == 255:
                 # NsJail probably failed to parse arguments so log output will still be in stdout
@@ -221,4 +230,4 @@ class NsJail:
 
         log.info(f"nsjail return code: {returncode}")
 
-        return CompletedProcess(args, returncode, output, None)
+        return CompletedProcess(args, returncode, output, None), attachments