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 {
"echo \"Motte: clearing old filesystem signatures\"\nwipefs -a \"$dev\"\n"
} else {
""
};
format!(
r#"set -eu
dev={drive}
iso={iso}
echo "Motte: preparing $dev"
{unmount_partitions}
{wipe}echo "Motte: writing $iso to $dev"
dd if="$iso" of="$dev" bs=4M conv=fsync status=progress
echo "Motte: syncing writes"
sync
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}
echo "Motte: preparing $dev"
{unmount_partitions}
echo "Motte: clearing old filesystem signatures"
wipefs -a "$dev"
echo "Motte: creating partition table"
parted -s "$dev" mklabel msdos
parted -s "$dev" mkpart primary 1MiB 100%
partprobe "$dev" || true
udevadm settle || true
sleep 1
part="$(lsblk -ln -o PATH "$dev" | sed -n '2p')"
if [ -z "$part" ]; then
echo "Motte: could not locate the new partition" >&2
exit 1
fi
echo "Motte: formatting $part as $fs"
case "$fs" in
fat32)
mkfs.vfat -F 32 -n "$label" "$part"
;;
exfat)
mkfs.exfat -n "$label" "$part"
;;
ext4)
mkfs.ext4 -F -L "$label" "$part"
;;
esac
sync
echo "Motte: format complete"
"#,
drive = drive,
label = label,
filesystem = shell_quote(filesystem.script_name()),
unmount_partitions = unmount_partitions_snippet(),
)
}
fn unmount_partitions_snippet() -> &'static str {
r#"while IFS= read -r part; do
[ "$part" = "$dev" ] && continue
if [ -n "$part" ]; then
echo "Motte: unmounting $part"
umount "$part" 2>/dev/null || true
fi
done <<MOTTE_PARTS
$(lsblk -ln -o PATH "$dev")
MOTTE_PARTS"#
}
#[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("dd if=\"$iso\" of=\"$dev\""));
assert!(script.contains("umount \"$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("mkfs.vfat -F 32"));
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%)");
}
}