313 lines
8.8 KiB
Python
313 lines
8.8 KiB
Python
import json
|
|
import os
|
|
import re
|
|
import subprocess
|
|
from pathlib import Path
|
|
|
|
|
|
FILESYSTEMS = {
|
|
"fat32": {
|
|
"mkfs": ["mkfs.fat", "-F", "32"],
|
|
"parted_type": "fat32",
|
|
"parted_flags": {"esp": "esp"},
|
|
},
|
|
"ext4": {
|
|
"mkfs": ["mkfs.ext4", "-F"],
|
|
"parted_type": "ext4",
|
|
"parted_flags": {},
|
|
},
|
|
}
|
|
|
|
|
|
def run(*args, check=True):
|
|
"""Run a system command and return stdout."""
|
|
result = subprocess.run(
|
|
[str(arg) for arg in args],
|
|
check=check,
|
|
text=True,
|
|
capture_output=True,
|
|
)
|
|
return result.stdout.strip()
|
|
|
|
|
|
def partition_device(disk, number):
|
|
"""Return the correct partition path for common Linux devices."""
|
|
disk = str(disk)
|
|
suffix = "p" if re.search(r"\d$", disk) else ""
|
|
return f"{disk}{suffix}{number}"
|
|
|
|
|
|
def parse_size(size, disk_bytes):
|
|
"""Convert a size such as 200M or 100% into bytes."""
|
|
if size == "100%":
|
|
return disk_bytes
|
|
|
|
match = re.fullmatch(
|
|
r"\s*(\d+(?:\.\d+)?)\s*([KMGT]?)B?\s*",
|
|
size.upper(),
|
|
)
|
|
if not match:
|
|
raise ValueError(f"Invalid partition size: {size!r}")
|
|
|
|
value = float(match.group(1))
|
|
unit = match.group(2)
|
|
multiplier = {
|
|
"": 1,
|
|
"K": 1024,
|
|
"M": 1024**2,
|
|
"G": 1024**3,
|
|
"T": 1024**4,
|
|
}[unit]
|
|
|
|
return int(value * multiplier)
|
|
|
|
|
|
def get_table(disk):
|
|
"""Read the partition table and return its partitions."""
|
|
try:
|
|
data = json.loads(run("sfdisk", "--json", disk))
|
|
except subprocess.CalledProcessError:
|
|
return None
|
|
|
|
table = data.get("partitiontable")
|
|
if not table:
|
|
return None
|
|
|
|
return table
|
|
|
|
|
|
def get_filesystem(device):
|
|
"""Return the detected filesystem type, or None."""
|
|
result = run(
|
|
"blkid", "-p", "-s", "TYPE", "-o", "value",
|
|
device, check=False,
|
|
)
|
|
return result.lower() if result else None
|
|
|
|
|
|
def apply_partition_config(
|
|
disk,
|
|
config,
|
|
*,
|
|
dry_run=True,
|
|
format_mismatches=False,
|
|
):
|
|
"""
|
|
Apply a partition configuration to a disk.
|
|
|
|
Safety behavior:
|
|
- Dry-run by default.
|
|
- Creates GPT only if no partition table exists.
|
|
- Does not automatically resize, delete, or reorder partitions.
|
|
- Refuses to format filesystems unless format_mismatches=True.
|
|
- Verifies the expected partition count and filesystem types.
|
|
|
|
Requires root privileges and: parted, sfdisk, blkid, blockdev,
|
|
mkfs.fat, mkfs.ext4, and partprobe.
|
|
"""
|
|
disk = str(Path(disk))
|
|
|
|
if os.geteuid() != 0:
|
|
raise PermissionError("Run this function as root.")
|
|
|
|
if not os.path.exists(disk):
|
|
raise FileNotFoundError(disk)
|
|
|
|
if run("lsblk", "-dn", "-o", "TYPE", disk) != "disk":
|
|
raise ValueError(f"{disk} is not a whole-disk device")
|
|
|
|
partitions = config["partitions"]
|
|
if not partitions:
|
|
raise ValueError("At least one partition is required")
|
|
|
|
for p in partitions:
|
|
fs = p["type"].lower()
|
|
if fs not in FILESYSTEMS:
|
|
raise ValueError(f"Unsupported filesystem: {fs}")
|
|
if not isinstance(p.get("flags", []), list):
|
|
raise ValueError("'flags' must be a list")
|
|
|
|
disk_bytes = int(run("blockdev", "--getsize64", disk))
|
|
sector_size = int(run("blockdev", "--getss", disk))
|
|
alignment = 1024 * 1024 # 1 MiB
|
|
|
|
def align_up(n):
|
|
return ((n + alignment - 1) // alignment) * alignment
|
|
|
|
# Calculate partition boundaries. Reserve 1 MiB at the start.
|
|
layout = []
|
|
start = alignment
|
|
|
|
for i, p in enumerate(partitions):
|
|
requested = p["size"]
|
|
if requested == "100%":
|
|
end = disk_bytes - alignment
|
|
else:
|
|
size_bytes = parse_size(requested, disk_bytes)
|
|
end = align_up(start + size_bytes)
|
|
|
|
if i == len(partitions) - 1 and requested == "100%":
|
|
end = disk_bytes - alignment
|
|
|
|
if end <= start:
|
|
raise ValueError(f"Invalid size for partition {i + 1}")
|
|
|
|
if end > disk_bytes - alignment:
|
|
raise ValueError(
|
|
f"Partition {i + 1} exceeds disk capacity"
|
|
)
|
|
|
|
layout.append({
|
|
"number": i + 1,
|
|
"start": start,
|
|
"end": end,
|
|
"type": p["type"].lower(),
|
|
"flags": p.get("flags", []),
|
|
})
|
|
start = end
|
|
|
|
def parted(*args):
|
|
if dry_run:
|
|
print("+ parted", "-s", disk, *args)
|
|
else:
|
|
run("parted", "-s", disk, *args)
|
|
|
|
table = get_table(disk)
|
|
|
|
# Create GPT only when the disk has no partition table.
|
|
if table is None:
|
|
print(f"No readable partition table found on {disk}.")
|
|
parted("mklabel", "gpt")
|
|
table = None
|
|
|
|
elif table.get("label") != "gpt":
|
|
raise RuntimeError(
|
|
f"{disk} has a {table.get('label')} table, not GPT. "
|
|
"Refusing to overwrite it."
|
|
)
|
|
|
|
existing = table.get("partitions", []) if table else []
|
|
|
|
# Never guess how to repair an incomplete or mismatched layout.
|
|
if existing and len(existing) != len(layout):
|
|
raise RuntimeError(
|
|
f"Found {len(existing)} partitions; "
|
|
f"configuration requires {len(layout)}. "
|
|
"Automatic resizing/deletion is disabled."
|
|
)
|
|
|
|
if not existing:
|
|
for item in layout:
|
|
fs = item["type"]
|
|
start_mib = item["start"] / (1024**2)
|
|
end_mib = item["end"] / (1024**2)
|
|
|
|
# parted accepts MiB boundaries.
|
|
parted(
|
|
"mkpart",
|
|
"primary",
|
|
FILESYSTEMS[fs]["parted_type"],
|
|
f"{start_mib:.6f}MiB",
|
|
f"{end_mib:.6f}MiB",
|
|
)
|
|
|
|
for flag in item["flags"]:
|
|
if flag not in FILESYSTEMS[fs]["parted_flags"]:
|
|
raise ValueError(
|
|
f"Unsupported flag {flag!r} for {fs}"
|
|
)
|
|
parted(
|
|
"set",
|
|
str(item["number"]),
|
|
FILESYSTEMS[fs]["parted_flags"][flag],
|
|
"on",
|
|
)
|
|
|
|
elif not dry_run:
|
|
# Re-read the on-disk table before inspecting filesystems.
|
|
table = get_table(disk)
|
|
existing = table.get("partitions", [])
|
|
|
|
# Validate the existing boundaries. Do not modify them.
|
|
for expected, actual in zip(layout, existing):
|
|
actual_start = int(actual["start"]) * sector_size
|
|
actual_end = (
|
|
int(actual["start"]) + int(actual["size"])
|
|
) * sector_size
|
|
|
|
if (
|
|
abs(actual_start - expected["start"]) > alignment
|
|
or abs(actual_end - expected["end"]) > alignment
|
|
):
|
|
raise RuntimeError(
|
|
f"Partition {expected['number']} has incorrect "
|
|
"boundaries. Automatic resizing is disabled."
|
|
)
|
|
|
|
# Ask the kernel to re-read the partition table.
|
|
if not dry_run:
|
|
run("partprobe", disk)
|
|
run("udevadm", "settle")
|
|
|
|
results = []
|
|
|
|
for item in layout:
|
|
device = partition_device(disk, item["number"])
|
|
expected_fs = item["type"]
|
|
|
|
if dry_run:
|
|
print(
|
|
f"Would inspect {device}; expected filesystem "
|
|
f"{expected_fs}"
|
|
)
|
|
if format_mismatches:
|
|
print(f"Would format {device} if needed")
|
|
results.append({
|
|
"device": device,
|
|
"expected": expected_fs,
|
|
"status": "dry-run",
|
|
})
|
|
continue
|
|
|
|
if not os.path.exists(device):
|
|
raise RuntimeError(
|
|
f"Partition device {device} did not appear"
|
|
)
|
|
|
|
actual_fs = get_filesystem(device)
|
|
|
|
if actual_fs != expected_fs:
|
|
if not format_mismatches:
|
|
raise RuntimeError(
|
|
f"{device}: expected {expected_fs}, "
|
|
f"found {actual_fs!r}. "
|
|
"Set format_mismatches=True to format it."
|
|
)
|
|
|
|
# Avoid accidentally formatting a mounted filesystem.
|
|
mounts = run(
|
|
"lsblk", "-n", "-o", "MOUNTPOINTS", device
|
|
)
|
|
if mounts.strip():
|
|
raise RuntimeError(
|
|
f"{device} appears to be mounted; refusing "
|
|
"to format it."
|
|
)
|
|
|
|
fs = FILESYSTEMS[expected_fs]
|
|
run(*fs["mkfs"], device)
|
|
|
|
actual_fs = get_filesystem(device)
|
|
if actual_fs != expected_fs:
|
|
raise RuntimeError(
|
|
f"Filesystem verification failed for {device}"
|
|
)
|
|
|
|
results.append({
|
|
"device": device,
|
|
"expected": expected_fs,
|
|
"actual": actual_fs,
|
|
"status": "ok",
|
|
})
|
|
|
|
return results |