Support for OS uninstall and partition grow

src/diskutil.py

@@ -9,6 +9,7 @@ class Partition:
     size: int
     free: bool
     type: str
+    free_space_after: int = 0
     uuid: str = None
     desc: str = None
     label: str = None
@@ -141,19 +142,26 @@ def get_partitions(self, dskname):
             parts.append(self.get_partition_info(dskpart["DeviceIdentifier"]))
         parts.sort(key=lambda i: i.offset)
 
+        prev_part = None
         prev_name = dskname
         parts2 = []
         for part in parts:
-            if (part.offset - p) > self.FREE_THRESHOLD:
+            free_space = part.offset - p
+            if free_space > self.FREE_THRESHOLD:
                 parts2.append(Partition(name=prev_name, free=True, type=None,
-                                        offset=p, size=(part.offset - p)))
+                                        offset=p, size=free_space))
+                if prev_part is not None:
+                    prev_part.free_space_after = free_space
             parts2.append(part)
+            prev_part = part
             prev_name = part.name
             p = part.offset + part.size
 
-        if (total_size - p) > self.FREE_THRESHOLD:
+        free_space = total_size - p
+        if free_space > self.FREE_THRESHOLD:
             parts2.append(Partition(name=prev_name, free=True, type=None,
-                                    offset=p, size=(total_size - p)))
+                                    offset=p, size=free_space))
+            prev_part.free_space_after = free_space
         return parts2
 
     def refresh_part(self, part):
@@ -197,6 +205,14 @@ def addPartition(self, after, fs, label, size):
 
         raise Exception("Could not find new partition")
 
+    def deletePartition(self, part):
+        if part.type == "Apple_APFS":
+            logging.info(f"Deleting APFS partition {part.name}")
+            self.action("apfs", "deleteContainer", part.name, verbose=True)
+        else:
+            logging.info(f"Deleting non-APFS partition {part.name}")
+            self.action("eraseVolume", "free", "free", part.name, verbose=True)
+
     def changeVolumeRole(self, volume, role):
         self.action("apfs", "changeVolumeRole", volume, role, verbose=True)
 

src/main.py

@@ -641,6 +641,141 @@ def action_resize(self, resizable):
 
         return True
 
+    def action_grow(self, growable):
+        p_message("The following partitions will be resized to their maximum size.")
+        for part in growable:
+            new_size = part.size + part.free_space_after
+            p_info(f" - {part.name}: {part.desc} to {ssize(new_size)}")
+
+        if self.yesno(f"Proceed?", default=True):
+            for part in growable:
+                p_progress(f"Resizing {part.name}...")
+                new_size = part.size + part.free_space_after
+                self.dutil.resizeContainer(part.name, new_size)
+
+        return True
+
+    def delete_partitions(self, partitions):
+        logging.info("Confirming deletion of partitions: {partitions}")
+
+        print()
+        p_warning("The following partitions will be DELETED:")
+        for part in partitions:
+            p_info(f" - {part.name}: {part.desc}")
+        p_warning("This operation cannot be reverted!")
+        if not self.yesno("Do you want to continue?"):
+            return
+
+        def order_key(part):
+            disk, idx = part.name.rsplit('s', 1)
+            return disk, -int(idx)
+
+        delete_order = sorted(partitions, key=order_key)
+        default_deleted = False
+        for part in delete_order:
+            if self.default_os in part.os:
+                default_deleted = True
+            p_progress(f"Deleting partition {part.name}...")
+            self.dutil.deletePartition(part)
+
+        if default_deleted:
+            p_warning("The default startup partition was deleted. Please use the Startup Disk")
+            p_warning("preference panel to select another one before rebooting.")
+            p_prompt("Press enter to continue.")
+            self.input()
+
+    def action_uninstall(self, m1n1_stubs):
+        print()
+        if len(m1n1_stubs) > 1 or self.expert:
+            choices = {str(i): s for i, s in enumerate(m1n1_stubs, 1)}
+            p_question("Choose an existing installation:")
+            idx = self.choice("Installation", choices)
+            stub_os = choices[idx]
+        else:
+            stub_os = m1n1_stubs[0]
+            p_plain(f"Selecting installation: {stub_os}")
+
+        stub_part = stub_os.partition
+
+        if len(stub_part.os) != 1:
+            p_error(f"This installation resides on partition ({stub_part.name}) that contains multiple ")
+            p_error(f"OSes; it cannot be automatically deleted.")
+            return False
+
+        efi_uuid = stub_os.m1n1_esp_uuid
+        efi_part = None
+        linux_parts = []
+
+        if efi_uuid:
+            efi_uuid = efi_uuid.upper()
+            for part in self.parts:
+                if efi_part:
+                    if part.type == "Linux Filesystem":
+                        linux_parts.append(part)
+                    else:
+                        break
+                elif part.uuid.upper() == efi_uuid:
+                    efi_part = part
+
+        if not efi_part:
+            p_error("Unable to find the EFI partition of this installation.")
+            p_error("Use the installer in expert mode to manually delete the partitions.")
+            return False
+
+        to_delete = [stub_part, efi_part]
+
+        if linux_parts:
+            p_info(f"{len(linux_parts)} Linux partition(s) seem to be part of this installation.")
+            if self.yesno("Delete them?", default=True):
+                to_delete.extend(linux_parts)
+            else:
+                p_info("You can delete these partitions later by re-running the installer in expert ")
+                p_info("mode.")
+        else:
+            p_warning("No Linux partitions attached to this installation were found.")
+            p_warning("If needed, you can delete any partition later by re-running the installer in ")
+            p_warning("expert mode.")
+
+        self.delete_partitions(to_delete)
+        return True
+
+    def action_delete_partitions(self):
+        choices = {}
+        for idx, part in enumerate(self.parts, 1):
+            if part.free:
+                continue
+            if part.type in {"Apple_APFS_ISC", "Apple_APFS_Recovery"}:
+                continue
+            if self.cur_os in part.os:
+                continue
+            choices[str(idx)] = part
+
+        if not choices:
+            p_info("No partitions can be deleted.")
+            return True
+
+        p_question("Choose partitions to delete:")
+        for idx, part in choices.items():
+            p_choice(f"  {col(BRIGHT)}{idx}{col(NORMAL)}: {part.desc}")
+
+        while True:
+            self.flush_input()
+            res = input_prompt("Partitions numbers to delete (separated by spaces): ")
+            res = [i for i in res.replace(',', ' ').split() if i]
+            if not res:
+                return
+
+            try:
+                partitions = [choices[idx] for idx in res]
+            except KeyError:
+                continue
+
+            break
+
+        self.delete_partitions(partitions)
+
+        return True
+
     def main(self):
         print()
         p_message("Welcome to the Asahi Linux installer!")
@@ -720,6 +855,8 @@ def main_loop(self):
         parts_free = []
         parts_empty_apfs = []
         parts_resizable = []
+        parts_growable = []
+        m1n1_stubs = []
 
         for i, p in enumerate(self.parts):
             if p.type in ("Apple_APFS_ISC",):
@@ -736,11 +873,16 @@ def main_loop(self):
                     p.desc += f" [{p.label}]"
                 vols = p.container["Volumes"]
                 p.desc += f" ({ssize(p.size)}, {len(vols)} volume{'s' if len(vols) != 1 else ''})"
+                if p.type == "Apple_APFS" and p.free_space_after > 0:
+                    parts_growable.append(p)
                 if self.can_resize(p):
                     parts_resizable.append(p)
                 else:
                     if p.size >= STUB_SIZE * 0.95:
                         parts_empty_apfs.append(p)
+                for part_os in p.os:
+                    if part_os.m1n1_ver:
+                        m1n1_stubs.append(part_os)
             else:
                 p.desc = f"{p.type} ({ssize(p.size)})"
 
@@ -749,7 +891,7 @@ def main_loop(self):
 
         self.cur_os = None
         self.is_sfr_recovery = self.sysinfo.boot_vgid in (osenum.UUID_SROS, osenum.UUID_FROS)
-        default_os = None
+        self.default_os = None
 
         r = col(YELLOW) + "R" + col()
         b = col(GREEN) + "B" + col()
@@ -776,7 +918,7 @@ def main_loop(self):
                 elif self.sysinfo.boot_vgid == os.vgid:
                     state = u
                 if self.sysinfo.default_boot == os.vgid:
-                    default_os = os
+                    self.default_os = os
                     state += d
                 else:
                     state += " "
@@ -788,7 +930,7 @@ def main_loop(self):
         print()
 
         if self.cur_os is None and self.sysinfo.boot_mode != "macOS":
-            self.cur_os = default_os
+            self.cur_os = self.default_os
         self.check_cur_os()
 
         actions = {}
@@ -802,8 +944,14 @@ def main_loop(self):
         if parts_resizable:
             actions["r"] = "Resize an existing partition to make space for a new OS"
             default = default or "r"
+        if parts_growable:
+            actions["g"] = "Grow a partition to reclaim free space after it"
+        if m1n1_stubs:
+            actions["u"] = "Uninstall an OS"
         if self.sysinfo.boot_mode == "one true recoveryOS" and False:
             actions["m"] = "Upgrade bootloader of an existing OS"
+        if self.expert:
+            actions["d"] = "Delete partitions"
 
         if not actions:
             p_error("No actions available on this system.")
@@ -822,9 +970,15 @@ def main_loop(self):
             return self.action_install_into_container(parts_empty_apfs)
         elif act == "r":
             return self.action_resize(parts_resizable)
+        elif act == "g":
+            return self.action_grow(parts_growable)
+        elif act == "u":
+            return self.action_uninstall(m1n1_stubs)
         elif act == "m":
             p_error("Unimplemented")
             sys.exit(1)
+        elif act == "d":
+            return self.action_delete_partitions()
         elif act == "q":
             return False
 

src/osenum.py

@@ -1,11 +1,12 @@
 # SPDX-License-Identifier: MIT
-import os, os.path, plistlib, subprocess, logging
+import os, os.path, plistlib, re, subprocess, logging
 from dataclasses import dataclass
 
 from util import *
 
 UUID_SROS = "3D3287DE-280D-4619-AAAB-D97469CA9C71"
 UUID_FROS = "C8858560-55AC-400F-BBB9-C9220A8DAC0D"
+M1N1_CHAINLOAD_RE = re.compile(rb'chainload=([a-zA-Z0-9-]+);([\x20-\x7E]+)(?:\x00|\n|$)')
 
 @dataclass
 class OSInfo:
@@ -17,6 +18,8 @@ class OSInfo:
     stub: bool = False
     version: str = None
     m1n1_ver: str = None
+    m1n1_esp_uuid: str = None
+    m1n1_next_obj: str = None
     system: object = None
     data: object = None
     preboot: object = None
@@ -207,6 +210,11 @@ def collect_os(self, part, volumes, vgid):
             if b"##m1n1_ver##" in fuos:
                 osi.m1n1_ver = fuos.split(b"##m1n1_ver##")[1].split(b"\0")[0].decode("ascii")
                 logging.info(f"  m1n1 version found: {osi.m1n1_ver}")
+                res = M1N1_CHAINLOAD_RE.search(fuos)
+                if res:
+                    osi.m1n1_esp_uuid = res.group(1).decode('ascii')
+                    osi.m1n1_next_obj = res.group(2).decode('ascii')
+                    logging.info(f"   m1n1 chainload found: ESP {osi.m1n1_esp_uuid} next {osi.m1n1_next_obj}")
 
         return osi