main.rs

extern crate ansi_term;
#[macro_use] extern crate lazy_static;
extern crate threadpool;
extern crate rand;

use std::io;
use std::io::prelude::*;
use std::io::BufReader;
use std::fs::{File};
use std::fmt::Display;
use std::sync::{Arc, Mutex};
use std::time::{Instant, Duration};

use ansi_term::Style;
use ansi_term::Colour::*;

use rand::Rng;

lazy_static! {
    static ref BOLD: Style = Style::new().bold();
    static ref UNDERLINE: Style = Style::new().underline();
    static ref HEADING: Style = Yellow.bold();
}

const WRITE_BUFFER_SIZE: usize = 4 * 1024;

fn main() {
    let model = get_file_contents("/sys/devices/virtual/dmi/id/product_version");
    println!("\n\n    {} running on {}\n", BOLD.paint("whipe buddy"), BOLD.paint(model));
    
    let disks = parse_disks();
    let pool = threadpool::ThreadPool::new(2);
    
    show_disk_status(&disks);
    loop {
        println!("\nPress {} to exit, {} for a list of available commands", BOLD.paint("q"), BOLD.paint("h"));
        
        let input = read_one_line_from_stdin();
        match input.trim() {
            "q" => return,
            "h" => print_help(),
            "l" => show_disk_status(&disks),
            "w" => select_disk_to_whipe(&disks),
            ""  => show_disk_status(&disks),
            com =>  println!("unknown command: {}", com),
        }
        
        update_jobs(&disks, &pool);
    }
}


fn show_disk_status(disks: &Vec<Arc<DiskInformation>>) {
    println!("{}", HEADING.paint("Status of found disks:") );

    let longest = disks.iter().fold(0, |b, e| std::cmp::max(b, e.name.len())) +1;

    for disk in disks {
        println!("{:6}) {}:{:width$} {} -- {}",
            disk.id,
            BOLD.paint(&*disk.name),
            " ",
            disk.size,
            *disk.state.lock().unwrap(),
            width = longest - disk.name.len()
        );
    }
}

fn parse_disks() -> Vec<Arc<DiskInformation>> {
    let partitions = get_file_contents("/proc/partitions");
    //println!("partitions: {:?}", partitions);
    
    let lines = partitions.split("\n").skip(2);
    //println!("lines: {:?}", lines);
    
    let mut disks = Vec::new();
    let mut i = 0;
    for line in lines {
        let line = line.trim();
        //println!("line: {:?}", line);
        
        let mut parts = line.split_whitespace();
        let (_major, minor, no_blocks, name) = (parts.next(), parts.next(), parts.next(), parts.next());
        
        if let Some(name) = name {
            if  minor.map_or(false, |input| input.parse::<u64>().map(|n| n % 16 == 0).unwrap_or(false))
                    && 
                (name.starts_with("sd") || name.starts_with("nvme") || name.starts_with("xvd") || name.starts_with("mmcblk")) {
                /// Contains the major and minor numbers of  each  partition  as  well  as  the  number  of 1024-byte blocks and the partition name.
                let no_blocks: u64 = no_blocks.unwrap_or("0").parse().unwrap_or(0);
                i += 1;
                disks.push(Arc::new(DiskInformation {
                    id: i,
                    name: name.into(),
                    size: KiloByte(no_blocks),
                    state: Mutex::new(Detected), // TODO erkenne wann eine Disk nicht überschrieben werden darf
                }));
            }
        }
    }
    
    disks
}

fn get_file_contents(path: &str) -> String {
    let f = match File::open(path) {
        Ok(f) => f,
        Err(_) => return format!("<unable to open {:?}", path),
    };
    let reader = BufReader::new(f);
    let mut buffer = String::new();

    let mut i = 0;
    for line in reader.lines() {
        i += 1;
        buffer += &* line.unwrap_or( format!("<read error on line {} - {}>\n", i, path) );
        buffer += "\n";
    }
    
    buffer
}

fn read_one_line_from_stdin() -> String {
    let mut input = String::new();
    io::stdin().read_line(&mut input).unwrap_or(0);
    input
}

#[derive(Debug)]
struct DiskInformation {
    id: u64,
    size: ByteUnit,
    name: String,
    state: Mutex<DiskState>,
}

#[derive(Debug,PartialEq,Eq)]
enum DiskState {
    Detected,
    Protected,
    QueueForWhipe,
    Wiping{ progress: ByteUnit, per_second: ByteUnit },
    Wiped{ took: Duration },
    WipeFailed,
}
use DiskState::*;

impl Display for DiskState {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        match self {
            &Detected => write!(f, "Detected"),
            &Protected => write!(f, "Protected"),
            &QueueForWhipe => write!(f, "Queued for Whipe"),
            &Wiping { ref progress, ref per_second } => write!(f, "Whiping {} ({}/s)", progress, per_second),
            &Wiped { ref took } => write!(f, "Wiped successfully in {} seconds", took.as_secs()),
            &WipeFailed=> write!(f, "Wipe failed"),
        }
    }
}


#[derive(Debug,Eq,Ord,Clone)]
enum ByteUnit {
    Byte(u64),
    KiloByte(u64),
}
use ByteUnit::*;

impl ByteUnit {
    pub fn to_byte(&self) -> ByteUnit {
        match self {
            &Byte(v) => Byte(v),
            &KiloByte(v) => Byte(1024 * v),
        }
    }
    
    pub fn byte_per_second(&self, periode: Duration) -> ByteUnit {
        match self {
            &Byte(v) => {
                let secs = periode.as_secs();
                
                Byte(if secs == 0 {
                    0
                } else {
                    v / secs
                })
            },
            something_else => something_else.to_byte().byte_per_second(periode),
        }
    }
}

impl std::cmp::PartialEq<Self> for ByteUnit {
    fn eq(&self, rhs: &Self) -> bool {
        match self {
            &Byte(l) => match rhs {
                &Byte(r) => l == r,
                &KiloByte(r) => l == 1024 * r,
            },
            &KiloByte(l) => match rhs {
                &Byte(r) => 1024 * l == r,
                &KiloByte(r) => l == r,
            }
        }
    }
}
impl std::cmp::PartialOrd<Self> for ByteUnit {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(match self {
            &Byte(l) => match other {
                &Byte(r) => l.cmp(&r),
                ref something_else => {
                    if let Byte(r) = something_else.to_byte() {
                        l.cmp(&r)
                    } else {
                        panic!("fn partial_cmp<ByteUnit>: this must never happen!");
                    }
                },
            },
            ref something_else => something_else.to_byte().cmp(other),
        })
    }
}
impl std::ops::Add for ByteUnit {
    type Output = ByteUnit;

    fn add(self, other: Self) -> Self {
        match self {
            Byte(l) => match other {
                Byte(r) => Byte(l + r),
                something_else => self + something_else.to_byte(),
            },
            something_else => something_else.to_byte() + other,
        }
    }
}
impl std::ops::AddAssign for ByteUnit {
    fn add_assign(&mut self, other: Self) {
        match self {
            &mut Byte(ref mut me) => match other {
                Byte(r) => { *me += r; return },
                KiloByte(r) => { *me += r * 1024; return; },
            },
            &mut KiloByte(ref mut me) => match other {
                Byte(_r) => { /* *me += r / 1024; würde zu einem Präzisions führen */ },
                KiloByte(r) => { *me += r; return },
            },
        }

        // Change Type of self
        let mut v = self.to_byte() + other;
        std::mem::swap(self, &mut v);
    }
}
impl std::ops::Sub for ByteUnit {
    type Output = ByteUnit;

    fn sub(self, other: Self) -> Self {
        match self {
            Byte(l) => match other {
                Byte(r) => Byte(l - r),
                something_else => self - something_else.to_byte(),
            },
            something_else => something_else.to_byte() - other,
        }
    }
}

impl Display for ByteUnit {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        match self {
            &Byte(size) => {
                if size < 1000 {
                    return write!(f, "{:>6.2} B ", size);
                }
                
                return KiloByte(size / 1024).fmt(f);
            }
            &KiloByte(size) => {   
                if size < 1000 {
                    return write!(f, "{:>6.2} KB", size);
                }
                
                let size = size as f64;                
                let size: f64 = size / 1024.0;
                if size < 1000.0 {
                    return write!(f, "{:>6.2} MB", size);
                }
                
                let size: f64 = size / 1024.0;
                if size < 1000.0 {
                    return write!(f, "{:>6.2} GB", size);
                }
                
                let size: f64 = size / 1024.0;
                write!(f, "{:>6.2} TB", size)
            }
        }
    }
}


fn print_help() {
    println!("\n{}", HEADING.paint("Available Commands:"));
    
    let options = [
        ("q", "quit"),
        ("h", "print this help"),
        ("l", "list found disks and status"),
        ("w", "start whiping a disk"),
    ];
    
    for option in options.iter() {
        println!("       {}        {}", BOLD.paint(option.0), option.1);
    }
}

fn select_disk_to_whipe(disks: &Vec<Arc<DiskInformation>>) {
    println!("{}", HEADING.paint("Enter the id of the disk you would like to be whiped:"));
    
    match read_one_line_from_stdin().trim().parse() {
        Ok(id) => {
            if let Some(disk) = disks.iter().find(|ref e| e.id == id && *e.state.lock().unwrap() == Detected) {
                while {
                    println!("\nWhiping is {}!\nDestroy any data on {} ({})\nEnter {} to continue or {} to abort?", UNDERLINE.paint("non-reversable"), BOLD.paint(&*disk.name), disk.size, BOLD.paint("y"), BOLD.paint("a"));
                    
                    match read_one_line_from_stdin().trim() {
                        "y" | "Y" | "j" | "J" => { false },
                        "a" => { return }
                        _ => { println!("{}", Red.paint("invalid response")); true }
                    }
                    
                } {}
                
                *disk.state.lock().unwrap() = QueueForWhipe;
                
                return;
            }
        },
        Err(e) => println!("parse_error: {:?}", e),
    }
    
    println!("{}", Red.paint("invalid selection"));
}

fn update_jobs(disks: &Vec<Arc<DiskInformation>>, pool: &threadpool::ThreadPool) {
    for disk in disks {
        if *disk.state.lock().unwrap() == QueueForWhipe {
            println!("Queued {} for whipe", BOLD.paint(&*disk.name));
            let disk = disk.clone();
            
            pool.execute(move || {
                println!("{}{}", Blue.paint("Wiping "), BOLD.paint(&*disk.name));

                match whipe_disk(&disk) {
                    Ok(duration) => {
                        *disk.state.lock().unwrap() = Wiped{ took: duration };
                        println!("{}{} in {} seconds", 
                            Green.dimmed().paint("Wiped "), 
                            BOLD.paint(&*disk.name),
                            duration.as_secs()
                        );
                    },
                    Err(e) => {
                        *disk.state.lock().unwrap() = WipeFailed;
                        println!("{}{}{}",
                            Red.paint("Wipe of "),
                            BOLD.paint(&*disk.name),
                            Red.paint(format!(" failed: {:?}", e))
                        );
                        return
                    }
                }

            })
        }
    }
}

fn whipe_disk(disk: &Arc<DiskInformation>) -> std::io::Result<std::time::Duration> {
    let mut rng = rand::thread_rng();

    *disk.state.lock().unwrap() = Wiping{ progress: Byte(0), per_second: Byte(0) };

    let mut buf = Box::new( [0; WRITE_BUFFER_SIZE] );
    let mut position = Byte(0);
    let mut disk_file = std::fs::OpenOptions::new()
                        .write(true)
                        .open(&*format!("/dev/{}", disk.name))?;

    let mut i: u64 = 0;
    let now = Instant::now();
    while position < disk.size {
        rng.fill_bytes(&mut *buf);
        disk_file.write_all(&*buf)?;

        position += Byte(WRITE_BUFFER_SIZE as u64);
        
        i += 1;
        if i % 1024*1024 == 0 {
            *disk.state.lock().unwrap() = Wiping{ 
                progress: position.clone(), 
                per_second: position.byte_per_second(now.elapsed()),
            };
        }
    }

    Ok( now.elapsed() )
}


#[cfg(test)]
mod test {
    use ::ByteUnit::*;
    
    #[test]
    fn unit_eq() {
        let one_k_bytes = Byte(1024);
        let one_kilobyte = KiloByte(1);
        
        assert_eq!(one_k_bytes, one_kilobyte);
        assert_eq!(one_kilobyte, one_k_bytes);
    }
    
    #[test]
    fn format_byte() {
        assert_eq!("    42 B ", format!("{}", Byte(42)));
    }
    #[test]
    fn format_kilobyte() {
        assert_eq!("    42 KB", format!("{}", KiloByte(42)));
    }
    #[test]
    fn format_megabyte() {
        assert_eq!(" 42.00 MB", format!("{}", KiloByte(42 * 1024)));
    }
    #[test]
    fn format_gigabyte() {
        assert_eq!(" 42.00 GB", format!("{}", KiloByte(42 * 1024 * 1024)));
    }
    #[test]
    fn format_terrabyte() {
        assert_eq!(" 42.00 TB", format!("{}", KiloByte(42 * 1024 * 1024 * 1024)));
    }


    #[test]
    fn ord_eqals() {
        let one_k_bytes = Byte(1024);
        let one_kilobyte = KiloByte(1);

        assert!(one_k_bytes <= one_kilobyte);
        assert!(one_kilobyte <= one_k_bytes);

        assert!(one_k_bytes >= one_kilobyte);
        assert!(one_kilobyte >= one_k_bytes);
    }
    #[test]
    fn ord_less() {
        let two_k_bytes = Byte(2048);
        let one_kilobyte = KiloByte(1);

        assert!(two_k_bytes > one_kilobyte);
        assert!(one_kilobyte < two_k_bytes);
    }
    #[test]
    fn ord_more() {
        let one_k_bytes = Byte(1024);
        let two_kilobytes = KiloByte(2);

        assert!(two_kilobytes > one_k_bytes);
        assert!(one_k_bytes < two_kilobytes);
    }


    #[test]
    fn add_bytes() {
        assert_eq!(Byte(4), Byte(2) + Byte(2));
    }
    #[test]
    fn add_kilobytes() {
        assert_eq!(KiloByte(4), KiloByte(2) + KiloByte(2));
    }
    #[test]
    fn add_bytes_mixed_up() {
        assert_eq!(Byte(2 + 2048), Byte(2) + KiloByte(2));
    }
    #[test]
    fn add_bytes_mixed_down() {
        assert_eq!(Byte(2 + 2048), KiloByte(2) + Byte(2));
    }

    #[test]
    fn add_assign_bytes() {
        let mut four = Byte(2);
        four += Byte(2);
        assert_eq!(Byte(4), four);
    }
    #[test]
    fn add_assign_kilobytes() {
        let mut four = KiloByte(2);
        four += KiloByte(2);
        assert_eq!(KiloByte(4), four);
    }
    #[test]
    fn add_assign_mixed_up() {
        let mut four = Byte(2);
        four += KiloByte(2);
        assert_eq!(Byte(2 + 2048), four);
    }
    #[test]
    fn add_assign_mixed_down() {
        let mut four = KiloByte(2);
        four += Byte(2);
        assert_eq!(Byte(2 + 2048), four);
    }
    
    
    #[test]
    fn sub_bytes() {
        assert_eq!(Byte(0), Byte(2) - Byte(2));
    }
    #[test]
    fn sub_kilobytes() {
        assert_eq!(KiloByte(0), KiloByte(2) - KiloByte(2));
    }
    #[test]
    fn sub_bytes_mixed_up() {
        assert_eq!(Byte(2), Byte(2048 + 2) - KiloByte(2));
    }
    #[test]
    fn sub_bytes_mixed_down() {
        assert_eq!(Byte(2048 - 2), KiloByte(2) - Byte(2));
    }
}