<commit_before>
// YAML files as iptables configuration sources

// Rust Core
use std::ascii::AsciiExt;
use std::fs::File;
use std::io::Read;
use std::path::Path;
use std::process::exit;

// Crate.io
extern crate clap;
extern crate yaml_rust;

use clap::{Arg, App};
use yaml_rust::{Yaml, YamlLoader};

fn main() {

    let app = App::new("iptables-compose")
        .version("1.0.0")
        .global_version(true)
        .unified_help_message(true)
        .arg_required_else_help(true)
        .about("\nYAML files as iptables configuration sources")
        .arg(Arg::with_name("CONFIG")
             .multiple(true)
             .short("c")
             .long("config")
             .help("yaml file as iptables configuration source\n")
             .conflicts_with("license")
             .takes_value(true))
        .args_from_usage("-l --license 'Prints License'");

    let matches = app.get_matches();

    if let Some(ref f_paths) = matches.values_of("CONFIG") {
        for f_path in f_paths.iter() {
            read_yaml(f_path);
        }
    }

    if matches.is_present("license") {
        print_license();
    }

}

fn print_license() {
    let s:String = "
The MIT License (MIT)

Copyright (c) 2015 Samuel <sam@infinitely.io>

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the \"Software\"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
".to_string();
    println!("{}", s);
}

fn read_yaml(f_path: &str) {

    // Create a path for yaml configuration file
    let path = Path::new(&f_path);
    let display = path.display();

    // Open file
    let mut file = match File::open(&path) {
        Err(_)   => {
            println!("Failed to open file \"{}\"", display);
            exit(1);
        },
        Ok(file) => file,
    };

    // Create buffer string and read content
    let mut buffer_str = String::new();
    match file.read_to_string(&mut buffer_str) {
        Err(_)  => {
            println!("Failed to read file \"{}\"", display);
            exit(1);
        },
        Ok(str) => str,
    };

    // Parse yaml file
    let yaml = match YamlLoader::load_from_str(&buffer_str) {
        Err(_)  => {
            println!("Failed to parse yaml data: \"{}\"", display);
            exit(1);
        },
        Ok(yaml) => yaml,
    };

    // Check if yaml file has no data, exits program if no data
    if yaml.is_empty() {
        println!("\"{}\" has no configuration(s)", display);
        exit(1);
    };

    // Validate configuration template format
    match yaml[0].as_hash() {
        Some(doc) => doc,
        None => {
            println!("Configuration template format is invalid: {}", display);
            exit(1);
        }
    };

    // Parse configuration from YAML configuration file
    parse_yaml(&yaml[0]);

}

fn parse_yaml(doc: &Yaml) {
    // Reset all rules
    reset_rules();
    // Read all rules from yaml
    match doc {
        // Parse if template data is a hash object
        &Yaml::Hash(ref h) => {
            for (k, v) in h {
                match k.as_str().unwrap() {
                    // Parse `filter` rules
                    "filter" | "FILTER" => parse_filter(k, v),
                    // Parse custom section
                    _ => parse_section(k, v)
                }
            }
        },
        _ => {
            println!("Configuration template format is invalid");
            exit(1);
        }
    }
}

fn reset_rules() {
    let mut s:String = "# reset all rules".to_string();
    s.push_str("\niptables -F");
    s.push_str("\niptables -t nat -F");
    s.push_str("\niptables -t mangle -F");
    s.push_str("\niptables -X");
    s.push_str("\n# defaults settings and accepts your current connection");
    s.push_str("\niptables -t nat -A POSTROUTING -j MASQUERADE");
    s.push_str("\niptables -A INPUT -i lo -j ACCEPT");
    s.push_str("\niptables -A INPUT -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT");
    println!("{}", s);
}

fn parse_section(id: &Yaml, doc: &Yaml) {
    match doc {
        &Yaml::Hash(ref h) => {
            println!("# {} rules", id.as_str().unwrap());
            for (k, v) in h {
                let k = k.as_str().unwrap();
                match k {
                    "ports" | "PORTS" => parse_ports(v),
                    _ => {
                        println!("Rules \"{}\" is not available", k);
                        exit(1);
                    }
                }
            }
        },
        _ => {
            println!("Configuration template format is invalid");
            exit(1);
        }
    }
}

fn parse_ports(doc: &Yaml) {
    match doc {
        &Yaml::Array(ref v) => {
            for x in v {
                parse_port_item(x);
            }
        },
        _ => {
            println!("Configuration \"ports\" format is invalid");
            exit(1);
        }
    }
}

fn parse_port_item(doc: &Yaml) {
    match doc {
        &Yaml::Hash(_) => {
            if doc["port"].is_badvalue() {
                println!("Port is not defined");
            }
            let port = doc["port"].as_i64();
            match port {
                Some(port) if port > -1 => (),
                _ => {
                    if port.unwrap() <= -1 {
                        println!("Port has to be greater or equals to 0");
                    } else {
                        println!("Port is not invalid");
                    }
                    exit(1);
                }
            }
            let mut cmd:String = "iptables".to_string();
            if ! doc["forward"].is_badvalue() {
                cmd.push_str(" -t nat -A PREROUTING");
            } else {
                let direction = doc["type"].as_str().unwrap_or("input");
                match direction {
                    "input" | "output" | "forward" | "INPUT" | "OUTPUT" | "FORWARD" => {
                        cmd.push_str(" -I ");
                        cmd.push_str(&direction.to_ascii_uppercase());
                    },
                    _ => {
                        println!("Direction value is invalid");
                        exit(1);
                    }
                }
            }
            match doc["subnet"] {
                Yaml::Array(ref v) => {
                    for x in v {
                        cmd.push_str(" -s ");
                        cmd.push_str(x.as_str().unwrap_or("0.0.0.0/0"));
                    }
                },
                _ => {}
            }
            let protocol = doc["protocol"].as_str().unwrap_or("tcp");
            match protocol {
                "tcp" | "udp" | "TCP" | "UDP" => {
                    cmd.push_str(" -p ");
                    cmd.push_str(protocol);
                    cmd.push_str(" -m ");
                    cmd.push_str(protocol);
                },
                _ => {
                    println!("Protocol value is invalid");
                    exit(1);
                }
            }
            cmd.push_str(" --dport ");
            cmd.push_str(&port.unwrap().to_string());
            if ! doc["forward"].is_badvalue() {
                let forward = doc["forward"].as_i64();
                match forward {
                    Some(forward) if forward > -1 => (),
                    _ => {
                        if forward.unwrap() <= -1 {
                            println!("forward port has to be greater or equals to 0");
                        } else {
                            println!("forward port is not invalid");
                        }
                        exit(1);
                    }
                }
                cmd.push_str(" -j REDIRECT --to-port ");
                cmd.push_str(&forward.unwrap().to_string());
            } else {
                let allow = doc["allow"].as_bool().unwrap_or(true);
                cmd.push_str(" -j ");
                match allow {
                    true => {
                        cmd.push_str("ACCEPT");
                    },
                    _ => {
                        cmd.push_str("DROP");
                    }
                }
            }
            println!("{}", cmd);
        },
        _ => {
            println!("Configuration \"ports\" item format is invalid");
            exit(1);
        }
    }
}

fn parse_filter(id: &Yaml, doc: &Yaml) {
    // Check if filter section exists
    match doc {
        &Yaml::Hash(ref h) => {
            println!("# {} rules", id.as_str().unwrap());
            for (k, v) in h {
                let k = k.as_str().unwrap();
                let v = v.as_str().unwrap();
                match k {
                    "input" | "output" | "forward" | "INPUT" | "OUTPUT" | "FORWARD" => {
                        match v {
                            "drop" | "reject" | "accept" | "DROP" | "REJECT" | "ACCEPT" => println!("iptables -P {} {}", k.to_ascii_uppercase(), v.to_ascii_uppercase()),
                            _ => {
                                println!("Rules \"{}\" only accept options of \"drop\",\"reject\" or \"accept\"", k);
                                exit(1);
                            }
                        }
                    },
                    _ => println!("iptables -P {}", k.to_ascii_uppercase())
                }
            }
        },
        _ => {
            println!("Configuration template format is invalid");
            exit(1);
        }
    }
}
<commit_msg>Removed unnecessary code, added optional reset rules<commit_after>
// YAML files as iptables configuration sources

// Rust Core
use std::ascii::AsciiExt;
use std::fs::File;
use std::io::Read;
use std::path::Path;
use std::process::exit;

// Crate.io
extern crate clap;
extern crate yaml_rust;

use clap::{Arg, App};
use yaml_rust::{Yaml, YamlLoader};

fn main() {

    let app = App::new("iptables-compose")
        .version("1.0.0")
        .global_version(true)
        .unified_help_message(true)
        .arg_required_else_help(true)
        .about("\nYAML files as iptables configuration sources")
        .arg(Arg::with_name("CONFIG")
             .multiple(false)
             .help("yaml file as iptables configuration source")
             .conflicts_with("license")
             .index(1))
        .arg(Arg::with_name("RESET")
             .short("r")
             .long("reset")
             .help("reset iptables rules\n")
             .requires("CONFIG"))
        .args_from_usage("-l --license 'Prints License'");

    let matches = app.get_matches();

    if let Some(ref f_path) = matches.value_of("CONFIG") {
        // Reset rules if "reset" argument is present
        if matches.is_present("RESET") {
            reset_rules();
        }
        read_yaml(f_path);
    }

    if matches.is_present("license") {
        print_license();
    }

}

fn print_license() {
    let s:String = "
The MIT License (MIT)

Copyright (c) 2015 Samuel <sam@infinitely.io>

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the \"Software\"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
".to_string();
    println!("{}", s);
}

fn read_yaml(f_path: &str) {

    // Create a path for yaml configuration file
    let path = Path::new(&f_path);
    let display = path.display();

    // Open file
    let mut file = match File::open(&path) {
        Err(_)   => {
            println!("Failed to open file \"{}\"", display);
            exit(1);
        },
        Ok(file) => file,
    };

    // Create buffer string and read content
    let mut buffer_str = String::new();
    match file.read_to_string(&mut buffer_str) {
        Err(_)  => {
            println!("Failed to read file \"{}\"", display);
            exit(1);
        },
        Ok(str) => str,
    };

    // Parse yaml file
    let yaml = match YamlLoader::load_from_str(&buffer_str) {
        Err(_)  => {
            println!("Failed to parse yaml data: \"{}\"", display);
            exit(1);
        },
        Ok(yaml) => yaml,
    };

    // Check if yaml file has no data, exits program if no data
    if yaml.is_empty() {
        println!("\"{}\" has no configuration(s)", display);
        exit(1);
    };

    // Validate configuration template format
    match yaml[0].as_hash() {
        Some(doc) => doc,
        None => {
            println!("Configuration template format is invalid: {}", display);
            exit(1);
        }
    };

    // Parse configuration from YAML configuration file
    parse_yaml(&yaml[0]);

}

fn parse_yaml(doc: &Yaml) {
    // Read all rules from yaml
    match doc {
        // Parse if template data is a hash object
        &Yaml::Hash(ref h) => {
            for (k, v) in h {
                match k.as_str().unwrap() {
                    // Parse `filter` rules
                    "filter" | "FILTER" => parse_filter(v),
                    // Parse custom section
                    _ => parse_section(v)
                }
            }
        },
        _ => {
            println!("Configuration template format is invalid");
            exit(1);
        }
    }
}

fn reset_rules() {
    let mut s:String = "iptables -F".to_string();
    s.push_str("\niptables -X");
    s.push_str("\niptables -t nat -F");
    s.push_str("\niptables -t nat -X");
    s.push_str("\niptables -t mangle -F");
    s.push_str("\niptables -t mangle -X");
    println!("{}", s);
}

fn parse_section(doc: &Yaml) {
    match doc {
        &Yaml::Hash(ref h) => {
            for (k, v) in h {
                let k = k.as_str().unwrap();
                match k {
                    "ports" | "PORTS" => parse_ports(v),
                    _ => {
                        println!("Rules \"{}\" is not available", k);
                        exit(1);
                    }
                }
            }
        },
        _ => {
            println!("Configuration template format is invalid");
            exit(1);
        }
    }
}

fn parse_ports(doc: &Yaml) {
    match doc {
        &Yaml::Array(ref v) => {
            for x in v {
                parse_port_item(x);
            }
        },
        _ => {
            println!("Configuration \"ports\" format is invalid");
            exit(1);
        }
    }
}

fn parse_port_item(doc: &Yaml) {
    match doc {
        &Yaml::Hash(_) => {
            if doc["port"].is_badvalue() {
                println!("Port is not defined");
            }
            let port = doc["port"].as_i64();
            match port {
                Some(port) if port > -1 => (),
                _ => {
                    if port.unwrap() <= -1 {
                        println!("Port has to be greater or equals to 0");
                    } else {
                        println!("Port is not invalid");
                    }
                    exit(1);
                }
            }
            let mut cmd:String = "iptables".to_string();
            if ! doc["forward"].is_badvalue() {
                cmd.push_str(" -t nat -A PREROUTING");
            } else {
                let direction = doc["type"].as_str().unwrap_or("input");
                match direction {
                    "input" | "output" | "forward" | "INPUT" | "OUTPUT" | "FORWARD" => {
                        cmd.push_str(" -I ");
                        cmd.push_str(&direction.to_ascii_uppercase());
                    },
                    _ => {
                        println!("Direction value is invalid");
                        exit(1);
                    }
                }
            }
            match doc["subnet"] {
                Yaml::Array(ref v) => {
                    for x in v {
                        cmd.push_str(" -s ");
                        cmd.push_str(x.as_str().unwrap_or("0.0.0.0/0"));
                    }
                },
                _ => {}
            }
            let protocol = doc["protocol"].as_str().unwrap_or("tcp");
            match protocol {
                "tcp" | "udp" | "TCP" | "UDP" => {
                    cmd.push_str(" -p ");
                    cmd.push_str(protocol);
                    cmd.push_str(" -m ");
                    cmd.push_str(protocol);
                },
                _ => {
                    println!("Protocol value is invalid");
                    exit(1);
                }
            }
            cmd.push_str(" --dport ");
            cmd.push_str(&port.unwrap().to_string());
            if ! doc["forward"].is_badvalue() {
                let forward = doc["forward"].as_i64();
                match forward {
                    Some(forward) if forward > -1 => (),
                    _ => {
                        if forward.unwrap() <= -1 {
                            println!("forward port has to be greater or equals to 0");
                        } else {
                            println!("forward port is not invalid");
                        }
                        exit(1);
                    }
                }
                cmd.push_str(" -j REDIRECT --to-port ");
                cmd.push_str(&forward.unwrap().to_string());
            } else {
                let allow = doc["allow"].as_bool().unwrap_or(true);
                cmd.push_str(" -j ");
                match allow {
                    true => {
                        cmd.push_str("ACCEPT");
                    },
                    _ => {
                        cmd.push_str("DROP");
                    }
                }
            }
            println!("{}", cmd);
        },
        _ => {
            println!("Configuration \"ports\" item format is invalid");
            exit(1);
        }
    }
}

fn parse_filter(doc: &Yaml) {
    // Check if filter section exists
    match doc {
        &Yaml::Hash(ref h) => {
            for (k, v) in h {
                let k = k.as_str().unwrap();
                let v = v.as_str().unwrap();
                match k {
                    "input" | "output" | "forward" | "INPUT" | "OUTPUT" | "FORWARD" => {
                        match v {
                            "drop" | "reject" | "accept" | "DROP" | "REJECT" | "ACCEPT" => println!("iptables -P {} {}", k.to_ascii_uppercase(), v.to_ascii_uppercase()),
                            _ => {
                                println!("Rules \"{}\" only accept options of \"drop\",\"reject\" or \"accept\"", k);
                                exit(1);
                            }
                        }
                    },
                    _ => println!("iptables -P {}", k.to_ascii_uppercase())
                }
            }
        },
        _ => {
            println!("Configuration template format is invalid");
            exit(1);
        }
    }
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate clap;

mod lib;

use clap::{App, Arg};
use lib::Petnames;
use rand::seq::IteratorRandom;
use std::collections::HashSet;
use std::str::FromStr;

fn main() {
    let matches = App::new("rust-petname")
        .version(crate_version!())
        .author(crate_authors!())
        .about("Generate human readable random names.")
        .after_help(concat!(
            "Based on Dustin Kirkland's petname project ",
            "<https://github.com/dustinkirkland/petname>."
        ))
        .arg(
            Arg::with_name("words")
                .short("w")
                .long("words")
                .value_name("WORDS")
                .default_value("2")
                .help("Number of words in name")
                .takes_value(true)
                .validator(can_be_parsed::<u8>),
        )
        .arg(
            Arg::with_name("separator")
                .short("s")
                .long("separator")
                .value_name("SEP")
                .default_value("-")
                .help("Separator between words")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("complexity")
                .short("c")
                .long("complexity")
                .value_name("COM")
                .possible_values(&["0", "1", "2"])
                .hide_possible_values(true)
                .default_value("0")
                .help("Use small words (0), medium words (1), or large words (2)")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("count")
                .long("count")
                .value_name("COUNT")
                .default_value("1")
                .help("Generate multiple names; pass 0 to produce infinite names!")
                .takes_value(true)
                .validator(can_be_parsed::<usize>),
        )
        .arg(
            Arg::with_name("letters")
                .short("l")
                .long("letters")
                .value_name("LETTERS")
                .default_value("0")
                .help("Maxiumum number of letters in each word; 0 for unlimited")
                .takes_value(true)
                .validator(can_be_parsed::<usize>),
        )
        .arg(
            Arg::with_name("alliterate")
                .short("a")
                .long("alliterate")
                .help("Generate names where each word begins with the same letter")
                .takes_value(false),
        )
        .arg(
            // For compatibility with upstream.
            Arg::with_name("ubuntu")
                .short("u")
                .long("ubuntu")
                .help("Alias; see --alliterate")
                .takes_value(false),
        )
        .get_matches();

    // Unwrapping is safe because these options have defaults.
    let opt_separator = matches.value_of("separator").unwrap();
    let opt_words = matches.value_of("words").unwrap();
    let opt_complexity = matches.value_of("complexity").unwrap();
    let opt_count = matches.value_of("count").unwrap();
    let opt_letters = matches.value_of("letters").unwrap();
    let opt_alliterate = matches.is_present("alliterate") || matches.is_present("ubuntu");

    // Parse numbers. Validated so unwrapping is okay.
    let opt_words: u8 = opt_words.parse().unwrap();
    let opt_count: usize = opt_count.parse().unwrap();
    let opt_letters: usize = opt_letters.parse().unwrap();

    // Select the appropriate word list.
    let mut petnames = match opt_complexity {
        "0" => Petnames::small(),
        "1" => Petnames::medium(),
        "2" => Petnames::large(),
        _ => Petnames::small(),
    };

    // If requested, limit the number of letters.
    if opt_letters != 0 {
        petnames.retain(|s| s.len() <= opt_letters);
    }

    // We're going to need a source of randomness.
    let mut rng = rand::thread_rng();

    // If requested, choose a random letter then discard all words that do not
    // begin with that letter.
    if opt_alliterate {
        // We choose the first letter from the intersection of the first letters
        // of each word list in `petnames`.
        let firsts =
            common_first_letters(&petnames.adjectives, &[&petnames.adverbs, &petnames.names]);
        // Choose the first letter at random; fails if there are no letters.
        match firsts.iter().choose(&mut rng) {
            Some(c) => petnames.retain(|s| s.chars().next() == Some(*c)),
            None => panic!("no letters in common"), // TODO: do this without a panic.
        };
    }

    // Stream if count is 0.
    if opt_count == 0 {
        loop {
            let petname = petnames.generate(&mut rng, opt_words, opt_separator);
            println!("{}", petname);
        }
    } else {
        for _ in 1..=opt_count {
            let petname = petnames.generate(&mut rng, opt_words, opt_separator);
            println!("{}", petname);
        }
    }
}

fn can_be_parsed<INTO>(value: String) -> Result<(), String>
where
    INTO: FromStr,
    <INTO as FromStr>::Err: std::fmt::Display,
{
    match value.parse::<INTO>() {
        Err(e) => Err(format!("{}", e)),
        Ok(_) => Ok(()),
    }
}

fn common_first_letters(init: &[&str], more: &[&[&str]]) -> HashSet<char> {
    let mut firsts = first_letters(init);
    let firsts_other: Vec<HashSet<char>> = more.iter().map(|list| first_letters(list)).collect();
    firsts.retain(|c| firsts_other.iter().all(|fs| fs.contains(c)));
    firsts
}

fn first_letters(names: &[&str]) -> HashSet<char> {
    names.iter().filter_map(|s| s.chars().next()).collect()
}
<commit_msg>Working implementation of the -d / --dir option.<commit_after>#[macro_use]
extern crate clap;

mod lib;
use lib::Petnames;

use std::collections::HashSet;
use std::fs;
use std::io;
use std::path;
use std::str::FromStr;

use clap::{App, Arg};
use rand::seq::IteratorRandom;

fn main() {
    let matches = App::new("rust-petname")
        .version(crate_version!())
        .author(crate_authors!())
        .about("Generate human readable random names.")
        .after_help(concat!(
            "Based on Dustin Kirkland's petname project ",
            "<https://github.com/dustinkirkland/petname>."
        ))
        .arg(
            Arg::with_name("words")
                .short("w")
                .long("words")
                .value_name("WORDS")
                .default_value("2")
                .help("Number of words in name")
                .takes_value(true)
                .validator(can_be_parsed::<u8>),
        )
        .arg(
            Arg::with_name("separator")
                .short("s")
                .long("separator")
                .value_name("SEP")
                .default_value("-")
                .help("Separator between words")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("complexity")
                .short("c")
                .long("complexity")
                .value_name("COM")
                .possible_values(&["0", "1", "2"])
                .hide_possible_values(true)
                .default_value("0")
                .help("Use small words (0), medium words (1), or large words (2)")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("directory")
                .short("d")
                .long("dir")
                .value_name("DIR")
                .help("Directory containing adjectives.txt, adverbs.txt, names.txt")
                .conflicts_with("complexity")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("count")
                .long("count")
                .value_name("COUNT")
                .default_value("1")
                .help("Generate multiple names; pass 0 to produce infinite names!")
                .takes_value(true)
                .validator(can_be_parsed::<usize>),
        )
        .arg(
            Arg::with_name("letters")
                .short("l")
                .long("letters")
                .value_name("LETTERS")
                .default_value("0")
                .help("Maxiumum number of letters in each word; 0 for unlimited")
                .takes_value(true)
                .validator(can_be_parsed::<usize>),
        )
        .arg(
            Arg::with_name("alliterate")
                .short("a")
                .long("alliterate")
                .help("Generate names where each word begins with the same letter")
                .takes_value(false),
        )
        .arg(
            // For compatibility with upstream.
            Arg::with_name("ubuntu")
                .short("u")
                .long("ubuntu")
                .help("Alias; see --alliterate")
                .takes_value(false),
        )
        .get_matches();

    // Unwrapping is safe because these options have defaults.
    let opt_separator = matches.value_of("separator").unwrap();
    let opt_words = matches.value_of("words").unwrap();
    let opt_complexity = matches.value_of("complexity").unwrap();
    let opt_count = matches.value_of("count").unwrap();
    let opt_letters = matches.value_of("letters").unwrap();
    let opt_alliterate = matches.is_present("alliterate") || matches.is_present("ubuntu");

    // Optional arguments without defaults.
    let opt_directory = matches.value_of("directory");

    // Parse numbers. Validated so unwrapping is okay.
    let opt_words: u8 = opt_words.parse().unwrap();
    let opt_count: usize = opt_count.parse().unwrap();
    let opt_letters: usize = opt_letters.parse().unwrap();

    // Load custom word lists, if specified.
    let (adjectives, adverbs, names) = match opt_directory {
        Some(dirname) => load_word_lists(dirname).unwrap(), // TODO: not with unwrap.
        None => ("".to_string(), "".to_string(), "".to_string()),
    };

    // Select the appropriate word list.
    let mut petnames = match (opt_directory, opt_complexity) {
        (None, "0") => Petnames::small(),
        (None, "1") => Petnames::medium(),
        (None, "2") => Petnames::large(),
        (None, _) => Petnames::small(),
        (Some(_), _) => Petnames {
            adjectives: adjectives.split_whitespace().collect(),
            adverbs: adverbs.split_whitespace().collect(),
            names: names.split_whitespace().collect(),
        },
    };

    // If requested, limit the number of letters.
    if opt_letters != 0 {
        petnames.retain(|s| s.len() <= opt_letters);
    }

    // We're going to need a source of randomness.
    let mut rng = rand::thread_rng();

    // If requested, choose a random letter then discard all words that do not
    // begin with that letter.
    if opt_alliterate {
        // We choose the first letter from the intersection of the first letters
        // of each word list in `petnames`.
        let firsts =
            common_first_letters(&petnames.adjectives, &[&petnames.adverbs, &petnames.names]);
        // Choose the first letter at random; fails if there are no letters.
        match firsts.iter().choose(&mut rng) {
            Some(c) => petnames.retain(|s| s.chars().next() == Some(*c)),
            None => panic!("no letters in common"), // TODO: do this without a panic.
        };
    }

    // Stream if count is 0.
    if opt_count == 0 {
        loop {
            let petname = petnames.generate(&mut rng, opt_words, opt_separator);
            println!("{}", petname);
        }
    } else {
        for _ in 1..=opt_count {
            let petname = petnames.generate(&mut rng, opt_words, opt_separator);
            println!("{}", petname);
        }
    }
}

fn can_be_parsed<INTO>(value: String) -> Result<(), String>
where
    INTO: FromStr,
    <INTO as FromStr>::Err: std::fmt::Display,
{
    match value.parse::<INTO>() {
        Err(e) => Err(format!("{}", e)),
        Ok(_) => Ok(()),
    }
}

fn common_first_letters(init: &[&str], more: &[&[&str]]) -> HashSet<char> {
    let mut firsts = first_letters(init);
    let firsts_other: Vec<HashSet<char>> = more.iter().map(|list| first_letters(list)).collect();
    firsts.retain(|c| firsts_other.iter().all(|fs| fs.contains(c)));
    firsts
}

fn first_letters(names: &[&str]) -> HashSet<char> {
    names.iter().filter_map(|s| s.chars().next()).collect()
}

// Load word lists from the given directory. This function expects to find three
// files in that directory: `adjectives.txt`, `adverbs.txt`, and `names.txt`.
// Each should be valid UTF-8, and contain words separated by whitespace.
fn load_word_lists<T: AsRef<path::Path>>(dirname: T) -> io::Result<(String, String, String)> {
    let dirname = dirname.as_ref();
    Ok((
        fs::read_to_string(dirname.join("adjectives.txt"))?,
        fs::read_to_string(dirname.join("adverbs.txt"))?,
        fs::read_to_string(dirname.join("names.txt"))?,
    ))
}
<|endoftext|>


<commit_before>extern crate libc;
extern crate termios;

use libc::{STDIN_FILENO, STDOUT_FILENO, TIOCGWINSZ, winsize};
use termios::{Termios, tcsetattr, TCSAFLUSH, VMIN, VTIME};
use termios::{BRKINT, ICRNL, INPCK, ISTRIP, IXON};
use termios::{OPOST, CS8};
use termios::{ECHO, ICANON, IEXTEN, ISIG};

use std::env;
use std::fs::File;
use std::io::{self, Stdin, Stdout, Read, BufRead, BufReader, Error, ErrorKind, Write};

const VERSION: &'static str = env!("CARGO_PKG_VERSION");

#[derive(PartialEq, Clone, Copy)]
enum Key {
    Character(u8),
    ArrowLeft,
    ArrowRight,
    ArrowUp,
    ArrowDown,
    DeleteKey,
    HomeKey,
    EndKey,
    PageUp,
    PageDown,
}

macro_rules! ctrl_key {
    ($k:expr) => ($k & 0x1f);
}

struct RawMode {
    orig_term: Termios
}

struct Row(String);

struct Editor {
    _mode: RawMode,
    cx: u16,
    cy: u16,
    screenrows: u16,
    screencols: u16,
    rows: Option<Vec<Row>>,
    stdin: Stdin,
    stdout: Stdout,
}

impl RawMode {
    fn enable_raw_mode() -> io::Result<RawMode> {
        let mut term = Termios::from_fd(STDIN_FILENO)?;
        let mode = RawMode { orig_term: term.clone() };

        term.c_iflag &= !(BRKINT | ICRNL | INPCK | ISTRIP | IXON);
        term.c_oflag &= !OPOST;
        term.c_cflag |= CS8;
        term.c_lflag &= !(ECHO | ICANON | IEXTEN | ISIG);
        term.c_cc[VMIN] = 0;
        term.c_cc[VTIME] = 1;

        tcsetattr(STDIN_FILENO, TCSAFLUSH, &term)?;
        Ok(mode)
    }
}

impl Drop for RawMode {
    fn drop(&mut self) {
        tcsetattr(STDIN_FILENO, TCSAFLUSH, &self.orig_term).unwrap()
    }
}

fn read_non_blocking<R : Read>(r : &mut R, buf: &mut [u8]) -> usize {
    r.read(buf)
     .or_else(|e| if e.kind() == ErrorKind::WouldBlock { Ok(0) } else { Err(e) })
     .expect("read_non_blocking")
}

fn truncate_bytes(s: &str, max_len: u16) -> &[u8] {
    if s.len() > max_len as usize {
        s[..(max_len as usize)].as_bytes()
    } else {
        s.as_bytes()
    }
}

fn editor_read_key(stdin: &mut Stdin) -> Key {
    let mut buf = [0; 1];
    loop {
        let n = read_non_blocking(stdin, &mut buf);
        if n == 1 {
            if buf[0] == b'\x1b' {
                let mut seq = [0; 2];
                let n = read_non_blocking(stdin, &mut seq);
                if n == 2 && seq[0] == b'[' {
                    if seq[1] >= b'0' && seq[1] <= b'9' {
                        let mut last = [0; 1];
                        let n = read_non_blocking(stdin, &mut last);
                        if n == 1 {
                            if last[0] == b'~' {
                                return match seq[1] {
                                    b'1' => Key::HomeKey,
                                    b'3' => Key::DeleteKey,
                                    b'4' => Key::EndKey,
                                    b'5' => Key::PageUp,
                                    b'6' => Key::PageDown,
                                    b'7' => Key::HomeKey,
                                    b'8' => Key::EndKey,
                                    _ => Key::Character(b'\x1b'),
                                }
                            }
                        }
                    } else {
                        return match seq[1] {
                            b'A' => Key::ArrowUp,
                            b'B' => Key::ArrowDown,
                            b'C' => Key::ArrowRight,
                            b'D' => Key::ArrowLeft,
                            b'H' => Key::HomeKey,
                            b'F' => Key::EndKey,
                            _ => Key::Character(b'\x1b'),
                        }
                    }
                } else if n == 2 && seq[0] == b'O' {
                    return match seq[1] {
                        b'H' => Key::HomeKey,
                        b'F' => Key::EndKey,
                        _ => Key::Character(b'\x1b'),
                    }
                }
                return Key::Character(b'\x1b');
            } else {
                return Key::Character(buf[0]);
            }
        }
    }
}

fn get_window_size() -> io::Result<(u16, u16)> {
    let ws = winsize { ws_col: 0, ws_row: 0, ws_xpixel: 0, ws_ypixel: 0 };
    unsafe {
        if libc::ioctl(STDOUT_FILENO, TIOCGWINSZ, &ws) == -1 || ws.ws_col == 0 {
            return Err(Error::new(ErrorKind::Other, "get_window_size: ioctl failed"))
        }
    }
    Ok((ws.ws_row, ws.ws_col))
}

impl Editor {
    fn new() -> io::Result<Editor> {
        let mode = RawMode::enable_raw_mode()?;
        let (rows, cols) = get_window_size()?;
        let stdin = io::stdin();
        let stdout = io::stdout();
        Ok(Editor {
            _mode: mode,
            cx: 0,
            cy: 0,
            screenrows: rows,
            screencols: cols,
            rows: None,
            stdin: stdin,
            stdout: stdout
        })
    }

    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.stdout.write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.stdout.flush()
    }

    fn draw_rows(&mut self) -> io::Result<()> {
        let numrows = self.rows.as_ref().map_or(0, |v| v.len());
        for y in 0..(self.screenrows) {
            if y as usize >= numrows {
                if self.rows.is_none() && y == self.screenrows / 3 {
                    let mut msg = format!("Kilo-rs editor -- version {}", VERSION);
                    msg.truncate(self.screencols as usize);
                    let padding = (self.screencols - msg.len() as u16) / 2;
                    if padding > 0 {
                        self.write(b"~")?;
                        for _ in 1..padding {
                            self.write(b" ")?;
                        }
                    }
                    self.write(msg.as_bytes())?;
                } else {
                    self.write(b"~")?;
                }
            } else {
                let rows = self.rows.as_ref().unwrap();
                self.stdout.write(truncate_bytes(&rows[y as usize].0, self.screencols))?;
            }

            self.write(b"\x1b[K")?;
            if y < self.screenrows - 1 {
                self.write(b"\r\n")?;
            }
        }
        Ok(())
    }

    fn refresh_screen(&mut self) -> io::Result<()> {
        self.write(b"\x1b[?25l")?;
        self.write(b"\x1b[H")?;

        self.draw_rows()?;

        let move_cursor = format!("\x1b[{};{}H", self.cy + 1, self.cx + 1).into_bytes();
        self.write(&move_cursor)?;

        self.write(b"\x1b[?25h")?;
        self.flush()
    }

    fn move_cursor(&mut self, k: Key) {
        match k {
            Key::ArrowUp => {
                if self.cy > 0 {
                    self.cy -= 1;
                }
            },
            Key::ArrowDown => {
                if self.cy < self.screenrows - 1 {
                    self.cy += 1;
                }
            },
            Key::ArrowLeft => {
                if self.cx > 0 {
                    self.cx -= 1;
                }
            },
            Key::ArrowRight => {
                if self.cx < self.screencols - 1 {
                    self.cx += 1;
                }
            },
            _ => (),
        }
    }

    fn process_keypress(&mut self) -> bool {
        let c = editor_read_key(&mut self.stdin);

        match c {
            Key::Character(k) if k == ctrl_key!(b'q') => {
                return false
            },
            Key::HomeKey => {
                self.cx = 0;
            },
            Key::EndKey => {
                self.cx = self.screencols - 1;
            },
            Key::PageUp | Key::PageDown => {
                let key = if c == Key::PageUp { Key::ArrowUp } else { Key::ArrowDown };
                for _ in 0..(self.screenrows) {
                    self.move_cursor(key);
                }
            },
            Key::ArrowUp | Key::ArrowDown | Key::ArrowLeft | Key::ArrowRight => {
                self.move_cursor(c);
            },
            _ => ()
        };
        true
    }

    fn open(&mut self, filename: &str) -> io::Result<()> {
        let f = File::open(filename)?;
        let file = BufReader::new(&f);
        let results: io::Result<Vec<Row>> = file.lines()
                        .map(|r| r.map(|l| Row(l)))
                        .collect();
        self.rows = Some(results?);
        Ok(())
    }
}

impl Drop for Editor {
    fn drop(&mut self) {
        // clear screen
        self.stdout.write(b"\x1b[2J").unwrap();
        self.stdout.write(b"\x1b[H").unwrap();
        self.stdout.flush().unwrap();
    }
}

fn main() {
    let mut editor = Editor::new().unwrap();
    if let Some(filename) = env::args().nth(1) {
        editor.open(&filename).unwrap();
    }

    editor.refresh_screen().unwrap();
    while editor.process_keypress() {
        editor.refresh_screen().unwrap();
    }
}
<commit_msg>Vertical scrolling<commit_after>extern crate libc;
extern crate termios;

use libc::{STDIN_FILENO, STDOUT_FILENO, TIOCGWINSZ, winsize};
use termios::{Termios, tcsetattr, TCSAFLUSH, VMIN, VTIME};
use termios::{BRKINT, ICRNL, INPCK, ISTRIP, IXON};
use termios::{OPOST, CS8};
use termios::{ECHO, ICANON, IEXTEN, ISIG};

use std::env;
use std::fs::File;
use std::io::{self, Stdin, Stdout, Read, BufRead, BufReader, Error, ErrorKind, Write};

const VERSION: &'static str = env!("CARGO_PKG_VERSION");

#[derive(PartialEq, Clone, Copy)]
enum Key {
    Character(u8),
    ArrowLeft,
    ArrowRight,
    ArrowUp,
    ArrowDown,
    DeleteKey,
    HomeKey,
    EndKey,
    PageUp,
    PageDown,
}

macro_rules! ctrl_key {
    ($k:expr) => ($k & 0x1f);
}

struct RawMode {
    orig_term: Termios
}

struct Row(String);

struct Editor {
    _mode: RawMode,
    cx: usize,
    cy: usize,
    rowoff: usize,
    screenrows: usize,
    screencols: usize,
    rows: Option<Vec<Row>>,
    stdin: Stdin,
    stdout: Stdout,
}

impl RawMode {
    fn enable_raw_mode() -> io::Result<RawMode> {
        let mut term = Termios::from_fd(STDIN_FILENO)?;
        let mode = RawMode { orig_term: term.clone() };

        term.c_iflag &= !(BRKINT | ICRNL | INPCK | ISTRIP | IXON);
        term.c_oflag &= !OPOST;
        term.c_cflag |= CS8;
        term.c_lflag &= !(ECHO | ICANON | IEXTEN | ISIG);
        term.c_cc[VMIN] = 0;
        term.c_cc[VTIME] = 1;

        tcsetattr(STDIN_FILENO, TCSAFLUSH, &term)?;
        Ok(mode)
    }
}

impl Drop for RawMode {
    fn drop(&mut self) {
        tcsetattr(STDIN_FILENO, TCSAFLUSH, &self.orig_term).unwrap()
    }
}

fn read_non_blocking<R : Read>(r : &mut R, buf: &mut [u8]) -> usize {
    r.read(buf)
     .or_else(|e| if e.kind() == ErrorKind::WouldBlock { Ok(0) } else { Err(e) })
     .expect("read_non_blocking")
}

fn truncate_bytes(s: &str, max_len: usize) -> &[u8] {
    if s.len() > max_len {
        s[..max_len].as_bytes()
    } else {
        s.as_bytes()
    }
}

fn editor_read_key(stdin: &mut Stdin) -> Key {
    let mut buf = [0; 1];
    loop {
        let n = read_non_blocking(stdin, &mut buf);
        if n == 1 {
            if buf[0] == b'\x1b' {
                let mut seq = [0; 2];
                let n = read_non_blocking(stdin, &mut seq);
                if n == 2 && seq[0] == b'[' {
                    if seq[1] >= b'0' && seq[1] <= b'9' {
                        let mut last = [0; 1];
                        let n = read_non_blocking(stdin, &mut last);
                        if n == 1 {
                            if last[0] == b'~' {
                                return match seq[1] {
                                    b'1' => Key::HomeKey,
                                    b'3' => Key::DeleteKey,
                                    b'4' => Key::EndKey,
                                    b'5' => Key::PageUp,
                                    b'6' => Key::PageDown,
                                    b'7' => Key::HomeKey,
                                    b'8' => Key::EndKey,
                                    _ => Key::Character(b'\x1b'),
                                }
                            }
                        }
                    } else {
                        return match seq[1] {
                            b'A' => Key::ArrowUp,
                            b'B' => Key::ArrowDown,
                            b'C' => Key::ArrowRight,
                            b'D' => Key::ArrowLeft,
                            b'H' => Key::HomeKey,
                            b'F' => Key::EndKey,
                            _ => Key::Character(b'\x1b'),
                        }
                    }
                } else if n == 2 && seq[0] == b'O' {
                    return match seq[1] {
                        b'H' => Key::HomeKey,
                        b'F' => Key::EndKey,
                        _ => Key::Character(b'\x1b'),
                    }
                }
                return Key::Character(b'\x1b');
            } else {
                return Key::Character(buf[0]);
            }
        }
    }
}

fn get_window_size() -> io::Result<(u16, u16)> {
    let ws = winsize { ws_col: 0, ws_row: 0, ws_xpixel: 0, ws_ypixel: 0 };
    unsafe {
        if libc::ioctl(STDOUT_FILENO, TIOCGWINSZ, &ws) == -1 || ws.ws_col == 0 {
            return Err(Error::new(ErrorKind::Other, "get_window_size: ioctl failed"))
        }
    }
    Ok((ws.ws_row, ws.ws_col))
}

impl Editor {
    fn new() -> io::Result<Editor> {
        let mode = RawMode::enable_raw_mode()?;
        let (rows, cols) = get_window_size()?;
        let stdin = io::stdin();
        let stdout = io::stdout();
        Ok(Editor {
            _mode: mode,
            cx: 0,
            cy: 0,
            rowoff: 0,
            screenrows: rows as usize,
            screencols: cols as usize,
            rows: None,
            stdin: stdin,
            stdout: stdout
        })
    }

    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.stdout.write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.stdout.flush()
    }

    fn scroll(&mut self) {
        if self.cy < self.rowoff {
            self.rowoff = self.cy;
        }
        if self.cy >= self.rowoff + self.screenrows {
            self.rowoff = self.cy - self.screenrows + 1;
        }
    }

    fn numrows(&self) -> usize {
        self.rows.as_ref().map_or(0, |v| v.len())
    }

    fn draw_rows(&mut self) -> io::Result<()> {
        let numrows = self.numrows();
        for y in 0..(self.screenrows) {
            let filerow = y + self.rowoff;
            if filerow >= numrows {
                if self.rows.is_none() && y == self.screenrows / 3 {
                    let mut msg = format!("Kilo-rs editor -- version {}", VERSION);
                    msg.truncate(self.screencols);
                    let padding = (self.screencols - msg.len()) / 2;
                    if padding > 0 {
                        self.write(b"~")?;
                        for _ in 1..padding {
                            self.write(b" ")?;
                        }
                    }
                    self.write(msg.as_bytes())?;
                } else {
                    self.write(b"~")?;
                }
            } else {
                let rows = self.rows.as_ref().unwrap();
                self.stdout.write(truncate_bytes(&rows[filerow].0, self.screencols))?;
            }

            self.write(b"\x1b[K")?;
            if y < self.screenrows - 1 {
                self.write(b"\r\n")?;
            }
        }
        Ok(())
    }

    fn refresh_screen(&mut self) -> io::Result<()> {
        self.scroll();

        self.write(b"\x1b[?25l")?;
        self.write(b"\x1b[H")?;

        self.draw_rows()?;

        let move_cursor = format!("\x1b[{};{}H", (self.cy - self.rowoff) + 1, self.cx + 1).into_bytes();
        self.write(&move_cursor)?;

        self.write(b"\x1b[?25h")?;
        self.flush()
    }

    fn move_cursor(&mut self, k: Key) {
        match k {
            Key::ArrowUp => {
                if self.cy > 0 {
                    self.cy -= 1;
                }
            },
            Key::ArrowDown => {
                if self.cy < self.numrows() {
                    self.cy += 1;
                }
            },
            Key::ArrowLeft => {
                if self.cx > 0 {
                    self.cx -= 1;
                }
            },
            Key::ArrowRight => {
                if self.cx < self.screencols - 1 {
                    self.cx += 1;
                }
            },
            _ => (),
        }
    }

    fn process_keypress(&mut self) -> bool {
        let c = editor_read_key(&mut self.stdin);

        match c {
            Key::Character(k) if k == ctrl_key!(b'q') => {
                return false
            },
            Key::HomeKey => {
                self.cx = 0;
            },
            Key::EndKey => {
                self.cx = self.screencols - 1;
            },
            Key::PageUp | Key::PageDown => {
                let key = if c == Key::PageUp { Key::ArrowUp } else { Key::ArrowDown };
                for _ in 0..(self.screenrows) {
                    self.move_cursor(key);
                }
            },
            Key::ArrowUp | Key::ArrowDown | Key::ArrowLeft | Key::ArrowRight => {
                self.move_cursor(c);
            },
            _ => ()
        };
        true
    }

    fn open(&mut self, filename: &str) -> io::Result<()> {
        let f = File::open(filename)?;
        let file = BufReader::new(&f);
        let results: io::Result<Vec<Row>> = file.lines()
                        .map(|r| r.map(|l| Row(l)))
                        .collect();
        self.rows = Some(results?);
        Ok(())
    }
}

impl Drop for Editor {
    fn drop(&mut self) {
        // clear screen
        self.stdout.write(b"\x1b[2J").unwrap();
        self.stdout.write(b"\x1b[H").unwrap();
        self.stdout.flush().unwrap();
    }
}

fn main() {
    let mut editor = Editor::new().unwrap();
    if let Some(filename) = env::args().nth(1) {
        editor.open(&filename).unwrap();
    }

    editor.refresh_screen().unwrap();
    while editor.process_keypress() {
        editor.refresh_screen().unwrap();
    }
}
<|endoftext|>


<commit_before>extern crate integrate;

use integrate::cquad;
use std::num::FloatMath;
use std::f64::consts;

fn main() {
    // Integrate sin(x)/2 from 0 to pi
    let (result, err, neval) = cquad(|x| 0.5f64 * FloatMath::sin(x), 0f64, consts::PI);
    println!("result: {}, err {}, neval {}", result, err, neval);
}

<commit_msg>Increase output precision in example<commit_after>extern crate integrate;

use integrate::cquad;
use std::num::FloatMath;
use std::f64::consts;

fn main() {
    // Integrate sin(x)/2 from 0 to pi
    let (result, err, neval) = cquad(|x| 0.5f64 * FloatMath::sin(x), 0f64, consts::PI);
    println!("result: {:.10e}, err {:.10e}, neval {}", result, err, neval);
}

<|endoftext|>


<commit_before>extern crate ansi_term;
extern crate getopts;
extern crate isatty;
extern crate regex;
extern crate walkdir;

use std::collections::HashMap;
use std::env;
use std::error::Error;
use std::ffi::OsStr;
use std::fs::File;
use std::io::{Write, BufReader,BufRead};
use std::path::{Path, Component};
use std::process;

use ansi_term::{Style, Colour};
use getopts::Options;
use isatty::stdout_isatty;
use regex::{Regex, RegexBuilder};
use walkdir::{WalkDir, DirEntry, WalkDirIterator};

/// Maps file extensions to ANSI colors / styles.
type ExtensionStyles = HashMap<String, ansi_term::Style>;

/// Configuration options for *fd*.
struct FdOptions {
    case_sensitive: bool,
    search_full_path: bool,
    search_hidden: bool,
    follow_links: bool,
    colored: bool,
    max_depth: usize,
    extension_styles: Option<ExtensionStyles>
}

/// The default maximum recursion depth.
const MAX_DEPTH_DEFAULT : usize = 25;

/// Print a search result to the console.
fn print_entry(path_root: &Path, path_entry: &Path, config: &FdOptions) {
    let path_full = path_root.join(path_entry);

    let path_str = match path_entry.to_str() {
        Some(p) => p,
        None    => return
    };

    if !config.colored {
        println!("{}", path_str);
    } else {
        let mut component_path = path_root.to_path_buf();

        // Traverse the path and colorize each component
        for component in path_entry.components() {
            let comp_str = match component {
                Component::Normal(p) => p,
                _                    => error("Unexpected path component")
            };

            component_path.push(Path::new(comp_str));

            let style =
                if component_path.symlink_metadata()
                                 .map(|md| md.file_type().is_symlink())
                                 .unwrap_or(false) {
                    Colour::Cyan.normal()
                } else if component_path.is_dir() {
                    Colour::Blue.bold()
                } else {
                    // Loop up file extension
                    if let Some(ref ext_styles) = config.extension_styles {
                        component_path.extension()
                                      .and_then(|e| e.to_str())
                                      .and_then(|e| ext_styles.get(e))
                                      .map(|r| r.clone())
                                      .unwrap_or(Style::new())
                    }
                    else {
                        Style::new()
                    }
                };

            print!("{}", style.paint(comp_str.to_str().unwrap()));

            if component_path.is_dir() && component_path != path_full {
                let sep = std::path::MAIN_SEPARATOR.to_string();
                print!("{}", style.paint(sep));
            }
        }
        println!();
    }
}

/// Check if filename of entry starts with a dot.
fn is_hidden(entry: &DirEntry) -> bool {
    entry.file_name()
         .to_str()
         .map(|s| s.starts_with("."))
         .unwrap_or(false)
}

/// Recursively scan the given root path and search for files / pathnames
/// matching the pattern.
fn scan(root: &Path, pattern: &Regex, config: &FdOptions) {
    let walker = WalkDir::new(root)
                     .follow_links(config.follow_links)
                     .max_depth(config.max_depth)
                     .into_iter()
                     .filter_entry(|e| config.search_hidden
                                       || !is_hidden(e)
                                       || e.path() == root)
                     .filter_map(|e| e.ok())
                     .filter(|e| e.path() != root);

    for entry in walker {
        let path_rel = match entry.path().strip_prefix(root) {
            Ok(p) => p,
            Err(_) => continue
        };

        let search_str =
            if config.search_full_path {
                path_rel.to_str()
            } else {
                if !path_rel.is_file() { continue }

                path_rel.file_name()
                        .and_then(OsStr::to_str)
            };

        search_str.and_then(|s| pattern.find(s))
                  .map(|_| print_entry(&root, path_rel, &config));
    }
}

/// Print error message to stderr and exit with status `1`.
fn error(message: &str) -> ! {
    writeln!(&mut std::io::stderr(), "{}", message)
        .expect("Failed writing to stderr");
    process::exit(1);
}

/// Parse `dircolors` file.
fn parse_dircolors(path: &Path) -> std::io::Result<ExtensionStyles> {
    let file = File::open(path)?;
    let mut ext_styles = HashMap::new();

    let pattern_ansi_256 =
        Regex::new(r"^\.([A-Za-z0-9]+)\s*(?:00;)?38;5;([0-9]+)\b")
            .unwrap();

    let pattern_ansi =
        Regex::new(r"^\.([A-Za-z0-9]+)\s*(?:([0-9]+);)?([0-9][0-9])\b")
            .unwrap();

    for line in BufReader::new(file).lines() {
        let line_s = line.unwrap();
        if let Some(caps) = pattern_ansi_256.captures(line_s.as_str()) {
            if let Some(ext) = caps.get(1).map(|m| m.as_str()) {
                let fg = caps.get(2)
                             .map(|m| m.as_str())
                             .and_then(|n| u8::from_str_radix(n, 10).ok())
                             .unwrap_or(7); // white
                ext_styles.insert(String::from(ext),
                                  Colour::Fixed(fg).normal());
            }
        } else if let Some(caps) = pattern_ansi.captures(line_s.as_str()) {
            if let Some(ext) = caps.get(1).map(|m| m.as_str()) {
                let color_s = caps.get(3)
                                  .map_or("", |m| m.as_str());
                let color = match color_s {
                    "31" => Colour::Red,
                    "32" => Colour::Green,
                    "33" => Colour::Yellow,
                    "34" => Colour::Blue,
                    "35" => Colour::Purple,
                    "36" => Colour::Cyan,
                    _    => Colour::White
                };
                let style_s = caps.get(2)
                                  .map_or("", |m| m.as_str());
                let style = match style_s {
                    "1"  => color.bold(),
                    "01" => color.bold(),
                    "3"  => color.italic(),
                    "03" => color.italic(),
                    "4"  => color.underline(),
                    "04" => color.underline(),
                    _    => color.normal()
                };
                ext_styles.insert(String::from(ext), style);
            }
        }
    }
    Ok(ext_styles)
}

fn main() {
    let args: Vec<String> = env::args().collect();

    let mut opts = Options::new();
    opts.optflag("h", "help", "print this help message");
    opts.optflag("s", "sensitive",
                      "case-sensitive search (default: smart case)");
    opts.optflag("f", "filename",
                      "search filenames only (default: full path)");
    opts.optflag("H", "hidden",
                      "search hidden files/directories (default: off)");
    opts.optflag("F", "follow", "follow symlinks (default: off)");
    opts.optflag("n", "no-color", "do not colorize output (default: on)");
    opts.optopt("d", "max-depth",
                     format!("maximum search depth (default: {})",
                             MAX_DEPTH_DEFAULT).as_str(),
                     "D");

    let matches = match opts.parse(&args[1..]) {
        Ok(m)  => m,
        Err(e) => error(e.description())
    };

    if matches.opt_present("h") {
        let brief = "Usage: fd [options] [PATTERN]";
        print!("{}", opts.usage(&brief));
        process::exit(1);
    }

    let empty = String::new();
    let pattern = matches.free.get(0).unwrap_or(&empty);

    let current_dir_buf = match env::current_dir() {
        Ok(cd) => cd,
        Err(_) => error("Could not get current directory!")
    };
    let current_dir = current_dir_buf.as_path();

    let colored_output = !matches.opt_present("no-color") &&
                         stdout_isatty();

    let ext_styles =
        if colored_output {
            env::home_dir()
                     .map(|h| h.join(".dir_colors"))
                     .and_then(|path| parse_dircolors(&path).ok())
        } else {
            None
        };

    let config = FdOptions {
        // The search will be case-sensitive if the command line flag is set or
        // if the pattern has an uppercase character (smart case).
        case_sensitive:    matches.opt_present("sensitive") ||
                           pattern.chars().any(char::is_uppercase),
        search_full_path: !matches.opt_present("filename"),
        search_hidden:     matches.opt_present("hidden"),
        colored:           colored_output,
        follow_links:      matches.opt_present("follow"),
        max_depth:
            matches.opt_str("max-depth")
                   .and_then(|ds| usize::from_str_radix(&ds, 10).ok())
                   .unwrap_or(MAX_DEPTH_DEFAULT),
        extension_styles:  ext_styles
    };

    match RegexBuilder::new(pattern)
              .case_insensitive(!config.case_sensitive)
              .build() {
        Ok(re)   => scan(&current_dir, &re, &config),
        Err(err) => error(err.description())
    }
}
<commit_msg>Search only filename by default, closes #22<commit_after>extern crate ansi_term;
extern crate getopts;
extern crate isatty;
extern crate regex;
extern crate walkdir;

use std::collections::HashMap;
use std::env;
use std::error::Error;
use std::ffi::OsStr;
use std::fs::File;
use std::io::{Write, BufReader,BufRead};
use std::path::{Path, Component};
use std::process;

use ansi_term::{Style, Colour};
use getopts::Options;
use isatty::stdout_isatty;
use regex::{Regex, RegexBuilder};
use walkdir::{WalkDir, DirEntry, WalkDirIterator};

/// Maps file extensions to ANSI colors / styles.
type ExtensionStyles = HashMap<String, ansi_term::Style>;

/// Configuration options for *fd*.
struct FdOptions {
    case_sensitive: bool,
    search_full_path: bool,
    search_hidden: bool,
    follow_links: bool,
    colored: bool,
    max_depth: usize,
    extension_styles: Option<ExtensionStyles>
}

/// The default maximum recursion depth.
const MAX_DEPTH_DEFAULT : usize = 25;

/// Print a search result to the console.
fn print_entry(path_root: &Path, path_entry: &Path, config: &FdOptions) {
    let path_full = path_root.join(path_entry);

    let path_str = match path_entry.to_str() {
        Some(p) => p,
        None    => return
    };

    if !config.colored {
        println!("{}", path_str);
    } else {
        let mut component_path = path_root.to_path_buf();

        // Traverse the path and colorize each component
        for component in path_entry.components() {
            let comp_str = match component {
                Component::Normal(p) => p,
                _                    => error("Unexpected path component")
            };

            component_path.push(Path::new(comp_str));

            let style =
                if component_path.symlink_metadata()
                                 .map(|md| md.file_type().is_symlink())
                                 .unwrap_or(false) {
                    Colour::Cyan.normal()
                } else if component_path.is_dir() {
                    Colour::Blue.bold()
                } else {
                    // Loop up file extension
                    if let Some(ref ext_styles) = config.extension_styles {
                        component_path.extension()
                                      .and_then(|e| e.to_str())
                                      .and_then(|e| ext_styles.get(e))
                                      .map(|r| r.clone())
                                      .unwrap_or(Style::new())
                    }
                    else {
                        Style::new()
                    }
                };

            print!("{}", style.paint(comp_str.to_str().unwrap()));

            if component_path.is_dir() && component_path != path_full {
                let sep = std::path::MAIN_SEPARATOR.to_string();
                print!("{}", style.paint(sep));
            }
        }
        println!();
    }
}

/// Check if filename of entry starts with a dot.
fn is_hidden(entry: &DirEntry) -> bool {
    entry.file_name()
         .to_str()
         .map(|s| s.starts_with("."))
         .unwrap_or(false)
}

/// Recursively scan the given root path and search for files / pathnames
/// matching the pattern.
fn scan(root: &Path, pattern: &Regex, config: &FdOptions) {
    let walker = WalkDir::new(root)
                     .follow_links(config.follow_links)
                     .max_depth(config.max_depth)
                     .into_iter()
                     .filter_entry(|e| config.search_hidden
                                       || !is_hidden(e)
                                       || e.path() == root)
                     .filter_map(|e| e.ok())
                     .filter(|e| e.path() != root);

    for entry in walker {
        let path_rel = match entry.path().strip_prefix(root) {
            Ok(p) => p,
            Err(_) => continue
        };

        let search_str =
            if config.search_full_path {
                path_rel.to_str()
            } else {
                if !path_rel.is_file() { continue }

                path_rel.file_name()
                        .and_then(OsStr::to_str)
            };

        search_str.and_then(|s| pattern.find(s))
                  .map(|_| print_entry(&root, path_rel, &config));
    }
}

/// Print error message to stderr and exit with status `1`.
fn error(message: &str) -> ! {
    writeln!(&mut std::io::stderr(), "{}", message)
        .expect("Failed writing to stderr");
    process::exit(1);
}

/// Parse `dircolors` file.
fn parse_dircolors(path: &Path) -> std::io::Result<ExtensionStyles> {
    let file = File::open(path)?;
    let mut ext_styles = HashMap::new();

    let pattern_ansi_256 =
        Regex::new(r"^\.([A-Za-z0-9]+)\s*(?:00;)?38;5;([0-9]+)\b")
            .unwrap();

    let pattern_ansi =
        Regex::new(r"^\.([A-Za-z0-9]+)\s*(?:([0-9]+);)?([0-9][0-9])\b")
            .unwrap();

    for line in BufReader::new(file).lines() {
        let line_s = line.unwrap();
        if let Some(caps) = pattern_ansi_256.captures(line_s.as_str()) {
            if let Some(ext) = caps.get(1).map(|m| m.as_str()) {
                let fg = caps.get(2)
                             .map(|m| m.as_str())
                             .and_then(|n| u8::from_str_radix(n, 10).ok())
                             .unwrap_or(7); // white
                ext_styles.insert(String::from(ext),
                                  Colour::Fixed(fg).normal());
            }
        } else if let Some(caps) = pattern_ansi.captures(line_s.as_str()) {
            if let Some(ext) = caps.get(1).map(|m| m.as_str()) {
                let color_s = caps.get(3)
                                  .map_or("", |m| m.as_str());
                let color = match color_s {
                    "31" => Colour::Red,
                    "32" => Colour::Green,
                    "33" => Colour::Yellow,
                    "34" => Colour::Blue,
                    "35" => Colour::Purple,
                    "36" => Colour::Cyan,
                    _    => Colour::White
                };
                let style_s = caps.get(2)
                                  .map_or("", |m| m.as_str());
                let style = match style_s {
                    "1"  => color.bold(),
                    "01" => color.bold(),
                    "3"  => color.italic(),
                    "03" => color.italic(),
                    "4"  => color.underline(),
                    "04" => color.underline(),
                    _    => color.normal()
                };
                ext_styles.insert(String::from(ext), style);
            }
        }
    }
    Ok(ext_styles)
}

fn main() {
    let args: Vec<String> = env::args().collect();

    let mut opts = Options::new();
    opts.optflag("h", "help", "print this help message");
    opts.optflag("s", "sensitive",
                      "case-sensitive search (default: smart case)");
    opts.optflag("p", "full-path",
                      "search full path (default: filename only)");
    opts.optflag("H", "hidden",
                      "search hidden files/directories (default: off)");
    opts.optflag("F", "follow", "follow symlinks (default: off)");
    opts.optflag("n", "no-color", "do not colorize output (default: on)");
    opts.optopt( "d", "max-depth",
                      format!("maximum search depth (default: {})",
                             MAX_DEPTH_DEFAULT).as_str(),
                     "D");

    let matches = match opts.parse(&args[1..]) {
        Ok(m)  => m,
        Err(e) => error(e.description())
    };

    if matches.opt_present("h") {
        let brief = "Usage: fd [options] [PATTERN]";
        print!("{}", opts.usage(&brief));
        process::exit(1);
    }

    let empty = String::new();
    let pattern = matches.free.get(0).unwrap_or(&empty);

    let current_dir_buf = match env::current_dir() {
        Ok(cd) => cd,
        Err(_) => error("Could not get current directory!")
    };
    let current_dir = current_dir_buf.as_path();

    let colored_output = !matches.opt_present("no-color") &&
                         stdout_isatty();

    let ext_styles =
        if colored_output {
            env::home_dir()
                     .map(|h| h.join(".dir_colors"))
                     .and_then(|path| parse_dircolors(&path).ok())
        } else {
            None
        };

    let config = FdOptions {
        // The search will be case-sensitive if the command line flag is set or
        // if the pattern has an uppercase character (smart case).
        case_sensitive:    matches.opt_present("sensitive") ||
                           pattern.chars().any(char::is_uppercase),
        search_full_path:  matches.opt_present("full-path"),
        search_hidden:     matches.opt_present("hidden"),
        colored:           colored_output,
        follow_links:      matches.opt_present("follow"),
        max_depth:
            matches.opt_str("max-depth")
                   .and_then(|ds| usize::from_str_radix(&ds, 10).ok())
                   .unwrap_or(MAX_DEPTH_DEFAULT),
        extension_styles:  ext_styles
    };

    match RegexBuilder::new(pattern)
              .case_insensitive(!config.case_sensitive)
              .build() {
        Ok(re)   => scan(&current_dir, &re, &config),
        Err(err) => error(err.description())
    }
}
<|endoftext|>


<commit_before>extern crate ansi_term;
extern crate getopts;
extern crate isatty;
extern crate regex;
extern crate walkdir;

use std::collections::HashMap;
use std::env;
use std::error::Error;
use std::ffi::OsStr;
use std::fs::File;
use std::io::{Write, BufReader,BufRead};
use std::path::{Path, Component};
use std::process;

use ansi_term::{Style, Colour};
use getopts::Options;
use isatty::stdout_isatty;
use regex::{Regex, RegexBuilder};
use walkdir::{WalkDir, DirEntry, WalkDirIterator};

/// Maps file extensions to ANSI colors / styles.
type ExtensionStyles = HashMap<String, ansi_term::Style>;

/// Configuration options for *fd*.
struct FdOptions {
    case_sensitive: bool,
    search_full_path: bool,
    search_hidden: bool,
    follow_links: bool,
    colored: bool,
    max_depth: usize,
    extension_styles: Option<ExtensionStyles>
}

/// The default maximum recursion depth.
const MAX_DEPTH_DEFAULT : usize = 25;

/// Print a search result to the console.
fn print_entry(path_root: &Path, path_entry: &Path, config: &FdOptions) {
    let path_full = path_root.join(path_entry);

    let path_str = match path_entry.to_str() {
        Some(p) => p,
        None    => return
    };

    if !config.colored {
        println!("{}", path_str);
    } else {
        let mut component_path = path_root.to_path_buf();

        // Traverse the path and colorize each component
        for component in path_entry.components() {
            let comp_str = match component {
                Component::Normal(p) => p,
                _                    => error("Unexpected path component")
            };

            component_path.push(Path::new(comp_str));

            let style =
                if component_path.symlink_metadata()
                                 .map(|md| md.file_type().is_symlink())
                                 .unwrap_or(false) {
                    Colour::Cyan.normal()
                } else if component_path.is_dir() {
                    Colour::Blue.bold()
                } else {
                    // Loop up file extension
                    if let Some(ref ext_styles) = config.extension_styles {
                        component_path.extension()
                                      .and_then(|e| e.to_str())
                                      .and_then(|e| ext_styles.get(e))
                                      .map(|r| r.clone())
                                      .unwrap_or(Style::new())
                    }
                    else {
                        Style::new()
                    }
                };

            print!("{}", style.paint(comp_str.to_str().unwrap()));

            if component_path.is_dir() && component_path != path_full {
                let sep = std::path::MAIN_SEPARATOR.to_string();
                print!("{}", style.paint(sep));
            }
        }
        println!();
    }
}

/// Check if filename of entry starts with a dot.
fn is_hidden(entry: &DirEntry) -> bool {
    entry.file_name()
         .to_str()
         .map(|s| s.starts_with("."))
         .unwrap_or(false)
}

/// Recursively scan the given root path and search for files / pathnames
/// matching the pattern.
fn scan(root: &Path, pattern: &Regex, config: &FdOptions) {
    let walker = WalkDir::new(root)
                     .follow_links(config.follow_links)
                     .max_depth(config.max_depth)
                     .into_iter()
                     .filter_entry(|e| config.search_hidden || !is_hidden(e))
                     .filter_map(|e| e.ok())
                     .filter(|e| e.path() != root);

    for entry in walker {
        let path_rel = match entry.path().strip_prefix(root) {
            Ok(p) => p,
            Err(_) => continue
        };

        let search_str =
            if config.search_full_path {
                path_rel.to_str()
            } else {
                if !path_rel.is_file() { continue }

                path_rel.file_name()
                        .and_then(OsStr::to_str)
            };

        search_str.and_then(|s| pattern.find(s))
                  .map(|_| print_entry(&root, path_rel, &config));
    }
}

/// Print error message to stderr and exit with status `1`.
fn error(message: &str) -> ! {
    writeln!(&mut std::io::stderr(), "{}", message)
        .expect("Failed writing to stderr");
    process::exit(1);
}

/// Parse `dircolors` file.
fn parse_dircolors(path: &Path) -> std::io::Result<ExtensionStyles> {
    let file = File::open(path)?;
    let mut ext_styles = HashMap::new();

    let pattern_ansi_256 =
        Regex::new(r"^\.([A-Za-z0-9]+)\s*(?:00;)?38;5;([0-9]+)\b")
            .unwrap();

    let pattern_ansi =
        Regex::new(r"^\.([A-Za-z0-9]+)\s*(?:([0-9][0-9]);)?([0-9][0-9])\b")
            .unwrap();

    for line in BufReader::new(file).lines() {
        let line_s = line.unwrap();
        if let Some(caps) = pattern_ansi_256.captures(line_s.as_str()) {
            if let Some(ext) = caps.get(1).map(|m| m.as_str()) {
                let fg = caps.get(2)
                             .map(|m| m.as_str())
                             .and_then(|n| u8::from_str_radix(n, 10).ok())
                             .unwrap_or(7); // white
                ext_styles.insert(String::from(ext),
                                  Colour::Fixed(fg).normal());
            }
        } else if let Some(caps) = pattern_ansi.captures(line_s.as_str()) {
            if let Some(ext) = caps.get(1).map(|m| m.as_str()) {
                let color_s = caps.get(3)
                                  .map_or("", |m| m.as_str());
                let color = match color_s {
                    "31" => Colour::Red,
                    "32" => Colour::Green,
                    "33" => Colour::Yellow,
                    "34" => Colour::Blue,
                    "35" => Colour::Purple,
                    "36" => Colour::Cyan,
                    _    => Colour::White
                };
                let style_s = caps.get(2)
                                  .map_or("", |m| m.as_str());
                let style = match style_s {
                    "01" => color.bold(),
                    "03" => color.italic(),
                    "04" => color.underline(),
                    _    => color.normal()
                };
                ext_styles.insert(String::from(ext), style);
            }
        }
    }
    Ok(ext_styles)
}

fn main() {
    let args: Vec<String> = env::args().collect();

    let mut opts = Options::new();
    opts.optflag("h", "help", "print this help message");
    opts.optflag("s", "sensitive",
                      "case-sensitive search (default: smart case)");
    opts.optflag("f", "filename",
                      "search filenames only (default: full path)");
    opts.optflag("H", "hidden",
                      "search hidden files/directories (default: off)");
    opts.optflag("F", "follow", "follow symlinks (default: off)");
    opts.optflag("n", "no-color", "do not colorize output (default: on)");
    opts.optopt("d", "max-depth",
                     format!("maximum search depth (default: {})",
                             MAX_DEPTH_DEFAULT).as_str(),
                     "D");

    let matches = match opts.parse(&args[1..]) {
        Ok(m)  => m,
        Err(e) => error(e.description())
    };

    if matches.opt_present("h") {
        let brief = "Usage: fd [options] [PATTERN]";
        print!("{}", opts.usage(&brief));
        process::exit(1);
    }

    let empty = String::new();
    let pattern = matches.free.get(0).unwrap_or(&empty);

    let current_dir_buf = match env::current_dir() {
        Ok(cd) => cd,
        Err(_) => error("Could not get current directory!")
    };
    let current_dir = current_dir_buf.as_path();

    let ext_styles = env::home_dir()
                         .map(|h| h.join(".dir_colors"))
                         .and_then(|path| parse_dircolors(&path).ok());

    let config = FdOptions {
        // The search will be case-sensitive if the command line flag is set or
        // if the pattern has an uppercase character (smart case).
        case_sensitive:    matches.opt_present("sensitive") ||
                           pattern.chars().any(char::is_uppercase),
        search_full_path: !matches.opt_present("filename"),
        search_hidden:     matches.opt_present("hidden"),
        colored:          !matches.opt_present("no-color") &&
                           stdout_isatty(),
        follow_links:      matches.opt_present("follow"),
        max_depth:
            matches.opt_str("max-depth")
                   .and_then(|ds| usize::from_str_radix(&ds, 10).ok())
                   .unwrap_or(MAX_DEPTH_DEFAULT),
        extension_styles:  ext_styles
    };

    match RegexBuilder::new(pattern)
              .case_insensitive(!config.case_sensitive)
              .build() {
        Ok(re)   => scan(&current_dir, &re, &config),
        Err(err) => error(err.description())
    }
}
<commit_msg>Another small fix for the ANSI regex<commit_after>extern crate ansi_term;
extern crate getopts;
extern crate isatty;
extern crate regex;
extern crate walkdir;

use std::collections::HashMap;
use std::env;
use std::error::Error;
use std::ffi::OsStr;
use std::fs::File;
use std::io::{Write, BufReader,BufRead};
use std::path::{Path, Component};
use std::process;

use ansi_term::{Style, Colour};
use getopts::Options;
use isatty::stdout_isatty;
use regex::{Regex, RegexBuilder};
use walkdir::{WalkDir, DirEntry, WalkDirIterator};

/// Maps file extensions to ANSI colors / styles.
type ExtensionStyles = HashMap<String, ansi_term::Style>;

/// Configuration options for *fd*.
struct FdOptions {
    case_sensitive: bool,
    search_full_path: bool,
    search_hidden: bool,
    follow_links: bool,
    colored: bool,
    max_depth: usize,
    extension_styles: Option<ExtensionStyles>
}

/// The default maximum recursion depth.
const MAX_DEPTH_DEFAULT : usize = 25;

/// Print a search result to the console.
fn print_entry(path_root: &Path, path_entry: &Path, config: &FdOptions) {
    let path_full = path_root.join(path_entry);

    let path_str = match path_entry.to_str() {
        Some(p) => p,
        None    => return
    };

    if !config.colored {
        println!("{}", path_str);
    } else {
        let mut component_path = path_root.to_path_buf();

        // Traverse the path and colorize each component
        for component in path_entry.components() {
            let comp_str = match component {
                Component::Normal(p) => p,
                _                    => error("Unexpected path component")
            };

            component_path.push(Path::new(comp_str));

            let style =
                if component_path.symlink_metadata()
                                 .map(|md| md.file_type().is_symlink())
                                 .unwrap_or(false) {
                    Colour::Cyan.normal()
                } else if component_path.is_dir() {
                    Colour::Blue.bold()
                } else {
                    // Loop up file extension
                    if let Some(ref ext_styles) = config.extension_styles {
                        component_path.extension()
                                      .and_then(|e| e.to_str())
                                      .and_then(|e| ext_styles.get(e))
                                      .map(|r| r.clone())
                                      .unwrap_or(Style::new())
                    }
                    else {
                        Style::new()
                    }
                };

            print!("{}", style.paint(comp_str.to_str().unwrap()));

            if component_path.is_dir() && component_path != path_full {
                let sep = std::path::MAIN_SEPARATOR.to_string();
                print!("{}", style.paint(sep));
            }
        }
        println!();
    }
}

/// Check if filename of entry starts with a dot.
fn is_hidden(entry: &DirEntry) -> bool {
    entry.file_name()
         .to_str()
         .map(|s| s.starts_with("."))
         .unwrap_or(false)
}

/// Recursively scan the given root path and search for files / pathnames
/// matching the pattern.
fn scan(root: &Path, pattern: &Regex, config: &FdOptions) {
    let walker = WalkDir::new(root)
                     .follow_links(config.follow_links)
                     .max_depth(config.max_depth)
                     .into_iter()
                     .filter_entry(|e| config.search_hidden || !is_hidden(e))
                     .filter_map(|e| e.ok())
                     .filter(|e| e.path() != root);

    for entry in walker {
        let path_rel = match entry.path().strip_prefix(root) {
            Ok(p) => p,
            Err(_) => continue
        };

        let search_str =
            if config.search_full_path {
                path_rel.to_str()
            } else {
                if !path_rel.is_file() { continue }

                path_rel.file_name()
                        .and_then(OsStr::to_str)
            };

        search_str.and_then(|s| pattern.find(s))
                  .map(|_| print_entry(&root, path_rel, &config));
    }
}

/// Print error message to stderr and exit with status `1`.
fn error(message: &str) -> ! {
    writeln!(&mut std::io::stderr(), "{}", message)
        .expect("Failed writing to stderr");
    process::exit(1);
}

/// Parse `dircolors` file.
fn parse_dircolors(path: &Path) -> std::io::Result<ExtensionStyles> {
    let file = File::open(path)?;
    let mut ext_styles = HashMap::new();

    let pattern_ansi_256 =
        Regex::new(r"^\.([A-Za-z0-9]+)\s*(?:00;)?38;5;([0-9]+)\b")
            .unwrap();

    let pattern_ansi =
        Regex::new(r"^\.([A-Za-z0-9]+)\s*(?:([0-9]+);)?([0-9][0-9])\b")
            .unwrap();

    for line in BufReader::new(file).lines() {
        let line_s = line.unwrap();
        if let Some(caps) = pattern_ansi_256.captures(line_s.as_str()) {
            if let Some(ext) = caps.get(1).map(|m| m.as_str()) {
                let fg = caps.get(2)
                             .map(|m| m.as_str())
                             .and_then(|n| u8::from_str_radix(n, 10).ok())
                             .unwrap_or(7); // white
                ext_styles.insert(String::from(ext),
                                  Colour::Fixed(fg).normal());
            }
        } else if let Some(caps) = pattern_ansi.captures(line_s.as_str()) {
            if let Some(ext) = caps.get(1).map(|m| m.as_str()) {
                let color_s = caps.get(3)
                                  .map_or("", |m| m.as_str());
                let color = match color_s {
                    "31" => Colour::Red,
                    "32" => Colour::Green,
                    "33" => Colour::Yellow,
                    "34" => Colour::Blue,
                    "35" => Colour::Purple,
                    "36" => Colour::Cyan,
                    _    => Colour::White
                };
                let style_s = caps.get(2)
                                  .map_or("", |m| m.as_str());
                let style = match style_s {
                    "1"  => color.bold(),
                    "01" => color.bold(),
                    "3"  => color.italic(),
                    "03" => color.italic(),
                    "4"  => color.underline(),
                    "04" => color.underline(),
                    _    => color.normal()
                };
                ext_styles.insert(String::from(ext), style);
            }
        }
    }
    Ok(ext_styles)
}

fn main() {
    let args: Vec<String> = env::args().collect();

    let mut opts = Options::new();
    opts.optflag("h", "help", "print this help message");
    opts.optflag("s", "sensitive",
                      "case-sensitive search (default: smart case)");
    opts.optflag("f", "filename",
                      "search filenames only (default: full path)");
    opts.optflag("H", "hidden",
                      "search hidden files/directories (default: off)");
    opts.optflag("F", "follow", "follow symlinks (default: off)");
    opts.optflag("n", "no-color", "do not colorize output (default: on)");
    opts.optopt("d", "max-depth",
                     format!("maximum search depth (default: {})",
                             MAX_DEPTH_DEFAULT).as_str(),
                     "D");

    let matches = match opts.parse(&args[1..]) {
        Ok(m)  => m,
        Err(e) => error(e.description())
    };

    if matches.opt_present("h") {
        let brief = "Usage: fd [options] [PATTERN]";
        print!("{}", opts.usage(&brief));
        process::exit(1);
    }

    let empty = String::new();
    let pattern = matches.free.get(0).unwrap_or(&empty);

    let current_dir_buf = match env::current_dir() {
        Ok(cd) => cd,
        Err(_) => error("Could not get current directory!")
    };
    let current_dir = current_dir_buf.as_path();

    let ext_styles = env::home_dir()
                         .map(|h| h.join(".dir_colors"))
                         .and_then(|path| parse_dircolors(&path).ok());

    let config = FdOptions {
        // The search will be case-sensitive if the command line flag is set or
        // if the pattern has an uppercase character (smart case).
        case_sensitive:    matches.opt_present("sensitive") ||
                           pattern.chars().any(char::is_uppercase),
        search_full_path: !matches.opt_present("filename"),
        search_hidden:     matches.opt_present("hidden"),
        colored:          !matches.opt_present("no-color") &&
                           stdout_isatty(),
        follow_links:      matches.opt_present("follow"),
        max_depth:
            matches.opt_str("max-depth")
                   .and_then(|ds| usize::from_str_radix(&ds, 10).ok())
                   .unwrap_or(MAX_DEPTH_DEFAULT),
        extension_styles:  ext_styles
    };

    match RegexBuilder::new(pattern)
              .case_insensitive(!config.case_sensitive)
              .build() {
        Ok(re)   => scan(&current_dir, &re, &config),
        Err(err) => error(err.description())
    }
}
<|endoftext|>


<commit_before>use std::env;
use std::io;
use std::io::Write;
use std::io::Read;
use std::io::BufRead;
use std::process;
use std::thread;

extern crate mio;
use mio::{Events, Ready, Poll, PollOpt, Token};
use mio::net::TcpStream;

//use std::collections::HashMap;

fn main() {
    // arg parsing
    let args: Vec<_> = env::args().collect();
    if args.len() != 3 {
        println!("Usage: {} <ip> <port>", args[0]);
        process::exit(1);
    }
    let port: u32 = match args[2].parse() {
        Ok(p) => p,
        Err(_) => { println!("invalid port number"); process::exit(1); }
    };
    let host = &args[1];
    let ending = "\x0A";
    ///////////////////////////////////////////////////////////////

    // Pick a token that will not be used by any other socket and
    // use that one for the listener.
    const SOCKREADER: Token = Token(0);

    // The `Poll` instance
    let poll = Poll::new().unwrap();

    let addr = match format!("{}:{}", host, port).parse() {
        Ok(a) => { a },
        //FIXME allow actual host names, not just IP addresses
        Err(e) => { println!("Couldn't parse address {}:{} ({}).", host, port, e); process::exit(1); }
    };

    let mut stream = TcpStream::connect(&addr).unwrap();
    let mut write_stream = stream.try_clone().expect("cloning failed, yikes");

    // Register the listener
    poll.register(&stream,
                  SOCKREADER,
                  Ready::readable(), // andere events wären toll, z.b. disconnect
                  PollOpt::edge()).unwrap();

    // Event storage
    let mut events = Events::with_capacity(1024);

    // Read buffer, this will never actually get filled
    let mut buf = [0; 256];

    thread::spawn(move || {
        loop {
            let input = std::io::stdin();
            for line in input.lock().lines() {
                let l = line.unwrap() + ending; // this could be an invalid string.
                let bytes = l.as_bytes();
                let byte_length = bytes.len();
                let written = write_stream.write(bytes).unwrap(); // dont unwrap, could be broken pipe
                assert_eq!(byte_length, written);
            }
        }
    });

    // The main event loop
    loop {
        // Wait for events
        poll.poll(&mut events, None).expect("Can't poll. Aborting.");

        for event in &events {
            match event.token() {
                SOCKREADER => {
                    // Continue reading in a loop until `WouldBlock` is
                    // encountered.
                    loop {
                        match stream.read(&mut buf) {
                            Ok(read_len) => {
                                match io::stdout().write(&buf[0..read_len]) {
                                    Ok(len) => { assert_eq!(len, read_len) },
                                    Err(e) => { panic!("{:?}", e); }
                                 }
                                break;
                            }
                            Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
                                // Socket is not ready anymore, stop reading
                                break;
                            },
                            Err(ref e) if e.kind() == io::ErrorKind::ConnectionRefused => {
                                println!("connect to {} port {} (tcp) failed: Connection refused", host, port);
                                process::exit(1);
                            },
                            Err(ref e)if e.raw_os_error().unwrap() == 113 => {
                                println!("connect to {} port {} (tcp) failed: No route to host", host, port);
                                process::exit(1);
                            },
                            e => panic!("err={:?}, ", e), // Other, unexpected error
                        }
                    }
                },
                x => panic!("unexpected token {:?}", x),
            }
        }
    }
}
<commit_msg>clean up some comments<commit_after>use std::env;
use std::io;
use std::io::Write;
use std::io::Read;
use std::io::BufRead;
use std::process;
use std::thread;

extern crate mio;
use mio::{Events, Ready, Poll, PollOpt, Token};
use mio::net::TcpStream;

fn main() {
    // arg parsing
    let args: Vec<_> = env::args().collect();
    if args.len() != 3 {
        println!("Usage: {} <ip> <port>", args[0]);
        process::exit(1);
    }
    let port: u32 = match args[2].parse() {
        Ok(p) => p,
        Err(_) => { println!("invalid port number"); process::exit(1); }
    };
    let host = &args[1];
    let ending = "\x0A";
    ///////////////////////////////////////////////////////////////

    // Pick a token that will not be used by any other socket and
    // use that one for the listener.
    const SOCKREADER: Token = Token(0);
    let poll = Poll::new().unwrap();

    let addr = match format!("{}:{}", host, port).parse() {
        Ok(a) => { a },
        //FIXME allow actual host names, not just IP addresses
        Err(e) => { println!("Couldn't parse address {}:{} ({}).", host, port, e); process::exit(1); }
    };

    let mut stream = TcpStream::connect(&addr).unwrap();
    let mut write_stream = stream.try_clone().expect("cloning failed, yikes");

    // Register the listener
    poll.register(&stream,
                  SOCKREADER,
                  Ready::readable(), // andere events wären toll, z.b. disconnect
                  PollOpt::edge()).unwrap();

    // Event storage
    let mut events = Events::with_capacity(1024);

    // Read buffer, this will never actually get filled
    let mut buf = [0; 256];

    thread::spawn(move || {
        loop {
            let input = std::io::stdin();
            for line in input.lock().lines() {
                let l = line.unwrap() + ending; // this could be an invalid string.
                let bytes = l.as_bytes();
                let byte_length = bytes.len();
                let written = write_stream.write(bytes).unwrap(); // dont unwrap, could be broken pipe
                assert_eq!(byte_length, written);
            }
        }
    });

    // The main event loop
    loop {
        // Wait for events
        poll.poll(&mut events, None).expect("Can't poll. Aborting.");

        for event in &events {
            match event.token() {
                SOCKREADER => {
                    // Continue reading in a loop until `WouldBlock` is
                    // encountered.
                    loop {
                        match stream.read(&mut buf) {
                            Ok(read_len) => {
                                match io::stdout().write(&buf[0..read_len]) {
                                    Ok(len) => { assert_eq!(len, read_len) },
                                    Err(e) => { panic!("{:?}", e); }
                                 }
                                break;
                            }
                            Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
                                // Socket is not ready anymore, stop reading
                                break;
                            },
                            Err(ref e) if e.kind() == io::ErrorKind::ConnectionRefused => {
                                println!("connect to {} port {} (tcp) failed: Connection refused", host, port);
                                process::exit(1);
                            },
                            Err(ref e)if e.raw_os_error().unwrap() == 113 => {
                                println!("connect to {} port {} (tcp) failed: No route to host", host, port);
                                process::exit(1);
                            },
                            e => panic!("err={:?}, ", e), // Other, unexpected error
                        }
                    }
                },
                x => panic!("unexpected token {:?}", x),
            }
        }
    }
}
<|endoftext|>