fix: extract the tar from the output device

Previously, the --artifact file was the whole output device.

Sponsored-by: author

1 files changed, 40 insertions, 2 deletions

diff --git a/ambient-run b/ambient-run
index be52226..f9b03e1 100755
--- a/ambient-run
+++ b/ambient-run
@@ -6,6 +6,7 @@ import os
 import shutil
 import subprocess
 import sys
+import tarfile
 import tempfile
 
 
@@ -18,11 +19,44 @@ def parse_args():
     p.add_argument("--verbose", action="store_true")
     p.add_argument("--image", required=True)
     p.add_argument("--log", required=True)
-    p.add_argument("--artifact", required=True)
+    p.add_argument("--artifact")
     p.add_argument("source_tree")
     return p.parse_args()
 
 
+def extract_tar_from_tar(filename, output):
+    N = 1024
+    BLOCK = 10 * 1024
+
+    with open(filename, "rb") as f:
+        tar = tarfile.open(fileobj=f, mode="r:")
+        while True:
+            info = tar.next()
+            if info is None:
+                break
+
+        # We've read to the end of the final tar file entry.
+        num_bytes = tar.fileobj.tell()
+
+        with open(output, "wb") as o:
+            # Copy first num_bytes of input file to output file.
+
+            f.seek(0)
+            n = 0
+            while n < num_bytes:
+                data = f.read(N)
+                if data is None:
+                    break
+                o.write(data)
+                n += len(data)
+
+            # Pad output file with zeroes until next block border.
+
+            while (n % BLOCK) != 0:
+                o.write(b"\0")
+                n += 1
+
+
 def main():
     args = parse_args()
 
@@ -40,7 +74,7 @@ def main():
         logging.debug(f"created {tmp}")
 
         image = os.path.join(tmp, "vm.qcow2")
-        output_drive = args.artifact
+        output_drive = os.path.join(tmp, "artifacts.tar")
         src_tar = os.path.join(tmp, "src.tar")
         vars = os.path.join(tmp, "vars.fd")
 
@@ -88,6 +122,10 @@ def main():
                 sys.stderr.write(p.stderr.decode())
                 raise Exception("qemu-system failed")
 
+        if args.artifact is not None:
+            logging.info("extract artifacts, if any")
+            extract_tar_from_tar(output_drive, args.artifact)
+
         logging.info("Build finished")
     except Exception as e:
         sys.stderr.write(f"{e}\n")