<commit_before>use std::fs;
use std::env;
use std::fs::File;
use std::io::Read;
use std::hash::Hasher;
use std::path::PathBuf;
use std::collections::HashMap;
//use std::os::unix::fs::MetadataExt;
use std::collections::hash_map::DefaultHasher;

#[macro_use]
extern crate log;
extern crate env_logger;

fn collect_files(dir: &String, h: &mut HashMap<u64, Vec<String>>) {

    let entries = match fs::read_dir(dir) {
        Ok(x) => x,
        _ => {
            warn!("{}: invalid directory", dir);
            return
        },
    };

    for path in  entries.filter( |x| x.is_ok() ).map( |x| x.unwrap().path() )
                .collect::<Vec<PathBuf>>() {
        if let (Some(path_str), Ok(metadata)) =
                (path.to_str(), path.symlink_metadata()) {
            let ft = metadata.file_type();
            if ft.is_symlink() {
                continue;
            } else if ft.is_file() {
                match metadata.len() {
                    0 => continue, //Ignore empty files.
                    i => h.entry(i).or_insert_with(Vec::new)
                          .push(String::from(path_str)),
                };
            } else if ft.is_dir() {
                collect_files(&(String::from(path_str)), h);
            }
        } else {
            warn!("Failed to retrieve metadata for {}", path.to_str().unwrap());
        }
    }
}

fn find_duplicates(duplicates: &mut HashMap<u64, Vec<String>>,
                   vec: &Vec<String>, thorough: bool ) {
    for file_entry in vec {
        if let Ok(mut file) = File::open(file_entry) {
            let mut hash: u64 = 0;
            let mut hasher = DefaultHasher::new();
            if thorough {
                let mut contents: Vec<u8> = Vec::new();
                if let Ok(_) = file.read_to_end(&mut contents) {
                    hasher.write(&contents[..]);
                    hash = hasher.finish();
                } else {
                    warn!("Failed to read {}", file_entry);
                }
            } else {
                let mut buf: [u8; 1024] = [0; 1024];
                if let Ok(nbytes) = file.read(&mut buf) {
                    hasher.write(&buf[..nbytes as usize]);
                    hash = hasher.finish();
                } else {
                    warn!("Failed to read {} chunk", file_entry);
                }
            }
            if hash != 0 {
                duplicates.entry(hash).or_insert_with(Vec::new)
                    .push(file_entry.to_string());
            }
        } else {
            warn!("Failed to open file {}", file_entry);
        }
    }
}

fn print_duplicates(vec: &Vec<String>, cluster: i32, key: u64, hash: u64) {
    println!("{} files in cluster {} (size: {}, digest: {})",
             vec.len(), cluster, key, hash);
    // for_each becomes stable v1.22.0 onwards. Should uncomment then.
    //vec.iter().for_each(|f| println!("{}", f.path));
    for file_entry in vec {
        println!("{}", file_entry);
    }
}

fn exhaustive_search(emap: &mut HashMap<Vec<u8>, Vec<String>>,
                     vec: &Vec<String>) {
    for file_entry in vec {
        if let Ok(mut file) = File::open(file_entry) {
            let mut contents: Vec<u8> = Vec::new();
            if let Ok(_) = file.read_to_end(&mut contents) {
                emap.entry(contents).or_insert_with(Vec::new)
                    .push(file_entry.to_string());
            } else {
                warn!("Failed to read {}", file_entry);
            }
        } else {
            warn!("Failed to open {}", file_entry);
        }
    }
}

fn print_usage_and_exit(binary_name: &String) {
        println!("Usage: {} [-t] [-e] <dir1> [dir2 [dir3 ...]]", binary_name);
        ::std::process::exit(0);
}

fn main() {
    env_logger::init().unwrap();

    let mut args: Vec<String> = env::args().collect();

    let binary_name = args.remove(0);

    // Print out usage if no directories are given.
    if args.len() == 0 {
        print_usage_and_exit(&binary_name);
    }

    let mut hmap: HashMap<u64, Vec<String> > = HashMap::new();

    let mut thorough = false;
    let mut exhaustive = false;
    let mut print_usage = true;

    for dir in &args[..] {
        if dir == "-t" {
            thorough = true;
            continue;
        } else if dir == "-e" {
            thorough = false;
            exhaustive = false;
            continue;
        }
        print_usage = false;
        collect_files(dir, &mut hmap);
    }

    if print_usage {
        print_usage_and_exit(&binary_name);
    }

    // Get rid of all the single entries.
    hmap.retain(|_, v| v.len() >= 2);

    let mut cluster = 1;

    for (key, val) in hmap {
        let mut duplicates: HashMap<u64, Vec<String> > = HashMap::new();
        find_duplicates(&mut duplicates, &val, thorough);

        for (hash, vec) in duplicates {
            if vec.len() >= 2 {
                if exhaustive {
                    let mut emap: HashMap<Vec<u8>, Vec<String>> = HashMap::new();
                    exhaustive_search(&mut emap, &vec);
                    for (_, v) in emap {
                        if v.len() >= 2 {
                            print_duplicates(&v, cluster, key, hash);
                            cluster += 1;
                        }
                    }
                } else {
                    print_duplicates(&vec, cluster, key, hash);
                    cluster += 1;
                }
            }
        }
    }
}
<commit_msg>Fix usage<commit_after>use std::fs;
use std::env;
use std::fs::File;
use std::io::Read;
use std::hash::Hasher;
use std::path::PathBuf;
use std::collections::HashMap;
//use std::os::unix::fs::MetadataExt;
use std::collections::hash_map::DefaultHasher;

#[macro_use]
extern crate log;
extern crate env_logger;

fn collect_files(dir: &String, h: &mut HashMap<u64, Vec<String>>) {

    let entries = match fs::read_dir(dir) {
        Ok(x) => x,
        _ => {
            warn!("{}: invalid directory", dir);
            return
        },
    };

    for path in  entries.filter( |x| x.is_ok() ).map( |x| x.unwrap().path() )
                .collect::<Vec<PathBuf>>() {
        if let (Some(path_str), Ok(metadata)) =
                (path.to_str(), path.symlink_metadata()) {
            let ft = metadata.file_type();
            if ft.is_symlink() {
                continue;
            } else if ft.is_file() {
                match metadata.len() {
                    0 => continue, //Ignore empty files.
                    i => h.entry(i).or_insert_with(Vec::new)
                          .push(String::from(path_str)),
                };
            } else if ft.is_dir() {
                collect_files(&(String::from(path_str)), h);
            }
        } else {
            warn!("Failed to retrieve metadata for {}", path.to_str().unwrap());
        }
    }
}

fn find_duplicates(duplicates: &mut HashMap<u64, Vec<String>>,
                   vec: &Vec<String>, thorough: bool ) {
    for file_entry in vec {
        if let Ok(mut file) = File::open(file_entry) {
            let mut hash: u64 = 0;
            let mut hasher = DefaultHasher::new();
            if thorough {
                let mut contents: Vec<u8> = Vec::new();
                if let Ok(_) = file.read_to_end(&mut contents) {
                    hasher.write(&contents[..]);
                    hash = hasher.finish();
                } else {
                    warn!("Failed to read {}", file_entry);
                }
            } else {
                let mut buf: [u8; 1024] = [0; 1024];
                if let Ok(nbytes) = file.read(&mut buf) {
                    hasher.write(&buf[..nbytes as usize]);
                    hash = hasher.finish();
                } else {
                    warn!("Failed to read {} chunk", file_entry);
                }
            }
            if hash != 0 {
                duplicates.entry(hash).or_insert_with(Vec::new)
                    .push(file_entry.to_string());
            }
        } else {
            warn!("Failed to open file {}", file_entry);
        }
    }
}

fn print_duplicates(vec: &Vec<String>, cluster: i32, key: u64, hash: u64) {
    println!("{} files in cluster {} (size: {}, digest: {})",
             vec.len(), cluster, key, hash);
    // for_each becomes stable v1.22.0 onwards. Should uncomment then.
    //vec.iter().for_each(|f| println!("{}", f.path));
    for file_entry in vec {
        println!("{}", file_entry);
    }
}

fn exhaustive_search(emap: &mut HashMap<Vec<u8>, Vec<String>>,
                     vec: &Vec<String>) {
    for file_entry in vec {
        if let Ok(mut file) = File::open(file_entry) {
            let mut contents: Vec<u8> = Vec::new();
            if let Ok(_) = file.read_to_end(&mut contents) {
                emap.entry(contents).or_insert_with(Vec::new)
                    .push(file_entry.to_string());
            } else {
                warn!("Failed to read {}", file_entry);
            }
        } else {
            warn!("Failed to open {}", file_entry);
        }
    }
}

fn main() {
    env_logger::init().unwrap();

    let mut args: Vec<String> = env::args().collect();

    let binary_name = args.remove(0);

    let mut hmap: HashMap<u64, Vec<String> > = HashMap::new();

    let mut thorough = false;
    let mut exhaustive = false;
    let mut print_usage = true;

    for dir in &args[..] {
        if dir == "-t" {
            thorough = true;
            continue;
        } else if dir == "-e" {
            thorough = false;
            exhaustive = false;
            continue;
        }
        print_usage = false;
        collect_files(dir, &mut hmap);
    }

    if print_usage {
        println!("Usage: {} [-t] [-e] <dir1> [dir2 [dir3 ...]]", binary_name);
        ::std::process::exit(0);
    }

    // Get rid of all the single entries.
    hmap.retain(|_, v| v.len() >= 2);

    let mut cluster = 1;

    for (key, val) in hmap {
        let mut duplicates: HashMap<u64, Vec<String> > = HashMap::new();
        find_duplicates(&mut duplicates, &val, thorough);

        for (hash, vec) in duplicates {
            if vec.len() >= 2 {
                if exhaustive {
                    let mut emap: HashMap<Vec<u8>, Vec<String>> = HashMap::new();
                    exhaustive_search(&mut emap, &vec);
                    for (_, v) in emap {
                        if v.len() >= 2 {
                            print_duplicates(&v, cluster, key, hash);
                            cluster += 1;
                        }
                    }
                } else {
                    print_duplicates(&vec, cluster, key, hash);
                    cluster += 1;
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>#![feature(std_misc, collections, core, env, os, libc)]

extern crate libc;
extern crate rand;

use libc::funcs::posix88::unistd;
use std::ffi::CString;
use std::cmp;
use std::env;
use std::os;
use std::ptr;
use rand::Rng;

static VERSION: &'static str = "0.1.0";

fn main() {
    let args: Vec<String> = env::args().skip(1).collect();
    match parse_args(args.as_slice()) {
        Err(Exit::Usage) => {
            die_usage();
        }
        Err(Exit::Help) =>
            print_usage(),
        Err(Exit::Version) =>
            print_version(),
        Ok(opts) => {
            if whack(&opts) {
                env::set_exit_status(1);
            } else {
                opts.exe.map_or((), |exe| {
                    execvp(exe.name, exe.args)
                });
            }
        }
    }
}

type ArgParse<T> = Result<T, Exit>;

#[derive(Debug)]
struct Options<'a> {
    chance: i32,
    exe: Option<Exe<'a>>,
}

#[derive(Debug)]
struct Exe<'a> {
    name: &'a str,
    args: &'a[String],
}

#[derive(Debug)]
enum Exit {
    Usage,
    Help,
    Version,
}

fn parse_args(args: &[String]) -> ArgParse<Options> {
    args.uncons().map_or(Err(Exit::Usage), |(s, args)| {
        match s.as_slice() {
            "--help"    | "-h" => Err(Exit::Help),
            "--version" | "-v" => Err(Exit::Version),
            "--chance"  | "-c" => {
                args.uncons().map_or(Err(Exit::Usage), |(s, args)| {
                    s.parse().ok().map_or(Err(Exit::Usage), |val| {
                        args.skip("--").uncons().map_or(
                            Ok(Options {
                                chance: val,
                                exe: None
                            }), |(exe, args)| {
                                Ok(Options {
                                    chance: val,
                                    exe: Some(Exe {
                                        name: exe.as_slice(),
                                        args: args,
                                    })
                                })
                            })
                        })
                    })
                }
            _ => Err(Exit::Usage),
        }
    })
}

trait Uncons<'a, A> { fn uncons(&self) -> Option<(&'a A, &'a [A])>; }
trait Skip<'a, A, B> { fn skip(&self, x: A) -> &'a [B]; }

impl<'a, A> Uncons<'a, A> for &'a [A] {
    fn uncons(&self) -> Option<(&'a A, &'a [A])> {
        self.first().map(|x| (x, self.tail()))
    }
}

impl<'a > Skip<'a, &'a str, String> for &'a [String] {
    fn skip(&self, x: &'a str) -> &'a [String] {
        self.uncons().map_or(*self, |(y, ys)| { if x == y.as_slice() { ys } else { *self } })
    }
}

fn whack(opts: &Options) -> bool {
    rand::thread_rng().gen_range(0, 100) > cmp::min(cmp::max(opts.chance, 0), 100)
}

fn execvp(name: &str, args: &[String]) {
    let c_name = CString::new(name).unwrap(); // should be safe, the string is genuine
    let c_args: Vec<CString> =
        args.iter().map(|tmp| CString::new(tmp.as_slice()).unwrap()).collect(); // likewise

    let mut xs: Vec<*const libc::c_char> = Vec::with_capacity(args.len() + 2);

    xs.push(c_name.as_ptr());
    xs.extend(c_args.iter().map(|tmp| tmp.as_ptr()));
    xs.push(ptr::null());

    let c_arg0 = xs[0];
    let c_argv = xs.as_mut_ptr();

    unsafe {
        unistd::execvp(c_arg0, c_argv);
    };

    panic!("execvp(3) failed with: {}", os::last_os_error());
}

fn die_usage() {
    env::set_exit_status(1);
    print_usage();
}

fn print_usage() {
    print_version();
    println!("\n\
        Usage: whacky (-c PERCENTAGE|--chance PERCENTAGE) [[--] COMMAND]\n\
          whacky, yet another randomly failing program\n\
        \n\
        Available options:\n\
          -h,--help           show this help text\n\
          -c,--chance         set success chance");
}

fn print_version() {
    println!("whacky {}", VERSION);
}
<commit_msg>Staying warning-free<commit_after>#![feature(std_misc, collections, core, exit_status, os, libc)]

extern crate libc;
extern crate rand;

use libc::funcs::posix88::unistd;
use std::ffi::CString;
use std::cmp;
use std::env;
use std::os;
use std::ptr;
use rand::Rng;

static VERSION: &'static str = "0.1.0";

fn main() {
    let args: Vec<String> = env::args().skip(1).collect();
    match parse_args(args.as_slice()) {
        Err(Exit::Usage) => {
            die_usage();
        }
        Err(Exit::Help) =>
            print_usage(),
        Err(Exit::Version) =>
            print_version(),
        Ok(opts) => {
            if whack(&opts) {
                env::set_exit_status(1);
            } else {
                opts.exe.map_or((), |exe| {
                    execvp(exe.name, exe.args)
                });
            }
        }
    }
}

type ArgParse<T> = Result<T, Exit>;

#[derive(Debug)]
struct Options<'a> {
    chance: i32,
    exe: Option<Exe<'a>>,
}

#[derive(Debug)]
struct Exe<'a> {
    name: &'a str,
    args: &'a[String],
}

#[derive(Debug)]
enum Exit {
    Usage,
    Help,
    Version,
}

fn parse_args(args: &[String]) -> ArgParse<Options> {
    args.uncons().map_or(Err(Exit::Usage), |(s, args)| {
        match s.as_slice() {
            "--help"    | "-h" => Err(Exit::Help),
            "--version" | "-v" => Err(Exit::Version),
            "--chance"  | "-c" => {
                args.uncons().map_or(Err(Exit::Usage), |(s, args)| {
                    s.parse().ok().map_or(Err(Exit::Usage), |val| {
                        args.skip("--").uncons().map_or(
                            Ok(Options {
                                chance: val,
                                exe: None
                            }), |(exe, args)| {
                                Ok(Options {
                                    chance: val,
                                    exe: Some(Exe {
                                        name: exe.as_slice(),
                                        args: args,
                                    })
                                })
                            })
                        })
                    })
                }
            _ => Err(Exit::Usage),
        }
    })
}

trait Uncons<'a, A> { fn uncons(&self) -> Option<(&'a A, &'a [A])>; }
trait Skip<'a, A, B> { fn skip(&self, x: A) -> &'a [B]; }

impl<'a, A> Uncons<'a, A> for &'a [A] {
    fn uncons(&self) -> Option<(&'a A, &'a [A])> {
        self.first().map(|x| (x, self.tail()))
    }
}

impl<'a > Skip<'a, &'a str, String> for &'a [String] {
    fn skip(&self, x: &'a str) -> &'a [String] {
        self.uncons().map_or(*self, |(y, ys)| { if x == y.as_slice() { ys } else { *self } })
    }
}

fn whack(opts: &Options) -> bool {
    rand::thread_rng().gen_range(0, 100) > cmp::min(cmp::max(opts.chance, 0), 100)
}

fn execvp(name: &str, args: &[String]) {
    let c_name = CString::new(name).unwrap(); // should be safe, the string is genuine
    let c_args: Vec<CString> =
        args.iter().map(|tmp| CString::new(tmp.as_slice()).unwrap()).collect(); // likewise

    let mut xs: Vec<*const libc::c_char> = Vec::with_capacity(args.len() + 2);

    xs.push(c_name.as_ptr());
    xs.extend(c_args.iter().map(|tmp| tmp.as_ptr()));
    xs.push(ptr::null());

    let c_arg0 = xs[0];
    let c_argv = xs.as_mut_ptr();

    unsafe {
        unistd::execvp(c_arg0, c_argv);
    };

    panic!("execvp(3) failed with: {}", os::last_os_error());
}

fn die_usage() {
    env::set_exit_status(1);
    print_usage();
}

fn print_usage() {
    print_version();
    println!("\n\
        Usage: whacky (-c PERCENTAGE|--chance PERCENTAGE) [[--] COMMAND]\n\
          whacky, yet another randomly failing program\n\
        \n\
        Available options:\n\
          -h,--help           show this help text\n\
          -c,--chance         set success chance");
}

fn print_version() {
    println!("whacky {}", VERSION);
}
<|endoftext|>


<commit_before>use macroquad::prelude::*;

use grid::Grid;
use graph::Graph;

mod colour;
mod grid;
mod graph;

#[macroquad::main("BasicShapes")]
async fn main() {
    let grid = Grid::generate(8);
    let mut graph = Graph::create(&grid);

    loop {
        clear_background(BLACK);

        let rendering_size = screen_width().min(screen_height());

        draw_rectangle((screen_width() - rendering_size) / 2.0, (screen_height() - rendering_size) / 2.0, rendering_size, rendering_size, GREEN);

        next_frame().await
    }
}
<commit_msg>Render the full grid<commit_after>use macroquad::prelude::*;

use colour::Colour;
use graph::{Graph, Position};
use grid::Grid;

mod colour;
mod graph;
mod grid;

fn macroquad_colour(colour: Colour) -> Color {
    match colour {
        Colour::Red => RED,
        Colour::Yellow => YELLOW,
        Colour::Green => GREEN,
        Colour::LightBrown => LIGHTGRAY,
        Colour::Purple => PURPLE,
        Colour::Cyan => PINK,
        Colour::Blue => BLUE,
        Colour::Fuchsia => MAGENTA,
    }
}

fn draw_cell(grid_x: f32, grid_y: f32, cell_size: f32, position: &Position, colour: Colour) {
    let color = macroquad_colour(colour);
    draw_rectangle(
        grid_x + cell_size * position.column as f32,
        grid_y + cell_size * position.row as f32,
        cell_size,
        cell_size,
        color,
    );
}

#[macroquad::main("BasicShapes")]
async fn main() {
    let size = 8;
    let grid = Grid::generate(size);
    let mut graph = Graph::create(&grid);

    loop {
        clear_background(BLACK);

        let grid_size = screen_width().min(screen_height());
        let grid_x = (screen_width() - grid_size) / 2.0;
        let grid_y = (screen_height() - grid_size) / 2.0;
        let cell_size = grid_size / size as f32;

        for component in graph.neighbours.keys() {
            for position in &component.cells {
                draw_cell(grid_x, grid_y, cell_size, position, component.colour);
            }
        }

        next_frame().await
    }
}
<|endoftext|>


<commit_before>/// Returns the next positional prefix array (Durbin's a_{k+1}) from the current one
///
/// # Arguments
///
/// * `pos_prefix_array` - An array slice holding the current positional prefix array (a_k)
/// * `haplotypes` - An array slice holding the haplotypes at the current (k-th) site
fn next_positional_prefix_array(pos_prefix_array: &[usize], haplotypes: &[u8]) -> Vec<usize> {
    let mut allele_0s: Vec<usize> = vec![];
    let mut allele_1s: Vec<usize> = vec![];

    for &i in pos_prefix_array {
        if haplotypes[i] == 0 {
            allele_0s.push(i);
        } else {
            allele_1s.push(i);
        }
    }
    let mut next_pos_prefix_array: Vec<usize> = vec![];
    next_pos_prefix_array.extend(&allele_0s);
    next_pos_prefix_array.extend(&allele_1s);
    next_pos_prefix_array
}

/// Returns the positional prefix array (Durbin's a_k)
///
/// # Arguments
///
/// * `haplotypes` - An array slice in which each row (of length n_haplo)
///                  corresponds to a single site, and contains 0s and 1s
///                  indicating allele
/// * `n_haplo` - The number of haplotypes
/// * `k_site` - The numbered site to build the array for
fn build_positional_prefix_array(haplotypes: &[u8], n_haplo: usize, k_site: usize) -> Vec<usize> {
    let mut pos_prefix_array: Vec<usize> = vec![0; n_haplo];
    for i in 0..n_haplo {
        pos_prefix_array[i] = i;
    }

    for k in 0..(k_site-1) {
        let haplotypes_at_k = &haplotypes[k*n_haplo..(k+1)*n_haplo];
        pos_prefix_array = next_positional_prefix_array(&pos_prefix_array, &haplotypes_at_k);
    }
    pos_prefix_array
}

fn main() {
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_positional_prefix_array() {
        let n_haplo: usize = 8;
        let n_sites: usize = 6;
        let haplotypes: Vec<u8> = vec![0, 1, 1, 0, 0, 1, 1, 0,
                                       1, 1, 1, 1, 0, 0, 1, 1,
                                       0, 0, 1, 1, 0, 0, 0, 0,
                                       1, 0, 1, 1, 0, 0, 0, 1,
                                       0, 0, 1, 1, 0, 1, 0, 1,
                                       1, 1, 1, 0, 0, 0, 1, 0];
        let expected_1: Vec<usize> = vec![0, 1, 2, 3, 4, 5, 6, 7];
        assert_eq!(&expected_1,
                   &build_positional_prefix_array(&haplotypes, n_haplo, 1));
        let expected_6: Vec<usize> = vec![4, 1, 6, 0, 5, 7, 3, 2];
        assert_eq!(&expected_6,
                   &build_positional_prefix_array(&haplotypes, n_haplo, n_sites));
    }
}
<commit_msg>Algoritm 2 and tests<commit_after>/// Derives the next positional prefix array and divergence array
/// (Durbin's a_{k+1}[i] and d_{k+1}[i]) from the current ones
///
/// # Arguments
///
/// * `pos_prefix_array` - An array slice holding the positional prefix array (a_k[i])
/// * `div_array` - An array slice holding the divergence array (d_k[i])
/// * `haplotypes` - An array slice holding the haplotypes at the current (k-th) site
fn next_pp_and_div_arrays(pos_prefix_array: &[usize], div_array: &[usize], haplotypes: &[u8], k: usize) -> (Vec<usize>, Vec<usize>) {
    let mut allele_0_pos: Vec<usize> = vec![];
    let mut allele_0_div: Vec<usize> = vec![];
    let mut allele_1_pos: Vec<usize> = vec![];
    let mut allele_1_div: Vec<usize> = vec![];

    let mut div_0_index: usize = k + 1;  // TODO: Naming. longest_common_prefix?
    let mut div_1_index: usize = k + 1;

    for (&haplo_index, &start_of_match) in pos_prefix_array.iter().zip(div_array.iter()) {
        if start_of_match > div_0_index {
            div_0_index = start_of_match;
        }
        if start_of_match > div_1_index {
            div_1_index = start_of_match;
        }
        if haplotypes[haplo_index] == 0 {
            allele_0_pos.push(haplo_index);
            allele_0_div.push(div_0_index);
            div_0_index = 0;
        } else {
            allele_1_pos.push(haplo_index);
            allele_1_div.push(div_1_index);
            div_1_index = 0;
        }
    }
    let mut next_pp_array: Vec<usize> = vec![];
    next_pp_array.extend(allele_0_pos);
    next_pp_array.extend(allele_1_pos);
    let mut next_div_array: Vec<usize> = vec![];
    next_div_array.extend(allele_0_div);
    next_div_array.extend(allele_1_div);
    (next_pp_array, next_div_array)
}

/// Return the positional prefix and divergence arrays (Durbin's a_k[i] and d_k[i])
///
/// # Arguments
///
/// * `haplotypes` - An array slice in which each row (of length n_haplo)
///                  corresponds to a single site, and contains 0s and 1s
///                  indicating allele
/// * `n_haplo` - The number of haplotypes
/// * `k_site` - The numbered site to build the array for
fn build_pp_and_div_arrays(haplotypes: &[u8], n_haplo: usize, k_site: usize) -> (Vec<usize>, Vec<usize>) {
    let mut pos_prefix_array: Vec<usize> = vec![0; n_haplo];
    for i in 0..n_haplo {
        pos_prefix_array[i] = i;
    }
    let mut div_array: Vec<usize> = vec![0; n_haplo];

    for k in 0..(k_site-1) {
        let haplotypes_at_k = &haplotypes[k*n_haplo..(k+1)*n_haplo];
        let (pos_prefix_array_tmp, div_array_tmp) = next_pp_and_div_arrays(&pos_prefix_array, &div_array, &haplotypes_at_k, k);
        pos_prefix_array = pos_prefix_array_tmp;
        div_array = div_array_tmp;
    }
    (pos_prefix_array, div_array)
}

fn main() {
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_pp_and_div_arrays() {
        let n_haplo: usize = 8;
        let n_sites: usize = 6;
        let haplotypes: Vec<u8> = vec![0, 1, 1, 0, 0, 1, 1, 0,
                                       1, 1, 1, 1, 0, 0, 1, 1,
                                       0, 0, 1, 1, 0, 0, 0, 0,
                                       1, 0, 1, 1, 0, 0, 0, 1,
                                       0, 0, 1, 1, 0, 1, 0, 1,
                                       1, 1, 1, 0, 0, 0, 1, 0];
        let expected_pp_1: Vec<usize> = (0..8).collect();
        let expected_div_1: Vec<usize> = vec![0; 8];
        let (pp_1, div_1) = build_pp_and_div_arrays(&haplotypes, n_haplo, 1);
        assert_eq!(&expected_pp_1, &pp_1);
        assert_eq!(&expected_div_1, &div_1);
        let expected_pp_6: Vec<usize> = vec![4, 1, 6, 0, 5, 7, 3, 2];
        let expected_div_6: Vec<usize> = vec![5, 2, 0, 4, 5, 4, 3, 1];
        let (pp_6, div_6) = build_pp_and_div_arrays(&haplotypes, n_haplo, n_sites);
        assert_eq!(&expected_pp_6, &pp_6);
        assert_eq!(&expected_div_6, &div_6);
    }
}
<|endoftext|>


<commit_before>extern crate image;
extern crate rand;
extern crate docopt;
extern crate rustc_serialize;
extern crate num_cpus;
extern crate scoped_threadpool;
extern crate clock_ticks;
extern crate tray_rust;

use std::vec::Vec;
use std::iter;
use std::path::PathBuf;
use std::io::ErrorKind;

use rand::StdRng;
use docopt::Docopt;

use tray_rust::film::{self, ImageSample, Colorf};
use tray_rust::geometry::{Geometry, Instance, Emitter};
use tray_rust::sampler::{self, Sampler};
use tray_rust::scene;
use tray_rust::integrator::Integrator;

static USAGE: &'static str = "
Usage: tray_rust <scenefile> [options]

Options:
  -o <file>     Specify the output file or directory to save the image or frames. Supported formats are
                PNG, JPG and PPM. Default is 'frame<#>.png'.
  -n <number>   Specify the number of threads to use for rendering. Defaults to the number of cores
                on the system.
  -h, --help    Show this message.
";

#[derive(RustcDecodable, Debug)]
struct Args {
    arg_scenefile: String,
    flag_o: Option<String>,
    flag_n: Option<u32>,
}

/// Threads are each sent a sender end of the channel that is
/// read from by the render target thread which then saves the
/// values recieved to the render target
fn thread_work(spp: usize, queue: &sampler::BlockQueue, scene: &scene::Scene,
               target: &film::RenderTarget, light_list: &Vec<&Emitter>) {
    let mut sampler = sampler::LowDiscrepancy::new(queue.block_dim(), spp);
    let mut sample_pos = Vec::with_capacity(sampler.max_spp());
    let mut time_samples: Vec<_> = iter::repeat(0.0).take(sampler.max_spp()).collect();
    let block_dim = queue.block_dim();
    let mut block_samples = Vec::with_capacity(sampler.max_spp() * (block_dim.0 * block_dim.1) as usize);
    let mut rng = match StdRng::new() {
        Ok(r) => r,
        Err(e) => { println!("Failed to get StdRng, {}", e); return }
    };
    // Grab a block from the queue and start working on it, submitting samples
    // to the render target thread after each pixel
    for b in queue.iter() {
        sampler.select_block(b);
        let mut pixel_samples = 0;
        while sampler.has_samples() {
            // Get samples for a pixel and render them
            sampler.get_samples(&mut sample_pos, &mut rng);
            sampler.get_samples_1d(&mut time_samples[..], &mut rng);
            for (s, t) in sample_pos.iter().zip(time_samples.iter()) {
                let mut ray = scene.camera.generate_ray(s, *t);
                if let Some(hit) = scene.intersect(&mut ray) {
                    let c = scene.integrator.illumination(scene, light_list, &ray,
                                                          &hit, &mut sampler, &mut rng).clamp();
                    block_samples.push(ImageSample::new(s.0, s.1, c));
                } else {
                    block_samples.push(ImageSample::new(s.0, s.1, Colorf::black()));
                }
            }
            // If the samples are ok the samples for the next pixel start at the end of the current
            // pixel's samples
            if sampler.report_results(&block_samples[pixel_samples..]) {
                pixel_samples = block_samples.len();
            }
        }
        target.write(&block_samples, sampler.get_region());
        block_samples.clear();
    }
}

/// Render the scene in parallel using `n` threads and write the result to the render target
fn render_parallel(rt: &mut film::RenderTarget, scene: &scene::Scene, n: u32, spp: usize){
    let mut pool = scoped_threadpool::Pool::new(n);
    let dim = rt.dimensions();
    let block_queue = sampler::BlockQueue::new((dim.0 as u32, dim.1 as u32), (8, 8));
    let light_list: Vec<_> = scene.bvh.into_iter().filter_map(|x| {
        match x {
            &Instance::Emitter(ref e) => Some(e),
            _ => None,
        }
    }).collect();
    assert!(!light_list.is_empty(), "At least one light is required");
    pool.scoped(|scope| {
        for _ in 0..n {
            let b = &block_queue;
            let ref r = rt;
            let l = &light_list;
            scope.execute(move || {
                thread_work(spp, b, scene, r, l);
            });
        }
    });
}

fn main() {
    let args: Args = Docopt::new(USAGE).and_then(|d| d.decode()).unwrap_or_else(|e| e.exit());
    let n = match args.flag_n {
        Some(n) => n,
        None => num_cpus::get() as u32,
    };
    let out_path = match &args.flag_o {
        &Some(ref f) => {
            let p = PathBuf::from(f);
            // If we're writing to a directory make sure it exists
            if p.extension() == None {
                match std::fs::create_dir(p.as_path()) {
                    Err(e) => {
                        if e.kind() != ErrorKind::AlreadyExists {
                            panic!("Failed to create output directory");
                        }
                    },
                    Ok(_) => {},
                }
            }
            p
        },
        &None => PathBuf::from("./"),
    };

    let (mut scene, mut rt, spp, frame_info) = scene::Scene::load_file(&args.arg_scenefile[..]);
    let image_dim = rt.dimensions();
    println!("Rendering using {} threads\n--------------------", n);
    let scene_start = clock_ticks::precise_time_s();

    let time_step = frame_info.time / frame_info.frames as f32;
    for i in frame_info.start..frame_info.end + 1 {
        let frame_start = i as f32 * time_step;
        let frame_end = (i as f32 + 1.0) * time_step;
        scene.camera.update_shutter(frame_start, frame_end);
        // TODO: How often to re-build the BVH?
        println!("Frame {}: re-building bvh for {} to {}", i, frame_start, frame_end);
        scene.bvh.rebuild(frame_start, frame_end);
        println!("Frame {}: rendering for {} to {}", i, frame_start, frame_end);
        let start = clock_ticks::precise_time_s();
        render_parallel(&mut rt, &scene, n, spp);
        let time = clock_ticks::precise_time_s() - start;
        println!("Frame {}: Rendering took {}s", i, time);

        let img = rt.get_render();
        let out_file = match out_path.extension() {
            Some(_) => out_path.clone(),
            None => out_path.join(PathBuf::from(format!("frame{:05}.png", i))),
        };
        match image::save_buffer(&out_file.as_path(), &img[..], image_dim.0 as u32, image_dim.1 as u32, image::RGB(8)) {
            Ok(_) => {},
            Err(e) => println!("Error saving image, {}", e),
        };
        rt.clear();
        println!("Frame {}: Rendered to '{}'\n--------------------", i, out_file.display());
    }
    let time = clock_ticks::precise_time_s() - scene_start;
    println!("Rendering entire sequence took {}s", time);
}

<commit_msg>Re-use threadpool instead of creating it when we call render parallel<commit_after>extern crate image;
extern crate rand;
extern crate docopt;
extern crate rustc_serialize;
extern crate num_cpus;
extern crate scoped_threadpool;
extern crate clock_ticks;
extern crate tray_rust;

use std::vec::Vec;
use std::iter;
use std::path::PathBuf;
use std::io::ErrorKind;

use rand::StdRng;
use docopt::Docopt;

use tray_rust::film::{self, ImageSample, Colorf};
use tray_rust::geometry::{Geometry, Instance, Emitter};
use tray_rust::sampler::{self, Sampler};
use tray_rust::scene;
use tray_rust::integrator::Integrator;

static USAGE: &'static str = "
Usage: tray_rust <scenefile> [options]

Options:
  -o <file>     Specify the output file or directory to save the image or frames. Supported formats are
                PNG, JPG and PPM. Default is 'frame<#>.png'.
  -n <number>   Specify the number of threads to use for rendering. Defaults to the number of cores
                on the system.
  -h, --help    Show this message.
";

#[derive(RustcDecodable, Debug)]
struct Args {
    arg_scenefile: String,
    flag_o: Option<String>,
    flag_n: Option<u32>,
}

/// Threads are each sent a sender end of the channel that is
/// read from by the render target thread which then saves the
/// values recieved to the render target
fn thread_work(spp: usize, queue: &sampler::BlockQueue, scene: &scene::Scene,
               target: &film::RenderTarget, light_list: &Vec<&Emitter>) {
    let mut sampler = sampler::LowDiscrepancy::new(queue.block_dim(), spp);
    let mut sample_pos = Vec::with_capacity(sampler.max_spp());
    let mut time_samples: Vec<_> = iter::repeat(0.0).take(sampler.max_spp()).collect();
    let block_dim = queue.block_dim();
    let mut block_samples = Vec::with_capacity(sampler.max_spp() * (block_dim.0 * block_dim.1) as usize);
    let mut rng = match StdRng::new() {
        Ok(r) => r,
        Err(e) => { println!("Failed to get StdRng, {}", e); return }
    };
    // Grab a block from the queue and start working on it, submitting samples
    // to the render target thread after each pixel
    for b in queue.iter() {
        sampler.select_block(b);
        let mut pixel_samples = 0;
        while sampler.has_samples() {
            // Get samples for a pixel and render them
            sampler.get_samples(&mut sample_pos, &mut rng);
            sampler.get_samples_1d(&mut time_samples[..], &mut rng);
            for (s, t) in sample_pos.iter().zip(time_samples.iter()) {
                let mut ray = scene.camera.generate_ray(s, *t);
                if let Some(hit) = scene.intersect(&mut ray) {
                    let c = scene.integrator.illumination(scene, light_list, &ray,
                                                          &hit, &mut sampler, &mut rng).clamp();
                    block_samples.push(ImageSample::new(s.0, s.1, c));
                } else {
                    block_samples.push(ImageSample::new(s.0, s.1, Colorf::black()));
                }
            }
            // If the samples are ok the samples for the next pixel start at the end of the current
            // pixel's samples
            if sampler.report_results(&block_samples[pixel_samples..]) {
                pixel_samples = block_samples.len();
            }
        }
        target.write(&block_samples, sampler.get_region());
        block_samples.clear();
    }
}

/// Render the scene in parallel using `n` threads and write the result to the render target
fn render_parallel(rt: &mut film::RenderTarget, scene: &scene::Scene, pool: &mut scoped_threadpool::Pool, spp: usize){
    let dim = rt.dimensions();
    let block_queue = sampler::BlockQueue::new((dim.0 as u32, dim.1 as u32), (8, 8));
    let light_list: Vec<_> = scene.bvh.into_iter().filter_map(|x| {
        match x {
            &Instance::Emitter(ref e) => Some(e),
            _ => None,
        }
    }).collect();
    assert!(!light_list.is_empty(), "At least one light is required");
    let n = pool.thread_count();
    pool.scoped(|scope| {
        for _ in 0..n {
            let b = &block_queue;
            let ref r = rt;
            let l = &light_list;
            scope.execute(move || {
                thread_work(spp, b, scene, r, l);
            });
        }
    });
}

fn main() {
    let args: Args = Docopt::new(USAGE).and_then(|d| d.decode()).unwrap_or_else(|e| e.exit());
    let n = match args.flag_n {
        Some(n) => n,
        None => num_cpus::get() as u32,
    };
    let out_path = match &args.flag_o {
        &Some(ref f) => {
            let p = PathBuf::from(f);
            // If we're writing to a directory make sure it exists
            if p.extension() == None {
                match std::fs::create_dir(p.as_path()) {
                    Err(e) => {
                        if e.kind() != ErrorKind::AlreadyExists {
                            panic!("Failed to create output directory");
                        }
                    },
                    Ok(_) => {},
                }
            }
            p
        },
        &None => PathBuf::from("./"),
    };

    let (mut scene, mut rt, spp, frame_info) = scene::Scene::load_file(&args.arg_scenefile[..]);
    let image_dim = rt.dimensions();
    println!("Rendering using {} threads\n--------------------", n);
    let mut pool = scoped_threadpool::Pool::new(n);
    let scene_start = clock_ticks::precise_time_s();

    let time_step = frame_info.time / frame_info.frames as f32;
    for i in frame_info.start..frame_info.end + 1 {
        let frame_start = i as f32 * time_step;
        let frame_end = (i as f32 + 1.0) * time_step;
        scene.camera.update_shutter(frame_start, frame_end);
        // TODO: How often to re-build the BVH?
        println!("Frame {}: re-building bvh for {} to {}", i, frame_start, frame_end);
        scene.bvh.rebuild(frame_start, frame_end);
        println!("Frame {}: rendering for {} to {}", i, frame_start, frame_end);
        let start = clock_ticks::precise_time_s();
        render_parallel(&mut rt, &scene, &mut pool, spp);
        let time = clock_ticks::precise_time_s() - start;
        println!("Frame {}: Rendering took {}s", i, time);

        let img = rt.get_render();
        let out_file = match out_path.extension() {
            Some(_) => out_path.clone(),
            None => out_path.join(PathBuf::from(format!("frame{:05}.png", i))),
        };
        match image::save_buffer(&out_file.as_path(), &img[..], image_dim.0 as u32, image_dim.1 as u32, image::RGB(8)) {
            Ok(_) => {},
            Err(e) => println!("Error saving image, {}", e),
        };
        rt.clear();
        println!("Frame {}: Rendered to '{}'\n--------------------", i, out_file.display());
    }
    let time = clock_ticks::precise_time_s() - scene_start;
    println!("Rendering entire sequence took {}s", time);
}

<|endoftext|>


<commit_before>extern crate mio;
extern crate log4rs;
#[macro_use]
extern crate log;
#[macro_use]
extern crate clap;
extern crate chain_repl_test;

use std::net::SocketAddr;
use std::time::Duration;
use clap::{Arg, App};

use chain_repl_test::{ChainRepl,Role, ConfigClient};

const LOG_FILE: &'static str = "log.toml";

fn main() {
    if let Err(e) = log4rs::init_file(LOG_FILE, Default::default()) {
        panic!("Could not init logger from file {}: {}", LOG_FILE, e);
    }
    let matches = App::new("chain-repl-test")
        .arg(Arg::with_name("bind").short("l").takes_value(true))
        .arg(Arg::with_name("peer").short("p").takes_value(true))
        .arg(Arg::with_name("next").short("n").takes_value(true))
        .arg(Arg::with_name("etcd").short("e").takes_value(true))
        .get_matches();

    let mut event_loop = mio::EventLoop::new().expect("Create event loop");
    let mut service = ChainRepl::new();

    if let Some(listen_addr) = matches.value_of("bind") {
        let listen_addr = listen_addr.parse::<std::net::SocketAddr>().expect("parse bind address");
        service.listen(&mut event_loop, listen_addr, Role::Client);
    }

    if let Some(listen_addr) = matches.value_of("peer") {
        let listen_addr = listen_addr.parse::<std::net::SocketAddr>().expect("peer listen address");
        service.listen(&mut event_loop, listen_addr, Role::Upstream);
    }

    if let Some(next_addr) = matches.value_of("next") {
        let next_addr = next_addr.parse::<std::net::SocketAddr>().expect("parse next address");
        info!("Forwarding to address {:?}", next_addr);
        service.set_downstream(&mut event_loop, Some(next_addr));
    } else {
        service.set_downstream(&mut event_loop, None);
    }

    let conf = if let Some(etcd) = matches.value_of("etcd") {
        info!("Etcd at: {:?}", etcd);
        let notifier = service.get_notifier(&mut event_loop);
        let config = ConfigClient::new(etcd, service.node_config(), Duration::new(5, 0),
            move |view| notifier.notify(view)).expect("Etcd");
        Some(config)
    } else {
        info!("No etcd");
        None
    };

    info!("running chain-repl-test listener");
    event_loop.run(&mut service).expect("Run loop");
}
<commit_msg>Actual name resolution.<commit_after>extern crate mio;
extern crate log4rs;
#[macro_use]
extern crate log;
#[macro_use]
extern crate clap;
extern crate chain_repl_test;

use std::net::SocketAddr;
use std::time::Duration;
use std::net::ToSocketAddrs;
use std::collections::HashSet;
use clap::{Arg, App};

use chain_repl_test::{ChainRepl,Role, ConfigClient};

const LOG_FILE: &'static str = "log.toml";

fn main() {
    if let Err(e) = log4rs::init_file(LOG_FILE, Default::default()) {
        panic!("Could not init logger from file {}: {}", LOG_FILE, e);
    }
    let matches = App::new("chain-repl-test")
        .arg(Arg::with_name("bind").short("l").takes_value(true))
        .arg(Arg::with_name("peer").short("p").takes_value(true))
        .arg(Arg::with_name("next").short("n").takes_value(true))
        .arg(Arg::with_name("etcd").short("e").takes_value(true))
        .get_matches();

    let mut event_loop = mio::EventLoop::new().expect("Create event loop");
    let mut service = ChainRepl::new();

    if let Some(listen_addr) = matches.value_of("bind") {
        let listen_addrs = listen_addr.to_socket_addrs().expect("parse bind address").collect::<HashSet<_>>();
        info!("Client addresses: {:?}", listen_addrs);
        for listen_addr in listen_addrs {
            service.listen(&mut event_loop, listen_addr, Role::Client);
        }
    }

    if let Some(listen_addr) = matches.value_of("peer") {
        let listen_addrs = listen_addr.to_socket_addrs().expect("peer listen address").collect::<HashSet<_>>();
        info!("Peer addresses: {:?}", listen_addrs);
        for listen_addr in listen_addrs {
            service.listen(&mut event_loop, listen_addr, Role::Upstream);
        }
    }

    if let Some(next_addr) = matches.value_of("next") {
        let next_addr = next_addr.to_socket_addrs().expect("address lookup").next().expect("parse next address");
        info!("Forwarding to address {:?}", next_addr);
        service.set_downstream(&mut event_loop, Some(next_addr));
    } else {
        service.set_downstream(&mut event_loop, None);
    }

    let conf = if let Some(etcd) = matches.value_of("etcd") {
        info!("Etcd at: {:?}", etcd);
        let notifier = service.get_notifier(&mut event_loop);
        let config = ConfigClient::new(etcd, service.node_config(), Duration::new(5, 0),
            move |view| notifier.notify(view)).expect("Etcd");
        Some(config)
    } else {
        info!("No etcd");
        None
    };

    info!("running chain-repl-test listener");
    event_loop.run(&mut service).expect("Run loop");
}
<|endoftext|>


<commit_before>extern crate alpm;
#[macro_use]
extern crate clap;
extern crate curl;
extern crate env_logger;
extern crate itertools;
#[macro_use]
extern crate log;
extern crate rustc_serialize;

use clap::App;
use curl::easy::Easy;
use itertools::Itertools;
use rustc_serialize::json::Json;
use std::cmp::Ordering;
use std::collections::BTreeMap;
use std::collections::btree_map::Entry::{Occupied, Vacant};
use std::default::Default;
use std::process::exit;
use std::str;
use std::str::FromStr;

#[derive(Clone, Debug, PartialEq, PartialOrd)]
enum Severity {
    Unknown,
    Low,
    Medium,
    High,
    Critical,
}

impl FromStr for Severity {
    type Err = ();

    fn from_str(s: &str) -> Result<Severity, ()> {
        match s {
            "Low" => Ok(Severity::Low),
            "Medium" => Ok(Severity::Medium),
            "High" => Ok(Severity::High),
            "Critical" => Ok(Severity::Critical),
            _ => Ok(Severity::Unknown),
        }
    }
}

#[derive(Clone, Debug, PartialEq)]
enum Status {
    Unknown,
    Vulnerable,
    Testing,
    Fixed,
    NotAffected,
}

impl FromStr for Status {
    type Err = ();

    fn from_str(s: &str) -> Result<Status, ()> {
        match s {
            "Vulnerable" => Ok(Status::Vulnerable),
            "Testing" => Ok(Status::Testing),
            "Fixed" => Ok(Status::Fixed),
            "Not affected" => Ok(Status::NotAffected),
            _ => Ok(Status::Unknown),
        }
    }
}

#[derive(Clone, Debug)]
struct AVG {
    issues: Vec<String>,
    fixed: Option<String>,
    severity: Severity,
    status: Status,
}

impl Default for AVG {
    fn default() -> AVG {
        AVG {
            issues: vec![],
            fixed: None,
            severity: Severity::Unknown,
            status: Status::Unknown,
        }
    }
}

struct Options {
    format: Option<String>,
    quiet: u64,
    upgradable_only: bool,
}

impl Default for Options {
    fn default() -> Options {
        Options {
            format: None,
            quiet: 0,
            upgradable_only: false,
        }
    }
}

fn main() {
    env_logger::init().unwrap();

    let yaml = load_yaml!("cli.yml");
    let args = App::from_yaml(yaml).get_matches();

    let options = Options {
        format: {
            match args.value_of("format") {
                Some(f) => Some(f.to_string()),
                None => None,
            }
        },
        quiet: args.occurrences_of("quiet"),
        upgradable_only: args.is_present("upgradable"),
    };

    let mut avgs = String::new();
    {
        info!("Downloading AVGs...");
        let avgs_url = "https://security.archlinux.org/json";

        let mut easy = Easy::new();
        easy.url(avgs_url).unwrap();
        let mut transfer = easy.transfer();
        transfer.write_function(|data| {
                avgs.push_str(str::from_utf8(data).unwrap());
                Ok(data.len())
            })
            .unwrap();
        match transfer.perform() {
            Ok(_) => {}
            Err(_) => {
                println!("Cannot fetch data, please check your network connection!");
                exit(1)
            }
        };
    }

    let pacman = match args.value_of("dbpath") {
        Some(path) => alpm::Alpm::with_dbpath(path.to_string()).unwrap(),
        None => alpm::Alpm::new().unwrap(),
    };

    let mut cves: BTreeMap<String, Vec<_>> = BTreeMap::new();
    {
        let json = Json::from_str(&avgs).unwrap();

        for avg in json.as_array().unwrap() {
            let packages = avg["packages"]
                .as_array()
                .unwrap()
                .iter()
                .map(|s| s.as_string().unwrap().to_string())
                .collect::<Vec<_>>();

            if !package_is_installed(&pacman, &packages) {
                continue;
            }

            let info = to_avg(avg);

            if info.status != Status::NotAffected {
                for p in packages {
                    match cves.entry(p) {
                            Occupied(c) => c.into_mut(),
                            Vacant(c) => c.insert(vec![]),
                        }
                        .push(info.clone());
                }
            }
        }
    }

    let mut affected_avgs: BTreeMap<String, Vec<_>> = BTreeMap::new();
    for (pkg, avgs) in cves {
        for avg in &avgs {
            if system_is_affected(&pacman, &pkg, &avg) {
                match affected_avgs.entry(pkg.clone()) {
                        Occupied(c) => c.into_mut(),
                        Vacant(c) => c.insert(vec![]),
                    }
                    .push(avg.clone());
            }
        }
    }

    let merged = merge_avgs(&pacman, &affected_avgs);
    print_avgs(&options, &merged);
}

/// Converts a JSON to an AVG
fn to_avg(data: &Json) -> AVG {
    AVG {
        issues: data["issues"]
            .as_array()
            .unwrap()
            .iter()
            .map(|s| s.as_string().unwrap().to_string())
            .collect(),
        fixed: match data["fixed"].as_string() {
            Some(s) => Some(s.to_string()),
            None => None,
        },
        severity: data["severity"]
            .as_string()
            .unwrap()
            .to_string()
            .parse::<Severity>()
            .unwrap(),
        status: data["status"]
            .as_string()
            .unwrap()
            .to_string()
            .parse::<Status>()
            .unwrap(),
    }
}

#[test]
fn test_to_avg() {
    let json = Json::from_str("{\"issues\": [\"CVE-1\", \"CVE-2\"], \"fixed\": \"1.0\", \
                               \"severity\": \"High\", \"status\": \"Not affected\"}")
        .unwrap();

    let avg1 = to_avg(&json);
    assert_eq!(2, avg1.issues.len());
    assert_eq!(Some("1.0".to_string()), avg1.fixed);
    assert_eq!(Severity::High, avg1.severity);
    assert_eq!(Status::NotAffected, avg1.status);

    let json = Json::from_str("{\"issues\": [\"CVE-1\"], \"fixed\": null, \
                               \"severity\": \"Low\", \"status\": \"Vulnerable\"}")
        .unwrap();

    let avg2 = to_avg(&json);
    assert_eq!(1, avg2.issues.len());
    assert_eq!(None, avg2.fixed);
    assert_eq!(Severity::Low, avg2.severity);
    assert_eq!(Status::Vulnerable, avg2.status);
}

/// Given a package and an AVG, returns true if the system is affected
fn system_is_affected(pacman: &alpm::Alpm, pkg: &String, avg: &AVG) -> bool {
    match pacman.query_package_version(pkg.clone()) {
        Ok(v) => {
            info!("Found installed version {} for package {}", v, pkg);
            match avg.fixed {
                Some(ref version) => {
                    info!("Comparing with fixed version {}", version);
                    match pacman.vercmp(v.clone(), version.clone()).unwrap() {
                        Ordering::Less => return true,
                        _ => {}
                    };
                }
                None => return true,
            };
        }
        Err(_) => debug!("Package {} not installed", pkg),
    }

    return false;
}

#[test]
fn test_system_is_affected() {
    let pacman = alpm::Alpm::new().unwrap();

    let avg1 = AVG {
        issues: vec!["CVE-1".to_string(), "CVE-2".to_string()],
        fixed: Some("1.0.0".to_string()),
        severity: Severity::Unknown,
        status: Status::Unknown,
    };

    assert_eq!(false,
               system_is_affected(&pacman, &"pacman".to_string(), &avg1));

    let avg2 = AVG {
        issues: vec!["CVE-1".to_string(), "CVE-2".to_string()],
        fixed: Some("7.0.0".to_string()),
        severity: Severity::Unknown,
        status: Status::Unknown,
    };

    assert!(system_is_affected(&pacman, &"pacman".to_string(), &avg2));
}

/// Given a list of package names, returns true when at least one is installed
fn package_is_installed(pacman: &alpm::Alpm, packages: &Vec<String>) -> bool {
    for pkg in packages {
        match pacman.query_package_version(pkg.as_str()) {
            Ok(_) => {
                info!("Package {} is installed", pkg);
                return true;
            }
            Err(_) => debug!("Package {} not installed", pkg),
        }
    }
    return false;
}

#[test]
fn test_package_is_installed() {
    let pacman = alpm::Alpm::new().unwrap();

    let packages = vec!["pacman".to_string(), "pac".to_string()];
    assert!(package_is_installed(&pacman, &packages));

    let packages = vec!["pac".to_string()];
    assert_eq!(false, package_is_installed(&pacman, &packages));
}

/// Merge a list of AVGs into a single AVG using major version as version
fn merge_avgs(pacman: &alpm::Alpm, cves: &BTreeMap<String, Vec<AVG>>) -> BTreeMap<String, AVG> {
    let mut avgs: BTreeMap<String, AVG> = BTreeMap::new();
    for (pkg, list) in cves.iter() {
        let mut avg_issues = vec![];
        let mut avg_fixed: Option<String> = None;
        let mut avg_severity = Severity::Unknown;
        let mut avg_status = Status::Unknown;

        for a in list.iter() {
            avg_issues.append(&mut a.issues.clone());

            match avg_fixed.clone() {
                Some(ref version) => {
                    match a.fixed {
                        Some(ref v) => {
                            match pacman.vercmp(version.to_string(), v.to_string()).unwrap() {
                                Ordering::Greater => avg_fixed = a.fixed.clone(),
                                _ => {}
                            }
                        }
                        None => {}
                    }
                }
                None => avg_fixed = a.fixed.clone(),
            }

            if a.severity > avg_severity {
                avg_severity = a.severity.clone();
            }

            // We only care about testing stuff
            if a.status == Status::Testing {
                avg_status = a.status.clone();
            }
        }

        let avg = AVG {
            issues: avg_issues,
            fixed: avg_fixed,
            severity: avg_severity,
            status: avg_status,
        };
        avgs.insert(pkg.to_string(), avg);
    }

    avgs
}

#[test]
fn test_merge_avgs() {
    let mut avgs: BTreeMap<String, Vec<_>> = BTreeMap::new();

    let avg1 = AVG {
        issues: vec!["CVE-1".to_string(), "CVE-2".to_string()],
        fixed: Some("1.0.0".to_string()),
        severity: Severity::Unknown,
        status: Status::Fixed,
    };

    let avg2 = AVG {
        issues: vec!["CVE-4".to_string(), "CVE-10".to_string()],
        fixed: Some("0.9.8".to_string()),
        severity: Severity::High,
        status: Status::Testing,
    };

    assert!(Severity::Critical > Severity::High);

    avgs.insert("package".to_string(), vec![avg1.clone(), avg2.clone()]);

    avgs.insert("package2".to_string(), vec![avg1, avg2]);

    let pacman = alpm::Alpm::new().unwrap();
    let merged = merge_avgs(&pacman, &avgs);

    assert_eq!(2, merged.len());
    assert_eq!(4, merged.get(&"package".to_string()).unwrap().issues.len());
    assert_eq!(Severity::High,
               merged.get(&"package".to_string()).unwrap().severity);
    assert_eq!(Status::Testing,
               merged.get(&"package".to_string()).unwrap().status);
}

/// Print a list of AVGs
fn print_avgs(options: &Options, avgs: &BTreeMap<String, AVG>) {
    for (pkg, avg) in avgs {
        let msg = format!("Package {} is affected by {:?}", pkg, avg.issues);

        match avg.fixed {
            Some(ref v) => {
                if options.quiet == 1 {
                    println!("{}>={}", pkg, v);
                } else if options.quiet >= 2 {
                    println!("{}", pkg);
                } else {
                    match options.format {
                        Some(ref f) => {
                            println!("{}",
                                     f.replace("%n", pkg.as_str())
                                         .replace("%c", avg.issues.iter().join(",").as_str()))
                        }
                        None => {
                            if avg.status == Status::Testing {
                                println!("{}. {:?} risk! Update to {} from testing repos!",
                                         msg,
                                         avg.severity,
                                         v)
                            } else {
                                println!("{}. {:?} risk! Update to {}!", msg, avg.severity, v)
                            }
                        }
                    }
                }
            }
            None => {
                if !options.upgradable_only {
                    if options.quiet > 0 {
                        println!("{}", pkg);
                    } else {
                        match options.format {
                            Some(ref f) => {
                                println!("{}",
                                         f.replace("%n", pkg.as_str())
                                             .replace("%c", avg.issues.iter().join(",").as_str()))
                            }
                            None => println!("{}. {:?} risk!", msg, avg.severity),
                        }
                    }
                }
            }
        }
    }
}
<commit_msg>replace all unwrap calls with expect<commit_after>extern crate alpm;
#[macro_use]
extern crate clap;
extern crate curl;
extern crate env_logger;
extern crate itertools;
#[macro_use]
extern crate log;
extern crate rustc_serialize;

use clap::App;
use curl::easy::Easy;
use itertools::Itertools;
use rustc_serialize::json::Json;
use std::cmp::Ordering;
use std::collections::BTreeMap;
use std::collections::btree_map::Entry::{Occupied, Vacant};
use std::default::Default;
use std::process::exit;
use std::str;
use std::str::FromStr;

#[derive(Clone, Debug, PartialEq, PartialOrd)]
enum Severity {
    Unknown,
    Low,
    Medium,
    High,
    Critical,
}

impl FromStr for Severity {
    type Err = ();

    fn from_str(s: &str) -> Result<Severity, ()> {
        match s {
            "Low" => Ok(Severity::Low),
            "Medium" => Ok(Severity::Medium),
            "High" => Ok(Severity::High),
            "Critical" => Ok(Severity::Critical),
            _ => Ok(Severity::Unknown),
        }
    }
}

#[derive(Clone, Debug, PartialEq)]
enum Status {
    Unknown,
    Vulnerable,
    Testing,
    Fixed,
    NotAffected,
}

impl FromStr for Status {
    type Err = ();

    fn from_str(s: &str) -> Result<Status, ()> {
        match s {
            "Vulnerable" => Ok(Status::Vulnerable),
            "Testing" => Ok(Status::Testing),
            "Fixed" => Ok(Status::Fixed),
            "Not affected" => Ok(Status::NotAffected),
            _ => Ok(Status::Unknown),
        }
    }
}

#[derive(Clone, Debug)]
struct AVG {
    issues: Vec<String>,
    fixed: Option<String>,
    severity: Severity,
    status: Status,
}

impl Default for AVG {
    fn default() -> AVG {
        AVG {
            issues: vec![],
            fixed: None,
            severity: Severity::Unknown,
            status: Status::Unknown,
        }
    }
}

struct Options {
    format: Option<String>,
    quiet: u64,
    upgradable_only: bool,
}

impl Default for Options {
    fn default() -> Options {
        Options {
            format: None,
            quiet: 0,
            upgradable_only: false,
        }
    }
}

fn main() {
    env_logger::init().expect("env_logger failed");

    let yaml = load_yaml!("cli.yml");
    let args = App::from_yaml(yaml).get_matches();

    let options = Options {
        format: {
            match args.value_of("format") {
                Some(f) => Some(f.to_string()),
                None => None,
            }
        },
        quiet: args.occurrences_of("quiet"),
        upgradable_only: args.is_present("upgradable"),
    };

    let mut avgs = String::new();
    {
        info!("Downloading AVGs...");
        let avgs_url = "https://security.archlinux.org/json";

        let mut easy = Easy::new();
        easy.url(avgs_url).expect("curl::Easy::url failed");
        let mut transfer = easy.transfer();
        transfer.write_function(|data| {
            avgs.push_str(str::from_utf8(data).expect("str conversion failed"));
            Ok(data.len())
        })
        .expect("write_function failed");
        match transfer.perform() {
            Ok(_) => {}
            Err(_) => {
                println!("Cannot fetch data, please check your network connection!");
                exit(1)
            }
        };
    }

    let pacman = match args.value_of("dbpath") {
        Some(path) => alpm::Alpm::with_dbpath(path.to_string())
                                 .expect("alpm::Alpm::with_dbpath failed"),
        None => alpm::Alpm::new().expect("alpm::Alpm::new failed"),
    };

    let mut cves: BTreeMap<String, Vec<_>> = BTreeMap::new();
    {
        let json = Json::from_str(&avgs).expect("Json::from_str failed");

        for avg in json.as_array().expect("Json::as_array failed") {
            let packages = avg["packages"]
                .as_array()
                .expect("Json::as_array failed")
                .iter()
                .map(|s| s.as_string().expect("Json::as_string failed").to_string())
                .collect::<Vec<_>>();

            if !package_is_installed(&pacman, &packages) {
                continue;
            }

            let info = to_avg(avg);

            if info.status != Status::NotAffected {
                for p in packages {
                    match cves.entry(p) {
                            Occupied(c) => c.into_mut(),
                            Vacant(c) => c.insert(vec![]),
                        }
                        .push(info.clone());
                }
            }
        }
    }

    let mut affected_avgs: BTreeMap<String, Vec<_>> = BTreeMap::new();
    for (pkg, avgs) in cves {
        for avg in &avgs {
            if system_is_affected(&pacman, &pkg, &avg) {
                match affected_avgs.entry(pkg.clone()) {
                        Occupied(c) => c.into_mut(),
                        Vacant(c) => c.insert(vec![]),
                    }
                    .push(avg.clone());
            }
        }
    }

    let merged = merge_avgs(&pacman, &affected_avgs);
    print_avgs(&options, &merged);
}

/// Converts a JSON to an AVG
fn to_avg(data: &Json) -> AVG {
    AVG {
        issues: data["issues"]
            .as_array()
            .expect("Json::as_array failed")
            .iter()
            .map(|s| s.as_string().expect("Json::as_string failed").to_string())
            .collect(),
        fixed: match data["fixed"].as_string() {
            Some(s) => Some(s.to_string()),
            None => None,
        },
        severity: data["severity"]
            .as_string()
            .expect("Json::as_string failed")
            .to_string()
            .parse::<Severity>()
            .expect("parse::<Severity> failed"),
        status: data["status"]
            .as_string()
            .expect("Json::as_string failed")
            .to_string()
            .parse::<Status>()
            .expect("parse::<Status> failed"),
    }
}

#[test]
fn test_to_avg() {
    let json = Json::from_str("{\"issues\": [\"CVE-1\", \"CVE-2\"], \"fixed\": \"1.0\", \
                               \"severity\": \"High\", \"status\": \"Not affected\"}")
        .expect("Json::from_str failed");

    let avg1 = to_avg(&json);
    assert_eq!(2, avg1.issues.len());
    assert_eq!(Some("1.0".to_string()), avg1.fixed);
    assert_eq!(Severity::High, avg1.severity);
    assert_eq!(Status::NotAffected, avg1.status);

    let json = Json::from_str("{\"issues\": [\"CVE-1\"], \"fixed\": null, \
                               \"severity\": \"Low\", \"status\": \"Vulnerable\"}")
        .expect("Json::from_str failed");

    let avg2 = to_avg(&json);
    assert_eq!(1, avg2.issues.len());
    assert_eq!(None, avg2.fixed);
    assert_eq!(Severity::Low, avg2.severity);
    assert_eq!(Status::Vulnerable, avg2.status);
}

/// Given a package and an AVG, returns true if the system is affected
fn system_is_affected(pacman: &alpm::Alpm, pkg: &String, avg: &AVG) -> bool {
    match pacman.query_package_version(pkg.clone()) {
        Ok(v) => {
            info!("Found installed version {} for package {}", v, pkg);
            match avg.fixed {
                Some(ref version) => {
                    info!("Comparing with fixed version {}", version);
                    match pacman.vercmp(v.clone(), version.clone()).expect("Alpm::vercmp failed") {
                        Ordering::Less => return true,
                        _ => {}
                    };
                }
                None => return true,
            };
        }
        Err(_) => debug!("Package {} not installed", pkg),
    }

    return false;
}

#[test]
fn test_system_is_affected() {
    let pacman = alpm::Alpm::new().expect("Alpm::new failed");

    let avg1 = AVG {
        issues: vec!["CVE-1".to_string(), "CVE-2".to_string()],
        fixed: Some("1.0.0".to_string()),
        severity: Severity::Unknown,
        status: Status::Unknown,
    };

    assert_eq!(false,
               system_is_affected(&pacman, &"pacman".to_string(), &avg1));

    let avg2 = AVG {
        issues: vec!["CVE-1".to_string(), "CVE-2".to_string()],
        fixed: Some("7.0.0".to_string()),
        severity: Severity::Unknown,
        status: Status::Unknown,
    };

    assert!(system_is_affected(&pacman, &"pacman".to_string(), &avg2));
}

/// Given a list of package names, returns true when at least one is installed
fn package_is_installed(pacman: &alpm::Alpm, packages: &Vec<String>) -> bool {
    for pkg in packages {
        match pacman.query_package_version(pkg.as_str()) {
            Ok(_) => {
                info!("Package {} is installed", pkg);
                return true;
            }
            Err(_) => debug!("Package {} not installed", pkg),
        }
    }
    return false;
}

#[test]
fn test_package_is_installed() {
    let pacman = alpm::Alpm::new().expect("Alpm::new failed");

    let packages = vec!["pacman".to_string(), "pac".to_string()];
    assert!(package_is_installed(&pacman, &packages));

    let packages = vec!["pac".to_string()];
    assert_eq!(false, package_is_installed(&pacman, &packages));
}

/// Merge a list of AVGs into a single AVG using major version as version
fn merge_avgs(pacman: &alpm::Alpm, cves: &BTreeMap<String, Vec<AVG>>) -> BTreeMap<String, AVG> {
    let mut avgs: BTreeMap<String, AVG> = BTreeMap::new();
    for (pkg, list) in cves.iter() {
        let mut avg_issues = vec![];
        let mut avg_fixed: Option<String> = None;
        let mut avg_severity = Severity::Unknown;
        let mut avg_status = Status::Unknown;

        for a in list.iter() {
            avg_issues.append(&mut a.issues.clone());

            match avg_fixed.clone() {
                Some(ref version) => {
                    match a.fixed {
                        Some(ref v) => {
                            match pacman.vercmp(version.to_string(), v.to_string())
                                        .expect("Alpm::vercmp failed") {
                                Ordering::Greater => avg_fixed = a.fixed.clone(),
                                _ => {}
                            }
                        }
                        None => {}
                    }
                }
                None => avg_fixed = a.fixed.clone(),
            }

            if a.severity > avg_severity {
                avg_severity = a.severity.clone();
            }

            // We only care about testing stuff
            if a.status == Status::Testing {
                avg_status = a.status.clone();
            }
        }

        let avg = AVG {
            issues: avg_issues,
            fixed: avg_fixed,
            severity: avg_severity,
            status: avg_status,
        };
        avgs.insert(pkg.to_string(), avg);
    }

    avgs
}

#[test]
fn test_merge_avgs() {
    let mut avgs: BTreeMap<String, Vec<_>> = BTreeMap::new();

    let avg1 = AVG {
        issues: vec!["CVE-1".to_string(), "CVE-2".to_string()],
        fixed: Some("1.0.0".to_string()),
        severity: Severity::Unknown,
        status: Status::Fixed,
    };

    let avg2 = AVG {
        issues: vec!["CVE-4".to_string(), "CVE-10".to_string()],
        fixed: Some("0.9.8".to_string()),
        severity: Severity::High,
        status: Status::Testing,
    };

    assert!(Severity::Critical > Severity::High);

    avgs.insert("package".to_string(), vec![avg1.clone(), avg2.clone()]);

    avgs.insert("package2".to_string(), vec![avg1, avg2]);

    let pacman = alpm::Alpm::new().expect("Alpm::new failed");
    let merged = merge_avgs(&pacman, &avgs);

    assert_eq!(2, merged.len());
    assert_eq!(4, merged.get(&"package".to_string()).expect("'package' key not found")
                                                    .issues.len());
    assert_eq!(Severity::High,
               merged.get(&"package".to_string()).expect("'package' key not found").severity);
    assert_eq!(Status::Testing,
               merged.get(&"package".to_string()).expect("'package' key not found").status);
}

/// Print a list of AVGs
fn print_avgs(options: &Options, avgs: &BTreeMap<String, AVG>) {
    for (pkg, avg) in avgs {
        let msg = format!("Package {} is affected by {:?}", pkg, avg.issues);

        match avg.fixed {
            Some(ref v) => {
                if options.quiet == 1 {
                    println!("{}>={}", pkg, v);
                } else if options.quiet >= 2 {
                    println!("{}", pkg);
                } else {
                    match options.format {
                        Some(ref f) => {
                            println!("{}",
                                     f.replace("%n", pkg.as_str())
                                         .replace("%c", avg.issues.iter().join(",").as_str()))
                        }
                        None => {
                            if avg.status == Status::Testing {
                                println!("{}. {:?} risk! Update to {} from testing repos!",
                                         msg,
                                         avg.severity,
                                         v)
                            } else {
                                println!("{}. {:?} risk! Update to {}!", msg, avg.severity, v)
                            }
                        }
                    }
                }
            }
            None => {
                if !options.upgradable_only {
                    if options.quiet > 0 {
                        println!("{}", pkg);
                    } else {
                        match options.format {
                            Some(ref f) => {
                                println!("{}",
                                         f.replace("%n", pkg.as_str())
                                             .replace("%c", avg.issues.iter().join(",").as_str()))
                            }
                            None => println!("{}. {:?} risk!", msg, avg.severity),
                        }
                    }
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>extern crate image;
extern crate clap;
extern crate rand;
extern crate itertools;

use clap::*;

mod codec;

mod rgba;
mod rgb;

mod utils;

fn main()
{
    let matches = App::new("stag")
        .version("0.0.1")
        .about("Image steganography")
        .subcommand(SubCommand::with_name("encode")
            .about("encodes a file")
            .arg(Arg::with_name("mode")
                 .short("m")
                 .long("mode")
                 .value_name("MODE")
                 .help("Set the encoding mode, default depends on SOURCE type")
                 .takes_value(true))
            .arg(Arg::with_name("SOURCE")
                 .help("The image source")
                 .index(1)
                 .required(true))
            .arg(Arg::with_name("OUTPUT")
                 .help("The output image")
                 .index(2)
                 .required(true)))
        .subcommand(SubCommand::with_name("decode")
            .about("decodes a file")
            .arg(Arg::with_name("mode")
                 .short("m")
                 .long("mode")
                 .value_name("MODE")
                 .help("Set the decoding mode, default depends on SOURCE type")
                 .takes_value(true))
            .arg(Arg::with_name("SOURCE")
                 .help("The image source")
                 .index(1)
                 .required(true))
            .arg(Arg::with_name("length")
                 .short("l")
                 .long("length")
                 .help("the amount of bytes to decode")
                 .value_name("LENGTH")
                 .takes_value(true)
                 .required(true)))
        .subcommand(SubCommand::with_name("estimate")
                .about("estimate how many bytes will fit into a file")
                .arg(Arg::with_name("SOURCE")
                     .help("The image source")
                     .index(1)
                     .takes_value(true))
                .arg(Arg::with_name("mode")
                     .short("m")
                     .long("mode")
                     .value_name("MODE")
                     .help("The encoding mode, default depends on SOURCE type")
                     .takes_value(true)))
        .get_matches();

    if let Some(matches) = matches.subcommand_matches("encode")
    {
        dispatch_encode(
            matches.value_of("mode"),
            matches.value_of("SOURCE").unwrap(),
            matches.value_of("OUTPUT").unwrap()
        );
    }

    if let Some(matches) = matches.subcommand_matches("decode")
    {
        dispatch_decode(
            matches.value_of("mode"),
            matches.value_of("SOURCE").unwrap(),
            matches.value_of("length").unwrap()
        );
    }

    if let Some(matches) = matches.subcommand_matches("estimate")
    {
        dispatch_estimate(
            matches.value_of("mode"),
            matches.value_of("SOURCE").unwrap()
        );
    }
}

use std::io::{stdin, stdout, Read, Write};
use std::str::FromStr;

use image::{open, DynamicImage};
use rand::StdRng;

use codec::Codec;
use rgba::RgbaCodec;

fn dispatch_encode(mode: Option<&str>, source: &str, output: &str)
{
    let dyimage = match open(source)
    {
        Ok(di) => di,
        Err(_) => error_out("Error opening source image for encoding")
    };

    match dyimage
    {
        DynamicImage::ImageRgba8(image) =>
        {
            match encode::<RgbaCodec>(image, mode).save(output)
            {
                Ok(_) => {},
                Err(_) => error_out("Error saving encoded output file"),
            };
        },
        _ => error_out("Unsupported filetype"),
    }
}

fn dispatch_decode(mode: Option<&str>, source: &str, len: &str)
{
    let len = match len.parse::<usize>()
    {
        Ok(l) => l,
        Err(_) => error_out("len argument to decode is not a number"),
    };

    let dyimage = match open(source)
    {
        Ok(di) => di,
        Err(_) => error_out("Error opening source image for decoding"),
    };

    match dyimage
    {
        DynamicImage::ImageRgba8(image) =>
        {
            decode::<RgbaCodec>(image, mode, len);
        },
        _ => error_out("Unsupported filetype"),
    }
}

fn dispatch_estimate(mode: Option<&str>, source: &str)
{
    let dyimage = match open(source)
    {
        Ok(di) => di,
        Err(_) => error_out("Error opening source image for estimating"),
    };

    match dyimage
    {
        DynamicImage::ImageRgba8(image) =>
        {
            estimate::<RgbaCodec>(image, mode);
        },
        _ => error_out("Unsupported filetype"),
    }
}

fn encode<C: Codec>(mut image: C::Input, mode: Option<&str>) -> C::Input
{
    let mode = if let Some(mode) = mode
    {
        match C::Mode::from_str(mode)
        {
            Ok(mode) => mode,
            Err(_) => error_out("Error parsing mode string"),
        }
    }
    else
    {
        C::Mode::default()
    };

    let rng = match StdRng::new()
    {
        Ok(r) => r,
        Err(_) => error_out("Error creating source of randomness"),
    };

    let mut payload = String::new();
    stdin().read_to_string(&mut payload).unwrap();

    C::encode(&mut image, payload.as_bytes(), mode, rng);

    image
}

fn decode<C: Codec>(image: C::Input, mode: Option<&str>, len: usize)
{
    let mode = if let Some(mode) = mode
    {
        match C::Mode::from_str(mode)
        {
            Ok(mode) => mode,
            Err(_) => error_out("Error parsing mode string"),
        }
    }
    else
    {
        C::Mode::default()
    };

    let mut buf = vec![0; len];

    C::decode(&image, &mut buf, len, mode);

    match stdout().write(&buf)
    {
        Ok(_) => {},
        Err(_) => error_out("Error writing decoded payload"),
    };
}

fn estimate<C: Codec>(image: C::Input, mode: Option<&str>)
{
    let mode = if let Some(mode) = mode
    {
        match C::Mode::from_str(mode)
        {
            Ok(mode) => mode,
            Err(_) => error_out("Error parsing mode string"),
        }
    }
    else
    {
        C::Mode::default()
    };

    match C::estimate(&image, mode)
    {
        Some(i) => println!("Estimate {} bytes", i),
        None => println!("Could not make an estimate"),
    }
}

fn error_out(msg: &str) -> !
{
    eprintln!("{}", msg);
    ::std::process::exit(1)
}
<commit_msg>Added Rgb codec to frontend<commit_after>extern crate image;
extern crate clap;
extern crate rand;
extern crate itertools;

use clap::*;

mod codec;

mod rgba;
mod rgb;

mod utils;

fn main()
{
    let matches = App::new("stag")
        .version("0.0.1")
        .about("Image steganography")
        .subcommand(SubCommand::with_name("encode")
            .about("encodes a file")
            .arg(Arg::with_name("mode")
                 .short("m")
                 .long("mode")
                 .value_name("MODE")
                 .help("Set the encoding mode, default depends on SOURCE type")
                 .takes_value(true))
            .arg(Arg::with_name("SOURCE")
                 .help("The image source")
                 .index(1)
                 .required(true))
            .arg(Arg::with_name("OUTPUT")
                 .help("The output image")
                 .index(2)
                 .required(true)))
        .subcommand(SubCommand::with_name("decode")
            .about("decodes a file")
            .arg(Arg::with_name("mode")
                 .short("m")
                 .long("mode")
                 .value_name("MODE")
                 .help("Set the decoding mode, default depends on SOURCE type")
                 .takes_value(true))
            .arg(Arg::with_name("SOURCE")
                 .help("The image source")
                 .index(1)
                 .required(true))
            .arg(Arg::with_name("length")
                 .short("l")
                 .long("length")
                 .help("the amount of bytes to decode")
                 .value_name("LENGTH")
                 .takes_value(true)
                 .required(true)))
        .subcommand(SubCommand::with_name("estimate")
                .about("estimate how many bytes will fit into a file")
                .arg(Arg::with_name("SOURCE")
                     .help("The image source")
                     .index(1)
                     .takes_value(true))
                .arg(Arg::with_name("mode")
                     .short("m")
                     .long("mode")
                     .value_name("MODE")
                     .help("The encoding mode, default depends on SOURCE type")
                     .takes_value(true)))
        .get_matches();

    if let Some(matches) = matches.subcommand_matches("encode")
    {
        dispatch_encode(
            matches.value_of("mode"),
            matches.value_of("SOURCE").unwrap(),
            matches.value_of("OUTPUT").unwrap()
        );
    }

    if let Some(matches) = matches.subcommand_matches("decode")
    {
        dispatch_decode(
            matches.value_of("mode"),
            matches.value_of("SOURCE").unwrap(),
            matches.value_of("length").unwrap()
        );
    }

    if let Some(matches) = matches.subcommand_matches("estimate")
    {
        dispatch_estimate(
            matches.value_of("mode"),
            matches.value_of("SOURCE").unwrap()
        );
    }
}

use std::io::{stdin, stdout, Read, Write};
use std::str::FromStr;

use image::{open, DynamicImage};
use rand::StdRng;

use codec::Codec;
use rgba::RgbaCodec;
use rgb::RgbCodec;

fn dispatch_encode(mode: Option<&str>, source: &str, output: &str)
{
    let dyimage = match open(source)
    {
        Ok(di) => di,
        Err(_) => error_out("Error opening source image for encoding")
    };

    match dyimage
    {
        DynamicImage::ImageRgba8(image) =>
        {
            match encode::<RgbaCodec>(image, mode).save(output)
            {
                Ok(_) => {},
                Err(_) => error_out("Error saving encoded output file"),
            };
        },
        DynamicImage::ImageRgb8(image) =>
        {
            match encode::<RgbCodec>(image, mode).save(output)
            {
                Ok(_) => {},
                Err(_) => error_out("Error saving encoded output file"),
            };
        },
        _ => error_out("Unsupported filetype"),
    }
}

fn dispatch_decode(mode: Option<&str>, source: &str, len: &str)
{
    let len = match len.parse::<usize>()
    {
        Ok(l) => l,
        Err(_) => error_out("len argument to decode is not a number"),
    };

    let dyimage = match open(source)
    {
        Ok(di) => di,
        Err(_) => error_out("Error opening source image for decoding"),
    };

    match dyimage
    {
        DynamicImage::ImageRgba8(image) =>
        {
            decode::<RgbaCodec>(image, mode, len);
        },
        DynamicImage::ImageRgb8(image) =>
        {
            decode::<RgbCodec>(image, mode, len);
        },
        _ => error_out("Unsupported filetype"),
    }
}

fn dispatch_estimate(mode: Option<&str>, source: &str)
{
    let dyimage = match open(source)
    {
        Ok(di) => di,
        Err(_) => error_out("Error opening source image for estimating"),
    };

    match dyimage
    {
        DynamicImage::ImageRgba8(image) =>
        {
            estimate::<RgbaCodec>(image, mode);
        },
        DynamicImage::ImageRgb8(image) =>
        {
            estimate::<RgbCodec>(image, mode);
        },
        _ => error_out("Unsupported filetype"),
    }
}

fn encode<C: Codec>(mut image: C::Input, mode: Option<&str>) -> C::Input
{
    let mode = if let Some(mode) = mode
    {
        match C::Mode::from_str(mode)
        {
            Ok(mode) => mode,
            Err(_) => error_out("Error parsing mode string"),
        }
    }
    else
    {
        C::Mode::default()
    };

    let rng = match StdRng::new()
    {
        Ok(r) => r,
        Err(_) => error_out("Error creating source of randomness"),
    };

    let mut payload = String::new();
    stdin().read_to_string(&mut payload).unwrap();

    C::encode(&mut image, payload.as_bytes(), mode, rng);

    image
}

fn decode<C: Codec>(image: C::Input, mode: Option<&str>, len: usize)
{
    let mode = if let Some(mode) = mode
    {
        match C::Mode::from_str(mode)
        {
            Ok(mode) => mode,
            Err(_) => error_out("Error parsing mode string"),
        }
    }
    else
    {
        C::Mode::default()
    };

    let mut buf = vec![0; len];

    C::decode(&image, &mut buf, len, mode);

    match stdout().write(&buf)
    {
        Ok(_) => {},
        Err(_) => error_out("Error writing decoded payload"),
    };
}

fn estimate<C: Codec>(image: C::Input, mode: Option<&str>)
{
    let mode = if let Some(mode) = mode
    {
        match C::Mode::from_str(mode)
        {
            Ok(mode) => mode,
            Err(_) => error_out("Error parsing mode string"),
        }
    }
    else
    {
        C::Mode::default()
    };

    match C::estimate(&image, mode)
    {
        Some(i) => println!("Estimate {} bytes", i),
        None => println!("Could not make an estimate"),
    }
}

fn error_out(msg: &str) -> !
{
    eprintln!("{}", msg);
    ::std::process::exit(1)
}
<|endoftext|>