main.rs 15.7 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");
Stefan Schindler's avatar
Stefan Schindler committed
29
    println!("\n\n    {} running on {}\n", BOLD.paint("wipe 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 {
Stefan Schindler's avatar
Stefan Schindler committed
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),
Stefan Schindler's avatar
Stefan Schindler committed
43
            "w" => select_disk_to_wipe(&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
fn read_one_line_from_stdin() -> String {
Stefan Schindler's avatar
Stefan Schindler committed
125
126
    print!("> ");
    io::stdout().flush().expect("Unable to flush stdout inside fn read_one_line_from_stdin()");
127
128
129
130
131
    let mut input = String::new();
    io::stdin().read_line(&mut input).unwrap_or(0);
    input
}

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

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


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

172
173
174
175
176
177
178
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
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
    
    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),
        }
    }
194
195
}

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

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

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

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


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

Stefan Schindler's avatar
Stefan Schindler committed
321
322
fn select_disk_to_wipe(disks: &Vec<Arc<DiskInformation>>) {
    println!("{}", HEADING.paint("Enter the id of the disk you would like to be wiped:"));
323
324
325
    
    match read_one_line_from_stdin().trim().parse() {
        Ok(id) => {
Stefan Schindler's avatar
Stefan Schindler committed
326
            if let Some(disk) = disks.iter().find(|ref e| e.id == id && *e.state.lock().unwrap() == Detected) {
327
                while {
Stefan Schindler's avatar
Stefan Schindler committed
328
                    println!("\nWiping 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"));
329
330
331
332
333
334
335
336
337
                    
                    match read_one_line_from_stdin().trim() {
                        "y" | "Y" | "j" | "J" => { false },
                        "a" => { return }
                        _ => { println!("{}", Red.paint("invalid response")); true }
                    }
                    
                } {}
                
Stefan Schindler's avatar
Stefan Schindler committed
338
                *disk.state.lock().unwrap() = QueueForWipe;
339
340
341
342
343
344
345
346
347
348
349
350
                
                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 {
Stefan Schindler's avatar
Stefan Schindler committed
351
352
        if *disk.state.lock().unwrap() == QueueForWipe {
            println!("Queued {} for wipe", BOLD.paint(&*disk.name));
353
            let disk = disk.clone();
Stefan Schindler's avatar
Stefan Schindler committed
354
            
355
            pool.execute(move || {
Stefan Schindler's avatar
Stefan Schindler committed
356
                println!("{}{}", Green.paint("Wiping "), BOLD.paint(&*disk.name));
Stefan Schindler's avatar
Stefan Schindler committed
357

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

378
379
380
381
382
            })
        }
    }
}

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

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

388
    let mut buf = Box::new( [0; WRITE_BUFFER_SIZE] );
Stefan Schindler's avatar
Stefan Schindler committed
389
390
391
392
393
    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
394
395
    let mut i: u64 = 0;
    let now = Instant::now();
Stefan Schindler's avatar
Stefan Schindler committed
396
    while position < disk.size {
397
398
        rng.fill_bytes(&mut *buf);
        disk_file.write_all(&*buf)?;
Stefan Schindler's avatar
Stefan Schindler committed
399

Stefan Schindler's avatar
Stefan Schindler committed
400
401
402
        position += Byte(WRITE_BUFFER_SIZE as u64);
        
        i += 1;
403
404
405
406
407
        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
408
        }
Stefan Schindler's avatar
Stefan Schindler committed
409
410
    }

411
    Ok( now.elapsed() )
Stefan Schindler's avatar
Stefan Schindler committed
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
446
447

#[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)));
    }
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
475
476


    #[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
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
518
519


    #[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
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
    
    
    #[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));
    }
Stefan Schindler's avatar
Stefan Schindler committed
538
539
540
541
542
543
544
545
546
547

    #[test]
    fn delta_byte_per_second_0() {
        let zero = Byte(0);
        let no_time = ::std::time::Duration::new(0, 0);

        let zero_bps = zero.byte_per_second( no_time );
        assert_eq!(zero, zero_bps)
    }

548
}