daniel/MontaukOS
1
Fork 0
Code Issues 3 Pull Requests Actions Packages Projects Releases Wiki Activity
Files
36291cd1ccf71ba710a28ab7e96e066b7e2f4138
MontaukOS/scripts/montauk-imager.py
T

1721 lines
69 KiB
Python
Executable File
Raw Blame History

#!/usr/bin/env python3
# montauk-imager.py - write and update MontaukOS SD cards for the Raspberry Pi 4/400
#
# Usage: ./scripts/montauk-imager.py [--image <file>] the graphical imager
# ./scripts/montauk-imager.py list list SD cards / USB drives
# sudo ./scripts/montauk-imager.py flash <image> <device> [--no-verify] [--expand] [--yes]
# sudo ./scripts/montauk-imager.py update <image> <device> [--keep-config] [--firmware]
# [--prune] [--expand] [--yes]
# sudo ./scripts/montauk-imager.py expand <device> [--yes]
#
# The image is what `make ARCH=aarch64 rpi4-sd` builds (montauk-rpi4.img, see
# scripts/mk-rpi4-sd.sh and docs/raspberry-pi.md), raw or compressed (.gz, .xz,
# .zip). The GUI runs as the user and does the work by starting this script
# again through pkexec with --machine, reading its progress from stdout.
#
# flash Writes the image over the whole card (cancellable), wipes the end of
# the card (a stale backup GPT from an earlier OS), reads it back to
# verify, and rereads the partition table.
#
# update Refreshes the installation on a card already written from an image,
# in place: the card's partition table is never touched, so a system
# partition that has grown to fill the card (first boot, or --expand)
# stays as it is. Both partitions are found by their labels.
# ESP (FAT32, MONTAUK) every file of the image's ESP is written;
# config.txt is backed up to config.txt.bak
# if it differs, and the UEFI firmware
# RPI_EFI.fd (which holds the firmware's own
# settings) is only replaced with --firmware.
# root (ext2, MONTAUKROOT)
# os/, apps/, sdk/, boot/ and any other system
# directory: new and changed files written.
# config/ and www/: only files the card lacks
# are added, except that the image's init.toml
# and capabilities.toml are merged in (below).
# users/ and state/ are never touched.
# init.toml becomes the image's, with each service's `enabled` taken
# from the card's copy and services only the card has appended;
# capabilities.toml becomes the image's with grants only the card has
# appended. The card's old copy is kept as <name>.old. --keep-config
# leaves both alone. os/firstboot.elf is only written if the card still
# has it: it deletes itself once the volume has grown. --prune deletes
# files of the system directories the image no longer ships (except
# os/btkeys.bin, the Bluetooth bonds the kernel keeps there).
# Unchanged files are not rewritten; each file is written to a temporary
# name, flushed and renamed over the old one.
#
# expand Grows the system partition to the end of the card (on a 1 MiB
# boundary, at most 1 TiB, the size mk-rpi4-sd.sh's resize_inode
# reserves for) and the ext2 volume with resize2fs, which is what
# 0:/os/firstboot.elf otherwise does on the Pi. firstboot.elf then
# finds nothing to grow and deletes itself. flash and update take
# --expand to do this as a last (flash) or first (update) step.
#
# The ext2 volume is checked with e2fsck before it is changed and afterwards,
# and its feature flags are checked against what Fs/Ext2.cpp mounts (no
# incompat feature but filetype).
import argparse
import contextlib
import errno
import fcntl
import filecmp
import gzip
import hashlib
import json
import lzma
import os
import queue
import re
import shutil
import signal
import stat
import struct
import subprocess
import sys
import tempfile
import threading
import time
import zipfile
ROOT_LABEL = "MONTAUKROOT" # RootLabel in kernel/src/Fs/Boot.cpp
ESP_LABEL = "MONTAUK" # mformat -v in scripts/mk-rpi4-sd.sh
SECTOR = 512
MIB = 1 << 20
ALIGN = MIB // SECTOR # partitions start and end on 1 MiB boundaries
MAX_ROOT_BYTES = (1 << 28) * 4096 # mke2fs -E resize=2^28 (4 KiB blocks); the kernel's grow cap too
CHUNK = 4 * MIB
FAT_TYPES = {0x01, 0x04, 0x06, 0x0B, 0x0C, 0x0E, 0xEF}
TYPE_ESP = 0xEF
TYPE_LINUX = 0x83
# ext2/3/4 incompat features (dumpe2fs names). Fs/Ext2.cpp refuses all but filetype.
EXT_INCOMPAT = {
"compression", "filetype", "needs_recovery", "journal_dev", "meta_bg",
"extent", "extents", "64bit", "mmp", "flex_bg", "ea_inode", "dirdata",
"metadata_csum_seed", "large_dir", "inline_data", "encrypt", "casefold",
}
SUPPORTED_INCOMPAT = {"filetype"}
# Mount points that mark a disk as the host's own: never offered, never written.
SYSTEM_MOUNTS = {"/", "/boot", "/boot/efi", "/efi", "/usr", "/var", "/home", "[SWAP]"}
# Update policy for the root volume (see the header).
MERGE_DIRS = {"config", "www"}
STATE_DIRS = {"users", "state", "lost+found"}
PRUNE_DIRS = {"os", "apps", "sdk", "boot"}
PRESERVE = {"os/btkeys.bin", "os/firstboot.elf"}
FIRSTBOOT = "os/firstboot.elf"
INIT_TOML = "config/init.toml"
CAPS_TOML = "config/capabilities.toml"
FIRMWARE_FILE = "RPI_EFI.fd"
ESP_BACKUP = {"config.txt"}
BLKRRPART = 0x125F
BLKFLSBUF = 0x1261
SCRIPT = os.path.realpath(__file__)
REPO = os.path.dirname(os.path.dirname(SCRIPT))
if not os.path.isfile(os.path.join(REPO, "scripts", "mk-rpi4-sd.sh")):
REPO = None
DEFAULT_IMAGE = os.path.join(REPO, "montauk-rpi4.img") if REPO else ""
class Fail(Exception):
pass
def human(n):
"""Decimal units, as SD cards are labelled."""
n = float(n)
for unit in ("B", "kB", "MB", "GB", "TB"):
if n < 1000 or unit == "TB":
return f"{n:.0f} {unit}" if unit == "B" else f"{n:.1f} {unit}"
n /= 1000
def tool_env():
env = dict(os.environ)
env["PATH"] = env.get("PATH", "/usr/bin:/bin") + ":/usr/sbin:/sbin"
env["LC_ALL"] = "C"
return env
# ==== Output ====
class Reporter:
"""Progress for the GUI (--machine: one '@tag text' line each) or a terminal."""
def __init__(self, machine):
self.machine = machine
self.cancel = threading.Event()
self.cancellable = False
self._last = 0.0
self._inline = False
self._tty = sys.stdout.isatty()
def _write(self, line, end="\n"):
try:
sys.stdout.write(line + end)
sys.stdout.flush()
except (BrokenPipeError, OSError):
# The GUI went away. Finish the job anyway: an update stopped
# halfway is worse than one nobody watches.
devnull = os.open(os.devnull, os.O_WRONLY)
os.dup2(devnull, sys.stdout.fileno())
def _emit(self, tag, text):
text = " ".join(str(text).splitlines()) if tag != "log" else str(text).rstrip("\n")
if self.machine:
for part in text.splitlines() or [""]:
self._write(f"@{tag} {part}")
return
if self._inline:
self._write("")
self._inline = False
prefix = {"stage": "==> ", "log": " ", "warn": "warning: ",
"error": "error: ", "done": "==> "}[tag]
self._write(prefix + text)
def stage(self, text): self._emit("stage", text)
def log(self, text): self._emit("log", text)
def warn(self, text): self._emit("warn", text)
def error(self, text): self._emit("error", text)
def done(self, text): self._emit("done", text)
def progress(self, done, total, text, force=False):
now = time.monotonic()
if not force and done < total and now - self._last < 0.2:
return
self._last = now
if self.machine:
self._write(f"@progress {int(done)} {int(max(total, 1))} {text}")
elif self._tty:
pct = 100 * done / max(total, 1)
self._write(f"\r[{pct:5.1f}%] {text}".ljust(78), end="")
self._inline = True
def check_cancel(self):
if self.cancellable and self.cancel.is_set():
raise Fail("cancelled")
def watch_stdin(rep):
"""The GUI cancels by writing 'cancel' (or by closing the pipe)."""
def run():
try:
for line in sys.stdin:
if line.strip() == "cancel":
rep.cancel.set()
except (OSError, ValueError):
pass
rep.cancel.set()
threading.Thread(target=run, daemon=True).start()
def run(rep, cmd, ok=(0,), quiet=False):
try:
p = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
stdin=subprocess.DEVNULL, text=True, errors="replace", env=tool_env())
except FileNotFoundError:
raise Fail(f"{cmd[0]} is not installed")
if not quiet or p.returncode not in ok:
for line in p.stdout.splitlines():
if line.strip():
rep.log(line)
if p.returncode not in ok:
raise Fail(f"{os.path.basename(cmd[0])} failed (exit {p.returncode})")
return p
# ==== Partition tables ====
class Part:
def __init__(self, number, ptype, start, sectors):
self.number, self.type, self.start, self.sectors = number, ptype, start, sectors
self.node = None
self.label = ""
self.fstype = ""
@property
def offset(self): return self.start * SECTOR
@property
def size(self): return self.sectors * SECTOR
@property
def end(self): return self.start + self.sectors
def parse_mbr(sector0):
"""The four primary entries of an MBR, or None. Raises on a GPT disk."""
if len(sector0) < 512 or sector0[510:512] != b"\x55\xaa":
return None
parts = []
for i in range(4):
e = sector0[446 + 16 * i: 462 + 16 * i]
ptype = e[4]
start, count = struct.unpack_from("<II", e, 8)
if ptype == 0xEE:
raise Fail("the card uses a GPT partition table, not the MBR a MontaukOS image has")
if ptype and count:
parts.append(Part(i + 1, ptype, start, count))
return parts
# ==== Images ====
class Image:
"""A raw, gzip, xz or zip image, streamed without unpacking it to disk."""
def __init__(self, path):
self.path = os.path.abspath(path)
if not os.path.isfile(self.path):
raise Fail(f"image {path} not found")
self.file_size = os.path.getsize(self.path)
with open(self.path, "rb") as f:
magic = f.read(6)
if magic[:2] == b"\x1f\x8b":
self.kind = "gzip"
elif magic == b"\xfd7zXZ\x00":
self.kind = "xz"
elif magic[:4] == b"PK\x03\x04":
self.kind = "zip"
else:
self.kind = "raw"
def _zip_member(self, zf):
infos = [i for i in zf.infolist() if not i.is_dir()]
imgs = [i for i in infos if i.filename.lower().endswith(".img")]
pick = imgs or infos
if not pick:
raise Fail("the zip file is empty")
return max(pick, key=lambda i: i.file_size)
def size(self):
"""The uncompressed size, or None where it is not cheap to find out."""
if self.kind == "raw":
return self.file_size
if self.kind == "zip":
with zipfile.ZipFile(self.path) as zf:
return self._zip_member(zf).file_size
if self.kind == "gzip" and self.file_size < (1 << 32):
with open(self.path, "rb") as f:
f.seek(-4, os.SEEK_END)
return struct.unpack("<I", f.read(4))[0] or None
return None
@contextlib.contextmanager
def open(self):
"""Yields (stream, raw) where raw.tell() / file_size is the progress."""
raw = open(self.path, "rb")
try:
if self.kind == "raw":
yield raw, raw
elif self.kind == "gzip":
with gzip.GzipFile(fileobj=raw) as s:
yield s, raw
elif self.kind == "xz":
with lzma.LZMAFile(raw) as s:
yield s, raw
else:
with zipfile.ZipFile(raw) as zf, zf.open(self._zip_member(zf)) as s:
yield s, raw
finally:
raw.close()
def materialize(self, rep, workdir):
"""A raw image file: this one, or a decompressed copy in workdir."""
if self.kind == "raw":
return self.path
out = os.path.join(workdir, "image.img")
rep.stage(f"Decompressing {os.path.basename(self.path)}")
with self.open() as (s, raw), open(out, "wb") as o:
while True:
buf = s.read(CHUNK)
if not buf:
break
o.write(buf)
rep.progress(raw.tell(), self.file_size, "Decompressing")
return out
def partitions(self, raw_path):
with open(raw_path, "rb") as f:
parts = parse_mbr(f.read(512))
if not parts:
raise Fail("the image has no MBR partition table; is it a MontaukOS Raspberry Pi image?")
return parts
# ==== Block devices ====
def disk_of(device):
"""The /sys/block name of a whole disk, or Fail."""
real = os.path.realpath(device)
name = os.path.basename(real)
if not real.startswith("/dev/") or not os.path.isdir(f"/sys/block/{name}"):
raise Fail(f"{device} is not a whole disk (give /dev/sdX or /dev/mmcblkN, not a partition)")
if not stat.S_ISBLK(os.stat(real).st_mode):
raise Fail(f"{device} is not a block device")
return name
def sysfs_int(path, default=0):
try:
with open(path) as f:
return int(f.read().strip())
except (OSError, ValueError):
return default
def disk_bytes(name):
return sysfs_int(f"/sys/block/{name}/size") * SECTOR
def sys_partitions(name):
"""{number: Part} as the kernel currently sees the disk's partitions."""
parts = {}
base = f"/sys/block/{name}"
for entry in os.listdir(base):
num = sysfs_int(f"{base}/{entry}/partition", 0)
if num:
p = Part(num, 0, sysfs_int(f"{base}/{entry}/start"), sysfs_int(f"{base}/{entry}/size"))
p.node = f"/dev/{entry}"
parts[num] = p
return parts
def holders(sysname):
"""Everything stacked on a block device (dm-crypt, LVM, md), recursively."""
out = []
for base in (f"/sys/block/{sysname}", f"/sys/class/block/{sysname}"):
hdir = os.path.join(base, "holders")
if os.path.isdir(hdir):
for h in os.listdir(hdir):
out.append(h)
out.extend(holders(h))
break
return out
def disk_mounts(name):
"""[(device node, mount point)] of the disk, its partitions and their holders."""
nodes = {f"/dev/{name}"} | {p.node for p in sys_partitions(name).values()}
for n in list(nodes):
nodes |= {f"/dev/{h}" for h in holders(os.path.basename(n))}
real = {os.path.realpath(n): n for n in nodes}
found = []
with open("/proc/self/mounts") as f:
for line in f:
dev, mnt = line.split()[:2]
mnt = mnt.replace("\\040", " ")
if dev.startswith("/dev/") and os.path.realpath(dev) in real:
found.append((dev, mnt))
with open("/proc/swaps") as f:
for line in f.readlines()[1:]:
dev = line.split()[0]
if os.path.realpath(dev) in real:
found.append((dev, "[SWAP]"))
return found
def check_target(rep, device):
"""The disk name of a device it is safe to write, or Fail."""
name = disk_of(device)
loop_ok = name.startswith("loop") and os.environ.get("MONTAUK_IMAGER_TEST_LOOP") == "1"
if name.startswith(("loop", "zram", "ram", "sr", "dm-", "md")) and not loop_ok:
raise Fail(f"{device} is not an SD card or USB drive")
for _, mnt in disk_mounts(name):
if mnt in SYSTEM_MOUNTS:
raise Fail(f"{device} holds this computer's own system ({mnt}); refusing to touch it")
if sysfs_int(f"/sys/block/{name}/ro"):
raise Fail(f"{device} is write-protected (check the lock switch on the side of the card)")
if disk_bytes(name) == 0:
raise Fail(f"{device} has no medium (is the card inserted?)")
lbs = sysfs_int(f"/sys/block/{name}/queue/logical_block_size", SECTOR)
if lbs != SECTOR:
raise Fail(f"{device} has {lbs}-byte sectors; the image's partition table assumes 512")
return name
def unmount_disk(rep, name):
for part in sys_partitions(name).values():
if holders(os.path.basename(part.node)):
raise Fail(f"{part.node} is in use by device mapper/LVM/RAID; close it first")
for attempt in range(3):
mounts = disk_mounts(name)
if not mounts:
return
for dev, mnt in sorted(mounts, key=lambda m: -len(m[1])):
if mnt == "[SWAP]":
raise Fail(f"{dev} is in use as swap")
rep.log(f"unmounting {dev} from {mnt}")
run(rep, ["umount", mnt], ok=(0, 32), quiet=True)
time.sleep(0.5)
if disk_mounts(name):
raise Fail(f"could not unmount /dev/{name}; close any program using it")
def settle():
with contextlib.suppress(FileNotFoundError, subprocess.SubprocessError):
subprocess.run(["udevadm", "settle", "--timeout=10"], env=tool_env(),
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, timeout=15)
def reread_partitions(rep, name, want):
"""Make the kernel load the new table; wait for the partitions we expect."""
dev = f"/dev/{name}"
for attempt in range(10):
settle()
unmount_disk(rep, name)
fd = os.open(dev, os.O_RDONLY | os.O_CLOEXEC)
try:
fcntl.ioctl(fd, BLKRRPART)
err = None
except OSError as e:
err = e
finally:
os.close(fd)
if err is None:
break
if err.errno != errno.EBUSY:
raise Fail(f"could not reread the partition table of {dev}: {err.strerror}")
time.sleep(1)
else:
# The kernel keeps the old table while a partition is open (a desktop
# probing it); partx updates the entries one by one instead.
run(rep, ["partx", "-u", dev], ok=(0, 1), quiet=True)
deadline = time.monotonic() + 10
while time.monotonic() < deadline:
settle()
parts = sys_partitions(name)
if all(n in parts and os.path.exists(parts[n].node) for n in want):
return parts
time.sleep(0.5)
raise Fail(f"the partitions of {dev} did not appear after writing the table")
@contextlib.contextmanager
def no_automount(rep, name):
"""Keep the desktop from mounting the card's partitions while we work.
Closing a disk opened for writing makes udev re-probe it, and a desktop
then mounts the partitions it finds (and Linux writes to an ext2 volume as
soon as it is mounted read-write). A runtime rule turns udisks' automount
off for this disk until we are done."""
rules = "/run/udev/rules.d"
path = os.path.join(rules, f"90-montauk-imager-{name}.rules")
sep = "p" if name[-1].isdigit() else ""
try:
os.makedirs(rules, exist_ok=True)
with open(path, "w") as f:
f.write("# Written by MontaukOS's montauk-imager.py while it writes this disk.\n"
f'KERNEL=="{name}", ENV{{UDISKS_AUTO}}="0"\n'
f'KERNEL=="{name}{sep}[0-9]*", ENV{{UDISKS_AUTO}}="0", ENV{{UDISKS_IGNORE}}="1"\n')
subprocess.run(["udevadm", "control", "--reload"], env=tool_env(),
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
except (OSError, subprocess.SubprocessError):
path = None
try:
yield
finally:
if path:
with contextlib.suppress(OSError, subprocess.SubprocessError):
os.unlink(path)
subprocess.run(["udevadm", "control", "--reload"], env=tool_env(),
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
# Let the desktop see the finished card.
subprocess.run(["udevadm", "trigger", "--action=change", f"/dev/{name}"],
env=tool_env(), stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
def probe(node):
"""blkid's LABEL and TYPE of a partition (low-level probe, no cache)."""
p = subprocess.run(["blkid", "-p", "-o", "export", node], stdout=subprocess.PIPE,
stderr=subprocess.DEVNULL, text=True, env=tool_env())
info = {}
for line in p.stdout.splitlines():
k, _, v = line.partition("=")
info[k] = v
return info.get("LABEL", ""), info.get("TYPE", "")
def card_layout(rep, name):
"""(esp, root) Parts of a MontaukOS card, root None for the ramdisk layout."""
with open(f"/dev/{name}", "rb") as f:
parts = parse_mbr(f.read(512))
if not parts:
raise Fail(f"/dev/{name} has no partition table; write a fresh image instead")
nodes = sys_partitions(name)
esp = root = None
for p in parts:
if p.number not in nodes:
continue
p.node = nodes[p.number].node
p.label, p.fstype = probe(p.node)
if p.label == ROOT_LABEL and p.fstype.startswith("ext"):
root = root or p
elif p.fstype == "vfat" and (p.label == ESP_LABEL or (p.type == TYPE_ESP and not esp)):
esp = p
if root and root.fstype != "ext2":
raise Fail(f"{root.node} is {root.fstype}, but MontaukOS needs ext2; write a fresh image")
if not esp:
raise Fail(f"/dev/{name} has no MontaukOS boot partition (FAT32, label {ESP_LABEL})")
return esp, root, parts
# ==== ext2 ====
def fs_info(rep, node):
p = run(rep, ["dumpe2fs", "-h", node], quiet=True)
info = {}
for line in p.stdout.splitlines():
k, _, v = line.partition(":")
info[k.strip()] = v.strip()
return (int(info.get("Block count", 0)), int(info.get("Block size", 4096)),
set(info.get("Filesystem features", "").split()))
def check_features(rep, node):
_, _, features = fs_info(rep, node)
bad = (features & EXT_INCOMPAT) - SUPPORTED_INCOMPAT
if bad:
raise Fail(f"{node} now has ext2 features MontaukOS cannot mount: {', '.join(sorted(bad))}")
def fsck(rep, node, check_only=False):
if check_only:
p = run(rep, ["e2fsck", "-f", "-n", node], ok=range(0, 256), quiet=True)
if p.returncode != 0:
for line in p.stdout.splitlines()[-20:]:
rep.log(line)
raise Fail(f"e2fsck reports problems on {node} after the update (exit {p.returncode})")
return
rep.log(f"checking {node}")
p = run(rep, ["e2fsck", "-f", "-p", node], ok=range(0, 256), quiet=True)
if p.returncode & ~3:
for line in p.stdout.splitlines()[-20:]:
rep.log(line)
raise Fail(f"{node} has errors e2fsck cannot fix by itself; run `sudo e2fsck -f {node}`")
if p.returncode:
rep.warn(f"e2fsck repaired {node} (it was probably not shut down cleanly)")
def expand_root(rep, name):
"""Grow the root partition to the end of the card, then the volume."""
unmount_disk(rep, name)
esp, root, parts = card_layout(rep, name)
if not root:
raise Fail("this card has no ext2 system partition to expand (ramdisk layout)")
if any(p.start > root.start for p in parts):
raise Fail(f"partition {root.number} is not the last one; cannot expand it")
rep.stage("Expanding the system partition")
total = disk_bytes(name) // SECTOR
end = total // ALIGN * ALIGN
count = min(end - root.start, MAX_ROOT_BYTES // SECTOR, 0xFFFFFFFF // ALIGN * ALIGN)
if count > root.sectors:
dev = f"/dev/{name}"
fd = os.open(dev, os.O_RDWR | os.O_EXCL | os.O_CLOEXEC)
try:
mbr = bytearray(os.pread(fd, SECTOR, 0))
struct.pack_into("<I", mbr, 446 + 16 * (root.number - 1) + 12, count)
os.pwrite(fd, bytes(mbr), 0)
os.fsync(fd)
finally:
os.close(fd)
rep.log(f"partition {root.number}: {human(root.size)} -> {human(count * SECTOR)}")
reread_partitions(rep, name, [p.number for p in parts])
root.sectors = count
else:
rep.log(f"partition {root.number} already fills the card ({human(root.size)})")
fsck(rep, root.node)
blocks, bs, _ = fs_info(rep, root.node)
target = min(root.size, MAX_ROOT_BYTES) // bs
if target > blocks:
rep.log(f"growing the volume: {human(blocks * bs)} -> {human(target * bs)}")
rep.progress(0, 1, "Growing the volume (resize2fs)", force=True)
run(rep, ["resize2fs", root.node, str(target)])
check_features(rep, root.node)
fsck(rep, root.node, check_only=True)
else:
rep.log(f"the volume already fills the partition ({human(blocks * bs)})")
# ==== Flash ====
def cmd_flash(rep, args):
img = Image(args.image)
name = check_target(rep, args.device)
dev = f"/dev/{name}"
card = disk_bytes(name)
known = img.size()
if known and known > card:
raise Fail(f"the image ({human(known)}) does not fit on the card ({human(card)})")
confirm(args, f"Everything on {dev} ({human(card)}) will be erased.")
rep.stage(f"Preparing {dev}")
with no_automount(rep, name):
unmount_disk(rep, name)
try:
fd = os.open(dev, os.O_RDWR | os.O_EXCL | os.O_CLOEXEC)
except OSError as e:
raise Fail(f"cannot open {dev}: {e.strerror} (is something still using it?)")
# One descriptor from the first write to the reread: udev re-probes
# the disk when a writer closes it, and a volume mounted before the
# read-back would be changed under the verification.
try:
written, digest = write_image(rep, img, fd, card, dev)
if not args.no_verify:
verify_image(rep, fd, written, digest)
parts = parse_mbr(os.pread(fd, SECTOR, 0)) or []
with contextlib.suppress(OSError):
fcntl.ioctl(fd, BLKRRPART)
finally:
os.close(fd)
reread_partitions(rep, name, [p.number for p in parts])
if args.expand:
expand_root(rep, name)
settle()
rep.done(f"MontaukOS is on {dev}. You can remove the card.")
def write_image(rep, img, fd, card, dev):
digest = hashlib.sha256()
written = 0
rep.stage(f"Writing {os.path.basename(img.path)} to {dev}")
rep.cancellable = True
try:
start = time.monotonic()
synced = 0
with img.open() as (stream, raw):
while True:
rep.check_cancel()
buf = stream.read(CHUNK)
if not buf:
break
if written == 0:
try:
ok = parse_mbr(buf[:SECTOR]) is not None
except Fail:
ok = False
if not ok:
raise Fail(f"{os.path.basename(img.path)} is not a MontaukOS Raspberry Pi image "
"(no MBR partition table); nothing was written")
if written + len(buf) > card:
raise Fail(f"the image does not fit on the card ({human(card)})")
view = memoryview(buf)
while view:
n = os.write(fd, view)
view = view[n:]
digest.update(buf)
written += len(buf)
if written - synced >= 32 * MIB:
os.fsync(fd) # keep the progress honest: page cache is not the card
synced = written
rate = written / max(time.monotonic() - start, 0.001)
rep.progress(raw.tell(), img.file_size,
f"Writing: {written // MIB} MiB ({rate / 1e6:.1f} MB/s)")
rep.progress(1, 1, "Flushing to the card", force=True)
os.fsync(fd)
except Fail as e:
if str(e) == "cancelled":
os.fsync(fd)
raise Fail("cancelled; the card is incomplete now, write it again before using it")
raise
finally:
rep.cancellable = False
# Clear the last MiB: a backup GPT left by an earlier OS would otherwise
# outlive the MBR this image brings.
if card - written > 2 * MIB:
os.pwrite(fd, bytes(MIB), card - MIB)
os.fsync(fd)
rep.log(f"wrote {human(written)}")
return written, digest.digest()
def verify_image(rep, fd, written, digest):
rep.stage("Verifying")
fcntl.ioctl(fd, BLKFLSBUF) # read the card, not the page cache
os.posix_fadvise(fd, 0, 0, os.POSIX_FADV_DONTNEED)
check = hashlib.sha256()
done = 0
start = time.monotonic()
while done < written:
if rep.cancel.is_set():
rep.warn("verification cancelled; the card was written completely but not checked")
return
buf = os.pread(fd, min(CHUNK, written - done), done)
if not buf:
break
check.update(buf)
done += len(buf)
rate = done / max(time.monotonic() - start, 0.001)
rep.progress(done, written, f"Verifying: {done // MIB} MiB ({rate / 1e6:.1f} MB/s)")
if done != written or check.digest() != digest:
raise Fail("verification failed: the card does not read back what was written "
"(a failing or counterfeit card?)")
rep.log("the card reads back correctly")
# ==== Update ====
def toml_sections(text):
"""[(table name or None, [lines])]: each table with the lines up to the next."""
header = re.compile(r"^\[([A-Za-z0-9_.\-\"]+)\]\s*(#.*)?$")
secs = [(None, [])]
for line in text.splitlines():
m = header.match(line)
if m:
secs.append((m.group(1), [line]))
else:
secs[-1][1].append(line)
return secs
def merge_init(new, old):
"""The image's init.toml, keeping the card's on/off choices and own services."""
enabled = re.compile(r"^\s*enabled\s*=\s*(true|false)\b")
old_secs = {n: l for n, l in toml_sections(old) if n}
new_secs = toml_sections(new)
names = {n for n, _ in new_secs}
out = []
for n, lines in new_secs:
lines = list(lines)
if n and n.startswith("services.") and n in old_secs:
choice = next((m.group(1) for m in map(enabled.match, old_secs[n]) if m), None)
if choice is not None:
idx = [i for i, l in enumerate(lines) if enabled.match(l)]
if idx:
lines[idx[0]] = f"enabled = {choice}"
else:
lines.insert(1, f"enabled = {choice}")
out.extend(lines)
extra = [(n, l) for n, l in toml_sections(old)
if n and n.startswith("services.") and n not in names]
if extra:
out += ["", "# Services from the previous installation."]
for _, lines in extra:
out.extend(lines)
return "\n".join(out).rstrip("\n") + "\n"
def merge_caps(new, old):
"""The image's capabilities.toml plus the grants only the card has."""
names = {n for n, _ in toml_sections(new)}
extra = [(n, l) for n, l in toml_sections(old)
if n and n.startswith("grant.") and n not in names]
out = new.rstrip("\n").splitlines()
if extra:
out += ["", "# ==== Grants from the previous installation ====", ""]
for _, lines in extra:
out.extend(lines)
return "\n".join(out).rstrip("\n") + "\n"
class Op:
def __init__(self, kind, rel, dst, src=None, data=None, size=0, replaced=0):
self.kind, self.rel, self.dst, self.src, self.data = kind, rel, dst, src, data
self.size, self.replaced = size, replaced
def alloc(size, bs, ext):
blocks = -(-size // bs)
if ext and blocks > 12:
blocks += -(-blocks // (bs // 4)) # indirect blocks
return blocks * bs
class Plan:
def __init__(self, what, dst_root, ext):
self.what, self.dst_root, self.ext = what, dst_root, ext
self.ops = []
self.unchanged = 0
st = os.statvfs(dst_root)
self.bs = st.f_frsize
self.free = st.f_bavail * st.f_frsize
def write(self, rel, src=None, data=None):
dst = os.path.join(self.dst_root, rel)
size = len(data) if data is not None else os.path.getsize(src)
replaced = os.path.getsize(dst) if os.path.isfile(dst) else 0
self.ops.append(Op("write", rel, dst, src, data, size, replaced))
def needed(self):
a = lambda n: alloc(n, self.bs, self.ext)
net = 0
biggest = 0
for op in self.ops:
if op.kind == "write":
net += a(op.size) - (a(op.replaced) if op.replaced else 0)
biggest = max(biggest, a(op.size))
elif op.kind == "backup":
net += a(op.size)
elif op.kind == "delete":
net -= a(op.size)
return net + biggest + MIB
def count(self, kind):
return sum(1 for op in self.ops if op.kind == kind)
def same_file(a, b):
try:
return os.path.getsize(a) == os.path.getsize(b) and filecmp.cmp(a, b, shallow=False)
except OSError:
return False
def plan_tree(rep, plan, src_root, decide, prune_dirs=(), preserve=()):
shipped = set()
files = []
for dirpath, dirnames, filenames in os.walk(src_root):
rel_dir = os.path.relpath(dirpath, src_root)
rel_dir = "" if rel_dir == "." else rel_dir
dirnames[:] = sorted(d for d in dirnames
if decide(os.path.join(rel_dir, d), True) != "skip")
for d in dirnames:
shipped.add(os.path.join(rel_dir, d))
for fn in sorted(filenames):
files.append(os.path.join(rel_dir, fn))
shipped.add(files[-1])
for i, rel in enumerate(files):
rep.progress(i, len(files), f"Comparing {plan.what} ({i}/{len(files)})")
src = os.path.join(src_root, rel)
dst = os.path.join(plan.dst_root, rel)
policy = decide(rel, False)
exists = os.path.lexists(dst)
if exists and os.path.isdir(dst):
raise Fail(f"{rel} is a directory on the card but a file in the image")
if policy == "skip" or (policy == "add" and exists):
continue
if policy == "firstboot" and not exists:
rep.log("not adding os/firstboot.elf: the card's system volume has already grown")
continue
if policy in ("merge-init", "merge-caps") and exists:
with open(src, encoding="utf-8", errors="replace") as f:
new = f.read()
with open(dst, encoding="utf-8", errors="replace") as f:
old = f.read()
merged = (merge_init if policy == "merge-init" else merge_caps)(new, old)
if merged == old:
plan.unchanged += 1
continue
plan.ops.append(Op("backup", rel, dst, size=len(old.encode())))
plan.write(rel, data=merged.encode())
continue
if exists and same_file(src, dst):
plan.unchanged += 1
continue
if policy == "backup" and exists:
plan.ops.append(Op("backup", rel, dst, size=os.path.getsize(dst)))
plan.write(rel, src=src)
for top in prune_dirs:
base = os.path.join(plan.dst_root, top)
if not os.path.isdir(base) or top not in shipped:
continue
for dirpath, dirnames, filenames in os.walk(base, topdown=False):
rel_dir = os.path.relpath(dirpath, plan.dst_root)
for fn in filenames:
rel = os.path.join(rel_dir, fn)
if rel not in shipped and rel not in preserve:
path = os.path.join(dirpath, fn)
plan.ops.append(Op("delete", rel, path, size=os.path.getsize(path)))
if rel_dir not in shipped:
plan.ops.append(Op("rmdir", rel_dir, dirpath))
def execute(rep, plan, done, total):
for op in plan.ops:
if op.kind == "delete":
rep.log(f"removing {op.rel}")
os.unlink(op.dst)
for op in plan.ops:
if op.kind == "rmdir":
with contextlib.suppress(OSError):
os.rmdir(op.dst) # only if it ended up empty
for op in plan.ops:
if op.kind == "backup":
suffix = ".bak" if op.rel in ESP_BACKUP else ".old"
rep.log(f"keeping the card's {op.rel} as {os.path.basename(op.rel)}{suffix}")
shutil.copyfile(op.dst, op.dst + suffix)
for op in plan.ops:
if op.kind != "write":
continue
os.makedirs(os.path.dirname(op.dst), exist_ok=True)
tmp = os.path.join(os.path.dirname(op.dst), ".imager-" + os.path.basename(op.dst))
with open(tmp, "wb") as out:
if op.data is not None:
out.write(op.data)
done += len(op.data)
else:
with open(op.src, "rb") as src:
while True:
buf = src.read(CHUNK)
if not buf:
break
out.write(buf)
done += len(buf)
rep.progress(done, total, f"Writing {op.rel}")
out.flush()
os.fsync(out.fileno())
if op.src:
with contextlib.suppress(OSError): # FAT has no modes
os.chmod(tmp, os.stat(op.src).st_mode & 0o777)
os.replace(tmp, op.dst)
rep.progress(done, total, f"Writing {op.rel}")
return done
@contextlib.contextmanager
def mounted(rep, workdir, label, source, fstype, opts):
mnt = os.path.join(workdir, label)
os.makedirs(mnt, exist_ok=True)
run(rep, ["mount", "-t", fstype, "-o", opts, source, mnt], quiet=True)
try:
yield mnt
finally:
os.sync()
p = subprocess.run(["umount", mnt], env=tool_env(),
stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT)
if p.returncode:
time.sleep(1)
subprocess.run(["umount", "-l", mnt], env=tool_env(),
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
def cmd_update(rep, args):
img = Image(args.image)
name = check_target(rep, args.device)
with no_automount(rep, name):
update(rep, args, img, name)
def update(rep, args, img, name):
dev = f"/dev/{name}"
workdir = tempfile.mkdtemp(prefix="montauk-imager-")
try:
raw = img.materialize(rep, workdir)
iparts = img.partitions(raw)
iesp = next((p for p in iparts if p.type in FAT_TYPES), None)
iroot = next((p for p in iparts if p.type == TYPE_LINUX), None)
if not iesp:
raise Fail("the image has no boot partition; is it a MontaukOS Raspberry Pi image?")
rep.stage(f"Inspecting {dev}")
unmount_disk(rep, name)
esp, root, _ = card_layout(rep, name)
if bool(iroot) != bool(root):
have = "an ext2 system partition" if root else "the single-partition (ramdisk) layout"
want = "an ext2 system partition" if iroot else "the single-partition (ramdisk) layout"
raise Fail(f"the card has {have} but the image uses {want}; write a fresh image instead")
rep.log(f"boot partition {esp.node} ({human(esp.size)})")
if root:
rep.log(f"system partition {root.node} ({human(root.size)}, "
f"{'fills the card' if root.end + ALIGN >= disk_bytes(name) // SECTOR else 'not expanded'})")
confirm(args, f"MontaukOS on {dev} will be updated; users, files and settings are kept.")
signal.signal(signal.SIGHUP, signal.SIG_IGN) # an update must not stop halfway
signal.signal(signal.SIGINT, signal.SIG_IGN)
if root:
rep.stage("Checking the system volume")
fsck(rep, root.node)
if args.expand:
expand_root(rep, name)
esp, root, _ = card_layout(rep, name)
with contextlib.ExitStack() as stack:
rep.stage("Mounting")
loop = lambda p: f"ro,loop,offset={p.offset},sizelimit={p.size}"
src_esp = stack.enter_context(mounted(rep, workdir, "image-esp", raw, "vfat", loop(iesp)))
dst_esp = stack.enter_context(mounted(rep, workdir, "card-esp", esp.node, "vfat", "rw,noatime"))
plans = []
def esp_rule(rel, is_dir):
if is_dir:
return "system"
if rel == FIRMWARE_FILE and not args.firmware:
return "add"
return "backup" if rel in ESP_BACKUP else "system"
rep.stage("Comparing the card with the image")
plans.append((Plan("the boot partition", dst_esp, False), src_esp, esp_rule, (), ()))
if root:
src_root = stack.enter_context(mounted(rep, workdir, "image-root", raw, "ext2", loop(iroot)))
dst_root = stack.enter_context(mounted(rep, workdir, "card-root", root.node, "ext2", "rw,noatime"))
def root_rule(rel, is_dir):
top = rel.split("/")[0]
if top in STATE_DIRS:
return "skip"
if is_dir:
return "system"
if rel == FIRSTBOOT:
return "firstboot"
if rel in (INIT_TOML, CAPS_TOML):
if args.keep_config:
return "add"
return "merge-init" if rel == INIT_TOML else "merge-caps"
return "add" if top in MERGE_DIRS else "system"
plans.append((Plan("the system partition", dst_root, True), src_root, root_rule,
PRUNE_DIRS if args.prune else (), PRESERVE))
for plan, src, rule, prune, keep in plans:
plan_tree(rep, plan, src, rule, prune, keep)
need = plan.needed()
rep.log(f"{plan.what}: {plan.count('write')} to write, {plan.unchanged} unchanged"
+ (f", {plan.count('delete')} to remove" if plan.count("delete") else ""))
if need > plan.free:
hint = ("enable Expand (--expand) or write a fresh image"
if plan.ext else "write a fresh image")
raise Fail(f"not enough space on {plan.what}: needs {human(need)}, "
f"{human(plan.free)} free; {hint}")
total = sum(op.size for plan, *_ in plans for op in plan.ops if op.kind == "write")
if total == 0 and not any(plan.ops for plan, *_ in plans):
rep.log("the card already matches the image")
rep.stage("Updating")
done = 0
for plan, *_ in plans:
done = execute(rep, plan, done, total)
rep.progress(1, 1, "Flushing to the card", force=True)
os.sync()
rep.stage("Unmounting")
if root:
rep.stage("Checking the system volume")
fsck(rep, root.node, check_only=True)
check_features(rep, root.node)
settle()
rep.done(f"MontaukOS on {dev} is up to date. You can remove the card.")
finally:
remove_workdir(rep, workdir)
def remove_workdir(rep, workdir):
"""rmtree crosses mount points: never let it into a card left mounted."""
for entry in os.listdir(workdir):
path = os.path.join(workdir, entry)
if os.path.ismount(path):
rep.warn(f"{path} is still mounted; unmount it by hand (sudo umount {path})")
elif os.path.isdir(path):
with contextlib.suppress(OSError):
os.rmdir(path)
else:
with contextlib.suppress(OSError):
os.unlink(path)
with contextlib.suppress(OSError):
os.rmdir(workdir)
def cmd_expand(rep, args):
name = check_target(rep, args.device)
confirm(args, f"The system partition on /dev/{name} will be grown to fill the card.")
signal.signal(signal.SIGINT, signal.SIG_IGN)
with no_automount(rep, name):
expand_root(rep, name)
settle()
rep.done(f"The system partition on /dev/{name} fills the card. You can remove the card.")
def confirm(args, what):
if args.yes:
return
if not sys.stdin.isatty():
raise Fail(f"{what} Pass --yes to go ahead.")
answer = input(f"{what} Type 'yes' to continue: ")
if answer.strip().lower() != "yes":
raise Fail("aborted")
# ==== Card discovery (unprivileged) ====
def truthy(v):
return v in (True, 1, "1", "true")
def list_cards(show_all=False):
"""Removable disks (or all non-system disks) as dicts, for the GUI and `list`."""
cols = "NAME,PATH,SIZE,TYPE,RM,HOTPLUG,TRAN,MODEL,VENDOR,FSTYPE,LABEL,MOUNTPOINT,RO"
try:
p = subprocess.run(["lsblk", "-J", "-b", "-o", cols], stdout=subprocess.PIPE,
stderr=subprocess.DEVNULL, text=True, env=tool_env(), timeout=10)
devices = json.loads(p.stdout or "{}").get("blockdevices", [])
except (OSError, ValueError, subprocess.SubprocessError):
return []
def mounts(d):
out = [d.get("mountpoint")] if d.get("mountpoint") else []
for c in d.get("children") or []:
out += mounts(c)
return out
cards = []
for d in devices:
name = d.get("name", "")
size = int(d.get("size") or 0)
if d.get("type") != "disk" or size == 0 or name.startswith(("zram", "loop", "ram")):
continue
if any(m in SYSTEM_MOUNTS for m in mounts(d)):
continue
removable = (truthy(d.get("rm")) or truthy(d.get("hotplug"))
or d.get("tran") in ("usb", "mmc", "sdio") or name.startswith("mmcblk"))
if not removable and not show_all:
continue
parts = []
for c in d.get("children") or []:
if c.get("type") != "part":
continue
cname = c.get("name", "")
parts.append({
"path": c.get("path") or f"/dev/{cname}",
"size": int(c.get("size") or 0),
"start": sysfs_int(f"/sys/block/{name}/{cname}/start") * SECTOR,
"label": c.get("label") or "",
"fstype": c.get("fstype") or "",
"mount": c.get("mountpoint") or "",
})
model = " ".join(x.strip() for x in (d.get("vendor") or "", d.get("model") or "") if x and x.strip())
card = {"path": d.get("path") or f"/dev/{name}", "name": name, "size": size,
"model": model or ("SD card" if name.startswith("mmcblk") else "Drive"),
"ro": truthy(d.get("ro")), "removable": removable, "parts": parts}
card.update(describe(card))
cards.append(card)
return cards
def describe(card):
"""What is on a card: {'kind': 'ext2'|'ramdisk'|'other'|'empty', 'contents', 'fills'}."""
parts = card["parts"]
root = next((p for p in parts if p["label"] == ROOT_LABEL), None)
esp = next((p for p in parts if p["label"] == ESP_LABEL and p["fstype"] == "vfat"), None)
if root:
fills = (root["start"] + root["size"] + 2 * MIB >= card["size"]
or root["size"] >= MAX_ROOT_BYTES - MIB)
text = f"MontaukOS, system partition {human(root['size'])}"
text += ", fills the card" if fills else ", not expanded yet"
return {"kind": "ext2", "contents": text, "fills": fills}
if esp:
return {"kind": "ramdisk", "contents": "MontaukOS (single partition, ramdisk)", "fills": True}
if parts:
names = [p["label"] or p["fstype"] or "unknown" for p in parts]
return {"kind": "other", "contents": ", ".join(names), "fills": False}
return {"kind": "empty", "contents": "empty", "fills": False}
def cmd_list(rep, args):
cards = list_cards(args.all)
if not cards:
print("No SD cards or USB drives found." + ("" if args.all else " (--all shows fixed disks too)"))
for c in cards:
print(f"{c['path']:<16} {human(c['size']):>9} {c['model']:<28} {c['contents']}"
+ (" [write-protected]" if c["ro"] else ""))
def image_summary(path):
"""One line about an image file for the GUI, and its layout if it can tell."""
try:
img = Image(path)
except Fail as e:
return str(e), None
when = time.strftime("%Y-%m-%d %H:%M", time.localtime(os.path.getmtime(img.path)))
if img.kind != "raw":
return f"{img.kind} compressed, {human(img.file_size)}, {when}", None
try:
parts = img.partitions(img.path)
except Fail as e:
return str(e), None
layout = "ext2" if any(p.type == TYPE_LINUX for p in parts) else "ramdisk"
desc = "boot + ext2 system partition" if layout == "ext2" else "single partition (ramdisk)"
return f"{human(img.file_size)}, {desc}, built {when}", layout
# ==== GUI ====
def gui_main(initial):
try:
import tkinter as tk
from tkinter import ttk, filedialog, messagebox
except ImportError:
sys.exit("montauk-imager: the GUI needs Tkinter (sudo apt install python3-tk)")
try:
root = tk.Tk()
except tk.TclError as e:
sys.exit(f"montauk-imager: cannot open a window ({e}); use the list/flash/update commands")
root.title("MontaukOS Imager")
root.minsize(680, 600)
style = ttk.Style(root)
if "clam" in style.theme_names():
style.theme_use("clam")
style.configure("Title.TLabel", font=("TkDefaultFont", 15, "bold"))
style.configure("Dim.TLabel", foreground="#666666")
style.configure("Accent.TButton", font=("TkDefaultFont", 10, "bold"))
S = {"busy": None, "proc": None, "cards": [], "sig": None, "layout": None}
events = queue.Queue()
image_var = tk.StringVar(value=initial or (DEFAULT_IMAGE if os.path.isfile(DEFAULT_IMAGE) else ""))
image_info = tk.StringVar()
show_all = tk.BooleanVar(value=False)
action = tk.StringVar(value="flash")
verify = tk.BooleanVar(value=True)
flash_expand = tk.BooleanVar(value=False)
merge_cfg = tk.BooleanVar(value=True)
firmware = tk.BooleanVar(value=False)
prune = tk.BooleanVar(value=False)
upd_expand = tk.BooleanVar(value=False)
status = tk.StringVar(value="Choose an image and an SD card.")
card_info = tk.StringVar()
main = ttk.Frame(root, padding=16)
main.pack(fill="both", expand=True)
main.columnconfigure(0, weight=1)
head = ttk.Frame(main)
head.grid(row=0, column=0, sticky="ew")
ttk.Label(head, text="MontaukOS Imager", style="Title.TLabel").pack(side="left")
ttk.Label(head, text="Raspberry Pi 4 / 400", style="Dim.TLabel").pack(side="left", padx=10, pady=(5, 0))
# Image
imgf = ttk.LabelFrame(main, text="Image", padding=10)
imgf.grid(row=1, column=0, sticky="ew", pady=(12, 0))
imgf.columnconfigure(0, weight=1)
img_entry = ttk.Entry(imgf, textvariable=image_var)
img_entry.grid(row=0, column=0, sticky="ew")
browse_btn = ttk.Button(imgf, text="Browse...", command=lambda: browse())
browse_btn.grid(row=0, column=1, padx=(6, 0))
build_btn = ttk.Button(imgf, text="Build", command=lambda: build())
if REPO:
build_btn.grid(row=0, column=2, padx=(6, 0))
ttk.Label(imgf, textvariable=image_info, style="Dim.TLabel").grid(row=1, column=0, columnspan=3, sticky="w", pady=(4, 0))
# Card
cardf = ttk.LabelFrame(main, text="SD card", padding=10)
cardf.grid(row=2, column=0, sticky="nsew", pady=(10, 0))
main.rowconfigure(2, weight=1)
cardf.columnconfigure(0, weight=1)
cardf.rowconfigure(0, weight=1)
tree = ttk.Treeview(cardf, columns=("device", "size", "model", "contents"), show="headings",
height=4, selectmode="browse")
for col, text, width, anchor in (("device", "Device", 110, "w"), ("size", "Size", 80, "e"),
("model", "Name", 170, "w"), ("contents", "Contents", 300, "w")):
tree.heading(col, text=text, anchor=anchor)
tree.column(col, width=width, anchor=anchor, stretch=(col == "contents"))
tree.grid(row=0, column=0, columnspan=3, sticky="nsew")
ttk.Label(cardf, textvariable=card_info, style="Dim.TLabel").grid(row=1, column=0, sticky="w", pady=(4, 0))
ttk.Checkbutton(cardf, text="Show all drives", variable=show_all,
command=lambda: refresh(force=True)).grid(row=1, column=1, padx=6, pady=(4, 0))
ttk.Button(cardf, text="Refresh", command=lambda: refresh(force=True)).grid(row=1, column=2, pady=(4, 0))
# Action
actf = ttk.LabelFrame(main, text="Action", padding=10)
actf.grid(row=3, column=0, sticky="ew", pady=(10, 0))
actf.columnconfigure(1, weight=1)
rb_flash = ttk.Radiobutton(actf, text="Write a fresh image (erases the card)", value="flash",
variable=action, command=lambda: update_ui())
rb_update = ttk.Radiobutton(actf, text="Update the installation on the card (keeps users, files and settings)",
value="update", variable=action, command=lambda: update_ui())
rb_flash.grid(row=0, column=0, columnspan=2, sticky="w")
rb_update.grid(row=1, column=0, columnspan=2, sticky="w")
optf = ttk.Frame(actf, padding=(22, 4, 0, 0))
optf.grid(row=2, column=0, columnspan=2, sticky="ew")
flash_opts = [
ttk.Checkbutton(optf, text="Verify the card after writing", variable=verify),
ttk.Checkbutton(optf, text="Expand the system partition to fill the card now "
"(otherwise MontaukOS does it on first boot)", variable=flash_expand),
]
update_opts = [
ttk.Checkbutton(optf, text="Merge new services and capability grants into the card's configuration",
variable=merge_cfg),
ttk.Checkbutton(optf, text="Update the UEFI firmware (RPI_EFI.fd; resets the firmware's settings)",
variable=firmware),
ttk.Checkbutton(optf, text="Remove system files the new image no longer ships", variable=prune),
ttk.Checkbutton(optf, text="Expand the system partition to fill the card", variable=upd_expand),
]
# Progress
progf = ttk.Frame(main)
progf.grid(row=4, column=0, sticky="ew", pady=(12, 0))
progf.columnconfigure(0, weight=1)
bar = ttk.Progressbar(progf, mode="determinate", maximum=1000)
bar.grid(row=0, column=0, columnspan=3, sticky="ew")
ttk.Label(progf, textvariable=status).grid(row=1, column=0, sticky="w", pady=(4, 0))
log_btn = ttk.Button(progf, text="Show log", command=lambda: toggle_log())
log_btn.grid(row=1, column=1, padx=6, pady=(4, 0))
logf = ttk.Frame(main)
logf.columnconfigure(0, weight=1)
logf.rowconfigure(0, weight=1)
log = tk.Text(logf, height=10, wrap="word", state="disabled", font=("TkFixedFont", 9),
relief="flat", borderwidth=1)
scroll = ttk.Scrollbar(logf, orient="vertical", command=log.yview)
log.configure(yscrollcommand=scroll.set)
log.grid(row=0, column=0, sticky="nsew")
scroll.grid(row=0, column=1, sticky="ns")
log.tag_configure("stage", font=("TkFixedFont", 9, "bold"))
log.tag_configure("warn", foreground="#a05a00")
log.tag_configure("error", foreground="#b00020", font=("TkFixedFont", 9, "bold"))
log.tag_configure("done", foreground="#0a7a2f", font=("TkFixedFont", 9, "bold"))
btns = ttk.Frame(main)
btns.grid(row=6, column=0, sticky="e", pady=(12, 0))
cancel_btn = ttk.Button(btns, text="Cancel", command=lambda: cancel())
start_btn = ttk.Button(btns, text="Write", style="Accent.TButton", command=lambda: start())
cancel_btn.pack(side="left", padx=(0, 6))
start_btn.pack(side="left")
def toggle_log(show=None):
visible = bool(logf.winfo_ismapped())
show = (not visible) if show is None else show
if show and not visible:
logf.grid(row=5, column=0, sticky="nsew", pady=(8, 0))
main.rowconfigure(5, weight=2)
log_btn.configure(text="Hide log")
elif not show and visible:
logf.grid_remove()
main.rowconfigure(5, weight=0)
log_btn.configure(text="Show log")
def append(text, tag=None):
log.configure(state="normal")
log.insert("end", text + "\n", (tag,) if tag else ())
log.see("end")
log.configure(state="disabled")
def selected():
sel = tree.selection()
return next((c for c in S["cards"] if c["path"] == (sel[0] if sel else None)), None)
def refresh(force=False):
if S["busy"] in ("flash", "update", "expand") and not force:
return
cards = list_cards(show_all.get())
sig = json.dumps(cards, sort_keys=True)
if sig == S["sig"] and not force:
return
S["sig"] = sig
keep = selected()
S["cards"] = cards
tree.delete(*tree.get_children())
for c in cards:
contents = c["contents"] + (" [write-protected]" if c["ro"] else "")
tree.insert("", "end", iid=c["path"],
values=(c["path"], human(c["size"]), c["model"], contents))
if keep and tree.exists(keep["path"]):
tree.selection_set(keep["path"])
elif len(cards) == 1:
tree.selection_set(cards[0]["path"])
update_ui()
def poll_cards():
if not S["busy"]:
refresh()
root.after(2000, poll_cards)
def refresh_image(*_):
path = image_var.get().strip()
if not path:
image_info.set("No image chosen." + (" Build one, or browse for montauk-rpi4.img." if REPO else ""))
S["layout"] = None
else:
info, S["layout"] = image_summary(path)
image_info.set(info)
update_ui()
def update_ui(*_):
card = selected()
installed = bool(card and card["kind"] in ("ext2", "ramdisk"))
rb_update.configure(state="normal" if installed and not S["busy"] else "disabled")
if action.get() == "update" and not installed:
action.set("flash")
for w in flash_opts + update_opts:
w.pack_forget()
opts = flash_opts if action.get() == "flash" else update_opts
for w in opts:
w.pack(anchor="w")
if card:
if card["kind"] == "ext2" and card["fills"]:
upd_expand.set(False)
card_info.set(f"{card['path']}: {card['model']}, {human(card['size'])}, {card['contents']}")
else:
card_info.set("Insert an SD card (directly or in a USB reader)." if not S["cards"]
else "Select the card to write.")
state = "disabled" if S["busy"] else "normal"
for w in opts + [img_entry, browse_btn, build_btn, rb_flash]:
w.configure(state=state)
if card and card["kind"] == "ext2" and card["fills"] and not S["busy"]:
update_opts[3].configure(state="disabled")
if card and card["kind"] == "ramdisk":
update_opts[3].configure(state="disabled")
start_btn.configure(text="Write" if action.get() == "flash" else "Update")
ready = card and image_var.get().strip() and not S["busy"] and not card["ro"]
start_btn.configure(state="normal" if ready else "disabled")
cancel_btn.configure(state="normal" if S["busy"] in ("flash", "build") else "disabled")
def browse():
start_dir = os.path.dirname(image_var.get()) or REPO or os.path.expanduser("~")
path = filedialog.askopenfilename(
parent=root, title="Choose a MontaukOS image", initialdir=start_dir,
filetypes=[("Disk images", "*.img *.img.gz *.img.xz *.gz *.xz *.zip"), ("All files", "*")])
if path:
image_var.set(path)
def launch(kind, cmd, cwd=None):
S["busy"] = kind
bar.configure(mode="determinate", value=0)
append("$ " + " ".join(cmd), "stage")
try:
proc = subprocess.Popen(cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE,
stderr=subprocess.STDOUT, text=True, bufsize=1,
errors="replace", cwd=cwd)
except OSError as e:
S["busy"] = None
fail(f"Could not start {cmd[0]}: {e}")
return
S["proc"] = proc
def reader():
for line in proc.stdout:
events.put(("line", line.rstrip("\n")))
events.put(("exit", proc.wait()))
threading.Thread(target=reader, daemon=True).start()
update_ui()
def fail(msg):
status.set(msg)
append(msg, "error")
toggle_log(True)
messagebox.showerror("MontaukOS Imager", msg, parent=root)
def pump():
try:
while True:
kind, val = events.get_nowait()
if kind == "line":
handle_line(val)
else:
finished(val)
except queue.Empty:
pass
root.after(50, pump)
def handle_line(line):
if S["busy"] == "build" or not line.startswith("@"):
append(line)
if S["busy"] == "build" and line.strip():
status.set(line.strip()[:90])
return
tag, _, text = line[1:].partition(" ")
if tag == "progress":
done, total, text = (text.split(" ", 2) + ["", ""])[:3]
bar.configure(mode="determinate", value=1000 * int(done) / max(int(total), 1))
status.set(text)
elif tag in ("stage", "log", "warn", "error", "done"):
append(("==> " + text) if tag == "stage" else text, tag if tag != "log" else None)
if tag in ("stage", "done"):
status.set(text)
if tag == "error":
S["error"] = text
if tag == "warn":
toggle_log(True)
def finished(rc):
kind = S["busy"]
S["busy"], S["proc"] = None, None
if kind == "build":
bar.stop()
bar.configure(mode="determinate", value=0)
if rc == 0 and os.path.isfile(DEFAULT_IMAGE):
status.set("Image built.")
image_var.set(DEFAULT_IMAGE)
refresh_image()
else:
fail(f"The build failed (make exit {rc}); see the log.")
elif rc == 0:
bar.configure(value=1000)
messagebox.showinfo("MontaukOS Imager", status.get(), parent=root)
else:
err = S.pop("error", None)
if rc == 126 and not err:
status.set("Authorization was cancelled.")
elif rc == 127 and not err:
fail("Could not get administrator rights through pkexec. "
"Run the imager with sudo instead.")
else:
fail(err or f"Failed (exit {rc}); see the log.")
S["sig"] = None
refresh(force=True)
def privileged(args):
base = [sys.executable, "-I", SCRIPT, "--machine"] + args + ["--yes"]
if os.geteuid() == 0:
return base
if not shutil.which("pkexec"):
return None
return ["pkexec"] + base
def start():
card = selected()
path = os.path.abspath(image_var.get().strip())
if not card:
return
if not os.path.isfile(path):
fail(f"The image {path} does not exist.")
return
what = f"{card['model']} ({card['path']}, {human(card['size'])})"
if action.get() == "flash":
if card["kind"] != "empty":
detail = f"It currently holds: {card['contents']}.\n\n"
else:
detail = ""
if not messagebox.askyesno(
"Erase the card?",
f"Everything on {what} will be erased and replaced with MontaukOS.\n\n{detail}"
"Continue?", icon="warning", default="no", parent=root):
return
args = ["flash", path, card["path"]]
if not verify.get():
args.append("--no-verify")
if flash_expand.get():
args.append("--expand")
else:
if S["layout"] and card["kind"] != S["layout"]:
fail("The card and the image use different layouts "
"(ext2 system partition vs single ramdisk partition). Write a fresh image instead.")
return
if not messagebox.askyesno(
"Update MontaukOS?",
f"MontaukOS on {what} will be updated from\n{os.path.basename(path)}.\n\n"
"Users, their files and settings are kept. Do not remove the card "
"until the update has finished.\n\nContinue?", parent=root):
return
args = ["update", path, card["path"]]
if not merge_cfg.get():
args.append("--keep-config")
if firmware.get():
args.append("--firmware")
if prune.get():
args.append("--prune")
if upd_expand.get():
args.append("--expand")
cmd = privileged(args)
if not cmd:
fail("Writing a card needs administrator rights, but pkexec is not installed. "
"Run the imager with sudo.")
return
S.pop("error", None)
status.set("Waiting for authorization..." if cmd[0] == "pkexec" else "Starting...")
launch(args[0], cmd)
def build():
if not REPO:
return
status.set("Building the image (make ARCH=aarch64 rpi4-sd)...")
toggle_log(True)
launch("build", ["make", "ARCH=aarch64", "rpi4-sd"], cwd=REPO)
bar.configure(mode="indeterminate")
bar.start(15)
def cancel():
proc = S["proc"]
if not proc:
return
if S["busy"] == "build":
proc.terminate()
bar.stop()
return
if S["busy"] == "flash" and messagebox.askyesno(
"Cancel writing?", "Stop writing the card? It will not boot until it is written again.",
parent=root):
with contextlib.suppress(OSError):
proc.stdin.write("cancel\n")
proc.stdin.flush()
status.set("Cancelling...")
def on_close():
if S["busy"] == "update" or S["busy"] == "expand":
messagebox.showwarning("MontaukOS Imager",
"The card is being updated. Wait until it has finished.", parent=root)
return
if S["busy"] == "flash":
if not messagebox.askyesno("Quit?", "The card is still being written. Stop and quit?",
parent=root):
return
with contextlib.suppress(OSError):
S["proc"].stdin.write("cancel\n")
S["proc"].stdin.flush()
if S["busy"] == "build":
S["proc"].terminate()
root.destroy()
tree.bind("<<TreeviewSelect>>", update_ui)
image_var.trace_add("write", refresh_image)
root.protocol("WM_DELETE_WINDOW", on_close)
refresh_image()
refresh(force=True)
root.after(50, pump)
root.after(2000, poll_cards)
root.mainloop()
# ==== Entry ====
def main():
ap = argparse.ArgumentParser(description="Write and update MontaukOS SD cards for the Raspberry Pi 4/400.")
ap.add_argument("--machine", action="store_true", help=argparse.SUPPRESS)
ap.add_argument("--image", help="image to preselect in the GUI")
sub = ap.add_subparsers(dest="cmd")
p = sub.add_parser("list", help="list SD cards and USB drives")
p.add_argument("--all", action="store_true", help="include fixed disks")
p = sub.add_parser("flash", help="write an image over the whole card")
p.add_argument("image")
p.add_argument("device")
p.add_argument("--no-verify", action="store_true", help="do not read the card back")
p.add_argument("--expand", action="store_true", help="grow the system partition to fill the card")
p.add_argument("--yes", action="store_true", help="do not ask for confirmation")
p = sub.add_parser("update", help="update the installation on a card, keeping user data")
p.add_argument("image")
p.add_argument("device")
p.add_argument("--keep-config", action="store_true", help="leave init.toml and capabilities.toml alone")
p.add_argument("--firmware", action="store_true", help="also replace RPI_EFI.fd (resets UEFI settings)")
p.add_argument("--prune", action="store_true", help="delete system files the image no longer ships")
p.add_argument("--expand", action="store_true", help="grow the system partition to fill the card first")
p.add_argument("--yes", action="store_true", help="do not ask for confirmation")
p = sub.add_parser("expand", help="grow the system partition to fill the card")
p.add_argument("device")
p.add_argument("--yes", action="store_true", help="do not ask for confirmation")
args = ap.parse_args()
if args.cmd is None:
gui_main(args.image)
return
rep = Reporter(args.machine)
if args.cmd == "list":
cmd_list(rep, args)
return
if os.geteuid() != 0:
rep.error("writing a card needs root: run with sudo (the GUI asks through pkexec)")
sys.exit(1)
if args.machine:
watch_stdin(rep)
try:
{"flash": cmd_flash, "update": cmd_update, "expand": cmd_expand}[args.cmd](rep, args)
except Fail as e:
rep.error(str(e))
sys.exit(1)
except KeyboardInterrupt:
rep.error("interrupted")
sys.exit(130)
except OSError as e:
rep.error(f"{e.filename + ': ' if e.filename else ''}{e.strerror or e}")
sys.exit(1)
if __name__ == "__main__":
main()
Reference in new issue View Git Blame Copy Permalink