main.rs 15.3 KB
Newer Older
1
2
3
extern crate ansi_term;
#[macro_use] extern crate lazy_static;
extern crate threadpool;
4
extern crate rand;
5

Stefan Schindler's avatar
Stefan Schindler committed
6
7
8
use std::io;
use std::io::prelude::*;
use std::io::BufReader;
9
use std::fs::{File};
Stefan Schindler's avatar
Stefan Schindler committed
10
use std::fmt::Display;
11
use std::sync::{Arc, Mutex};
Stefan Schindler's avatar
Stefan Schindler committed
12
use std::time::{Instant, Duration};
13
14
15

use ansi_term::Style;
use ansi_term::Colour::*;
Stefan Schindler's avatar
Stefan Schindler committed
16

Stefan Schindler's avatar
Stefan Schindler committed
17
18
use rand::Rng;

19
20
21
22
23
lazy_static! {
    static ref BOLD: Style = Style::new().bold();
    static ref UNDERLINE: Style = Style::new().underline();
    static ref HEADING: Style = Yellow.bold();
}
Stefan Schindler's avatar
Stefan Schindler committed
24

25
const WRITE_BUFFER_SIZE: usize = 4 * 1024;
Stefan Schindler's avatar
Stefan Schindler committed
26
27
28

fn main() {
    let model = get_file_contents("/sys/devices/virtual/dmi/id/product_version");
29
    println!("\n\n    {} running on {}\n", BOLD.paint("whipe buddy"), BOLD.paint(model));
Stefan Schindler's avatar
Stefan Schindler committed
30
31
    
    let disks = parse_disks();
32
    let pool = threadpool::ThreadPool::new(2);
Stefan Schindler's avatar
Stefan Schindler committed
33
    
34
    show_disk_status(&disks);
Stefan Schindler's avatar
Stefan Schindler committed
35
    loop {
36
        println!("\nPress {} to exit, {} for a list of available commands", BOLD.paint("q"), BOLD.paint("h"));
Stefan Schindler's avatar
Stefan Schindler committed
37
        
38
        let input = read_one_line_from_stdin();
Stefan Schindler's avatar
Stefan Schindler committed
39
40
        match input.trim() {
            "q" => return,
41
            "h" => print_help(),
42
            "l" => show_disk_status(&disks),
43
            "w" => select_disk_to_whipe(&disks),
44
            ""  => show_disk_status(&disks),
45
            com =>  println!("unknown command: {}", com),
Stefan Schindler's avatar
Stefan Schindler committed
46
        }
47
48
        
        update_jobs(&disks, &pool);
Stefan Schindler's avatar
Stefan Schindler committed
49
50
51
    }
}

52
53
54
55
56
57

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;

Stefan Schindler's avatar
Stefan Schindler committed
58
    for disk in disks {
Stefan Schindler's avatar
Stefan Schindler committed
59
        println!("{:6}) {}:{:width$} {} -- {}",
60
61
62
63
64
65
66
            disk.id,
            BOLD.paint(&*disk.name),
            " ",
            disk.size,
            *disk.state.lock().unwrap(),
            width = longest - disk.name.len()
        );
Stefan Schindler's avatar
Stefan Schindler committed
67
68
69
    }
}

70
fn parse_disks() -> Vec<Arc<DiskInformation>> {
Stefan Schindler's avatar
Stefan Schindler committed
71
72
73
74
75
76
77
    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();
78
    let mut i = 0;
Stefan Schindler's avatar
Stefan Schindler committed
79
80
81
82
83
84
85
86
    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 {
Stefan Schindler's avatar
Stefan Schindler committed
87
            if  minor.map_or(false, |input| input.parse::<u64>().map(|n| n % 16 == 0).unwrap_or(false))
Stefan Schindler's avatar
Stefan Schindler committed
88
89
                    && 
                (name.starts_with("sd") || name.starts_with("nvme") || name.starts_with("xvd") || name.starts_with("mmcblk")) {
Stefan Schindler's avatar
Stefan Schindler committed
90
91
                /// 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);
92
93
94
                i += 1;
                disks.push(Arc::new(DiskInformation {
                    id: i,
Stefan Schindler's avatar
Stefan Schindler committed
95
96
                    name: name.into(),
                    size: KiloByte(no_blocks),
Stefan Schindler's avatar
Stefan Schindler committed
97
                    state: Mutex::new(Detected), // TODO erkenne wann eine Disk nicht überschrieben werden darf
98
                }));
Stefan Schindler's avatar
Stefan Schindler committed
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
            }
        }
    }
    
    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;
117
        buffer += &* line.unwrap_or( format!("<read error on line {} - {}>\n", i, path) );
Stefan Schindler's avatar
Stefan Schindler committed
118
119
120
121
122
123
        buffer += "\n";
    }
    
    buffer
}

124
125
126
127
128
129
fn read_one_line_from_stdin() -> String {
    let mut input = String::new();
    io::stdin().read_line(&mut input).unwrap_or(0);
    input
}

Stefan Schindler's avatar
Stefan Schindler committed
130
131
#[derive(Debug)]
struct DiskInformation {
132
    id: u64,
Stefan Schindler's avatar
Stefan Schindler committed
133
134
    size: ByteUnit,
    name: String,
135
    state: Mutex<DiskState>,
Stefan Schindler's avatar
Stefan Schindler committed
136
137
}

Stefan Schindler's avatar
Stefan Schindler committed
138
139
140
141
142
#[derive(Debug,PartialEq,Eq)]
enum DiskState {
    Detected,
    Protected,
    QueueForWhipe,
Stefan Schindler's avatar
Stefan Schindler committed
143
    Wiping{ progress: ByteUnit, per_second: ByteUnit },
Stefan Schindler's avatar
Stefan Schindler committed
144
    Wiped{ took: Duration },
Stefan Schindler's avatar
Stefan Schindler committed
145
    WipeFailed,
Stefan Schindler's avatar
Stefan Schindler committed
146
147
148
149
150
151
152
153
154
}
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"),
Stefan Schindler's avatar
Stefan Schindler committed
155
            &Wiping { ref progress, ref per_second } => write!(f, "Whiping {} ({}/s)", progress, per_second),
Stefan Schindler's avatar
Stefan Schindler committed
156
            &Wiped { ref took } => write!(f, "Wiped successfully in {} seconds", took.as_secs()),
Stefan Schindler's avatar
Stefan Schindler committed
157
            &WipeFailed=> write!(f, "Wipe failed"),
Stefan Schindler's avatar
Stefan Schindler committed
158
159
160
161
162
163
        }
    }
}


#[derive(Debug,Eq,Ord,Clone)]
Stefan Schindler's avatar
Stefan Schindler committed
164
enum ByteUnit {
165
    Byte(u64),
Stefan Schindler's avatar
Stefan Schindler committed
166
167
168
169
    KiloByte(u64),
}
use ByteUnit::*;

170
171
172
173
174
175
176
impl ByteUnit {
    pub fn to_byte(&self) -> ByteUnit {
        match self {
            &Byte(v) => Byte(v),
            &KiloByte(v) => Byte(1024 * v),
        }
    }
Stefan Schindler's avatar
Stefan Schindler committed
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
    
    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),
        }
    }
192
193
}

194
195
196
197
198
199
200
201
202
203
204
205
206
207
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,
            }
        }
    }
}
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
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),
        })
    }
}
Stefan Schindler's avatar
Stefan Schindler committed
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
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 {
242
243
                Byte(r) => { *me += r; return },
                KiloByte(r) => { *me += r * 1024; return; },
Stefan Schindler's avatar
Stefan Schindler committed
244
245
            },
            &mut KiloByte(ref mut me) => match other {
Stefan Schindler's avatar
Stefan Schindler committed
246
                Byte(_r) => { /* *me += r / 1024; würde zu einem Präzisions führen */ },
247
                KiloByte(r) => { *me += r; return },
Stefan Schindler's avatar
Stefan Schindler committed
248
249
            },
        }
250
251
252
253

        // Change Type of self
        let mut v = self.to_byte() + other;
        std::mem::swap(self, &mut v);
Stefan Schindler's avatar
Stefan Schindler committed
254
255
    }
}
Stefan Schindler's avatar
Stefan Schindler committed
256
257
impl std::ops::Sub for ByteUnit {
    type Output = ByteUnit;
258

Stefan Schindler's avatar
Stefan Schindler committed
259
260
261
262
263
264
265
266
267
    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,
        }
    }
268
269
}

Stefan Schindler's avatar
Stefan Schindler committed
270
impl Display for ByteUnit {
Stefan Schindler's avatar
Stefan Schindler committed
271
272
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        match self {
273
274
275
276
277
278
279
            &Byte(size) => {
                if size < 1000 {
                    return write!(f, "{:>6.2} B ", size);
                }
                
                return KiloByte(size / 1024).fmt(f);
            }
Stefan Schindler's avatar
Stefan Schindler committed
280
            &KiloByte(size) => {   
281
282
                if size < 1000 {
                    return write!(f, "{:>6.2} KB", size);
Stefan Schindler's avatar
Stefan Schindler committed
283
284
285
286
                }
                
                let size = size as f64;                
                let size: f64 = size / 1024.0;
287
288
                if size < 1000.0 {
                    return write!(f, "{:>6.2} MB", size);
Stefan Schindler's avatar
Stefan Schindler committed
289
290
291
                }
                
                let size: f64 = size / 1024.0;
292
293
                if size < 1000.0 {
                    return write!(f, "{:>6.2} GB", size);
Stefan Schindler's avatar
Stefan Schindler committed
294
295
296
                }
                
                let size: f64 = size / 1024.0;
297
                write!(f, "{:>6.2} TB", size)
Stefan Schindler's avatar
Stefan Schindler committed
298
299
300
301
            }
        }
    }
}
302
303
304
305
306
307
308
309


fn print_help() {
    println!("\n{}", HEADING.paint("Available Commands:"));
    
    let options = [
        ("q", "quit"),
        ("h", "print this help"),
310
        ("l", "list found disks and status"),
311
312
313
314
        ("w", "start whiping a disk"),
    ];
    
    for option in options.iter() {
315
        println!("       {}        {}", BOLD.paint(option.0), option.1);
316
317
318
319
320
321
322
323
    }
}

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) => {
Stefan Schindler's avatar
Stefan Schindler committed
324
            if let Some(disk) = disks.iter().find(|ref e| e.id == id && *e.state.lock().unwrap() == Detected) {
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
                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();
Stefan Schindler's avatar
Stefan Schindler committed
352
            
353
            pool.execute(move || {
Stefan Schindler's avatar
Stefan Schindler committed
354
355
                println!("{}{}", Blue.paint("Wiping "), BOLD.paint(&*disk.name));

Stefan Schindler's avatar
Stefan Schindler committed
356
                match whipe_disk(&disk) {
357
                    Ok(duration) => {
Stefan Schindler's avatar
Stefan Schindler committed
358
359
                        *disk.state.lock().unwrap() = Wiped{ took: duration };
                        println!("{}{} in {} seconds", 
360
361
                            Green.dimmed().paint("Wiped "), 
                            BOLD.paint(&*disk.name),
Stefan Schindler's avatar
Stefan Schindler committed
362
                            duration.as_secs()
363
                        );
Stefan Schindler's avatar
Stefan Schindler committed
364
365
                    },
                    Err(e) => {
Stefan Schindler's avatar
Stefan Schindler committed
366
                        *disk.state.lock().unwrap() = WipeFailed;
Stefan Schindler's avatar
Stefan Schindler committed
367
                        println!("{}{}{}",
Stefan Schindler's avatar
Stefan Schindler committed
368
                            Red.paint("Wipe of "),
Stefan Schindler's avatar
Stefan Schindler committed
369
370
371
372
373
374
                            BOLD.paint(&*disk.name),
                            Red.paint(format!(" failed: {:?}", e))
                        );
                        return
                    }
                }
375

376
377
378
379
380
            })
        }
    }
}

381
fn whipe_disk(disk: &Arc<DiskInformation>) -> std::io::Result<std::time::Duration> {
Stefan Schindler's avatar
Stefan Schindler committed
382
383
    let mut rng = rand::thread_rng();

Stefan Schindler's avatar
Stefan Schindler committed
384
    *disk.state.lock().unwrap() = Wiping{ progress: Byte(0), per_second: Byte(0) };
Stefan Schindler's avatar
Stefan Schindler committed
385

386
    let mut buf = Box::new( [0; WRITE_BUFFER_SIZE] );
Stefan Schindler's avatar
Stefan Schindler committed
387
388
389
390
391
    let mut position = Byte(0);
    let mut disk_file = std::fs::OpenOptions::new()
                        .write(true)
                        .open(&*format!("/dev/{}", disk.name))?;

Stefan Schindler's avatar
Stefan Schindler committed
392
393
    let mut i: u64 = 0;
    let now = Instant::now();
Stefan Schindler's avatar
Stefan Schindler committed
394
    while position < disk.size {
395
396
        rng.fill_bytes(&mut *buf);
        disk_file.write_all(&*buf)?;
Stefan Schindler's avatar
Stefan Schindler committed
397

Stefan Schindler's avatar
Stefan Schindler committed
398
399
400
        position += Byte(WRITE_BUFFER_SIZE as u64);
        
        i += 1;
401
402
403
404
405
        if i % 1024*1024 == 0 {
            *disk.state.lock().unwrap() = Wiping{ 
                progress: position.clone(), 
                per_second: position.byte_per_second(now.elapsed()),
            };
Stefan Schindler's avatar
Stefan Schindler committed
406
        }
Stefan Schindler's avatar
Stefan Schindler committed
407
408
    }

409
    Ok( now.elapsed() )
Stefan Schindler's avatar
Stefan Schindler committed
410
411
}

412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445

#[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)));
    }
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474


    #[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);
    }
Stefan Schindler's avatar
Stefan Schindler committed
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517


    #[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);
    }
Stefan Schindler's avatar
Stefan Schindler committed
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
    
    
    #[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));
    }
536
}