Use timeout and update TransactionProgress for APT calls

onionkit/container.py

@@ -3,9 +3,11 @@ from contextlib import contextmanager
 from pathlib import Path
 import shutil
 import time
+import subprocess
 
 import sh
 
+from onionkit import _
 from onionkit import CONTAINER_BASE, MACHINE_DIR, SYSTEMD_UNIT_DIR
 
 from typing import List
@@ -15,12 +17,18 @@ logger = getLogger(__name__)
 
 START_TIMEOUT = 5
 STOP_TIMEOUT = 5
+APT_UPDATE_TIMEOUT = 60
+APT_INSTALL_TIMEOUT = 60
 
 
 class ContainerError(Exception):
     pass
 
 
+class APTError(Exception):
+    pass
+
+
 UNIT_FILE_CONTENT = str(
     "[Service]\n"
     "ExecStart=/usr/bin/systemd-nspawn --quiet "
@@ -162,10 +170,10 @@ class ContainerManager(object):
     def install_packages(self, packages: List[str]):
         logger.debug("Installing packages in %r", self.name)
         # XXX: Use `apt-get -o APT::Status-Fd` to get progress
-        self.status_callback("updating package cache")
-        self.execute_command("/usr/bin/apt", "update")
-        self.status_callback("installing packages")
-        self.execute_command("/usr/bin/apt", "install", "-y", *packages)
+        self.status_callback(_("Updating package cache"))
+        self.execute_apt("update", timeout=APT_UPDATE_TIMEOUT)
+        self.status_callback(_("Downloading packages"))
+        self.execute_apt("install", "-y", *packages, timeout=APT_INSTALL_TIMEOUT)
 
     def bind(self, source: str, target: str):
         logger.debug("Bind mounting %r in %r", target, self.name)
@@ -193,10 +201,53 @@ class ContainerManager(object):
             return False
 
     def execute_command(self, *args):
-        # Could also use systemd-run instead
-        # XXX: Implement timeout
         sh.machinectl("shell", self.name, *args)
 
+    def execute_apt(self, *args, timeout):
+
+        def handle_line(line):
+            line = line.strip()
+            if not line:
+                return
+            logger.debug("APT stderr: %r", line)
+
+            # Parse APT error
+            if line.startswith('E:'):
+                error_message = line[2:]
+                raise APTError("APT command '%s' failed: %s" % (' '.join(command), error_message))
+
+            # Parse APT status information
+            # https://github.com/mvo5/apt/blob/master/README.progress-reporting
+            info = line.split(':')
+            if info[0] == "dlstatus":
+                self.progress_callback(int(info[2].split('.')[0]))
+            elif info[0] == "pmstatus":
+                self.status_callback(info[3])
+
+        exit_code = None
+        command = ["/usr/bin/apt-get", "--option=APT::Status-Fd=2"] + list(args)
+        p = subprocess.Popen(["systemd-run", "--pipe", "--wait", "-M", self.name] + command,
+                             stdout=subprocess.PIPE,
+                             stderr=subprocess.PIPE,
+                             universal_newlines=True,
+                             bufsize=1)
+        start_time = time.perf_counter()
+        while time.perf_counter() - start_time < timeout:
+            exit_code = p.poll()
+            if exit_code is not None:
+                break
+            handle_line(p.stderr.readline())
+            time.sleep(0.1)
+
+        if exit_code is None:
+            raise ContainerError("Failed to run APT command '%s' (timeout %s)" % (' '.join(command), timeout))
+
+        for l in p.stderr.readlines():
+            handle_line(l)
+
+        if exit_code != 0:
+            raise APTError("APT command '%s' failed with exit code %s" % (' '.join(command), exit_code))
+
     @contextmanager
     def run(self):
         self.start()