harakit/src/mm.rs

/*
 * Copyright (c) 2024 Emma Tebibyte <emma@tebibyte.media>
 * Copyright (c) 2024 DTB <trinity@trinity.moe>
 * SPDX-License-Identifier: AGPL-3.0-or-later
 *
 * This program is free software: you can redistribute it and/or modify it under
 * the terms of the GNU Affero General Public License as published by the Free
 * Software Foundation, either version 3 of the License, or (at your option) any
 * later version.
 * 
 * This program is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
 * FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more
 * details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program. If not, see https://www.gnu.org/licenses/.
 */

use std::{
	env::args,
	fs::File,
	io::{ Error, BufWriter, Read, Write, stderr, stdin, stdout },
	os::fd::{ AsRawFd, FromRawFd },
	process::{ ExitCode, exit},
};

extern crate getopt;
extern crate strerror;
extern crate sysexits;

use getopt::GetOpt;
use strerror::StrError;
use sysexits::{ EX_IOERR, EX_USAGE };

#[cfg(target_os="openbsd")] use sysexits::EX_OSERR;
#[cfg(target_os="openbsd")] extern crate openbsd;
#[cfg(target_os="openbsd")]
use openbsd::{
	Promises,
	UnveilPerms,
	pledge,
	unveil,
};

use ArgMode::*;

enum ArgMode { In, Out }

fn err(argv0: &String, e: Error, code: Option<u8>) -> ExitCode {
	eprintln!("{}: {}", argv0, e.strerror());
	ExitCode::from(code.unwrap_or(1 /* unknown error */))
}

fn usage(argv0: &String) -> ExitCode {
	eprintln!("Usage: {} [-aetu] [-i input] [-o output]", argv0);
	ExitCode::from(EX_USAGE)
}

fn main() -> ExitCode {
	let argv = args().collect::<Vec<_>>();

	#[cfg(target_os="openbsd")] {
		let promises = Promises::new("cpath rpath stdio unveil wpath");
		if let Err(e) = pledge(Some(promises), Some(Promises::default())) {
			return err(&argv[0], e, Some(EX_OSERR));
		}
	}

	let mut a = false; /* append to the file */
	let mut e = false; /* use stderr as an output */
	let mut t = true; /* do not truncate the file before writing */
	let mut u = false; /* unbuffer i/o */
	let mut ins = Vec::new(); /* initial input file path vector */
	let mut outs = Vec::new(); /* initial output file path vector */
	let mut mode: Option<ArgMode> = None; /* mode set by last-used option */
	let mut optind = 0;

	while let Some(opt) = argv.getopt("aei:o:tu") {
		match opt.opt() {
			Ok("a") => a = true,
			Ok("e") => e = true, 
			Ok("u") => u = true,
			Ok("t") => t = false,
			Ok("i") => { /* add inputs */
				let input = opt.arg().unwrap();
				ins.push(input);
				mode = Some(In); /* latest argument == -i */
			},
			Ok("o") => { /* add output */
				let output = opt.arg().unwrap();
				outs.push(output);
				mode = Some(Out); /* latest argument == -o */
			},
			Err(_) | Ok(_) => {
				return usage(&argv[0]);
			},
		};

		optind = opt.ind();
	}

	let remaining = argv.iter().skip(optind);

	/* check the last flag specified */
	if let Some(m) = mode {
		for arg in remaining {
			/* move the subsequent arguments to the list of inputs or outputs */
			match m {
				In => ins.push(arg.to_string()),
				Out => outs.push(arg.to_string()),
			};
		}
	}

	#[cfg(target_os="openbsd")] {
		for input in &ins {
			let perms = UnveilPerms::new(vec!['r']);

			if let Err(e) = unveil(Some(&input), Some(perms)) {
				return err(&argv[0], e, Some(EX_OSERR));
			}
		}

		for output in &outs {
			let perms = UnveilPerms::new(vec!['c', 'w']);

			if let Err(e) = unveil(Some(&output), Some(perms)) {
				return err(&argv[0], e, Some(EX_OSERR));
			}
		}

		if let Err(e) = unveil(None, None) {
			return err(&argv[0], e, Some(EX_OSERR));
		}
	}

	if ins.is_empty() && outs.is_empty() && argv.len() > optind {
		return usage(&argv[0]);
	}

	/* use stdin if no inputs are specified */
	if ins.is_empty() { ins.push("-".to_string()); }

	/* use stdout if no outputs are specified */
	if outs.is_empty() && !e { outs.push("-".to_string()); }
	
	/* map all path strings to files */
	let inputs = ins.iter().map(|file| {
		/* if a file is “-”, it is stdin */
		if *file == "-" {
			/* portable way to access stdin as a file */
			return unsafe { File::from_raw_fd(stdin().as_raw_fd()) };
		}

		match File::open(file) {
			Ok(f) => f,
			Err(e) => {
				let _ = err(&(argv[0].clone() + ": " + file), e, None);
				exit(EX_IOERR.into());
			},
		}
	}).collect::<Vec<_>>();

	/* map all path strings to files */
	let mut outputs = outs.iter().map(|file| {
		/* of a file is “-”, it is stdout */
		if *file == "-" {
			/* portable way to access stdout as a file */
			return unsafe { File::from_raw_fd(stdout().as_raw_fd()) }; 
		}

		let options = File::options()
			/* don’t truncate if -t is specified, append if -a is specified */
			.truncate(t)
			.append(a)
			/* enable the ability to create and write to files */
			.create(true)
			.write(true)
			/* finally, open the file! */
			.open(file);

		match options {
			Ok(f) => return f,
			Err(e) => {
				let _ = err(&(argv[0].clone() + ": " + file), e, None);
				exit(EX_IOERR.into());
			},
		};
	}).collect::<Vec<_>>();

	/* if -e is specified, use stderr */
	if e {
		/* portable way to access stderr as a file */
		outputs.push(unsafe { File::from_raw_fd(stderr().as_raw_fd()) });
	}

	let mut outputs = outputs.iter().map(|o| {
		if u {
			/* unbuffered writing through a buffer of capacity 0 */
			BufWriter::with_capacity(0, o)
		} else {
			/* theoretically buffered writing */
			BufWriter::new(o)
		}
	}).collect::<Vec<_>>();

	for file in inputs {
		for byte in file.bytes().map(|b| {
			b.unwrap_or_else(|e| {
				let _ = err(&argv[0], e, None);
				exit(EX_IOERR.into());
			})
		}) {
			for out in &mut outputs {
				if let Err(e) = out.write(&[byte]) {
					return err(&argv[0], e, Some(EX_IOERR));
				}

				if u {
					/* immediately flush the output for -u */
					if let Err(e) = out.flush() { 
						return err(&argv[0], e, Some(EX_IOERR));
					}
				}
			}
		}
	}

	ExitCode::SUCCESS
}