From 32d7cd44ffe8c681f884136c15b0a9d4b895d8a3 Mon Sep 17 00:00:00 2001
From: 2Shirt <2xShirt@gmail.com>
Date: Tue, 14 Dec 2021 15:08:01 -0700
Subject: [PATCH] Save extra image list to UFD to support --update

Addresses #176
---
 scripts/wk/kit/ufd.py | 21 ++++++++++++++++++++-
 1 file changed, 20 insertions(+), 1 deletion(-)

diff --git a/scripts/wk/kit/ufd.py b/scripts/wk/kit/ufd.py
index a29bfcd6..49ec21e5 100644
--- a/scripts/wk/kit/ufd.py
+++ b/scripts/wk/kit/ufd.py
@@ -52,6 +52,7 @@ Options:
   -U --update           Don't format device, just update
 '''
 LOG = logging.getLogger(__name__)
+EXTRA_IMAGES_LIST = '/mnt/UFD/arch/extra_images.list'
 MIB = 1024 ** 2
 ISO_LABEL = f'{KIT_NAME_SHORT}_LINUX'
 UFD_LABEL = f'{KIT_NAME_SHORT}_UFD'
@@ -78,6 +79,7 @@ def build_ufd():
     log.enable_debug_mode()
   if args['--update'] and args['EXTRA_IMAGES']:
     std.print_warning('Extra images are ignored when updating')
+    args['EXTRA_IMAGES'] = []
   log.update_log_path(dest_name='build-ufd', timestamp=True)
   try_print = std.TryAndPrint()
   try_print.add_error('FileNotFoundError')
@@ -137,6 +139,14 @@ def build_ufd():
     read_write=True,
     )
 
+  # Load extra images if updating
+  if args['--update'] and os.path.exists(EXTRA_IMAGES_LIST):
+    with open(EXTRA_IMAGES_LIST, 'r', encoding='utf-8') as _f:
+      extra_images = [
+        io.get_path_obj(image.strip(), resolve=False)
+        for image in _f.readlines()
+        ]
+
   # Remove Arch folder
   if args['--update']:
     try_print.run(
@@ -172,6 +182,11 @@ def build_ufd():
         image_path=image_path,
         )
 
+  # Save extra image list
+  if extra_images:
+    with open(EXTRA_IMAGES_LIST, 'w', encoding='utf-8') as _f:
+      _f.write('\n'.join([image.name for image in extra_images]))
+
   # Update boot entries
   std.print_standard(' ')
   std.print_info('Boot Setup')
@@ -539,8 +554,12 @@ def update_boot_entries(ufd_dev, images=None):
       ]
     run_program(cmd, check=False)
 
-  # Bail early
+  # Check if we're working with extra images
+  if not images and os.path.exists(EXTRA_IMAGES_LIST):
+    with open(EXTRA_IMAGES_LIST, 'r', encoding='utf-8') as _f:
+      images = [image.strip() for image in _f.readlines()]
   if not images:
+    # No extra images detected
     return
 
   # Get PARTUUID values