arch-setup/shared/gen-disk-config.py

#!/usr/bin/env python3
# Genera el disk_config de archinstall (>=4.x) dimensionado al disco REAL y lo
# inyecta en el JSON base. archinstall 4.x quitó la unidad "Percent", así que un
# JSON con tamaños fijos no es portable entre la VM (20G/10G) y el hardware
# (250G/1TB): por eso se calcula en tiempo de instalación con los propios modelos
# de archinstall (válido ante cambios de schema entre versiones).
#
#   gen-disk-config.py BASE.json [SYS_DEV] [HOME_DEV] > merged.json
#
# Los discos se autodetectan: no se puede contar discos para decidir el layout
# porque el USB instalador tambien aparece como disco (p.ej. sdb). Filtramos por
# NO removible (RM=0): el USB es removible y queda fuera; el ISO de la VM es 'rom'
# y tampoco cuenta. Asi da igual que el disco de sistema sea nvme0n1 o un SSD
# PCIe que enumera como sda. Pasa SYS_DEV/HOME_DEV para forzar nombres a mano.
#
# Layout (replica el hardware): disco de sistema = ESP 512M + /boot 1G ext4 +
# / btrfs (resto); segundo disco = /home ext4 (todo). swap = zram ("swap": true
# del base). Si solo hay 1 disco (portatil), /home queda dentro del / btrfs.
import json, subprocess, sys
from pathlib import Path
# archinstall solo existe en el entorno live; el --selftest corre en cualquier
# sitio porque pick_fixed no lo necesita. ponytail: se importa al usar main().
try:
    from archinstall.lib.disk.device_handler import device_handler
    from archinstall.lib.models.device import (
        DiskLayoutConfiguration, DiskLayoutType, DeviceModification,
        PartitionModification, ModificationStatus, PartitionType,
        FilesystemType, PartitionFlag, Size, Unit,
    )
except ImportError:
    device_handler = None

MiB = 1024 ** 2
END_MARGIN = MiB  # deja 1 MiB al final para la cabecera GPT de respaldo

def usable_end(total):
    # Final alineado a la rejilla de 1 MiB. Los discos reales tienen tamanos que
    # no son multiplo de MiB (p.ej. 250_059_350_016 B), asi que terminar en
    # total-END_MARGIN deja la particion fuera de rejilla y parted la marca como
    # "misaligned". Redondeamos el final hacia abajo a MiB (se pierden <2 MiB).
    return ((int(total) - END_MARGIN) // MiB) * MiB

def part(status, ptype, start_b, len_b, ss, fs=None, mnt=None, flags=()):
    return PartitionModification(
        status=status, type=ptype,
        start=Size(start_b, Unit.B, ss),
        length=Size(len_b, Unit.B, ss),
        fs_type=fs, mountpoint=Path(mnt) if mnt else None,
        flags=list(flags),
    )

def nvme_mods(dev):
    ss = dev.device_info.sector_size
    total = dev.device_info.total_size._normalize()
    esp_start, esp_len = 1 * MiB, 511 * MiB           # 1 -> 512 MiB
    boot_start, boot_len = 512 * MiB, 1024 * MiB        # 512 -> 1536 MiB
    root_start = boot_start + boot_len                  # 1536 MiB
    root_len = usable_end(total) - root_start
    return DeviceModification(device=dev, wipe=True, partitions=[
        part(ModificationStatus.CREATE, PartitionType.PRIMARY, esp_start, esp_len, ss,
             FilesystemType.FAT32, "/boot/efi", (PartitionFlag.BOOT, PartitionFlag.ESP)),
        part(ModificationStatus.CREATE, PartitionType.PRIMARY, boot_start, boot_len, ss,
             FilesystemType.EXT4, "/boot"),
        part(ModificationStatus.CREATE, PartitionType.PRIMARY, root_start, root_len, ss,
             FilesystemType.BTRFS, "/"),
    ])

def sda_mods(dev):
    ss = dev.device_info.sector_size
    total = dev.device_info.total_size._normalize()
    start = 1 * MiB
    return DeviceModification(device=dev, wipe=True, partitions=[
        part(ModificationStatus.CREATE, PartitionType.PRIMARY, start,
             usable_end(total) - start, ss, FilesystemType.EXT4, "/home"),
    ])

def pick_fixed(lsblk_out):
    # Discos NO removibles, en orden de kernel. Excluye el USB instalador (RM=1)
    # y el ISO de la VM (TYPE=rom). nvme va antes que sda si ambos existen.
    disks = []
    for line in lsblk_out.splitlines():
        name, rm, typ = line.split()
        if typ == "disk" and rm == "0":
            disks.append(f"/dev/{name}")
    return sorted(disks, key=lambda d: (0 if "nvme" in d else 1, d))

def fixed_disks():
    return pick_fixed(subprocess.check_output(
        ["lsblk", "-dnro", "NAME,RM,TYPE"], text=True))

def main():
    if device_handler is None:
        sys.exit("archinstall no disponible (¿estás en el entorno live?)")
    base_path = sys.argv[1]
    override = sys.argv[2:]  # SYS_DEV [HOME_DEV] para forzar a mano
    disks = override if override else fixed_disks()
    if not disks:
        sys.exit("No encuentro ningun disco fijo no removible (revisa lsblk)")

    dev_sys = device_handler.get_device(Path(disks[0]))
    if not dev_sys:
        sys.exit(f"No encuentro el disco de sistema {disks[0]} (revisa lsblk)")

    # >=2 discos fijos -> sistema + /home aparte. 1 disco (portatil) -> todo
    # en el de sistema; /home vive dentro del btrfs raiz.
    mods = [nvme_mods(dev_sys)]
    if len(disks) > 1:
        dev_home = device_handler.get_device(Path(disks[1]))
        if dev_home:
            mods.append(sda_mods(dev_home))

    layout = DiskLayoutConfiguration(
        config_type=DiskLayoutType.Manual,
        device_modifications=mods,
    )
    cfg = json.loads(Path(base_path).read_text())
    cfg["disk_config"] = layout.json()
    json.dump(cfg, sys.stdout, indent=2)

def _selftest():
    # portatil: SSD PCIe = sda (fijo), USB instalador = sdb (RM=1), ISO = sr0
    laptop = "sda 0 disk\nsdb 1 disk\nsr0 1 rom"
    assert pick_fixed(laptop) == ["/dev/sda"], pick_fixed(laptop)
    # hardware/VM: nvme = sistema, sda = /home; nvme primero aunque ordene despues
    two = "sda 0 disk\nnvme0n1 0 disk"
    assert pick_fixed(two) == ["/dev/nvme0n1", "/dev/sda"], pick_fixed(two)
    # solo USB conectado -> sin discos fijos (main aborta)
    assert pick_fixed("sdb 1 disk\nsr0 1 rom") == []
    # final alineado: disco "feo" -> end multiplo de MiB, dentro del disco
    real = 250_059_350_016  # SSD 250G tipico, no multiplo de MiB
    e = usable_end(real)
    assert e % MiB == 0 and 0 < e <= real - END_MARGIN, e
    # disco MiB-exacto (VM 40G) -> pierde justo END_MARGIN, nada mas
    assert usable_end(40 * 1024 * MiB) == 40 * 1024 * MiB - MiB
    print("ok")

if __name__ == "__main__":
    if sys.argv[1:2] == ["--selftest"]:
        _selftest()
    else:
        main()