<commit_before>extern crate crypto;
extern crate base64;
extern crate byteorder;
extern crate docopt;
extern crate rustc_serialize;
extern crate rpassword;
extern crate ring;
extern crate untrusted;
#[macro_use]
extern crate error_chain;

use std::process;
use std::io::prelude::*;
use std::io::BufReader;
use std::fs::{OpenOptions, File};
use std::path::Path;

use ring::rand::SystemRandom;
use ring::signature::Ed25519KeyPair;
use ring::digest;
use crypto::bcrypt_pbkdf::bcrypt_pbkdf;

use docopt::Docopt;

mod structs;
mod errors;

use structs::*;
use errors::*;

const USAGE: &'static str = "
signify-rs

Usage:
  signify -h
  signify -G [-n] [-c <comment>] -p <pubkey> -s <seckey>
  signify -S [-x <sigfile>] -s <seckey> -m <message>
  signify -V [-x <sigfile>] -p <pubkey> -m <message>

Options:
  -h --help     Show this screen.
  -c <comment>  Specify the comment to be added during key generation.
  -m <message>  When signing, the file containing the message to sign.  When verifying, the file containing the
                message to verify.  When verifying with -e, the file to create.
  -n            Do not ask for a passphrase during key generation. Otherwise, signify will prompt the user for a
                passphrase to protect the secret key.
  -p <pubkey>   Public key produced by -G, and used by -V to check a signature.
  -s <seckey>   Secret (private) key produced by -G, and used by -S to sign a message.
  -x <sigfile>  The signature file to create or verify.  The default is <message>.sig.
";

#[allow(non_snake_case)]
#[derive(Debug, RustcDecodable)]
struct Args {
    flag_G: bool,
    flag_S: bool,
    flag_V: bool,

    flag_x: Option<String>,
    flag_c: Option<String>,

    flag_p: String,
    flag_s: String,
    flag_m: String,
    flag_n: bool,
}

fn write_base64_file<P: AsRef<Path>>(file: P, comment: &str, buf: &[u8]) -> Result<()> {
    let mut f = try!(OpenOptions::new().write(true).create_new(true).open(file));

    try!(write!(f, "{}", COMMENTHDR));
    try!(write!(f, "{}\n", comment));
    let out = base64::encode(buf);
    try!(write!(f, "{}\n", out));

    Ok(())
}

fn read_base64_file<P: AsRef<Path>>(file: P) -> Result<Vec<u8>> {
    let file_display = format!("{}", file.as_ref().display());
    let f = try!(File::open(file));
    let mut reader = BufReader::new(f);

    let mut comment_line = String::new();
    let len = try!(reader.read_line(&mut comment_line));

    if len == 0 || len < COMMENTHDRLEN || !comment_line.starts_with(COMMENTHDR) {
        return Err(format!("invalid comment in {}; must start with '{}'", file_display, COMMENTHDR).into());
    }

    if &comment_line[len-1..len] != "\n" {
        return Err(format!("missing new line after comment in {}", file_display).into());
    }

    if len > COMMENTHDRLEN + COMMENTMAXLEN {
        return Err("comment too long".into());
    }

    let mut base64_line = String::new();
    let len = try!(reader.read_line(&mut base64_line));

    if len == 0 {
        return Err(format!("missing line in {}", file_display).into());
    }

    if &base64_line[len-1..len] != "\n" {
        return Err(format!("missing new line after comment in {}", file_display).into());
    }

    let base64_line = &base64_line[0..len-1];

    let data = try!(base64::decode(base64_line)
                    .chain_err(|| format!("invalid base64 encoding in {}", file_display)));

    if &data[0..2] != PKGALG {
        return Err(format!("unsupported file {}", file_display).into());
    }

    Ok(data)
}

fn verify(pubkey_path: String, msg_path: String, signature_path: Option<String>) -> Result<()> {
    // TODO: Better error message?
    let serialized_pkey = try!(read_base64_file(&pubkey_path));
    let pkey = try!(PublicKey::from_buf(&serialized_pkey));

    let signature_path = match signature_path {
        Some(path) => path,
        None => format!("{}.sig", msg_path)
    };

    // TODO: Better error message?
    let serialized_signature = try!(read_base64_file(&signature_path));
    let signature = try!(Signature::from_buf(&serialized_signature));

    let mut msgfile = try!(File::open(&msg_path).chain_err(|| read_error(&msg_path)));
    let mut msg = vec![];
    try!(msgfile.read_to_end(&mut msg).chain_err(|| read_error(&msg_path)));

    if signature.keynum != pkey.keynum {
        return Err("signature verification failed: checked against wrong key".into());
    }

    if signature.verify(&msg, &pkey) {
        println!("Signature Verified");
        Ok(())
    } else {
        Err("signature verification failed".into())
    }
}

fn sign(seckey_path: String, msg_path: String, signature_path: Option<String>) -> Result<()> {
    let serialized_skey = try!(read_base64_file(&seckey_path));
    let mut skey = try!(PrivateKey::from_buf(&serialized_skey));

    let rounds = skey.kdfrounds;
    let xorkey = try!(kdf(&skey.salt, rounds, false, SECRETBYTES));

    for (prv, xor) in skey.seckey.iter_mut().zip(xorkey.iter()) {
        *prv = *prv ^ xor;
    }
    let skey = skey;

    let mut msgfile = try!(File::open(&msg_path)
                           .chain_err(|| read_error(&msg_path)).into());
    let mut msg = vec![];
    try!(msgfile.read_to_end(&mut msg).chain_err(|| read_error(&msg_path)));

    let signature_path = match signature_path {
        Some(path) => path,
        None => format!("{}.sig", msg_path)
    };

    let sig = try!(skey.sign(&msg).chain_err(|| "Failed to sign message"));

    let mut out = vec![];
    try!(sig.write(&mut out).chain_err(|| "Can't write to internal buffer"));

    let sig_comment = "signature from signify secret key";

    write_base64_file(&signature_path, sig_comment, &out)
        .chain_err(|| "Failed to write signature file")
}

fn read_password(prompt: &str) -> Result<String> {
    let mut stdout = std::io::stdout();
    try!(stdout.write_all(prompt.as_bytes()).chain_err(|| "Write to stdout failed"));
    try!(stdout.flush().chain_err(|| "Flushing stdout failed"));

    rpassword::read_password()
        .chain_err(|| "unable to read passphrase")
}

fn kdf(salt: &[u8], rounds: u32, confirm: bool, keylen: usize) -> Result<Vec<u8>> {
    let mut result = vec![0; keylen];
    if rounds == 0 {
        return Ok(result);
    }

    let passphrase = try!(read_password("passphrase: "));

    if confirm {
        let confirm_passphrase = try!(read_password("confirm passphrase: "));

        if passphrase != confirm_passphrase {
            return Err("passwords don't match".into());
        }
    }

    bcrypt_pbkdf(passphrase.as_bytes(), salt, rounds, &mut result);
    Ok(result)
}

fn generate(pubkey_path: String, privkey_path: String, comment: Option<String>, kdfrounds: u32) -> Result<()> {
    let comment = match comment {
        Some(s) => s,
        None    => "signify".into()
    };

    let mut keynum = [0; KEYNUMLEN];
    try!(SystemRandom.fill(&mut keynum).chain_err(|| "Can't fill keynum randomly"));

    let (_, keypair_bytes) = try!(Ed25519KeyPair::generate_serializable(&SystemRandom));
    let mut skey = keypair_bytes.private_key;
    let pkey = keypair_bytes.public_key;

    let mut salt = [0; 16];
    try!(SystemRandom.fill(&mut salt).chain_err(|| "Can't fill salt randomly"));

    let xorkey = try!(kdf(&salt, kdfrounds, true, SECRETBYTES));

    for (prv, xor) in skey.iter_mut().zip(xorkey.iter()) {
        *prv = *prv ^ xor;
    }

    // signify stores the extended key as the private key,
    // that is the 32 byte of the secret key, followed by the 32 byte of the public key,
    // summing up to 64 byte.
    //
    //  *ring* separates them, so we need to stick them together again.
    let mut complete_key = [0; 64];
    complete_key[0..32].copy_from_slice(&skey[0..32]);
    complete_key[32..].copy_from_slice(&pkey);

    // Store private key
    let digest = digest::digest(&digest::SHA512, &complete_key);
    let mut checksum = [0; 8];
    checksum.copy_from_slice(&digest.as_ref()[0..8]);

    let private_key = PrivateKey {
        pkgalg: PKGALG,
        kdfalg: KDFALG,
        kdfrounds: kdfrounds,
        salt: salt,
        checksum: checksum,
        keynum: keynum,
        seckey: complete_key,
    };

    let mut out = vec![];
    try!(private_key.write(&mut out).chain_err(|| "Can't write to internal buffer"));

    let priv_comment = format!("{} secret key", comment);
    try!(write_base64_file(&privkey_path, &priv_comment, &out)
        .chain_err(|| write_error(&privkey_path))
    );

    // Store public key
    let public_key = PublicKey::with_key_and_keynum(pkey, keynum);

    let mut out = vec![];
    try!(public_key.write(&mut out).chain_err(|| "Can't write to internal buffer"));

    let pub_comment = format!("{} public key", comment);
    write_base64_file(&pubkey_path, &pub_comment, &out)
        .chain_err(|| write_error(&pubkey_path))
}

fn human(res: Result<()>) {
    match res {
        Err(e) => {
            let mut it = e.iter();
            let e = it.next().unwrap(); // We have an error, so we can definitely unwrap it.
            let next_err = it.next();
            if let Some(next_err) = next_err {
                println!("{}: {}", e, next_err);
            } else {
                println!("{}", e);
            }

            process::exit(1);
        },
        Ok(()) => {}
    }
}

fn main() {
    let args: Args = Docopt::new(USAGE)
        .and_then(|d| d.decode())
        .unwrap_or_else(|e| e.exit());

    if args.flag_V {
        human(verify(args.flag_p, args.flag_m, args.flag_x));
    } else if args.flag_G {
        let rounds = if args.flag_n { 0 } else { 42 };
        human(generate(args.flag_p, args.flag_s, args.flag_c, rounds));
    } else if args.flag_S {
        human(sign(args.flag_s, args.flag_m, args.flag_x));
    }
}
<commit_msg>Generate embedded signature and verify it<commit_after>extern crate crypto;
extern crate base64;
extern crate byteorder;
extern crate docopt;
extern crate rustc_serialize;
extern crate rpassword;
extern crate ring;
extern crate untrusted;
#[macro_use]
extern crate error_chain;

use std::process;
use std::io::prelude::*;
use std::io::BufReader;
use std::fs::{OpenOptions, File};

use ring::rand::SystemRandom;
use ring::signature::Ed25519KeyPair;
use ring::digest;
use crypto::bcrypt_pbkdf::bcrypt_pbkdf;

use docopt::Docopt;

mod structs;
mod errors;

use structs::*;
use errors::*;

const USAGE: &'static str = "
signify-rs

Usage:
  signify -h
  signify -G [-n] [-c <comment>] -p <pubkey> -s <seckey>
  signify -S [-e] [-x <sigfile>] -s <seckey> -m <message>
  signify -V [-e] [-x <sigfile>] -p <pubkey> -m <message>

Options:
  -h --help     Show this screen.
  -c <comment>  Specify the comment to be added during key generation.
  -e            When signing, embed the message after the signature. When verifying extract the message from
                the signature.
  -m <message>  When signing, the file containing the message to sign.  When verifying, the file containing the
                message to verify.  When verifying with -e, the file to create.
  -n            Do not ask for a passphrase during key generation. Otherwise, signify will prompt the user for a
                passphrase to protect the secret key.
  -p <pubkey>   Public key produced by -G, and used by -V to check a signature.
  -s <seckey>   Secret (private) key produced by -G, and used by -S to sign a message.
  -x <sigfile>  The signature file to create or verify.  The default is <message>.sig.
";

#[allow(non_snake_case)]
#[derive(Debug, RustcDecodable)]
struct Args {
    flag_G: bool,
    flag_S: bool,
    flag_V: bool,

    flag_x: Option<String>,
    flag_c: Option<String>,
    flag_e: bool,

    flag_p: String,
    flag_s: String,
    flag_m: String,
    flag_n: bool,
}

fn write_base64_file(file: &mut File, comment: &str, buf: &[u8]) -> Result<()> {
    try!(write!(file, "{}", COMMENTHDR));
    try!(write!(file, "{}\n", comment));
    let out = base64::encode(buf);
    try!(write!(file, "{}\n", out));

    Ok(())
}

fn read_base64_file<R: Read>(file_display: &str, reader: &mut BufReader<R>) -> Result<Vec<u8>> {
    let mut comment_line = String::new();
    let len = try!(reader.read_line(&mut comment_line));

    if len == 0 || len < COMMENTHDRLEN || !comment_line.starts_with(COMMENTHDR) {
        return Err(format!("invalid comment in {}; must start with '{}'", file_display, COMMENTHDR).into());
    }

    if &comment_line[len-1..len] != "\n" {
        return Err(format!("missing new line after comment in {}", file_display).into());
    }

    if len > COMMENTHDRLEN + COMMENTMAXLEN {
        return Err("comment too long".into());
    }

    let mut base64_line = String::new();
    let len = try!(reader.read_line(&mut base64_line));

    if len == 0 {
        return Err(format!("missing line in {}", file_display).into());
    }

    if &base64_line[len-1..len] != "\n" {
        return Err(format!("missing new line after comment in {}", file_display).into());
    }

    let base64_line = &base64_line[0..len-1];

    let data = try!(base64::decode(base64_line)
                    .chain_err(|| format!("invalid base64 encoding in {}", file_display)));

    if &data[0..2] != PKGALG {
        return Err(format!("unsupported file {}", file_display).into());
    }

    Ok(data)
}

fn verify(pubkey_path: String, msg_path: String, signature_path: Option<String>, embed: bool) -> Result<()> {
    // TODO: Better error message?

    let pubkey_file = try!(File::open(&pubkey_path));
    let mut pubkey = BufReader::new(pubkey_file);
    let serialized_pkey = try!(read_base64_file(&pubkey_path, &mut pubkey));
    let pkey = try!(PublicKey::from_buf(&serialized_pkey));

    let signature_path = match signature_path {
        Some(path) => path,
        None => format!("{}.sig", msg_path)
    };

    let signature_file = try!(File::open(&signature_path));
    let mut sig_data = BufReader::new(signature_file);

    // TODO: Better error message?
    let serialized_signature = try!(read_base64_file(&signature_path, &mut sig_data));
    let signature = try!(Signature::from_buf(&serialized_signature));

    let mut msg = vec![];

    if embed {
        try!(sig_data.read_to_end(&mut msg).chain_err(|| read_error(&msg_path)));
    } else {
        let mut msgfile = try!(File::open(&msg_path).chain_err(|| read_error(&msg_path)));
        try!(msgfile.read_to_end(&mut msg).chain_err(|| read_error(&msg_path)));
    }

    if signature.keynum != pkey.keynum {
        return Err("signature verification failed: checked against wrong key".into());
    }

    if signature.verify(&msg, &pkey) {
        println!("Signature Verified");
        Ok(())
    } else {
        Err("signature verification failed".into())
    }
}

fn sign(seckey_path: String, msg_path: String, signature_path: Option<String>, embed: bool) -> Result<()> {
    let seckey_file = try!(File::open(&seckey_path));
    let mut seckey = BufReader::new(seckey_file);

    let serialized_skey = try!(read_base64_file(&seckey_path, &mut seckey));
    let mut skey = try!(PrivateKey::from_buf(&serialized_skey));

    let rounds = skey.kdfrounds;
    let xorkey = try!(kdf(&skey.salt, rounds, false, SECRETBYTES));

    for (prv, xor) in skey.seckey.iter_mut().zip(xorkey.iter()) {
        *prv = *prv ^ xor;
    }
    let skey = skey;

    let mut msgfile = try!(File::open(&msg_path)
                           .chain_err(|| read_error(&msg_path)).into());
    let mut msg = vec![];
    try!(msgfile.read_to_end(&mut msg).chain_err(|| read_error(&msg_path)));

    let signature_path = match signature_path {
        Some(path) => path,
        None => format!("{}.sig", msg_path)
    };

    let sig = try!(skey.sign(&msg).chain_err(|| "Failed to sign message"));

    let mut out = vec![];
    try!(sig.write(&mut out).chain_err(|| "Can't write to internal buffer"));

    let sig_comment = "signature from signify secret key";

    let mut file = try!(OpenOptions::new().write(true).create_new(true).open(&signature_path));
    try!(write_base64_file(&mut file, sig_comment, &out)
        .chain_err(|| "Failed to write signature file"));

    if embed {
        try!(file.write(&msg));
    }

    Ok(())
}

fn read_password(prompt: &str) -> Result<String> {
    let mut stdout = std::io::stdout();
    try!(stdout.write_all(prompt.as_bytes()).chain_err(|| "Write to stdout failed"));
    try!(stdout.flush().chain_err(|| "Flushing stdout failed"));

    rpassword::read_password()
        .chain_err(|| "unable to read passphrase")
}

fn kdf(salt: &[u8], rounds: u32, confirm: bool, keylen: usize) -> Result<Vec<u8>> {
    let mut result = vec![0; keylen];
    if rounds == 0 {
        return Ok(result);
    }

    let passphrase = try!(read_password("passphrase: "));

    if confirm {
        let confirm_passphrase = try!(read_password("confirm passphrase: "));

        if passphrase != confirm_passphrase {
            return Err("passwords don't match".into());
        }
    }

    bcrypt_pbkdf(passphrase.as_bytes(), salt, rounds, &mut result);
    Ok(result)
}

fn generate(pubkey_path: String, privkey_path: String, comment: Option<String>, kdfrounds: u32) -> Result<()> {
    let comment = match comment {
        Some(s) => s,
        None    => "signify".into()
    };

    let mut keynum = [0; KEYNUMLEN];
    try!(SystemRandom.fill(&mut keynum).chain_err(|| "Can't fill keynum randomly"));

    let (_, keypair_bytes) = try!(Ed25519KeyPair::generate_serializable(&SystemRandom));
    let mut skey = keypair_bytes.private_key;
    let pkey = keypair_bytes.public_key;

    let mut salt = [0; 16];
    try!(SystemRandom.fill(&mut salt).chain_err(|| "Can't fill salt randomly"));

    let xorkey = try!(kdf(&salt, kdfrounds, true, SECRETBYTES));

    for (prv, xor) in skey.iter_mut().zip(xorkey.iter()) {
        *prv = *prv ^ xor;
    }

    // signify stores the extended key as the private key,
    // that is the 32 byte of the secret key, followed by the 32 byte of the public key,
    // summing up to 64 byte.
    //
    //  *ring* separates them, so we need to stick them together again.
    let mut complete_key = [0; 64];
    complete_key[0..32].copy_from_slice(&skey[0..32]);
    complete_key[32..].copy_from_slice(&pkey);

    // Store private key
    let digest = digest::digest(&digest::SHA512, &complete_key);
    let mut checksum = [0; 8];
    checksum.copy_from_slice(&digest.as_ref()[0..8]);

    let private_key = PrivateKey {
        pkgalg: PKGALG,
        kdfalg: KDFALG,
        kdfrounds: kdfrounds,
        salt: salt,
        checksum: checksum,
        keynum: keynum,
        seckey: complete_key,
    };

    let mut out = vec![];
    try!(private_key.write(&mut out).chain_err(|| "Can't write to internal buffer"));

    let priv_comment = format!("{} secret key", comment);
    let mut file = try!(OpenOptions::new().write(true).create_new(true).open(&privkey_path));
    try!(write_base64_file(&mut file, &priv_comment, &out)
        .chain_err(|| write_error(&privkey_path)));

    // Store public key
    let public_key = PublicKey::with_key_and_keynum(pkey, keynum);

    let mut out = vec![];
    try!(public_key.write(&mut out).chain_err(|| "Can't write to internal buffer"));

    let pub_comment = format!("{} public key", comment);
    let mut file = try!(OpenOptions::new().write(true).create_new(true).open(&pubkey_path));
    write_base64_file(&mut file, &pub_comment, &out)
        .chain_err(|| write_error(&pubkey_path))
}

fn human(res: Result<()>) {
    match res {
        Err(e) => {
            let mut it = e.iter();
            let e = it.next().unwrap(); // We have an error, so we can definitely unwrap it.
            let next_err = it.next();
            if let Some(next_err) = next_err {
                println!("{}: {}", e, next_err);
            } else {
                println!("{}", e);
            }

            process::exit(1);
        },
        Ok(()) => {}
    }
}

fn main() {
    let args: Args = Docopt::new(USAGE)
        .and_then(|d| d.decode())
        .unwrap_or_else(|e| e.exit());

    if args.flag_V {
        human(verify(args.flag_p, args.flag_m, args.flag_x, args.flag_e));
    } else if args.flag_G {
        let rounds = if args.flag_n { 0 } else { 42 };
        human(generate(args.flag_p, args.flag_s, args.flag_c, rounds));
    } else if args.flag_S {
        human(sign(args.flag_s, args.flag_m, args.flag_x, args.flag_e));
    }
}
<|endoftext|>


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

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

mod avg;
mod enums;

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

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

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

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

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

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

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

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

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

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

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

            let info = to_avg(avg);

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

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

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

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

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

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

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

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

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

    false
}

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

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

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

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

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

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

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

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

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

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

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

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

            if a.severity > avg_severity {
                avg_severity = a.severity;
            }

            if a.status > avg_status {
                avg_status = a.status;
            }
        }

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

    avgs
}

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

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

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

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

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

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

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

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

/// Print a list of `avg::AVG`
fn print_avgs(options: &Options, avgs: &BTreeMap<String, avg::AVG>) {
    for (pkg, avg) in avgs {
        match avg.fixed {
            Some(ref v) if avg.status != enums::Status::Vulnerable => {
                if options.quiet >= 2 {
                    println!("{}", pkg);
                } else if options.quiet == 1 {
                    println!("{}>={}", pkg, v);
                } else {
                    match options.format {
                        Some(ref f) => println!(
                            "{}",
                            f.replace("%n", pkg.as_str())
                                .replace("%c", avg.issues.iter().join(",").as_str(),)
                        ),
                        None => {
                            let msg = format!(
                                "Package {} is affected by {}. {}!",
                                pkg,
                                avg.issues.join(", "),
                                avg.severity
                            );

                            if avg.status == enums::Status::Testing {
                                println!("{}. Update to {} from testing repos!", msg, v)
                            } else if avg.status == enums::Status::Fixed {
                                println!("{}. Update to {}!", msg, v)
                            } else {
                                println!("{}", msg)
                            }
                        }
                    }
                }
            }
            _ => {
                if !options.upgradable_only {
                    if options.quiet > 0 {
                        println!("{}", pkg);
                    } else {
                        match options.format {
                            Some(ref f) => println!(
                                "{}",
                                f.replace("%n", pkg.as_str())
                                    .replace("%c", avg.issues.iter().join(",").as_str(),)
                            ),
                            None => {
                                println!(
                                    "Package {} is affected by {}. {}!",
                                    pkg,
                                    avg.issues.join(", "),
                                    avg.severity
                                );
                            }
                        }
                    }
                }
            }
        }
    }
}
<commit_msg>Report the failure when the server is down<commit_after>extern crate alpm;
#[macro_use]
extern crate clap;
extern crate curl;
extern crate env_logger;
extern crate itertools;
#[macro_use]
extern crate log;
extern crate serde_json;

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

mod avg;
mod enums;

const WEBSITE: &str = "https://security.archlinux.org";

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

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

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

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

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

    let mut avgs = String::new();
    {
        info!("Downloading AVGs...");
        let avgs_url = format!("{}/issues/all.json", WEBSITE);

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

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

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

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

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

            let info = to_avg(avg);

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

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

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

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

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

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

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

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

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

    false
}

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

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

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

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

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

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

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

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

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

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

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

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

            if a.severity > avg_severity {
                avg_severity = a.severity;
            }

            if a.status > avg_status {
                avg_status = a.status;
            }
        }

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

    avgs
}

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

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

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

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

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

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

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

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

/// Print a list of `avg::AVG`
fn print_avgs(options: &Options, avgs: &BTreeMap<String, avg::AVG>) {
    for (pkg, avg) in avgs {
        match avg.fixed {
            Some(ref v) if avg.status != enums::Status::Vulnerable => {
                if options.quiet >= 2 {
                    println!("{}", pkg);
                } else if options.quiet == 1 {
                    println!("{}>={}", pkg, v);
                } else {
                    match options.format {
                        Some(ref f) => println!(
                            "{}",
                            f.replace("%n", pkg.as_str())
                                .replace("%c", avg.issues.iter().join(",").as_str(),)
                        ),
                        None => {
                            let msg = format!(
                                "Package {} is affected by {}. {}!",
                                pkg,
                                avg.issues.join(", "),
                                avg.severity
                            );

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


<commit_before>#![deny(warnings)]

extern crate chan;
extern crate crypto;
extern crate walkdir;

use self::crypto::digest::Digest;
use self::crypto::sha1::Sha1;
use std::env;
use std::fs::File;
use std::io::Read;
use std::path::Path;
use std::io;
use walkdir::WalkDir;
use std::thread;
use std::sync::Arc;
use walkdir::DirEntry;

fn calc_hash(p: &Path, hasher: &mut Sha1, buf: &mut [u8]) -> io::Result<String> {
    hasher.reset();
    let mut f = File::open(p)?;
    loop {
        let num_read = f.read(buf)?;
        if num_read == 0 {
            break;
        }
        hasher.input(&buf[..num_read]);
    }
    Ok(hasher.result_str())
}

type Result<T> = std::result::Result<T, TreesumError>;

#[derive(Debug)]
enum TreesumError {
    Io(io::Error),
    Ignored(DirEntry),
    WalkDir(walkdir::Error),
}

fn process_root(root: &Path) -> io::Result<()> {
    let pb = root.to_path_buf();
    let producer_ctor = || WalkDir::new(pb);
    let xform_ctor = || {
        let mut hasher = Sha1::new();
        let mut buf = [0; 1024 * 8];
        let f = move |e: walkdir::Result<DirEntry>| -> Result<String> {
            let e = e.map_err(TreesumError::WalkDir)?;
            if !e.file_type().is_file() {
                return Err(TreesumError::Ignored(e));
            }
            calc_hash(e.path(), &mut hasher, &mut buf).map_err(TreesumError::Io)
        };
        f
    };
    let results = scatter_gather(producer_ctor, xform_ctor);
    for r in results {
        println!("{:?}", r);
    }
    Ok(())
}

fn main() {
    let root = env::args().nth(1).unwrap_or(".".to_string());
    let root = Path::new(root.as_str());
    process_root(root).expect("Hello error 2!");
}

/// The purpose of this newtype is to hide the channel implementation.
struct GatherIter<T>(chan::Iter<T>);
impl<T> Iterator for GatherIter<T> {
    type Item = T;
    fn next(&mut self) -> Option<T> {
        self.0.next()
    }
}

/// producer_ctor and xform_xtor is needed to allow construction
/// in the correct producer / worker thread.
fn scatter_gather<PC, XC, P, X, J, R>(producer_ctor: PC, xform_ctor: XC) -> GatherIter<R>
    where PC: 'static + std::marker::Send + FnOnce() -> P,
          XC: 'static + std::marker::Send + std::marker::Sync + Fn() -> X,
          X: FnMut(J) -> R,
          J: 'static + std::marker::Send,
          R: 'static + std::marker::Send,
          P: IntoIterator<Item = J>
{
    let jobs_rx = {
        let (tx, rx) = chan::sync(0);
        thread::spawn(move || {
            for e in producer_ctor().into_iter() {
                tx.send(e);
            }
        });
        rx
    };
    let results_rx = {
        let (tx, rx) = chan::sync(0);
        let xform_ctor = Arc::new(xform_ctor); // TODO: Investigate why this is needed.
        for _ in 0..8 {
            let tx = tx.clone();
            let jobs_rx = jobs_rx.clone();
            let xform_ctor = xform_ctor.clone();
            thread::spawn(move || {
                let mut xform = xform_ctor();
                for e in jobs_rx {
                    tx.send(xform(e));
                }
            });
        }
        rx
    };
    GatherIter(results_rx.iter())
}<commit_msg>Code cleanup.<commit_after>#![deny(warnings)]

extern crate chan;
extern crate crypto;
extern crate walkdir;

use self::crypto::digest::Digest;
use self::crypto::sha1::Sha1;
use std::env;
use std::fs::File;
use std::io::Read;
use std::path::Path;
use std::io;
use walkdir::WalkDir;
use std::thread;
use std::sync::Arc;
use walkdir::DirEntry;

/// The purpose of this newtype is to hide the channel implementation.
struct GatherIter<T>(chan::Iter<T>);
impl<T> Iterator for GatherIter<T> {
    type Item = T;
    fn next(&mut self) -> Option<T> {
        self.0.next()
    }
}

/// producer_ctor and xform_ctor is needed to allow construction
/// in the correct producer / worker thread.
/// TODO: Figure out how to simplify this trait mess :)
fn scatter_gather<PC, XC, P, X, J, R>(producer_ctor: PC, xform_ctor: XC) -> GatherIter<R>
    where PC: 'static + std::marker::Send + FnOnce() -> P,
          XC: 'static + std::marker::Send + std::marker::Sync + Fn() -> X,
          X: FnMut(J) -> R,
          J: 'static + std::marker::Send,
          R: 'static + std::marker::Send,
          P: IntoIterator<Item = J>
{
    let jobs_rx = {
        let (tx, rx) = chan::sync(0);
        thread::spawn(move || {
            for e in producer_ctor().into_iter() {
                tx.send(e);
            }
        });
        rx
    };
    let results_rx = {
        let (tx, rx) = chan::sync(0);
        let xform_ctor = Arc::new(xform_ctor); // TODO: Investigate why this is needed.
        for _ in 0..8 {
            let tx = tx.clone();
            let jobs_rx = jobs_rx.clone();
            let xform_ctor = xform_ctor.clone();
            thread::spawn(move || {
                let mut xform = xform_ctor();
                for e in jobs_rx {
                    tx.send(xform(e));
                }
            });
        }
        rx
    };
    GatherIter(results_rx.iter())
}

fn calc_hash(p: &Path, hasher: &mut Sha1, buf: &mut [u8]) -> io::Result<String> {
    hasher.reset();
    let mut f = File::open(p)?;
    loop {
        let num_read = f.read(buf)?;
        if num_read == 0 {
            break;
        }
        hasher.input(&buf[..num_read]);
    }
    Ok(hasher.result_str())
}

type Result<T> = std::result::Result<T, Error>;

#[derive(Debug)]
enum Error {
    Io(io::Error),
    Ignored(DirEntry),
    WalkDir(walkdir::Error),
}

fn process_root(root: &Path) -> io::Result<()> {
    let pb = root.to_path_buf();
    let producer_ctor = || WalkDir::new(pb);
    let xform_ctor = || {
        let mut hasher = Sha1::new();
        let mut buf = [0; 1024 * 8];
        let f = move |e: walkdir::Result<DirEntry>| -> Result<String> {
            let e = e.map_err(Error::WalkDir)?;
            if !e.file_type().is_file() {
                return Err(Error::Ignored(e));
            }
            calc_hash(e.path(), &mut hasher, &mut buf).map_err(Error::Io)
        };
        f
    };
    let results = scatter_gather(producer_ctor, xform_ctor);
    for r in results {
        println!("{:?}", r);
    }
    Ok(())
}

fn main() {
    let root = env::args().nth(1).unwrap_or(".".to_string());
    let root = Path::new(root.as_str());
    process_root(root).expect("Hello error 2!");
}
<|endoftext|>


<commit_before>
extern crate civet;
extern crate conduit;
extern crate conduit_router;
extern crate cql_ffi;

use std::collections::HashMap;
use std::io::Cursor;
use std::sync::mpsc::channel;

use cql_ffi::{CassCluster,CassSession,CassError,CassPrepared};
use civet::{Config, response, Server};
use conduit::{Request, Response};
use conduit_router::{RouteBuilder, RequestParams};

const CONTACT_POINTS:&'static str = "127.0.0.1";

fn main() {
	let cluster = &CassCluster::new()
                        .set_contact_points(CONTACT_POINTS).unwrap()
                        .set_load_balance_round_robin().unwrap();

    let session_future = CassSession::new().connect(cluster).wait();

    let all_prepared:Vec<CassPrepared> = Vec::new();

    match session_future {
   		Err(err) => panic!("{:?}",err),
        Ok(session) => {
			println!("session established");

		    let mut router = RouteBuilder::new();
    		router.get("/prepare/:statement", |r: &mut Request| prepare(r, &session, &all_prepared));
    		router.get("/execute/:statement_id", |r: &mut Request| execute(r, &session, &all_prepared));
    		let _server = Server::start(Config { port: 8888, threads: 1 }, router);
    		let (_tx, rx) = channel::<()>();
    		rx.recv().unwrap();
		}
    }
}

fn prepare(req: &mut Request, session: &CassSession, all_prepared:&Vec<CassPrepared>) -> Result<Response,CassError> {
	let statement = req.params().find("statement").unwrap();
	let prepared:CassPrepared = try!(try!(session.prepare(statement)).wait());
    let bytes = format!("Preparing: {}!", statement).into_bytes();
     Ok(response(200, HashMap::new(), Cursor::new(bytes)))
}

fn execute(req: &mut Request, session: &CassSession, all_prepared:&Vec<CassPrepared>) -> Result<Response,CassError> {
	let statement = req.params().find("statement").unwrap();
	let prepared:CassPrepared = try!(try!(session.prepare(statement)).wait());
	//prepared.
    let bytes = format!("Preparing: {}!", statement).into_bytes();
     Ok(response(200, HashMap::new(), Cursor::new(bytes)))
}

<commit_msg>working method of passing session into multiple closures<commit_after>
extern crate civet;
extern crate conduit;
extern crate conduit_router;
extern crate cql_ffi;

use std::collections::HashMap;
use std::io::Cursor;
use std::sync::Arc;
use std::sync::mpsc::channel;

use cql_ffi::{CassCluster,CassSession,CassError,CassPrepared};
use civet::{Config, response, Server};
use conduit::{Request, Response};
use conduit_router::{RouteBuilder, RequestParams};

const CONTACT_POINTS:&'static str = "127.0.0.1";

fn main() {
	let cluster = &CassCluster::new()
                        .set_contact_points(CONTACT_POINTS).unwrap()
                        .set_load_balance_round_robin().unwrap();

    let session_future = CassSession::new().connect(cluster).wait();

    match session_future {
   		Err(err) => panic!("{:?}",err),
        Ok(session) => {
		    let all_prepared1:Arc<Vec<CassPrepared>> = Arc::new(Vec::new());
		    let all_prepared2:Arc<Vec<CassPrepared>> = all_prepared1.clone();

			println!("session established");
			let session = Arc::new(session);
			let prepare_session = session.clone();
			let execute_session = session.clone();
		    let mut router = RouteBuilder::new();

		    router.get("/prepare/:statement", move |r: &mut Request| prepare(r, &prepare_session, &all_prepared1));
    		router.get("/execute/:statement_id", move |r: &mut Request| execute(r, &execute_session, &all_prepared2));

    		let _server = Server::start(Config { port: 8888, threads: 1 }, router);
    		let (_tx, rx) = channel::<()>();
    		rx.recv().unwrap();
		}
    }
}

fn prepare(req: &mut Request, session: &Arc<CassSession>, all_prepared: &Arc<Vec<CassPrepared>>) -> Result<Response,CassError> {
	let statement = req.params().find("statement").unwrap();
	let prepared:CassPrepared = try!(try!(session.prepare(statement)).wait());
    let bytes = format!("Preparing: {}!", statement).into_bytes();
     Ok(response(200, HashMap::new(), Cursor::new(bytes)))
}

fn execute(req: &mut Request, session: &CassSession, all_prepared: &Arc<Vec<CassPrepared>>) -> Result<Response,CassError> {
	let statement = req.params().find("statement").unwrap();
	let prepared:CassPrepared = try!(try!(session.prepare(statement)).wait());
	//prepared.
    let bytes = format!("Preparing: {}!", statement).into_bytes();
     Ok(response(200, HashMap::new(), Cursor::new(bytes)))
}

<|endoftext|>


<commit_before>#![feature(core)]
#![feature(path)]

extern crate git2;
extern crate "rustc-serialize" as rustc_serialize;
extern crate docopt;

use git2::{Repository, Branch, BranchType, DiffLine,
    Commit, Diff, DiffFormat, DiffDelta, DiffHunk, Error, Oid};
use docopt::Docopt;
use std::old_io as io;
use std::old_io::BufferedReader;
use std::old_io::File;
use std::str;

// Write the Docopt usage string.
static USAGE: &'static str = "
Usage: transit <repo> [<old> <new>]

If no commits are given, transit will revwalk from latest to oldest.
";

#[derive(RustcDecodable, Debug)]
struct Args {
    arg_repo: String,
    arg_old: Option<String>,
    arg_new: Option<String>,
}

fn main() {
    // Parse the args above or die.
    let args: Args = Docopt::new(USAGE)
        .and_then(|d| d.decode())
        .unwrap_or_else(|e| e.exit());
    // Pull up the revwalk.
    let path = Path::new(args.arg_repo);
    let repo = Repository::discover(&path)
        .ok().expect("Unable to find repo.");
    // If we can destructures the two optional args into real things.
    if let (Some(old_string), Some(new_string)) = (args.arg_old, args.arg_new) {
        // Compare a specific commit pair.
        // Fist, get the commits. (Error checked)
        let old = Oid::from_str(&old_string[]).and_then(|oid| repo.find_commit(oid));
        let new = Oid::from_str(&new_string[]).and_then(|oid| repo.find_commit(oid));
        if old.is_ok() && new.is_ok() {
            find_moves(&repo, &old.unwrap(), &new.unwrap()).unwrap();
        } else {
            panic!("Commit ids were not valid.");
        }
    } else {
        // Revwalk.
        let mut revwalk = repo.revwalk()
            .ok().expect("Unable to get revwalk.");
        // Setup some options.
        revwalk.simplify_first_parent(); // TODO: Maybe remove?
        let mut flags = git2::Sort::empty();
        flags.insert(git2::SORT_TIME);
        flags.insert(git2::SORT_TOPOLOGICAL);
        revwalk.set_sorting(flags);
        // Push HEAD to the revwalk.
        revwalk.push_head()
            .ok().expect("Unable to push HEAD.");
        // We sadly must collect here to use `.windows()`
        let history = revwalk.filter_map(|id| repo.find_commit(id).ok())
            .collect::<Vec<Commit>>();
        // Walk through each pair of commits.
        for pair in history.windows(2) {let (old, new) = (&pair[1], &pair[0]);
			let output = find_moves(&repo, old.clone(), new.clone()).unwrap();
			// mangle the vectors here to make pretty output
			make_output(output);
		}
	}
}

fn make_output(output: Vec<Output>) {
	println!("some json output");

}

fn find_moves(repo: &Repository, old: &Commit, new: &Commit) -> Result<Vec<Output>, Error> {
	let old_tree = try!(old.tree());
	let new_tree = try!(new.tree());
	// Build up a diff of the two trees.
	let diff = try!(Diff::tree_to_tree(repo, Some(&old_tree), Some(&new_tree), None));
	// State. TODO: Don't just use a u64.
	let moved = 0;
	let mut current_hunk: Vec<String> = Vec::new();
	// Read about this function in http://alexcrichton.com/git2-rs/git2/struct.Diff.html#method.print
	// It's a bit weird, but I think it will provide the necessary information.
	diff.print(DiffFormat::Patch, |delta, maybe_hunk, line| -> bool {
		// Thinking:
		//  * If is not a hunk, keep going.
		//  * If it's a hunk, do regex magic.
		//  * Stick regex output into a hashmap as a hash.
		//  * Later, we will iterate through and output pased on the values.
		// Filter out all the boring context lines.
		// If we're not interested in this line just return since it will iterate to the next.
		if maybe_hunk.is_none() { return true }; // Return early.
		match line.origin() {
			// Context
			' ' | '=' => true,
			// Headers
			'F' | 'H' => true,
			// Additions
			'+' | '>' => {
				println!("+ {}", str::from_utf8(line.content()).unwrap());
				true
			},
			// Deletions
			'-' | '<' => {
				println!("- {}", str::from_utf8(line.content()).unwrap());
				true
			},
			// Other (We don't care about these.)
			_         => true,
		}
			// Look at that match statement, what a hunk. So hunky.
			// match maybe_hunk {
			//     Some(hunk) => unimplemented!(),
			//     None => unimplemented!(),
			// }
	});
	Ok(vec![Output {
		old_commit: old.id(),
		new_commit: new.id(),
		origin_line: 0,
		destination_line: 0,
		num_lines: 0
	}])
}

#[derive[Show]]
struct Output {
    old_commit: Oid,
    new_commit: Oid,
    origin_line: u32,
    destination_line: u32,
    num_lines: u32,
}
<commit_msg>fixed syntax error.<commit_after>#![feature(core)]
#![feature(path)]

extern crate git2;
extern crate "rustc-serialize" as rustc_serialize;
extern crate docopt;

use git2::{Repository, Branch, BranchType, DiffLine,
    Commit, Diff, DiffFormat, DiffDelta, DiffHunk, Error, Oid};
use docopt::Docopt;
use std::old_io as io;
use std::old_io::BufferedReader;
use std::old_io::File;
use std::str;

// Write the Docopt usage string.
static USAGE: &'static str = "
Usage: transit <repo> [<old> <new>]

If no commits are given, transit will revwalk from latest to oldest.
";

#[derive(RustcDecodable, Debug)]
struct Args {
    arg_repo: String,
    arg_old: Option<String>,
    arg_new: Option<String>,
}

fn main() {
    // Parse the args above or die.
    let args: Args = Docopt::new(USAGE)
        .and_then(|d| d.decode())
        .unwrap_or_else(|e| e.exit());
    // Pull up the revwalk.
    let path = Path::new(args.arg_repo);
    let repo = Repository::discover(&path)
        .ok().expect("Unable to find repo.");
    // If we can destructures the two optional args into real things.
    if let (Some(old_string), Some(new_string)) = (args.arg_old, args.arg_new) {
        // Compare a specific commit pair.
        // Fist, get the commits. (Error checked)
        let old = Oid::from_str(&old_string[]).and_then(|oid| repo.find_commit(oid));
        let new = Oid::from_str(&new_string[]).and_then(|oid| repo.find_commit(oid));
        if old.is_ok() && new.is_ok() {
            find_moves(&repo, &old.unwrap(), &new.unwrap()).unwrap();
        } else {
            panic!("Commit ids were not valid.");
        }
    } else {
        // Revwalk.
        let mut revwalk = repo.revwalk()
            .ok().expect("Unable to get revwalk.");
        // Setup some options.
        revwalk.simplify_first_parent(); // TODO: Maybe remove?
        let mut flags = git2::Sort::empty();
        flags.insert(git2::SORT_TIME);
        flags.insert(git2::SORT_TOPOLOGICAL);
        revwalk.set_sorting(flags);
        // Push HEAD to the revwalk.
        revwalk.push_head()
            .ok().expect("Unable to push HEAD.");
        // We sadly must collect here to use `.windows()`
        let history = revwalk.filter_map(|id| repo.find_commit(id).ok())
            .collect::<Vec<Commit>>();
        // Walk through each pair of commits.
        for pair in history.windows(2) {let (old, new) = (&pair[1], &pair[0]);
			let output = find_moves(&repo, old.clone(), new.clone()).unwrap();
			// mangle the vectors here to make pretty output
			make_output(output);
		}
	}
}

fn make_output(output: Vec<Output>) {
	println!("some json output");

}

fn find_moves(repo: &Repository, old: &Commit, new: &Commit) -> Result<Vec<Output>, Error> {
	let old_tree = try!(old.tree());
	let new_tree = try!(new.tree());
	// Build up a diff of the two trees.
	let diff = try!(Diff::tree_to_tree(repo, Some(&old_tree), Some(&new_tree), None));
	// State. TODO: Don't just use a u64.
	let moved = 0;
	let mut current_hunk: Vec<String> = Vec::new();
	// Read about this function in http://alexcrichton.com/git2-rs/git2/struct.Diff.html#method.print
	// It's a bit weird, but I think it will provide the necessary information.
	diff.print(DiffFormat::Patch, |delta, maybe_hunk, line| -> bool {
		// Thinking:
		//  * If is not a hunk, keep going.
		//  * If it's a hunk, do regex magic.
		//  * Stick regex output into a hashmap as a hash.
		//  * Later, we will iterate through and output pased on the values.
		// Filter out all the boring context lines.
		// If we're not interested in this line just return since it will iterate to the next.
		if maybe_hunk.is_none() { return true }; // Return early.
		match line.origin() {
			// Context
			' ' | '=' => true,
			// Headers
			'F' | 'H' => true,
			// Additions
			'+' | '>' => {
				println!("+ {}", str::from_utf8(line.content()).unwrap());
				true
			},
			// Deletions
			'-' | '<' => {
				println!("- {}", str::from_utf8(line.content()).unwrap());
				true
			},
			// Other (We don't care about these.)
			_         => true,
		}
			// Look at that match statement, what a hunk. So hunky.
			// match maybe_hunk {
			//     Some(hunk) => unimplemented!(),
			//     None => unimplemented!(),
			// }
	});
	Ok(vec![Output {
		old_commit: old.id(),
		new_commit: new.id(),
		origin_line: 0,
		destination_line: 0,
		num_lines: 0
	}])
}

#[derive(Show)]
struct Output {
    old_commit: Oid,
    new_commit: Oid,
    origin_line: u32,
    destination_line: u32,
    num_lines: u32,
}
<|endoftext|>


<commit_before>extern crate yaml_rust;
extern crate regex;
extern crate rusqlite;

use std::fs::File;
use std::io;
use std::io::prelude::*;
use std::net::TcpListener;
use std::time::Duration;
use std::sync::Arc;
use std::thread;
use std::process::exit;
use yaml_rust::YamlLoader;
use regex::bytes::RegexSetBuilder;
use rusqlite::Connection;

const VERSION: &'static str = env!("CARGO_PKG_VERSION");
const BINARY_MATCHES: [(&str, &str);24] = [ // Global array, so needs an explicit length
  ("SSL3.0 Record Protocol", r"^\x16\x03\x00..\x01"),
  ("TLS1.0 Record Protocol", r"^\x16\x03\x01..\x01"),
  ("TLS1.1 Record Protocol", r"^\x16\x03\x02..\x01"),
  ("TLS1.2 Record Protocol", r"^\x16\x03\x03..\x01"),
  ("SSL3.0 CLIENT_HELLO", r"^\x16....\x01...\x03\x00"),
  ("TLS1.0 CLIENT_HELLO", r"^\x16....\x01...\x03\x01"),
  ("TLS1.1 CLIENT_HELLO", r"^\x16....\x01...\x03\x02"),
  ("TLS1.2 CLIENT_HELLO", r"^\x16....\x01...\x03\x03"),
  ("SMB1 COMMAND NEGOTIATE", r"^....\xffSMB\x72"),
  ("SMB1 NT_STATUS Success", r"^....\xffSMB.[\x00-\x0f]"),
  ("SMB1 NT_STATUS Information", r"^....\xffSMB.[\x40-\x4f]"),
  ("SMB1 NT_STATUS Warning", r"^....\xffSMB.[\x80-\x8f]"),
  ("SMB1 NT_STATUS Error", r"^....\xffSMB.[\xc0-\xcf]"),
  ("SMB2 COMMAND NEGOTIATE", r"^\x00...\xfeSMB........\x00\x00"),
  ("SMB2 NT_STATUS Success", r"^\x00...\xfeSMB....[\x00-\x0f]"),
  ("SMB2 NT_STATUS Information", r"^\x00...\xfeSMB....[\x40-\x4f]"),
  ("SMB2 NT_STATUS Warning", r"^\x00...\xfeSMB....[\x80-\x8f]"),
  ("SMB2 NT_STATUS Error", r"^\x00...\xfeSMB....[\xc0-\xcf]"),
  ("MS-TDS PRELOGIN Request", r"^\x12\x01\x00.\x00\x00"),
  ("MS-TDS LOGIN Request", r"^\x10\x01\x00.\x00\x00"),
  ("SOCKS4 NOAUTH Request", r"^\x04\x01\x00\x50"),
  ("SOCKS5 NOAUTH Request", r"^\x05\x01\x00$"), // Tested ok-ish
  ("SOCKS5 USER/PASS Request", r"^\x05\x02\x00\x02$"), // possibly broken
  ("Bitcoin main chain magic number", r"\xf9\xbe\xb4\xd9")
];

#[derive(Copy, Clone)]
struct App {
  print_ascii: bool,
  print_binary: bool,
  sql_logging: bool
}

#[derive(Debug)]
struct LoggedConnection {
  remoteip: String,
  remoteport: u16,
  localport: i64
}

fn to_hex(bytes: &[u8]) -> String {
  let mut count = 0;
  let mut result = String::with_capacity(67);
  let mut dotline = String::with_capacity(18);

  for byte in bytes.iter() {
    if count != 0 && count%8 == 0 {
      result.push(' ');
      if count%16 == 0 {
        dotline.push('\n');
        result.push_str(&dotline);
        result.reserve(67);
        dotline.truncate(0);
      }
      else { dotline.push(' '); }
    }
    result.push_str(&format!("{:02X} ", byte));
    if *byte > 31 && *byte < 127 { dotline.push(*byte as char) }
    else if *byte == 0 { dotline.push('-'); }
    else { dotline.push('.'); }
    count += 1;
  }
  while count%16 != 0 {
    if count%8 == 0 { result.push(' '); }
    result.push_str("   ");
    count += 1;
  }
  if dotline.len() != 0 {
    result.push(' ');
    result.push_str(&dotline);
  }
  result
}

fn setup() -> App {
  let authorstring: String = str::replace(env!("CARGO_PKG_AUTHORS"), ":", "\n");
  println!("Portlurker v{}", VERSION);
  println!("{}", authorstring);
  println!("-----------------------------------------");

  let mut app = App { print_ascii: false, print_binary: false, sql_logging: false };

  let mut config_str = String::new();
  let mut file = match File::open("config.yml") {
      Ok(file) => file,
      Err(e) => { println!("Unable to open configuration file: {}", e.to_string()); exit(-1); },
  };

  file.read_to_string(&mut config_str).unwrap();
  let docs = YamlLoader::load_from_str(&config_str).unwrap();
  let config = &docs[0];
  //println!("{:?}", config);
  if config["general"].is_badvalue() {
    println!("No 'general' section found in configuration file");
    exit(-1);
  }
  if config["ports"].is_badvalue() {
    println!("No 'ports' section found in configuration file");
    exit(-1);
  }
  if !config["general"]["print_ascii"].is_badvalue() {
    if config["general"]["print_ascii"].as_bool().unwrap() {
      app.print_ascii = true;
      println!("Printing ASCII");
    }
  }
  if !config["general"]["print_binary"].is_badvalue() {
    if config["general"]["print_binary"].as_bool().unwrap() {
      app.print_binary = true;
      println!("Printing binary in hexadecimal");
    }
  }
  if !config["general"]["sql_logging"].is_badvalue() {
    if config["general"]["sql_logging"].as_bool().unwrap() {
      app.sql_logging = true;
      println!("Logging events to portlurker.sqlite");
      match Connection::open("portlurker.sqlite") {
        Ok(conn) => { conn.execute("CREATE TABLE IF NOT EXISTS connections (
                                    id         INTEGER PRIMARY KEY,
                                    time       INTEGER NOT NULL DEFAULT (strftime('%s', 'now')),
                                    remoteip   TEXT NOT NULL,
                                    remoteport INTEGER NOT NULL,
                                    localport  INTEGER NOT NULL
                                  )", &[]).expect("Failed to create table inside database! Logging may not function correctly!"); },
        Err(e) => {
          println!("Failed to open or create database: {}\nContinuing without logging", e.to_string());
        },
      };
    }
  }
  return app; // send our config to the main function
}

fn main() {
  let app: App = setup(); // Print initial UI stuff, then parse the config file, and store the config
  let io_timeout = Duration::new(300, 0); // 5 minutes

  let mut patterns = Vec::with_capacity(BINARY_MATCHES.len());
  for &(_, pattern) in BINARY_MATCHES.into_iter() {
    patterns.push(pattern);
  }
  let regexset = RegexSetBuilder::new(patterns)
    .unicode(false)
    .dot_matches_new_line(false)
    .build().unwrap();

  println!("\nStarting listeners on the following ports:");

  // Have to reload config file here - can we improve this?
  let mut config_str = String::new();
  let mut file = match File::open("config.yml") {
      Ok(file) => file,
      Err(e) => { println!("Unable to open configuration file: {}", e.to_string()); exit(-1); },
  };

  file.read_to_string(&mut config_str).unwrap();
  let docs = YamlLoader::load_from_str(&config_str).unwrap();
  let config = &docs[0];
  for port in config["ports"].as_vec().unwrap() {
    if !port["tcp"].is_badvalue() {
      let portno = port["tcp"].as_i64().unwrap();
      println!("TCP port {}", portno);
      let mut banner = Arc::new(String::new());
      if let Some(x) = port["banner"].as_str() {
        Arc::get_mut(& mut banner).unwrap().push_str(x);
        println!("  with banner: {}", *banner);
      }
      let regexset = regexset.clone();
      thread::spawn(move || {
        let server = TcpListener::bind(("0.0.0.0", portno as u16)).expect("Port can't be bound (is it in use?)"); // Add more error checking here
        for res in server.incoming() {
          let mut stream = match res {
            Ok(stream) => stream,
            Err(e) => { println!("ACCEPT ERROR TCP {}: {}", portno, e.to_string()); continue; }
          };
          stream.set_read_timeout(Some(io_timeout)).expect("Failed to set read timeout on TcpStream");
          stream.set_write_timeout(Some(io_timeout)).expect("Failed to set write timeout on TcpStream");
          let addr = stream.peer_addr().unwrap();
          println!("CONNECT TCP {} from {}", portno, addr);

          if app.sql_logging {
            let newdbentry = LoggedConnection {
              remoteip: addr.ip().to_string(),
              remoteport: addr.port(),
              localport: portno
            };

            match Connection::open("portlurker.sqlite") {
              Ok(conn) => { conn.execute("INSERT INTO connections (
                              remoteip, remoteport, localport) VALUES (
                              ?1, ?2, ?3)", &[&newdbentry.remoteip, &newdbentry.remoteport, &newdbentry.localport]
                            ).expect("Can't write new row into table! Subsequent logging may also fail.");},
              Err(e) => {
                println!("Failed to open database: {} - Continuing without logging", e.to_string());
                let mut app = app;
                app.sql_logging = false;
              },
            }
          }

          let regexset = regexset.clone();
          let banner = banner.clone();
          thread::spawn(move || {
            if banner.len() > 0 {
              match stream.write((*banner).as_bytes()) {
                Ok(_) => {
                  for line in (*banner).lines() {
                    println!("> {}", line);
                  }
                }
                Err(e) => {
                  if e.kind() == io::ErrorKind::WouldBlock { println!("WRITE TIMEOUT TCP {} from {}", portno, addr); }
                  else { println!("WRITE ERROR TCP {} from {}: {}", portno, addr, e.to_string()); }
                  return;
                }
              }
            }
            let mut buf: [u8; 2048] = [0; 2048];
            loop {
              match stream.read(&mut buf) {
                Ok(c) => {
                  if c == 0 {
                    println!("CLOSE TCP {} from {}", portno, addr);
                    break;
                  }
                  let mut printables = Vec::new();
                  let mut found = false;
                  let mut mbfound = false;
                  let mut mbstring = String::new();
                  let mut start = 0;
                  for i in 0..c {
                    if (buf[i] > 31 && buf[i] < 127) || buf[i] == 10 || buf[i] == 13 {
                      if !found {
                        start = i;
                        found = true;
                      }
                    }
                    else {
                      if found {
                        if i-start == 1 {
                          if (start > 0) && (buf[start-1] == 0) {
                            mbstring.push(buf[i-1] as char);
                            if !mbfound { mbfound = true; }
                          }
                        }
                        else { printables.push(&buf[start..i]); }
                        found = false;
                      }
                      else if mbfound {
                        mbstring.push('\n');
                        mbfound = false;
                      }
                    }
                  }
                  if found { printables.push(&buf[start..c]); }
                  if printables.len() == 1 && printables[0].len() == c {
                    if app.print_ascii {
                      let data = String::from_utf8_lossy(printables[0]);
                      for line in data.lines() {
                        println!("| {}", line);
                      }
                    }
                    else { println!("! Read {} bytes of printable ASCII", c); }
                  }
                  else {
                    println!("! Read {} bytes of binary", c);
                    for id in regexset.matches(&buf[..c]).into_iter() {
                      println!("^ Matches pattern {}", BINARY_MATCHES[id].0);
                    }
                    for printable in printables {
                      if printable.len() > 3 {
                        let data = String::from_utf8_lossy(printable);
                        for line in data.lines() {
                          println!("$ {}", line);
                        }
                      }
                    }
                    for line in mbstring.lines() {
                      if line.len() > 3 { println!("% {}", line); }
                    }
                    if app.print_binary {
                      let hex = to_hex(&buf[..c]);
                      for line in hex.lines() { println!(". {}", line); }
                    }
                  }
                }
                Err(e) => {
                  if e.kind() == io::ErrorKind::WouldBlock { println!("READ TIMEOUT TCP {} from {}", portno, addr); }
                  else { println!("READ ERROR TCP {} from {}: {}", portno, addr, e.to_string()); }
                  break;
                }
              }
              match stream.take_error() {
                Ok(opt) => {
                  if opt.is_some() {
                    println!("ERROR TCP {} from {}: {}", portno, addr, opt.unwrap().to_string());
                    break;
                  }
                }
                Err(_) => {
                  println!("This shouldn't happen...");
                  break;
                }
              }
            }
          });
        }
      });
    }
    else if !port["udp"].is_badvalue() {
      println!("UDP port {}", port["udp"].as_i64().unwrap());
    }
    else {
      println!("Invalid port specification in configuration file");
    }
  }

  loop { thread::sleep(Duration::new(60, 0)); }
}
<commit_msg>Fixed #8: sql logging config bug<commit_after>extern crate yaml_rust;
extern crate regex;
extern crate rusqlite;

use std::fs::File;
use std::io;
use std::io::prelude::*;
use std::net::TcpListener;
use std::time::Duration;
use std::sync::Arc;
use std::thread;
use std::process::exit;
use yaml_rust::YamlLoader;
use regex::bytes::RegexSetBuilder;
use rusqlite::Connection;

const VERSION: &'static str = env!("CARGO_PKG_VERSION");
const BINARY_MATCHES: [(&str, &str);24] = [ // Global array, so needs an explicit length
  ("SSL3.0 Record Protocol", r"^\x16\x03\x00..\x01"),
  ("TLS1.0 Record Protocol", r"^\x16\x03\x01..\x01"),
  ("TLS1.1 Record Protocol", r"^\x16\x03\x02..\x01"),
  ("TLS1.2 Record Protocol", r"^\x16\x03\x03..\x01"),
  ("SSL3.0 CLIENT_HELLO", r"^\x16....\x01...\x03\x00"),
  ("TLS1.0 CLIENT_HELLO", r"^\x16....\x01...\x03\x01"),
  ("TLS1.1 CLIENT_HELLO", r"^\x16....\x01...\x03\x02"),
  ("TLS1.2 CLIENT_HELLO", r"^\x16....\x01...\x03\x03"),
  ("SMB1 COMMAND NEGOTIATE", r"^....\xffSMB\x72"),
  ("SMB1 NT_STATUS Success", r"^....\xffSMB.[\x00-\x0f]"),
  ("SMB1 NT_STATUS Information", r"^....\xffSMB.[\x40-\x4f]"),
  ("SMB1 NT_STATUS Warning", r"^....\xffSMB.[\x80-\x8f]"),
  ("SMB1 NT_STATUS Error", r"^....\xffSMB.[\xc0-\xcf]"),
  ("SMB2 COMMAND NEGOTIATE", r"^\x00...\xfeSMB........\x00\x00"),
  ("SMB2 NT_STATUS Success", r"^\x00...\xfeSMB....[\x00-\x0f]"),
  ("SMB2 NT_STATUS Information", r"^\x00...\xfeSMB....[\x40-\x4f]"),
  ("SMB2 NT_STATUS Warning", r"^\x00...\xfeSMB....[\x80-\x8f]"),
  ("SMB2 NT_STATUS Error", r"^\x00...\xfeSMB....[\xc0-\xcf]"),
  ("MS-TDS PRELOGIN Request", r"^\x12\x01\x00.\x00\x00"),
  ("MS-TDS LOGIN Request", r"^\x10\x01\x00.\x00\x00"),
  ("SOCKS4 NOAUTH Request", r"^\x04\x01\x00\x50"),
  ("SOCKS5 NOAUTH Request", r"^\x05\x01\x00$"), // Tested ok-ish
  ("SOCKS5 USER/PASS Request", r"^\x05\x02\x00\x02$"), // possibly broken
  ("Bitcoin main chain magic number", r"\xf9\xbe\xb4\xd9")
];

#[derive(Copy, Clone)]
struct App {
  print_ascii: bool,
  print_binary: bool,
  sql_logging: bool
}

#[derive(Debug)]
struct LoggedConnection {
  remoteip: String,
  remoteport: u16,
  localport: i64
}

fn to_hex(bytes: &[u8]) -> String {
  let mut count = 0;
  let mut result = String::with_capacity(67);
  let mut dotline = String::with_capacity(18);

  for byte in bytes.iter() {
    if count != 0 && count%8 == 0 {
      result.push(' ');
      if count%16 == 0 {
        dotline.push('\n');
        result.push_str(&dotline);
        result.reserve(67);
        dotline.truncate(0);
      }
      else { dotline.push(' '); }
    }
    result.push_str(&format!("{:02X} ", byte));
    if *byte > 31 && *byte < 127 { dotline.push(*byte as char) }
    else if *byte == 0 { dotline.push('-'); }
    else { dotline.push('.'); }
    count += 1;
  }
  while count%16 != 0 {
    if count%8 == 0 { result.push(' '); }
    result.push_str("   ");
    count += 1;
  }
  if dotline.len() != 0 {
    result.push(' ');
    result.push_str(&dotline);
  }
  result
}

fn setup() -> App {
  let authorstring: String = str::replace(env!("CARGO_PKG_AUTHORS"), ":", "\n");
  println!("Portlurker v{}", VERSION);
  println!("{}", authorstring);
  println!("-----------------------------------------");

  let mut app = App { print_ascii: false, print_binary: false, sql_logging: false };

  let mut config_str = String::new();
  let mut file = match File::open("config.yml") {
      Ok(file) => file,
      Err(e) => { println!("Unable to open configuration file: {}", e.to_string()); exit(-1); },
  };

  file.read_to_string(&mut config_str).unwrap();
  let docs = YamlLoader::load_from_str(&config_str).unwrap();
  let config = &docs[0];
  //println!("{:?}", config);
  if config["general"].is_badvalue() {
    println!("No 'general' section found in configuration file");
    exit(-1);
  }
  if config["ports"].is_badvalue() {
    println!("No 'ports' section found in configuration file");
    exit(-1);
  }
  if !config["general"]["print_ascii"].is_badvalue() {
    if config["general"]["print_ascii"].as_bool().unwrap() {
      app.print_ascii = true;
      println!("Printing ASCII");
    }
  }
  if !config["general"]["print_binary"].is_badvalue() {
    if config["general"]["print_binary"].as_bool().unwrap() {
      app.print_binary = true;
      println!("Printing binary in hexadecimal");
    }
  }
  if !config["general"]["sql_logging"].is_badvalue() {
    if config["general"]["sql_logging"].as_bool().unwrap() {
      println!("Logging events to portlurker.sqlite");
      match Connection::open("portlurker.sqlite") {
        Ok(conn) => { app.sql_logging = true;
                      conn.execute("CREATE TABLE IF NOT EXISTS connections (
                                    id         INTEGER PRIMARY KEY,
                                    time       INTEGER NOT NULL DEFAULT (strftime('%s', 'now')),
                                    remoteip   TEXT NOT NULL,
                                    remoteport INTEGER NOT NULL,
                                    localport  INTEGER NOT NULL
                                  )", &[]).expect("Failed to create table inside database! Logging may not function correctly!"); },
        Err(e) => {
          println!("Failed to open or create database: {}\nContinuing without logging", e.to_string());
        },
      };
    }
  }
  return app; // send our config to the main function
}

fn main() {
  let app: App = setup(); // Print initial UI stuff, then parse the config file, and store the config
  let io_timeout = Duration::new(300, 0); // 5 minutes

  let mut patterns = Vec::with_capacity(BINARY_MATCHES.len());
  for &(_, pattern) in BINARY_MATCHES.into_iter() {
    patterns.push(pattern);
  }
  let regexset = RegexSetBuilder::new(patterns)
    .unicode(false)
    .dot_matches_new_line(false)
    .build().unwrap();

  println!("\nStarting listeners on the following ports:");

  // Have to reload config file here - can we improve this?
  let mut config_str = String::new();
  let mut file = match File::open("config.yml") {
      Ok(file) => file,
      Err(e) => { println!("Unable to open configuration file: {}", e.to_string()); exit(-1); },
  };

  file.read_to_string(&mut config_str).unwrap();
  let docs = YamlLoader::load_from_str(&config_str).unwrap();
  let config = &docs[0];
  for port in config["ports"].as_vec().unwrap() {
    if !port["tcp"].is_badvalue() {
      let portno = port["tcp"].as_i64().unwrap();
      println!("TCP port {}", portno);
      let mut banner = Arc::new(String::new());
      if let Some(x) = port["banner"].as_str() {
        Arc::get_mut(& mut banner).unwrap().push_str(x);
        println!("  with banner: {}", *banner);
      }
      let regexset = regexset.clone();
      thread::spawn(move || {
        let server = TcpListener::bind(("0.0.0.0", portno as u16)).expect("Port can't be bound (is it in use?)"); // Add more error checking here
        for res in server.incoming() {
          let mut stream = match res {
            Ok(stream) => stream,
            Err(e) => { println!("ACCEPT ERROR TCP {}: {}", portno, e.to_string()); continue; }
          };
          stream.set_read_timeout(Some(io_timeout)).expect("Failed to set read timeout on TcpStream");
          stream.set_write_timeout(Some(io_timeout)).expect("Failed to set write timeout on TcpStream");
          let addr = stream.peer_addr().unwrap();
          println!("CONNECT TCP {} from {}", portno, addr);

          if app.sql_logging {
            let newdbentry = LoggedConnection {
              remoteip: addr.ip().to_string(),
              remoteport: addr.port(),
              localport: portno
            };

            match Connection::open("portlurker.sqlite") {
              Ok(conn) => { conn.execute("INSERT INTO connections (
                              remoteip, remoteport, localport) VALUES (
                              ?1, ?2, ?3)", &[&newdbentry.remoteip, &newdbentry.remoteport, &newdbentry.localport]
                            ).expect("Can't write new row into table! Subsequent logging may also fail.");},
              Err(e) => {
                println!("Failed to open database: {} - Continuing without logging", e.to_string());
                let mut app = app;
                app.sql_logging = false;
              },
            }
          }

          let regexset = regexset.clone();
          let banner = banner.clone();
          thread::spawn(move || {
            if banner.len() > 0 {
              match stream.write((*banner).as_bytes()) {
                Ok(_) => {
                  for line in (*banner).lines() {
                    println!("> {}", line);
                  }
                }
                Err(e) => {
                  if e.kind() == io::ErrorKind::WouldBlock { println!("WRITE TIMEOUT TCP {} from {}", portno, addr); }
                  else { println!("WRITE ERROR TCP {} from {}: {}", portno, addr, e.to_string()); }
                  return;
                }
              }
            }
            let mut buf: [u8; 2048] = [0; 2048];
            loop {
              match stream.read(&mut buf) {
                Ok(c) => {
                  if c == 0 {
                    println!("CLOSE TCP {} from {}", portno, addr);
                    break;
                  }
                  let mut printables = Vec::new();
                  let mut found = false;
                  let mut mbfound = false;
                  let mut mbstring = String::new();
                  let mut start = 0;
                  for i in 0..c {
                    if (buf[i] > 31 && buf[i] < 127) || buf[i] == 10 || buf[i] == 13 {
                      if !found {
                        start = i;
                        found = true;
                      }
                    }
                    else {
                      if found {
                        if i-start == 1 {
                          if (start > 0) && (buf[start-1] == 0) {
                            mbstring.push(buf[i-1] as char);
                            if !mbfound { mbfound = true; }
                          }
                        }
                        else { printables.push(&buf[start..i]); }
                        found = false;
                      }
                      else if mbfound {
                        mbstring.push('\n');
                        mbfound = false;
                      }
                    }
                  }
                  if found { printables.push(&buf[start..c]); }
                  if printables.len() == 1 && printables[0].len() == c {
                    if app.print_ascii {
                      let data = String::from_utf8_lossy(printables[0]);
                      for line in data.lines() {
                        println!("| {}", line);
                      }
                    }
                    else { println!("! Read {} bytes of printable ASCII", c); }
                  }
                  else {
                    println!("! Read {} bytes of binary", c);
                    for id in regexset.matches(&buf[..c]).into_iter() {
                      println!("^ Matches pattern {}", BINARY_MATCHES[id].0);
                    }
                    for printable in printables {
                      if printable.len() > 3 {
                        let data = String::from_utf8_lossy(printable);
                        for line in data.lines() {
                          println!("$ {}", line);
                        }
                      }
                    }
                    for line in mbstring.lines() {
                      if line.len() > 3 { println!("% {}", line); }
                    }
                    if app.print_binary {
                      let hex = to_hex(&buf[..c]);
                      for line in hex.lines() { println!(". {}", line); }
                    }
                  }
                }
                Err(e) => {
                  if e.kind() == io::ErrorKind::WouldBlock { println!("READ TIMEOUT TCP {} from {}", portno, addr); }
                  else { println!("READ ERROR TCP {} from {}: {}", portno, addr, e.to_string()); }
                  break;
                }
              }
              match stream.take_error() {
                Ok(opt) => {
                  if opt.is_some() {
                    println!("ERROR TCP {} from {}: {}", portno, addr, opt.unwrap().to_string());
                    break;
                  }
                }
                Err(_) => {
                  println!("This shouldn't happen...");
                  break;
                }
              }
            }
          });
        }
      });
    }
    else if !port["udp"].is_badvalue() {
      println!("UDP port {}", port["udp"].as_i64().unwrap());
    }
    else {
      println!("Invalid port specification in configuration file");
    }
  }

  loop { thread::sleep(Duration::new(60, 0)); }
}
<|endoftext|>