snap-seed-parse.py: Update to allow parsing uc20-style seeds. (LP: #2028984)

debian/changelog

@@ -3,6 +3,8 @@ livecd-rootfs (23.10.24) UNRELEASED; urgency=medium
   * update-source-catalog: Fix case where a variaton does not point at the
     base layer (i.e. most builds) (LP: #2033168) 
   * Configure universe sources in canary ISO. (LP: #2033109)
+  * snap-seed-parse.py: Update to allow parsing uc20-style seeds.
+    (LP: #2028984)
 
  -- Michael Hudson-Doyle <michael.hudson@ubuntu.com>  Mon, 28 Aug 2023 10:39:52 +1200
 

live-build/snap-seed-parse.py

@@ -10,6 +10,7 @@ The $chroot_dir argument is optional and will default to the empty string.
 """
 
 import argparse
+import glob
 import os.path
 import re
 import yaml
@@ -32,29 +33,25 @@ CHROOT_ROOT = ARGS.chroot
 FNAME = ARGS.file
 
 # Trim any trailing slashes for correct appending
+CHROOT_ROOT = CHROOT_ROOT.rstrip('/')
 log("CHROOT_ROOT: {}".format(CHROOT_ROOT))
-if len(CHROOT_ROOT) > 0 and CHROOT_ROOT[-1] == '/':
-    CHROOT_ROOT = CHROOT_ROOT[:-1]
-
-# This is where we expect to find the seed.yaml file
-YAML_PATH = CHROOT_ROOT + '/var/lib/snapd/seed/seed.yaml'
 
 # Snaps are prepended with this string in the manifest
 LINE_PREFIX = 'snap:'
 
+# This is where we expect to find the seed.yaml file
+YAML_PATH = CHROOT_ROOT + '/var/lib/snapd/seed/seed.yaml'
+
 log("yaml path: {}".format(YAML_PATH))
-if not os.path.isfile(YAML_PATH):
-    log("WARNING: yaml path not found; no seeded snaps found.")
-    exit(0)
-else:
-    log("yaml path found.")
 
-with open(YAML_PATH, 'r') as fh:
+
+def make_manifest_from_seed_yaml(path):
+    with open(YAML_PATH, 'r') as fh:
         yaml_lines = yaml.safe_load(fh)['snaps']
 
-log('Writing manifest to {}'.format(FNAME))
+    log('Writing manifest to {}'.format(FNAME))
 
-with open(FNAME, 'a+') as fh:
+    with open(FNAME, 'a+') as fh:
         for item in yaml_lines:
             filestring = item['file']
             # Pull the revision number off the file name
@@ -65,4 +62,96 @@ with open(FNAME, 'a+') as fh:
                                              item['channel'],
                                              revision,
                                              ))
+
+
+def look_for_uc20_model(chroot):
+    modeenv = f"{chroot}/var/lib/snapd/modeenv"
+    system_name = None
+    if os.path.isfile(modeenv):
+        log(f"found modeenv file at {modeenv}")
+        with open(modeenv) as fh:
+            for line in fh:
+                if line.startswith("recovery_system="):
+                    system_name = line.split('=', 1)[1].strip()
+                    log(f"read system name {system_name!r} from modeenv")
+                    break
+    if system_name is None:
+        system_names = os.listdir(f"{chroot}/var/lib/snapd/seed/systems")
+        if len(system_names) == 0:
+            log("no systems found")
+            return None
+        elif len(system_names) > 1:
+            log("multiple systems found, refusing to guess which to parse")
+            return None
+        else:
+            system_name = system_names[0]
+            log(f"parsing only system found {system_name}")
+    system_dir = f"{chroot}/var/lib/snapd/seed/systems/{system_name}"
+    if not os.path.isdir(system_dir):
+        log(f"could not find system called {system_name}")
+        return None
+    return system_dir
+
+
+def parse_assertion_file(asserts, filename):
+    # Parse the snapd assertions file 'filename' and store the
+    # assertions found in 'asserts'.
+    with open(filename) as fp:
+        text = fp.read()
+
+    k = ''
+
+    for block in text.split('\n\n'):
+        if block.startswith('type:'):
+            this_assert = {}
+            for line in block.split('\n'):
+                if line.startswith(' '):
+                    this_assert[k.strip()] += '\n' + line
+                    continue
+                k, v = line.split(':', 1)
+                this_assert[k.strip()] = v.strip()
+            asserts.setdefault(this_assert['type'], []).append(this_assert)
+
+
+def make_manifest_from_system(system_dir):
+    files = [f"{system_dir}/model"] + glob.glob(f"{system_dir}/assertions/*")
+
+    asserts = {}
+    for filename in files:
+        parse_assertion_file(asserts, filename)
+
+    [model] = asserts['model']
+    snaps = yaml.safe_load(model['snaps'])
+
+    snap_names = []
+    for snap in snaps:
+        snap_names.append(snap['name'])
+    snap_names.sort()
+
+    snap_name_to_id = {}
+    snap_id_to_rev = {}
+    for decl in asserts['snap-declaration']:
+        snap_name_to_id[decl['snap-name']] = decl['snap-id']
+    for rev in asserts['snap-revision']:
+        snap_id_to_rev[rev['snap-id']] = rev['snap-revision']
+
+    log('Writing manifest to {}'.format(FNAME))
+
+    with open(FNAME, 'a+') as fh:
+        for snap_name in snap_names:
+            channel = snap['default-channel']
+            rev = snap_id_to_rev[snap_name_to_id[snap_name]]
+            fh.write(f"{LINE_PREFIX}{snap_name}\t{channel}\t{rev}\n")
+
+
+if os.path.isfile(YAML_PATH):
+    log(f"seed.yaml found at {YAML_PATH}")
+    make_manifest_from_seed_yaml(YAML_PATH)
+else:
+    system_dir = look_for_uc20_model(CHROOT_ROOT)
+    if system_dir is None:
+        log("WARNING: could not find seed.yaml or uc20-style seed")
+        exit(0)
+    make_manifest_from_system(system_dir)
+
 log('Manifest output finished.')