Interactive device selection

src/args.rs

@@ -33,7 +33,7 @@ pub enum Command {
 pub struct CreateCommand {
     /// Either a path to a removable block device or a nonexiting file if --image is specified
     #[structopt(parse(from_os_str))]
-    pub path: PathBuf,
+    pub path: Option<PathBuf>,
 
     /// Additional pacakges to install
     #[structopt(short = "p", long = "extra-packages", value_name = "package")]
@@ -55,7 +55,8 @@ pub struct CreateCommand {
     #[structopt(
         long = "image",
         parse(try_from_str = "parse_bytes"),
-        value_name = "size"
+        value_name = "size",
+        requires = "path"
     )]
     pub image: Option<Byte>,
 }

src/error.rs

@@ -91,6 +91,15 @@ pub enum ErrorKind {
 
     #[fail(display = "Error setting up a loop device: {}", _0)]
     Losetup(String),
+
+    #[fail(display = "Error querying removeable devices")]
+    RemoveableDevicesQuery,
+
+    #[fail(display = "There are no removable devices")]
+    NoRemovableDevices,
+
+    #[fail(display = "Error selecing device")]
+    DeviceSelection,
 }
 
 impl Fail for Error {

src/main.rs

@@ -13,14 +13,15 @@ use crate::tool::Tool;
 use byte_unit::Byte;
 use failure::{Fail, ResultExt};
 use log::{debug, error, info, warn};
-use nix::sys::signal;
 use simplelog::*;
 use std::collections::HashSet;
 use std::fs;
-use std::io::Write;
+use std::io::{stdin, stdout, BufRead, Write};
 use std::os::unix::{fs::PermissionsExt, process::CommandExt as UnixCommandExt};
-use std::path::Path;
+use std::path::{Path, PathBuf};
 use std::process::exit;
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::Arc;
 use std::thread;
 use std::time::Duration;
 use structopt::StructOpt;
@@ -91,7 +92,45 @@ fn create_image(path: &Path, size: Byte) -> Result<LoopDevice, Error> {
     LoopDevice::create(path)
 }
 
-fn create(command: CreateCommand) -> Result<(), Error> {
+fn select_block_device(running: Arc<AtomicBool>) -> Result<PathBuf, Error> {
+    let devices = get_removable_devices()?;
+
+    if devices.is_empty() {
+        Err(ErrorKind::NoRemovableDevices)?
+    }
+
+    devices
+        .iter()
+        .enumerate()
+        .for_each(|(i, d)| println!("{}) {}", i + 1, d));
+
+    print!("\nSelect a removable device by Typing its number: ");
+    stdout().lock().flush().ok();
+
+    let mut buffer = String::new();
+    loop {
+        stdin().lock().read_line(&mut buffer).unwrap();
+
+        if buffer.is_empty() || !running.load(Ordering::SeqCst) {
+            println!();
+            Err(ErrorKind::DeviceSelection)?;
+        }
+
+        let choice = buffer
+            .trim()
+            .parse::<usize>()
+            .ok()
+            .filter(|n| 0 < *n && *n <= devices.len());
+
+        if let Some(choice) = choice {
+            return Ok(PathBuf::from("/dev").join(&devices[choice - 1].name));
+        } else {
+            error!("Bad choice");
+        }
+    }
+}
+
+fn create(command: CreateCommand, running: Arc<AtomicBool>) -> Result<(), Error> {
     let presets = presets::Presets::load(&command.presets)?;
 
     let sgdisk = Tool::find("sgdisk")?;
@@ -111,8 +150,14 @@ fn create(command: CreateCommand) -> Result<(), Error> {
         None
     };
 
+    let storage_device_path = if let Some(path) = command.path {
+        path
+    } else {
+        select_block_device(running)?
+    };
+
     let image_loop = if let Some(size) = command.image {
-        Some(create_image(&command.path, size)?)
+        Some(create_image(&storage_device_path, size)?)
     } else {
         None
     };
@@ -124,8 +169,9 @@ fn create(command: CreateCommand) -> Result<(), Error> {
                 info!("Using loop device at {}", loop_dev.path().display());
                 loop_dev.path()
             })
-            .unwrap_or(&command.path),
+            .unwrap_or(&storage_device_path),
     )?;
+
     let mount_point = tempdir().context(ErrorKind::TmpDirError)?;
     let disk_path = storage_device.path();
 
@@ -403,10 +449,6 @@ fn qemu(command: QemuCommand) -> Result<(), Error> {
     Err(err).context(ErrorKind::Qemu)?
 }
 
-extern "C" fn handle_sigint(_: i32) {
-    warn!("Interrupted");
-}
-
 fn main() {
     let app = App::from_args();
 
@@ -417,19 +459,17 @@ fn main() {
     };
     CombinedLogger::init(vec![TermLogger::new(log_level, Config::default()).unwrap()]).unwrap();
 
-    let sig_action = signal::SigAction::new(
-        signal::SigHandler::Handler(handle_sigint),
-        signal::SaFlags::empty(),
-        signal::SigSet::empty(),
-    );
-    unsafe {
-        signal::sigaction(signal::SIGINT, &sig_action).unwrap();
-        signal::sigaction(signal::SIGTERM, &sig_action).unwrap();
-        signal::sigaction(signal::SIGQUIT, &sig_action).unwrap();
-    }
+    let running = Arc::new(AtomicBool::new(true));
+    let r = running.clone();
+
+    ctrlc::set_handler(move || {
+        warn!("Interrupted");
+        r.store(false, Ordering::SeqCst);
+    })
+    .expect("Error setting Ctrl-C handler");
 
     let result = match app.cmd {
-        Command::Create(command) => create(command),
+        Command::Create(command) => create(command, running),
         Command::Chroot(command) => chroot(command),
         Command::Qemu(command) => qemu(command),
     };

src/storage/mod.rs

@@ -4,6 +4,7 @@ mod loop_device;
 mod markers;
 mod mount_stack;
 mod partition;
+mod removeable_devices;
 mod storage_device;
 
 pub use crypt::{is_encrypted_device, EncryptedDevice};
@@ -11,4 +12,5 @@ pub use filesystem::{Filesystem, FilesystemType};
 pub use loop_device::LoopDevice;
 pub use markers::BlockDevice;
 pub use mount_stack::MountStack;
+pub use removeable_devices::get_removable_devices;
 pub use storage_device::StorageDevice;

src/storage/removeable_devices.rs

@@ -0,0 +1,85 @@
+use crate::error::{Error, ErrorKind};
+use byte_unit::Byte;
+use failure::ResultExt;
+use std::{fmt, fs};
+
+#[derive(Debug)]
+pub struct Device {
+    model: String,
+    vendor: String,
+    size: Byte,
+    pub name: String,
+}
+
+impl fmt::Display for Device {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(
+            f,
+            "{} {} ({})",
+            self.vendor,
+            self.model,
+            self.size.get_appropriate_unit(true)
+        )
+    }
+}
+
+fn trimmed(source: String) -> String {
+    String::from(source.trim_end())
+}
+
+pub fn get_removable_devices() -> Result<Vec<Device>, Error> {
+    let mut result = Vec::new();
+
+    for entry in fs::read_dir("/sys/block").context(ErrorKind::RemoveableDevicesQuery)? {
+        let entry = entry.context(ErrorKind::RemoveableDevicesQuery)?;
+
+        let removable = fs::read_to_string(entry.path().join("removable"))
+            .context(ErrorKind::RemoveableDevicesQuery)?;
+
+        if removable != "1\n" {
+            continue;
+        }
+
+        let model = fs::read_to_string(entry.path().join("device/model"))
+            .map(trimmed)
+            .context(ErrorKind::RemoveableDevicesQuery)?;
+
+        if model == "CD-ROM" {
+            continue;
+        }
+
+        result.push(Device {
+            name: entry
+                .path()
+                .file_name()
+                .unwrap()
+                .to_string_lossy()
+                .into_owned(),
+            model,
+            vendor: fs::read_to_string(entry.path().join("device/vendor"))
+                .map(trimmed)
+                .context(ErrorKind::RemoveableDevicesQuery)?,
+            size: Byte::from_bytes(
+                fs::read_to_string(entry.path().join("size"))
+                    .context(ErrorKind::RemoveableDevicesQuery)?
+                    .trim()
+                    .parse::<u128>()
+                    .unwrap()
+                    * 512,
+            ),
+        })
+    }
+
+    Ok(result)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn sanity() {
+        let devices = get_removable_devices().unwrap();
+        println!("{:?}", devices);
+    }
+}