Improve volume matching on macOS and fix snapshot-based volume handling on macOS

README.md

@@ -40,6 +40,11 @@ Some observations that lead to the development of this tool:
 Note: in CI, it is recommended to use `--volume-initialization-rate 300` for
 best performance and consistency (additional costs apply).
 
+## Disclaimer
+
+This script may wipe, or otherwise mishandle volumes. Risks go up if
+`--mountpoint` is used.
+
 ## Goals and scope
 
 * Generic/flexible
@@ -54,12 +59,6 @@ best performance and consistency (additional costs apply).
 * Runs unprivileged (sudo to format/(u)mount)
 * Under 1k lines
 
-## Disclaimer
-
-When the `--mountpoint` option is used, volumes determined to be unformatted
-are formatted with no additional confirmation. There is a risk of data loss if
-that detection is (or ever becomes) inaccurate, especially on macOS.
-
 ## Usage
 
 ```console
@@ -294,7 +293,7 @@ EOF
 
 ```bash
 sudo install -Tm440 /dev/stdin /etc/sudoers.d/99-vol <<-'EOF'
-Cmnd_Alias VOL_FORMAT = /usr/sbin/diskutil ^eraseDisk -noEFI APFS vol-[0-9a-fA-F]+ /dev/disk[1-9][0-9]*$
+Cmnd_Alias VOL_FORMAT = /usr/sbin/diskutil ^eraseDisk -noEFI APFS [a-zA-Z]+ /dev/disk[1-9][0-9]*$
 user ALL=(root) NOPASSWD: VOL_FORMAT
 EOF
 ```
@@ -303,9 +302,9 @@ EOF
 
 ```bash
 sudo install -Tm440 /dev/stdin /etc/sudoers.d/99-vol <<-'EOF'
-Cmnd_Alias VOL_FORMAT = /usr/sbin/diskutil ^eraseDisk -noEFI APFS vol-[0-9a-fA-F]+ /dev/disk[1-9][0-9]*$
-Cmnd_Alias VOL_MOUNT = /usr/sbin/diskutil ^mount -mountPoint /mnt/point vol-[0-9a-fA-F]+$
-Cmnd_Alias VOL_UMOUNT = /usr/sbin/diskutil ^umount /mnt/point$, /usr/sbin/diskutil ^umount /Volumes/vol-[0-9a-fA-F]+$
+Cmnd_Alias VOL_FORMAT = /usr/sbin/diskutil ^eraseDisk -noEFI APFS [a-zA-Z]+ /dev/disk[1-9][0-9]*$
+Cmnd_Alias VOL_MOUNT = /usr/sbin/diskutil ^mount -mountPoint /mnt/point /dev/disk[1-9][0-9]*$
+Cmnd_Alias VOL_UMOUNT = /usr/sbin/diskutil ^umount /mnt/point$, /usr/sbin/diskutil ^umount /Volumes/[a-zA-Z]+$
 user ALL=(root) NOPASSWD: VOL_FORMAT, VOL_MOUNT, VOL_UMOUNT
 EOF
 ```

vol

@@ -19,6 +19,7 @@ from operator import itemgetter
 from platform import system
 from os import getgid, getuid, makedirs
 from os.path import exists, ismount
+from string import ascii_letters
 from textwrap import dedent
 from time import sleep
 from urllib.request import Request, urlopen
@@ -228,9 +229,7 @@ def err(template = None, text = None):
             'ATTACH_REQUIRES_REAL_TAGS': 'The attach action may create volumes, so it requires non-wildcard tags',
             'AZ_DOES_NOT_MATCH_INSTANCE': 'The specified availability zone does not match the instance',
             'MOUNTPOINT_WO_FSTYPE': 'If a mountpoint is specified, the filesystem type also needs to be specified',
-            'MULTIPLE_EMPTY_EBS': 'Found multiple empty EBS volumes attached, expected one',
             'NO_AVAILABLE_DEVICES': 'All available device names are in use on the target instance',
-            'NO_EMPTY_EBS': 'Could not find a matching EBS volume (neither formatted, nor empty)',
             'NO_ROOT': 'Please don\'t run this script as root, it will call sudo if necessary ((u)mount/format)',
             'NO_TAGS_SPECIFIED': 'This action requires --tags, or --snapshot-tags, or both',
             'NO_VOLUMES': 'No volumes found with the specified tags',
@@ -281,20 +280,23 @@ def get_tag_dict(tags):
 def get_disk(volume_id, fail = True):
     match system():
         case 'Darwin':
-            disk_info = run(['diskutil', 'info', '-plist', volume_id], check=False)
-            match disk_info.returncode:
-                case 0:
-                    return run(['plutil', '-extract', 'DeviceNode', 'raw', '-'], input=disk_info.stdout).stdout.strip()
+            nvme_all = json.loads(run(['system_profiler', '-json', 'SPNVMeDataType']).stdout)['SPNVMeDataType']
+            nvme_items = sum((i['_items'] for i in nvme_all if i['_items']), [])
+            nvme_matching = next((i for i in nvme_items if i['device_serial'] == volume_id.replace('-', '')), None)
+            match nvme_matching, fail:
+                case None, True:
+                    err(text = f'Could not find disk for volume {volume_id}')
+                case None, _:
+                    return None
                 case _:
-                    # TODO: filter by diskutil activity for a tighter match
-                    candidates = run2json(['diskutil', 'list', '-plist', 'external', 'physical'])['WholeDisks']
-                    candidates_info = [run2json(['diskutil', 'info', '-plist', d]) for d in candidates]
-                    empty_ebs_disks = [d['DeviceNode'] for d in candidates_info if is_empty_ebs(d)]
-                    match empty_ebs_disks, fail:
-                        case [], True: err('NO_EMPTY_EBS')
-                        case [], False: return None
-                        case [dev], _: return dev
-                        case _: err('MULTIPLE_EMPTY_EBS')
+                    # If unformatted, returns the block device, otherwise returns the apfs device
+                    if nvme_matching.get('volumes'):
+                        containers = run2json(['diskutil', 'apfs', 'list', '-plist'])['Containers']
+                        device = [c['Volumes'][0]['DeviceIdentifier'] for c in containers
+                            if c['PhysicalStores'][0]['DeviceIdentifier'] == nvme_matching['volumes'][0]['bsd_name']][0]
+                        return f'/dev/{device}'
+                    else:
+                        return f'/dev/{nvme_matching['bsd_name']}'
         case _:
             return f'/dev/disk/by-id/nvme-Amazon_Elastic_Block_Store_{volume_id.replace("-", "")}'
 
@@ -323,14 +325,6 @@ def is_formatted(disk):
     else:
         err(text = f'Disk {disk} does not exist')
 
-def is_empty_ebs(disk_info):
-    return ( # TODO: could also match on TotalSize
-        disk_info['VolumeName'] == '' and
-        disk_info['VolumeSize'] == 0 and
-        disk_info['Content'] == '' and
-        disk_info['MediaName'] == 'Amazon Elastic Block Store'
-    )
-
 def get_volumes(tags, availability_zone = None, only_active = True):
     dbg(f'Getting volumes with tags {get_tag_string(tags)}' +
        (f' in availability zone {availability_zone}' if availability_zone else ""))
@@ -554,8 +548,8 @@ def format_volume(volume_id, fstype):
     msg(f'Formatting {disk} ({volume_id})')
     match system():
         case 'Darwin':
-            # On macOS, we rely on the disk name matching the EBS volume ID
-            run(['diskutil', 'eraseDisk', '-noEFI', fstype, volume_id, disk], sudo=True)
+            volume_name = ''.join(random.choices(ascii_letters, k=16))
+            run(['diskutil', 'eraseDisk', '-noEFI', fstype, volume_name, disk], sudo=True)
         case _:
             run(['mkfs', '-t', fstype, disk], sudo=True)
 
@@ -566,17 +560,19 @@ def mount(volume_id, mountpoint, fstype, retry = 5):
     disk = wait_for_disk(volume_id)
     if not is_formatted(disk):
         format_volume(volume_id, fstype)
+        disk = get_disk(volume_id)
 
     makedirs(mountpoint, exist_ok=True)
     match system():
         case 'Darwin':
-            existing_mountpoint = run2json(['diskutil', 'info', '-plist', volume_id])['MountPoint'].strip()
+            # TODO: only autoumount if it's under /Volumes, to handle the automounter
+            existing_mountpoint = run2json(['diskutil', 'info', '-plist', disk])['MountPoint'].strip()
             if existing_mountpoint:
                 umount(existing_mountpoint)
             msg(f'Mounting {disk} ({volume_id}) at {mountpoint}')
-            run(['diskutil', 'mount', '-mountPoint', mountpoint, volume_id], sudo=not is_mac_admin())
+            run(['diskutil', 'mount', '-mountPoint', mountpoint, disk], sudo=not is_mac_admin())
             if not ismount(mountpoint):
-                msg(f'Path {mountpoint} is not a mountpoint (likely automounted at /Volumes/{volume_id})')
+                msg(f'Path {mountpoint} is not a mountpoint (likely automounted under /Volumes)')
                 sleep(2)
                 mount(volume_id, mountpoint, fstype, retry - 1)
         case _: