use crate::drives::human_size;
use std::fmt;
use std::io::Read;
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio};
use std::sync::mpsc::{self, Receiver, Sender};
use std::thread;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum OperationKind {
WipeAndFlash,
FlashIso,
FormatOnly,
}
impl OperationKind {
pub const ALL: [OperationKind; 3] = [
OperationKind::WipeAndFlash,
OperationKind::FlashIso,
OperationKind::FormatOnly,
];
pub fn title(self) -> &'static str {
match self {
OperationKind::WipeAndFlash => "Wipe + Flash",
OperationKind::FlashIso => "Flash ISO",
OperationKind::FormatOnly => "Format Drive",
}
}
pub fn description(self) -> &'static str {
match self {
OperationKind::WipeAndFlash => {
"Clear old signatures, then write the Arch ISO byte-for-byte."
}
OperationKind::FlashIso => "Write the selected ISO directly to the USB device.",
OperationKind::FormatOnly => {
"Create one clean partition and filesystem for normal storage."
}
}
}
pub fn needs_iso(self) -> bool {
matches!(self, OperationKind::WipeAndFlash | OperationKind::FlashIso)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FileSystem {
Fat32,
Exfat,
Ext4,
}
impl FileSystem {
pub const ALL: [FileSystem; 3] = [FileSystem::Fat32, FileSystem::Exfat, FileSystem::Ext4];
pub fn label(self) -> &'static str {
match self {
FileSystem::Fat32 => "FAT32",
FileSystem::Exfat => "exFAT",
FileSystem::Ext4 => "ext4",
}
}
fn script_name(self) -> &'static str {
match self {
FileSystem::Fat32 => "fat32",
FileSystem::Exfat => "exfat",
FileSystem::Ext4 => "ext4",
}
}
}
#[derive(Debug, Clone)]
pub struct OperationRequest {
pub kind: OperationKind,
pub drive_path: String,
pub iso_path: Option<PathBuf>,
pub filesystem: FileSystem,
pub label: String,
}
impl OperationRequest {
pub fn expected_image_bytes(&self) -> Option<u64> {
let path = self.iso_path.as_ref()?;
std::fs::metadata(path).ok().map(|metadata| metadata.len())
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CommandSpec {
pub program: String,
pub args: Vec<String>,
pub script: String,
}
impl fmt::Display for CommandSpec {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(formatter, "{}", self.program)?;
for arg in &self.args {
write!(formatter, " {}", shell_quote(arg))?;
}
Ok(())
}
}
#[derive(Debug, thiserror::Error)]
pub enum OperationError {
#[error("select a removable drive first")]
MissingDrive,
#[error("drive path must be an absolute /dev path")]
InvalidDrivePath,
#[error("select an ISO or disk image first")]
MissingIso,
#[error("image does not exist: {0}")]
MissingIsoFile(String),
#[error("image must be a regular file: {0}")]
IsoNotFile(String),
#[error("image should be an .iso or .img file: {0}")]
InvalidImageExtension(String),
#[error("filesystem label must contain letters, numbers, dashes, or underscores")]
InvalidLabel,
}
pub fn build_command(request: &OperationRequest) -> Result<CommandSpec, OperationError> {
validate_request(request)?;
let script = build_script(request)?;
Ok(privileged_command(script))
}
pub fn build_script(request: &OperationRequest) -> Result<String, OperationError> {
validate_request(request)?;
let drive = shell_quote(&request.drive_path);
let script = match request.kind {
OperationKind::WipeAndFlash => {
let iso = request
.iso_path
.as_ref()
.ok_or(OperationError::MissingIso)?;
flash_script(&drive, &shell_quote_path(iso), true)
}
OperationKind::FlashIso => {
let iso = request
.iso_path
.as_ref()
.ok_or(OperationError::MissingIso)?;
flash_script(&drive, &shell_quote_path(iso), false)
}
OperationKind::FormatOnly => format_script(
&drive,
&shell_quote(&sanitize_label(&request.label)?),
request.filesystem,
),
};
Ok(script)
}
pub fn validate_request(request: &OperationRequest) -> Result<(), OperationError> {
validate_drive_path(&request.drive_path)?;
if request.kind.needs_iso() {
let path = request
.iso_path
.as_ref()
.ok_or(OperationError::MissingIso)?;
validate_image_path(path)?;
}
if matches!(request.kind, OperationKind::FormatOnly) {
sanitize_label(&request.label)?;
}
Ok(())
}
pub fn validate_drive_path(path: &str) -> Result<(), OperationError> {
if path.trim().is_empty() {
return Err(OperationError::MissingDrive);
}
let valid = path.starts_with("/dev/")
&& path.len() > "/dev/".len()
&& path
.chars()
.all(|ch| ch.is_ascii_alphanumeric() || matches!(ch, '/' | '_' | '-'));
valid.then_some(()).ok_or(OperationError::InvalidDrivePath)
}
pub fn validate_image_path(path: &Path) -> Result<(), OperationError> {
let display = path.display().to_string();
let metadata =
std::fs::metadata(path).map_err(|_| OperationError::MissingIsoFile(display.clone()))?;
if !metadata.is_file() {
return Err(OperationError::IsoNotFile(display));
}
let extension = path
.extension()
.and_then(|extension| extension.to_str())
.unwrap_or_default()
.to_ascii_lowercase();
if extension != "iso" && extension != "img" {
return Err(OperationError::InvalidImageExtension(display));
}
Ok(())
}
pub fn sanitize_label(label: &str) -> Result<String, OperationError> {
let label = label.trim();
if label.is_empty() || label.len() > 11 {
return Err(OperationError::InvalidLabel);
}
let valid = label
.chars()
.all(|ch| ch.is_ascii_alphanumeric() || matches!(ch, '_' | '-'));
if !valid {
return Err(OperationError::InvalidLabel);
}
Ok(label.to_ascii_uppercase())
}
pub fn shell_quote(value: &str) -> String {
if value.is_empty() {
return "''".to_owned();
}
format!("'{}'", value.replace('\'', "'\"'\"'"))
}
fn shell_quote_path(path: &Path) -> String {
shell_quote(&path.display().to_string())
}
fn privileged_command(script: String) -> CommandSpec {
let shell = "/bin/sh";
let mut args = Vec::new();
if std::env::var_os("MOTTE_NO_PKEXEC").is_some() || !Path::new("/usr/bin/pkexec").exists() {
args.push("-c".to_owned());
args.push(script.clone());
return CommandSpec {
program: shell.to_owned(),
args,
script,
};
}
args.push(shell.to_owned());
args.push("-c".to_owned());
args.push(script.clone());
CommandSpec {
program: "/usr/bin/pkexec".to_owned(),
args,
script,
}
}
fn flash_script(drive: &str, iso: &str, wipe_first: bool) -> String {
let wipe = if wipe_first {
format!(
"echo \"Motte: clearing old filesystem signatures\"\n{}\nwipefs -a \"$dev\"\nblockdev --rereadpt \"$dev\" || true\npartprobe \"$dev\" || true\nudevadm settle || true\n",
wipe_partition_signatures_snippet()
)
} else {
String::new()
};
format!(
r#"set -eu
dev={drive}
iso={iso}
echo "Motte: preparing $dev"
{unmount_partitions}
{wipe}echo "Motte: writing $iso to $dev"
echo "Motte: dd is writing and flushing the target device"
dd if="$iso" of="$dev" bs=4M iflag=fullblock conv=fsync status=progress
echo "Motte: refreshing target partition table"
blockdev --rereadpt "$dev" || true
partprobe "$dev" || true
udevadm settle || true
echo "Motte: target flush complete"
echo "Motte: flash complete"
"#,
drive = drive,
iso = iso,
unmount_partitions = unmount_partitions_snippet(),
wipe = wipe,
)
}
fn format_script(drive: &str, label: &str, filesystem: FileSystem) -> String {
format!(
r#"set -eu
dev={drive}
label={label}
fs={filesystem}
{format_tools}
echo "Motte: checking formatter tools"
case "$fs" in
fat32)
ensure_any_command "mkfs.vfat mkfs.fat" "dosfstools" "FAT32 formatter not found"
;;
exfat)
ensure_any_command "mkfs.exfat" "exfatprogs" "exFAT formatter not found"
;;
ext4)
ensure_any_command "mkfs.ext4" "e2fsprogs" "ext4 formatter not found"
;;
esac
echo "Motte: preparing $dev"
{unmount_partitions}
echo "Motte: clearing old filesystem signatures"
{wipe_partition_signatures}
if ! wipefs -a "$dev"; then
echo "Motte: retrying signature clear after device settle"
udevadm settle || true
sleep 1
wipefs -a "$dev"
fi
blockdev --rereadpt "$dev" || true
partprobe "$dev" || true
udevadm settle || true
echo "Motte: creating partition table"
parted -s "$dev" mklabel msdos
case "$fs" in
fat32|exfat)
parted -s "$dev" mkpart primary fat32 1MiB 100%
parted -s "$dev" set 1 lba on || true
;;
ext4)
parted -s "$dev" mkpart primary ext4 1MiB 100%
;;
esac
blockdev --rereadpt "$dev" || true
partprobe "$dev" || true
udevadm settle || true
echo "Motte: waiting for the new partition"
attempt=0
part=""
while [ "$attempt" -lt 15 ]; do
part=""
part_count=0
while read -r path type; do
if [ "$type" = "part" ]; then
part_count=$((part_count + 1))
if [ -z "$part" ]; then
part="$path"
fi
fi
done <<MOTTE_PARTITION_SCAN
$(lsblk -rno PATH,TYPE "$dev" 2>/dev/null || true)
MOTTE_PARTITION_SCAN
if [ "$part_count" -eq 1 ] && [ -n "$part" ] && [ -b "$part" ]; then
break
fi
blockdev --rereadpt "$dev" || true
partprobe "$dev" || true
udevadm settle || true
sleep 1
attempt=$((attempt + 1))
done
if [ -z "$part" ] || [ ! -b "$part" ]; then
echo "Motte: could not locate the new partition" >&2
exit 1
fi
echo "Motte: clearing old signatures from $part"
wipefs -a "$part" || true
echo "Motte: formatting $part as $fs"
case "$fs" in
fat32)
fat_mkfs="$(first_available_command mkfs.vfat mkfs.fat)"
"$fat_mkfs" -F 32 -n "$label" "$part"
;;
exfat)
mkfs.exfat -n "$label" "$part"
;;
ext4)
mkfs.ext4 -F -L "$label" "$part"
;;
esac
echo "Motte: flushing $dev"
blockdev --flushbufs "$dev" || true
echo "Motte: format complete"
"#,
drive = drive,
label = label,
filesystem = shell_quote(filesystem.script_name()),
format_tools = format_tool_install_snippet(),
unmount_partitions = unmount_partitions_snippet(),
wipe_partition_signatures = wipe_partition_signatures_snippet(),
)
}
fn unmount_partitions_snippet() -> &'static str {
r#"unmount_partition() {
part="$1"
if command -v udisksctl >/dev/null 2>&1; then
udisksctl unmount -b "$part" >/dev/null 2>&1 || true
fi
findmnt -rn --source "$part" --output TARGET 2>/dev/null | while IFS= read -r mountpoint; do
[ -n "$mountpoint" ] || continue
echo "Motte: unmounting $part from $mountpoint"
umount "$mountpoint" 2>/dev/null || umount -l "$mountpoint" 2>/dev/null || umount "$part" 2>/dev/null || umount -l "$part" 2>/dev/null || true
done
}
wait_for_partition_unmounted() {
part="$1"
attempt=0
while [ "$attempt" -lt 10 ]; do
if ! findmnt -rn --source "$part" >/dev/null 2>&1; then
return 0
fi
unmount_partition "$part"
udevadm settle || true
sleep 1
attempt=$((attempt + 1))
done
echo "Motte: $part is still mounted after unmount attempts" >&2
return 1
}
while IFS= read -r part; do
[ "$part" = "$dev" ] && continue
[ -n "$part" ] || continue
if findmnt -rn --source "$part" >/dev/null 2>&1; then
echo "Motte: releasing mounted $part"
unmount_partition "$part"
wait_for_partition_unmounted "$part"
fi
done <<MOTTE_PARTS
$(lsblk -ln -o PATH "$dev")
MOTTE_PARTS"#
}
fn wipe_partition_signatures_snippet() -> &'static str {
r#"while IFS= read -r part; do
[ "$part" = "$dev" ] && continue
if [ -n "$part" ]; then
wipefs -a "$part" || true
fi
done <<MOTTE_EXISTING_PARTS
$(lsblk -ln -o PATH "$dev")
MOTTE_EXISTING_PARTS"#
}
fn format_tool_install_snippet() -> &'static str {
r#"first_available_command() {
for command_name in "$@"; do
if command -v "$command_name" >/dev/null 2>&1; then
command -v "$command_name"
return 0
fi
done
return 1
}
install_package() {
package_name="$1"
if command -v pacman >/dev/null 2>&1; then
echo "Motte: installing $package_name with pacman"
pacman -Sy --noconfirm --needed "$package_name"
elif command -v apt-get >/dev/null 2>&1; then
echo "Motte: installing $package_name with apt-get"
export DEBIAN_FRONTEND=noninteractive
apt-get update
apt-get install -y "$package_name"
elif command -v dnf >/dev/null 2>&1; then
echo "Motte: installing $package_name with dnf"
dnf install -y "$package_name"
elif command -v zypper >/dev/null 2>&1; then
echo "Motte: installing $package_name with zypper"
zypper --non-interactive install "$package_name"
else
echo "Motte: no supported package manager found for $package_name" >&2
return 127
fi
}
ensure_any_command() {
command_names="$1"
package_name="$2"
missing_message="$3"
for command_name in $command_names; do
if command -v "$command_name" >/dev/null 2>&1; then
return 0
fi
done
echo "Motte: $missing_message; attempting to install $package_name"
install_package "$package_name"
for command_name in $command_names; do
if command -v "$command_name" >/dev/null 2>&1; then
return 0
fi
done
echo "Motte: $missing_message; install $package_name manually" >&2
return 127
}
"#
}
#[derive(Debug, Clone, PartialEq)]
pub enum OperationEvent {
Started(String),
Log(String),
Progress { copied_bytes: u64, total_bytes: u64 },
Finished(Result<(), String>),
}
pub fn start_operation(
command: CommandSpec,
expected_bytes: Option<u64>,
) -> Receiver<OperationEvent> {
let (tx, rx) = mpsc::channel();
thread::spawn(move || run_operation(command, expected_bytes, tx));
rx
}
fn run_operation(command: CommandSpec, expected_bytes: Option<u64>, tx: Sender<OperationEvent>) {
let _ = tx.send(OperationEvent::Started(command.to_string()));
let mut child = match Command::new(&command.program)
.args(&command.args)
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.spawn()
{
Ok(child) => child,
Err(error) => {
let _ = tx.send(OperationEvent::Finished(Err(format!(
"failed to start operation: {error}"
))));
return;
}
};
let stdout = child.stdout.take();
let stderr = child.stderr.take();
let stdout_thread = stdout.map(|stream| forward_stream(stream, tx.clone(), expected_bytes));
let stderr_thread = stderr.map(|stream| forward_stream(stream, tx.clone(), expected_bytes));
let result = match child.wait() {
Ok(status) if status.success() => Ok(()),
Ok(status) => Err(format!("operation exited with status {status}")),
Err(error) => Err(format!("failed to wait for operation: {error}")),
};
if let Some(thread) = stdout_thread {
let _ = thread.join();
}
if let Some(thread) = stderr_thread {
let _ = thread.join();
}
let _ = tx.send(OperationEvent::Finished(result));
}
fn forward_stream<R>(
mut stream: R,
tx: Sender<OperationEvent>,
expected_bytes: Option<u64>,
) -> thread::JoinHandle<()>
where
R: Read + Send + 'static,
{
thread::spawn(move || {
let mut buffer = [0; 1024];
let mut pending = String::new();
loop {
let bytes_read = match stream.read(&mut buffer) {
Ok(0) => break,
Ok(bytes_read) => bytes_read,
Err(error) => {
let _ = tx.send(OperationEvent::Log(format!("stream read failed: {error}")));
break;
}
};
let chunk = String::from_utf8_lossy(&buffer[..bytes_read]);
for character in chunk.chars() {
if character == '\n' || character == '\r' {
flush_pending(&tx, &mut pending, expected_bytes);
} else {
pending.push(character);
}
}
}
flush_pending(&tx, &mut pending, expected_bytes);
})
}
fn flush_pending(tx: &Sender<OperationEvent>, pending: &mut String, expected_bytes: Option<u64>) {
let line = pending.trim().to_owned();
pending.clear();
if line.is_empty() {
return;
}
if let (Some(copied_bytes), Some(total_bytes)) =
(parse_dd_progress_bytes(&line), expected_bytes)
{
let _ = tx.send(OperationEvent::Progress {
copied_bytes: copied_bytes.min(total_bytes),
total_bytes,
});
}
let _ = tx.send(OperationEvent::Log(line));
}
pub fn parse_dd_progress_bytes(line: &str) -> Option<u64> {
let line = line.trim();
let (bytes, rest) = line.split_once(" bytes")?;
if !rest.contains("copied") {
return None;
}
bytes.trim().replace(',', "").parse().ok()
}
pub fn progress_label(copied_bytes: u64, total_bytes: u64) -> String {
let percent = if total_bytes == 0 {
0.0
} else {
copied_bytes as f32 / total_bytes as f32 * 100.0
};
format!(
"{} / {} ({percent:.1}%)",
human_size(copied_bytes),
human_size(total_bytes)
)
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::Write;
use tempfile::NamedTempFile;
#[test]
fn quotes_shell_values_safely() {
assert_eq!(shell_quote(""), "''");
assert_eq!(shell_quote("/dev/sdb"), "'/dev/sdb'");
assert_eq!(shell_quote("arch user's.iso"), "'arch user'\"'\"'s.iso'");
}
#[test]
fn validates_iso_paths() {
let mut file = NamedTempFile::with_suffix(".iso").expect("temp iso");
writeln!(file, "fake iso").expect("write temp iso");
validate_image_path(file.path()).expect("valid iso path");
}
#[test]
fn rejects_non_images() {
let file = NamedTempFile::new().expect("temp file");
let err = validate_image_path(file.path()).expect_err("reject extension");
assert!(matches!(err, OperationError::InvalidImageExtension(_)));
}
#[test]
fn builds_wipe_and_flash_script() {
let mut file = NamedTempFile::with_suffix(".iso").expect("temp iso");
writeln!(file, "fake iso").expect("write temp iso");
let request = OperationRequest {
kind: OperationKind::WipeAndFlash,
drive_path: "/dev/sdb".to_owned(),
iso_path: Some(file.path().to_path_buf()),
filesystem: FileSystem::Fat32,
label: "MOTTE".to_owned(),
};
let script = build_script(&request).expect("build script");
assert!(script.contains("wipefs -a \"$dev\""));
assert!(script.contains("MOTTE_EXISTING_PARTS"));
assert!(script.contains("dd if=\"$iso\" of=\"$dev\""));
assert!(script.contains("iflag=fullblock"));
assert!(script.contains("conv=fsync"));
assert!(script.contains("Motte: refreshing target partition table"));
assert!(script.contains("blockdev --rereadpt \"$dev\""));
assert!(script.contains("partprobe \"$dev\""));
assert!(script.contains("Motte: target flush complete"));
assert!(!script.contains("\nsync\n"));
assert!(script.contains("udisksctl unmount -b \"$part\""));
assert!(script.contains("findmnt -rn --source \"$part\""));
assert!(script.contains("umount \"$mountpoint\""));
assert!(script.contains("wait_for_partition_unmounted \"$part\""));
}
#[test]
fn builds_format_script_for_fat32() {
let request = OperationRequest {
kind: OperationKind::FormatOnly,
drive_path: "/dev/sdc".to_owned(),
iso_path: None,
filesystem: FileSystem::Fat32,
label: "arch_usb".to_owned(),
};
let script = build_script(&request).expect("build script");
assert!(script.contains("parted -s \"$dev\" mklabel msdos"));
assert!(script.contains("Motte: checking formatter tools"));
assert!(script.contains("ensure_any_command \"mkfs.vfat mkfs.fat\" \"dosfstools\""));
assert!(script.contains("parted -s \"$dev\" mkpart primary fat32"));
assert!(script.contains("parted -s \"$dev\" set 1 lba on"));
assert!(script.contains("blockdev --rereadpt \"$dev\""));
assert!(script.contains("Motte: waiting for the new partition"));
assert!(script.contains("while [ \"$attempt\" -lt 15 ]"));
assert!(script.contains("lsblk -rno PATH,TYPE \"$dev\""));
assert!(script.contains("first_available_command mkfs.vfat mkfs.fat"));
assert!(script.contains("pacman -Sy --noconfirm --needed \"$package_name\""));
assert!(!script.contains("sed -n '2p'"));
assert!(script.contains("label='ARCH_USB'"));
}
#[test]
fn parses_dd_progress_lines() {
assert_eq!(
parse_dd_progress_bytes("1,048,576 bytes (1.0 MB, 1.0 MiB) copied, 1 s, 1 MB/s"),
Some(1_048_576)
);
assert_eq!(parse_dd_progress_bytes("512 records in"), None);
}
#[test]
fn labels_progress() {
assert_eq!(progress_label(512, 1024), "512 B / 1.0 KiB (50.0%)");
}
}