rust-dms/src/main.rs
2025-03-15 19:28:47 +04:00

125 lines
3.2 KiB
Rust

#![doc = include_str!("../README.md")]
pub mod cli;
pub mod error;
#[cfg(target_os = "linux")]
pub mod signal_handler;
use std::{
process::{Command, exit},
thread::sleep,
time::Duration,
};
use clap::Parser;
use cli::Cli;
use error::*;
use log::{info, trace};
use snafu::ResultExt;
pub fn check_if_plugged(pid: u16, vid: u16, first_run: bool) -> Result<bool, MyError> {
let device_iter = nusb::list_devices();
if device_iter.is_err() {
if first_run {
let _ = device_iter.context(USBSnafu {
message: "Getting the device list for the first time",
})?;
unreachable!();
} else {
dbg!(&device_iter.err());
return Ok(false);
}
}
Ok(device_iter
.unwrap()
.any(|dev| dev.vendor_id() == vid && dev.product_id() == pid))
}
pub fn execute_actions(cli: &Cli, actions: impl Iterator<Item = impl AsRef<str>>) {
for action in actions {
let action = action.as_ref();
if cli.dry_run {
info!("[Dry run] {}", action);
continue;
}
let mut cmd = Command::new(
std::env::var_os("SHELL")
.map(|el| el.into_string().unwrap())
.unwrap_or_else(|| "/bin/sh".to_string()),
);
cmd.args(["-c", action]);
let mut handle = cmd.spawn().expect("Can't start the shell");
let _ = handle.wait().expect("Error waiting for the action");
}
}
#[snafu::report]
fn main() -> Result<(), MyError> {
let cli = Cli::parse();
if cfg!(target_os = "linux") {
use signal_handler::handle_signals;
let _signals_handler = handle_signals(&cli)?;
}
let (pid, vid) = (
hex::decode(cli.pid.clone()).context(FromHexSnafu {
message: "Invalid PID",
})?,
hex::decode(cli.vid.clone()).context(FromHexSnafu {
message: "Invalid VID",
})?,
);
let (pid, vid) = (
u16::from_be_bytes(
pid.as_slice()
.try_into()
.with_context(|_| TryFromSliceSnafu {
expected_len: 2usize,
got: pid.len(),
})?,
),
u16::from_be_bytes(
vid.as_slice()
.try_into()
.with_context(|_| TryFromSliceSnafu {
expected_len: 2usize,
got: vid.len(),
})?,
),
);
let first_check = check_if_plugged(pid, vid, true);
if let Err(ref _first_check_err) = first_check {
if !cli.not_plugged_on_boot_actions.is_empty() {
execute_actions(&cli, cli.not_plugged_on_boot_actions.iter());
exit(1);
} else {
let _ = first_check?;
unreachable!();
}
}
info!("Successfully got through the first iteration of the loop");
loop {
sleep(Duration::from_micros(
(cli.sleep_time * 1000.).round() as u64
));
if !(check_if_plugged(pid, vid, false)?) {
execute_actions(&cli, cli.actions.iter());
if !cli.continue_on_unplug {
exit(1);
}
} else {
trace!("Device plugged in")
}
}
}