<commit_before>extern crate combine;
extern crate nix; // As in *nix
extern crate rand;
extern crate signal;

mod commander;

// Reference: https://github.com/NixOS/nixpkgs/blob/b07051ce6c3ba3039c89b6755da279002b0c3ace/nixos/modules/virtualisation/nixos-container.pl

use nix::sys::signal::{kill, SIGTERM, SIGINT};
use std::io;
use std::io::Write;
use std::process::Command;

#[derive(Debug)]
enum Error {
    Usage(String),
    NonZero(String, i32),
    ControlError(String),
    IoError(io::Error),
    NixError(nix::Error),
    Interrupted(commander::Signal)
}

use Error::*;

impl From<io::Error> for Error {
    fn from(e: io::Error) -> Self { IoError(e) }
}

impl From<nix::Error> for Error {
    fn from(e: nix::Error) -> Self { NixError(e) }
}

impl From<commander::NonZero> for Error {
    fn from(e: commander::NonZero) -> Self { match e { commander::NonZero::NonZero(s, o) => NonZero(s, o), commander::NonZero::Interrupted(n) => Interrupted(n) } }
}

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

// TODO: why can't this be inferred away
fn run<T>(cmd: &mut Command) -> Result<T> where T: commander::CommandOut { commander::run(cmd) }

fn probably_unique_name(length: usize) -> String {
    use rand::distributions::{IndependentSample, Range};
    let between = Range::new(0, 36);
    let mut rng = rand::thread_rng();
    unsafe {
        std::str::from_utf8_unchecked(&(0..length).map(|_| between.ind_sample(&mut rng)).map(|v| if v < 10 { v + 48 } else { v + 87 }).collect::<Vec<u8>>()).to_string()
    }
}

// Result is of the form "10.233.x"
fn unused_ip_block() -> Result<String> {
    // Can bail early anywhere here because we've got the lock
    let mut used = std::collections::HashSet::new();
    for entry in try!(std::fs::read_dir("/etc/containers")) {
        let entry = try!(entry);
        if try!(entry.file_type()).is_file() {
            let mut str = String::new();
            use std::io::Read;
            try!(try!(std::fs::File::open(entry.path())).read_to_string(&mut str));
            for line in str.lines() {
                use combine::{string, many1, choice, digit, token, eof, Parser, ParserExt};
                let ip_component = || { string("10.233.").with(many1::<String, _>(digit())).skip(token('.')).skip(many1::<String, _>(digit())) };
                if let Ok((c, _)) = choice([string("HOST_ADDRESS="), string("LOCAL_ADDRESS=")]).with(ip_component()).skip(eof()).parse(line) {
                    used.insert(c);
                }
            }
        }
    }
    (0..255).map(|p| p.to_string()).find(|s| !used.contains(s)).map(|s| format!("10.233.{}", s)).ok_or(ControlError("Out of IP addresses".to_string()))
}

fn mkpath(mode : u32, path : &str) -> Result<()> {
    use std::os::unix::fs::DirBuilderExt;
    match std::fs::DirBuilder::new().mode(mode).create(path) {
        Err(e) => if let std::io::ErrorKind::AlreadyExists = e.kind() { Ok(()) } else { Err(IoError(e)) },
        Ok(()) => Ok(())
    }
}

fn system_init() -> Result<()> {
    try!(mkpath(0o0755, "/etc/containers"));
    try!(mkpath(0o0700, "/var/lib/containers"));
    Ok(())
}

fn profile_dir(name : &str) -> String {
    format!("/nix/var/nix/profiles/per-container/{}", name)
}

fn container_root(name : &str) -> String {
    format!("/var/lib/containers/{}", name)
}

fn conf_file(name : &str) -> String {
    format!("/etc/containers/{}.conf", name)
}

fn create(config_file : &str) -> Result<String> {
    fn init_container_unsynced(name : &str) -> Result<()> {
        let ip_block = try!(unused_ip_block());
        // Use create_new to avoid clobbering existing container
        let mut f = try!(std::fs::OpenOptions::new().create_new(true).write(true).open(conf_file(name)));
        write!(&mut f, "PRIVATE_NETWORK=1\nHOST_ADDRESS={0}.1\nLOCAL_ADDRESS={0}.2\nAUTO_START=0\n", ip_block).map_err(IoError)
    }

    fn populate_container(name : &str, config_file : &str) -> Result<()> {
        use std::os::unix::fs::DirBuilderExt;
        // The per-container directory is restricted to prevent users on
        // the host from messing with guest users who happen to have the
        // same uid.
        try!(mkpath(0o0700, "/nix/var/nix/profiles/per-container"));
        let profile_dir = profile_dir(name);
        let mut dir_builder = std::fs::DirBuilder::new();
        try!(dir_builder.mode(0o0755).create(&profile_dir));

        let container_root = container_root(name);
        dir_builder.recursive(true);
        try!(dir_builder.mode(0o0755).create(format!("{}/etc/nixos", container_root)));

        // Write configuration.nix into the container fs
        let configuration_nix = format!("{}/etc/nixos/configuration.nix", container_root);
        {
            let mut f = try!(std::fs::File::create(format!("{}/etc/nixos/configuration.nix", container_root)));
            try!(write!(&mut f, "
    {{ config, lib, pkgs, ... }}:
    with lib;

    {{
        boot.isContainer = true;
        networking.hostName = mkDefault \"{}\";
        networking.useDHCP = false;
        imports = [ {} ];
    }}", name, config_file));
            // TODO: guaranteed sync at end of block?
        }

        // Build the container config
        run(Command::new("nix-env")
            .arg("-p").arg(format!("{}/system", profile_dir))
            .arg("-I").arg(format!("nixos-config={}", configuration_nix))
            .arg("-f").arg("<nixpkgs/nixos>")
            .arg("--set")
            .arg("-A").arg("system")
            .arg("--show-trace"))
    }

    let name = format!("fr{}", probably_unique_name(9));

    // Acquire lock
    use std::os::unix::io::IntoRawFd;
    let lock_fd = try!(std::fs::OpenOptions::new().create(true).append(true).open("/run/lock/nixos-container")).into_raw_fd();
    try!(nix::fcntl::flock(lock_fd, nix::fcntl::FlockArg::LockExclusive));

    let res = init_container_unsynced(&name);

    // Drop lock
    unsafe {
        let _ = nix::libc::close(lock_fd);  // If closing fails, ¯\_(ツ)_/¯
    }

    match res.and_then(|()| populate_container(&name, config_file)) {
        Ok(()) => Ok(name),
        // Don't leave a half-baked container behind
        Err(e) => { let _ = destroy(&name); Err(e) }
    }
}

fn run_test(container_name : &str, args : &[String]) -> Result<i32> {
    try!(run(Command::new("systemctl").arg("start").arg(format!("container@{}", container_name))));
    let leader_pid = try!(run(Command::new("machinectl").arg("show").arg(container_name).arg("-p").arg("Leader")).and_then(|o : String| {
        use combine::{string, many1, digit, spaces, eof, Parser, ParserExt};
        string("Leader=").with(many1::<String, _>(digit())).skip(spaces()).skip(eof()).parse(o.as_str())
            .map(|(ds, _)| ds)
            .map_err(|_| ControlError(format!("Bad machinectl output {}", o)))
    }));
    let su_cmd = args.iter().map(|s| format!("'{}'", s.replace("'", "'\\''"))).collect::<Vec<String>>().join(" ");
    match run(Command::new("nsenter")
            .arg("-t").arg(leader_pid)
            .arg("-m").arg("-u").arg("-i").arg("-n").arg("-p")
            .arg("--").arg("/var/setuid-wrappers/su").arg("root").arg("-l").arg("-c").arg(format!("exec {}", su_cmd))) {
        Ok(()) => Ok(0),
        Err(NonZero(_, code)) => Ok(code),
        Err(e) => Err(e)
    }
}

// Remove a directory while recursively unmounting all mounted filesystems within
// that directory and unmounting/removing that directory afterwards as well.
//
// Specified path shouldn't be a mountpoint.
fn safe_remove_tree(path : &str) -> Result<()> {
    if !std::path::Path::new(path).is_dir() {
        return Ok(());
    }
    try!(run(Command::new("find").arg(path)
        .arg("-mindepth").arg("1")
        .arg("-xdev")
        .arg("(")
            .arg("-type").arg("d")
            .arg("-exec").arg("mountpoint").arg("-q").arg("{}").arg(";")
        .arg(")")
        .arg("-exec").arg("umount").arg("-fR").arg("{}").arg("+")));
    try!(run(Command::new("rm").arg("--one-file-system").arg("-rf").arg(path)));

    Ok(())
}

fn stop(container_name : &str) -> Result<()> {
    match run(Command::new("systemctl").arg("stop").arg(format!("container@{}", container_name))) {
        Ok(()) => Ok(()),
        Err(Interrupted(_)) => {
            return run(Command::new("systemctl").arg("kill").arg(format!("container@{}", container_name)))
        },
        Err(e) => Err(e)
    }
}

fn destroy(container_name : &str) -> Result<()> {
    fn log_err(e: Error) {
        println!("Error while destroying container: {:?}", e);
    }

    safe_remove_tree(&profile_dir(container_name)).unwrap_or_else(log_err);
    safe_remove_tree(&format!("/nix/var/nix/gcroots/per-container/{}", container_name)).unwrap_or_else(log_err);
    safe_remove_tree(&container_root(container_name)).unwrap_or_else(log_err);

    std::fs::remove_file(conf_file(container_name)).or_else(|e| match e.kind() {
        std::io::ErrorKind::NotFound => Ok(()),
        _ => Err(IoError(e))
    })
}

fn go() -> Result<i32> {
    // Skip program name
    let mut args = std::env::args().skip(1).peekable();

    let no_destroy = args.peek().map_or(false, |arg| arg == "--no-destroy");
    if no_destroy { args.next(); }

    let config_file = try!(args.next().ok_or(Error::Usage("Missing config file argument".to_string())));

    let test_args : Vec<String> = args.collect();
    if test_args.len() == 0 {
        return Err(Usage("Missing test args".to_string()));
    }

    let config_file = {
        let p = try!(std::fs::canonicalize(config_file).map_err(|e| Usage(format!("For config file: {}", e))));
        try!(p.to_str().ok_or(Usage("Config file path is invalid UTF-8".to_string()))).to_string()
    };

    try!(system_init());
    // We suppress signals in this process and allow the commander module to deal with child processes
    let trap = signal::trap::Trap::trap(&[SIGTERM, SIGINT]);
    let container_name = try!(create(&config_file));
    let res = match trap.wait(std::time::Instant::now()) {
        None => run_test(&container_name, &test_args),
        Some(n) => Err(Interrupted(n))
    };

    let _ = stop(&container_name);
    if !no_destroy {
        try!(destroy(&container_name));
    }

    res
}

fn main() {
    std::process::exit(go().unwrap_or_else(|e| match e {
        NonZero(cmd, code) => { writeln!(io::stderr(), "Command {:?} failed with code {:?}", cmd, code).unwrap(); 2 },
        Usage(msg) => { writeln!(io::stderr(), "Usage error: {}", msg).unwrap(); 1 },
        IoError(err) => { writeln!(io::stderr(), "{}", err).unwrap(); 2 },
        NixError(err) => { writeln!(io::stderr(), "{}", err).unwrap(); 2 },
        ControlError(msg) => { writeln!(io::stderr(), "{}", msg).unwrap(); 2 },
        Interrupted(SIGINT) => { kill(nix::unistd::getpid(), SIGINT).unwrap(); loop {} }
        Interrupted(n) => 128 + n
    }));
}
<commit_msg>Call stop before destroy<commit_after>extern crate combine;
extern crate nix; // As in *nix
extern crate rand;
extern crate signal;

mod commander;

// Reference: https://github.com/NixOS/nixpkgs/blob/b07051ce6c3ba3039c89b6755da279002b0c3ace/nixos/modules/virtualisation/nixos-container.pl

use nix::sys::signal::{kill, SIGTERM, SIGINT};
use std::io;
use std::io::Write;
use std::process::Command;

#[derive(Debug)]
enum Error {
    Usage(String),
    NonZero(String, i32),
    ControlError(String),
    IoError(io::Error),
    NixError(nix::Error),
    Interrupted(commander::Signal)
}

use Error::*;

impl From<io::Error> for Error {
    fn from(e: io::Error) -> Self { IoError(e) }
}

impl From<nix::Error> for Error {
    fn from(e: nix::Error) -> Self { NixError(e) }
}

impl From<commander::NonZero> for Error {
    fn from(e: commander::NonZero) -> Self { match e { commander::NonZero::NonZero(s, o) => NonZero(s, o), commander::NonZero::Interrupted(n) => Interrupted(n) } }
}

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

// TODO: why can't this be inferred away
fn run<T>(cmd: &mut Command) -> Result<T> where T: commander::CommandOut { commander::run(cmd) }

fn probably_unique_name(length: usize) -> String {
    use rand::distributions::{IndependentSample, Range};
    let between = Range::new(0, 36);
    let mut rng = rand::thread_rng();
    unsafe {
        std::str::from_utf8_unchecked(&(0..length).map(|_| between.ind_sample(&mut rng)).map(|v| if v < 10 { v + 48 } else { v + 87 }).collect::<Vec<u8>>()).to_string()
    }
}

// Result is of the form "10.233.x"
fn unused_ip_block() -> Result<String> {
    // Can bail early anywhere here because we've got the lock
    let mut used = std::collections::HashSet::new();
    for entry in try!(std::fs::read_dir("/etc/containers")) {
        let entry = try!(entry);
        if try!(entry.file_type()).is_file() {
            let mut str = String::new();
            use std::io::Read;
            try!(try!(std::fs::File::open(entry.path())).read_to_string(&mut str));
            for line in str.lines() {
                use combine::{string, many1, choice, digit, token, eof, Parser, ParserExt};
                let ip_component = || { string("10.233.").with(many1::<String, _>(digit())).skip(token('.')).skip(many1::<String, _>(digit())) };
                if let Ok((c, _)) = choice([string("HOST_ADDRESS="), string("LOCAL_ADDRESS=")]).with(ip_component()).skip(eof()).parse(line) {
                    used.insert(c);
                }
            }
        }
    }
    (0..255).map(|p| p.to_string()).find(|s| !used.contains(s)).map(|s| format!("10.233.{}", s)).ok_or(ControlError("Out of IP addresses".to_string()))
}

fn mkpath(mode : u32, path : &str) -> Result<()> {
    use std::os::unix::fs::DirBuilderExt;
    match std::fs::DirBuilder::new().mode(mode).create(path) {
        Err(e) => if let std::io::ErrorKind::AlreadyExists = e.kind() { Ok(()) } else { Err(IoError(e)) },
        Ok(()) => Ok(())
    }
}

fn system_init() -> Result<()> {
    try!(mkpath(0o0755, "/etc/containers"));
    try!(mkpath(0o0700, "/var/lib/containers"));
    Ok(())
}

fn profile_dir(name : &str) -> String {
    format!("/nix/var/nix/profiles/per-container/{}", name)
}

fn container_root(name : &str) -> String {
    format!("/var/lib/containers/{}", name)
}

fn conf_file(name : &str) -> String {
    format!("/etc/containers/{}.conf", name)
}

fn create(config_file : &str) -> Result<String> {
    fn init_container_unsynced(name : &str) -> Result<()> {
        let ip_block = try!(unused_ip_block());
        // Use create_new to avoid clobbering existing container
        let mut f = try!(std::fs::OpenOptions::new().create_new(true).write(true).open(conf_file(name)));
        write!(&mut f, "PRIVATE_NETWORK=1\nHOST_ADDRESS={0}.1\nLOCAL_ADDRESS={0}.2\nAUTO_START=0\n", ip_block).map_err(IoError)
    }

    fn populate_container(name : &str, config_file : &str) -> Result<()> {
        use std::os::unix::fs::DirBuilderExt;
        // The per-container directory is restricted to prevent users on
        // the host from messing with guest users who happen to have the
        // same uid.
        try!(mkpath(0o0700, "/nix/var/nix/profiles/per-container"));
        let profile_dir = profile_dir(name);
        let mut dir_builder = std::fs::DirBuilder::new();
        try!(dir_builder.mode(0o0755).create(&profile_dir));

        let container_root = container_root(name);
        dir_builder.recursive(true);
        try!(dir_builder.mode(0o0755).create(format!("{}/etc/nixos", container_root)));

        // Write configuration.nix into the container fs
        let configuration_nix = format!("{}/etc/nixos/configuration.nix", container_root);
        {
            let mut f = try!(std::fs::File::create(format!("{}/etc/nixos/configuration.nix", container_root)));
            try!(write!(&mut f, "
    {{ config, lib, pkgs, ... }}:
    with lib;

    {{
        boot.isContainer = true;
        networking.hostName = mkDefault \"{}\";
        networking.useDHCP = false;
        imports = [ {} ];
    }}", name, config_file));
            // TODO: guaranteed sync at end of block?
        }

        // Build the container config
        run(Command::new("nix-env")
            .arg("-p").arg(format!("{}/system", profile_dir))
            .arg("-I").arg(format!("nixos-config={}", configuration_nix))
            .arg("-f").arg("<nixpkgs/nixos>")
            .arg("--set")
            .arg("-A").arg("system")
            .arg("--show-trace"))
    }

    let name = format!("fr{}", probably_unique_name(9));

    // Acquire lock
    use std::os::unix::io::IntoRawFd;
    let lock_fd = try!(std::fs::OpenOptions::new().create(true).append(true).open("/run/lock/nixos-container")).into_raw_fd();
    try!(nix::fcntl::flock(lock_fd, nix::fcntl::FlockArg::LockExclusive));

    let res = init_container_unsynced(&name);

    // Drop lock
    unsafe {
        let _ = nix::libc::close(lock_fd);  // If closing fails, ¯\_(ツ)_/¯
    }

    match res.and_then(|()| populate_container(&name, config_file)) {
        Ok(()) => Ok(name),
        // Don't leave a half-baked container behind
        Err(e) => { let (_, _) = (stop(&name), destroy(&name)); Err(e) }
    }
}

fn run_test(container_name : &str, args : &[String]) -> Result<i32> {
    try!(run(Command::new("systemctl").arg("start").arg(format!("container@{}", container_name))));
    let leader_pid = try!(run(Command::new("machinectl").arg("show").arg(container_name).arg("-p").arg("Leader")).and_then(|o : String| {
        use combine::{string, many1, digit, spaces, eof, Parser, ParserExt};
        string("Leader=").with(many1::<String, _>(digit())).skip(spaces()).skip(eof()).parse(o.as_str())
            .map(|(ds, _)| ds)
            .map_err(|_| ControlError(format!("Bad machinectl output {}", o)))
    }));
    let su_cmd = args.iter().map(|s| format!("'{}'", s.replace("'", "'\\''"))).collect::<Vec<String>>().join(" ");
    match run(Command::new("nsenter")
            .arg("-t").arg(leader_pid)
            .arg("-m").arg("-u").arg("-i").arg("-n").arg("-p")
            .arg("--").arg("/var/setuid-wrappers/su").arg("root").arg("-l").arg("-c").arg(format!("exec {}", su_cmd))) {
        Ok(()) => Ok(0),
        Err(NonZero(_, code)) => Ok(code),
        Err(e) => Err(e)
    }
}

// Remove a directory while recursively unmounting all mounted filesystems within
// that directory and unmounting/removing that directory afterwards as well.
//
// Specified path shouldn't be a mountpoint.
fn safe_remove_tree(path : &str) -> Result<()> {
    if !std::path::Path::new(path).is_dir() {
        return Ok(());
    }
    try!(run(Command::new("find").arg(path)
        .arg("-mindepth").arg("1")
        .arg("-xdev")
        .arg("(")
            .arg("-type").arg("d")
            .arg("-exec").arg("mountpoint").arg("-q").arg("{}").arg(";")
        .arg(")")
        .arg("-exec").arg("umount").arg("-fR").arg("{}").arg("+")));
    try!(run(Command::new("rm").arg("--one-file-system").arg("-rf").arg(path)));

    Ok(())
}

fn stop(container_name : &str) -> Result<()> {
    match run(Command::new("systemctl").arg("stop").arg(format!("container@{}", container_name))) {
        Ok(()) => Ok(()),
        Err(Interrupted(_)) => {
            return run(Command::new("systemctl").arg("kill").arg(format!("container@{}", container_name)))
        },
        Err(e) => Err(e)
    }
}

fn destroy(container_name : &str) -> Result<()> {
    fn log_err(e: Error) {
        println!("Error while destroying container: {:?}", e);
    }

    safe_remove_tree(&profile_dir(container_name)).unwrap_or_else(log_err);
    safe_remove_tree(&format!("/nix/var/nix/gcroots/per-container/{}", container_name)).unwrap_or_else(log_err);
    safe_remove_tree(&container_root(container_name)).unwrap_or_else(log_err);

    std::fs::remove_file(conf_file(container_name)).or_else(|e| match e.kind() {
        std::io::ErrorKind::NotFound => Ok(()),
        _ => Err(IoError(e))
    })
}

fn go() -> Result<i32> {
    // Skip program name
    let mut args = std::env::args().skip(1).peekable();

    let no_destroy = args.peek().map_or(false, |arg| arg == "--no-destroy");
    if no_destroy { args.next(); }

    let config_file = try!(args.next().ok_or(Error::Usage("Missing config file argument".to_string())));

    let test_args : Vec<String> = args.collect();
    if test_args.len() == 0 {
        return Err(Usage("Missing test args".to_string()));
    }

    let config_file = {
        let p = try!(std::fs::canonicalize(config_file).map_err(|e| Usage(format!("For config file: {}", e))));
        try!(p.to_str().ok_or(Usage("Config file path is invalid UTF-8".to_string()))).to_string()
    };

    try!(system_init());
    // We suppress signals in this process and allow the commander module to deal with child processes
    let trap = signal::trap::Trap::trap(&[SIGTERM, SIGINT]);
    let container_name = try!(create(&config_file));
    let res = match trap.wait(std::time::Instant::now()) {
        None => run_test(&container_name, &test_args),
        Some(n) => Err(Interrupted(n))
    };

    let _ = stop(&container_name);
    if !no_destroy {
        try!(destroy(&container_name));
    }

    res
}

fn main() {
    std::process::exit(go().unwrap_or_else(|e| match e {
        NonZero(cmd, code) => { writeln!(io::stderr(), "Command {:?} failed with code {:?}", cmd, code).unwrap(); 2 },
        Usage(msg) => { writeln!(io::stderr(), "Usage error: {}", msg).unwrap(); 1 },
        IoError(err) => { writeln!(io::stderr(), "{}", err).unwrap(); 2 },
        NixError(err) => { writeln!(io::stderr(), "{}", err).unwrap(); 2 },
        ControlError(msg) => { writeln!(io::stderr(), "{}", msg).unwrap(); 2 },
        Interrupted(SIGINT) => { kill(nix::unistd::getpid(), SIGINT).unwrap(); loop {} }
        Interrupted(n) => 128 + n
    }));
}
<|endoftext|>


<commit_before>use std::fs::File;
use std::io::prelude::*;

extern crate advent;

use advent::Keypad;

fn main() {
    let mut keypad = Keypad::new("123\n456\n789");

    let mut file = File::open("./src/input.txt")
                        .expect("Could not open input.txt");
    let mut input = String::new();
    file.read_to_string(&mut input).expect("Could not read input.txt");

    let answer = advent::puzzle(&mut keypad, &input);
    println!("The answer is {}", answer);
}
<commit_msg>Wooo finished day 2!!<commit_after>use std::fs::File;
use std::io::prelude::*;

extern crate advent;

use advent::Keypad;

fn main() {
    let keypad_input = [
        "  1  ",
        " 234 ",
        "56789",
        " ABC ",
        "  D  ",
    ].join("\n");
    let mut keypad = Keypad::new(&keypad_input);

    let mut file = File::open("./src/input.txt")
                        .expect("Could not open input.txt");
    let mut input = String::new();
    file.read_to_string(&mut input).expect("Could not read input.txt");

    let answer = advent::puzzle(&mut keypad, &input);
    println!("The answer is {}", answer);
}
<|endoftext|>


<commit_before>extern crate yaml_rust;
extern crate regex;

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;

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

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

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_str(&dotline); }
  result
}

fn main() {
  println!("PortLurker v{}", VERSION);
  println!("{}", AUTHORS);

  let mut app = App { print_ascii: false, print_binary: false };
  let binary_matches = [
    ("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 Requst", r"^\x10\x01\x00.\x00\x00")
  ];
  let io_timeout = Duration::new(300, 0); // 5 minutes

  let mut file = File::open("config.yml").unwrap();
  let mut config_str = String::new();
  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");
    }
  }
  let app = app; // Revert to immutable

  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!("Starting listeners on the following ports:");

  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)).unwrap();
        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);
          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>Update main.rs<commit_after>extern crate yaml_rust;
extern crate regex;

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;

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

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

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_str(&dotline); }
  result
}

fn main() {
  println!("Portlurker v{}", VERSION);
  println!("{}", AUTHORS);

  let mut app = App { print_ascii: false, print_binary: false };
  let binary_matches = [
    ("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 Requst", r"^\x10\x01\x00.\x00\x00")
  ];
  let io_timeout = Duration::new(300, 0); // 5 minutes

  let mut file = File::open("config.yml").unwrap();
  let mut config_str = String::new();
  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");
    }
  }
  let app = app; // Revert to immutable

  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!("Starting listeners on the following ports:");

  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)).unwrap();
        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);
          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|>


<commit_before>#[macro_use]
extern crate futures;
extern crate tokio_proto;
extern crate tokio_service;
extern crate tokio_core;
extern crate tokio_io;
extern crate bytes;

use std::str;
use std::io::{self, ErrorKind, Write};
use std::net::ToSocketAddrs;

use futures::{future, Future, BoxFuture};
use futures::{Stream, Poll, Async};
use futures::{Sink, AsyncSink, StartSend};
//use tokio_proto::TcpServer;
use tokio_proto::TcpClient;
//use tokio_proto::pipeline::ServerProto;
use tokio_proto::pipeline::ClientProto;
use tokio_service::Service;
use tokio_core::reactor::Core;
use tokio_io::codec::{Encoder, Decoder};
use tokio_io::codec::Framed;
use tokio_io::{AsyncRead, AsyncWrite};
use bytes::{BytesMut, BufMut};

#[derive(Default)]
pub struct LineCodec;

impl Decoder for LineCodec {
    type Item = String;
    type Error = io::Error;

    fn decode(&mut self, buf: &mut BytesMut) -> io::Result<Option<String>> {
        if let Some(i) = buf.iter().position(|&b| b == b'\n') {
            let line = buf.split_to(i);
            buf.split_to(1);
            match str::from_utf8(&line) {
                Ok(s) => Ok(Some(s.to_string())),
                Err(_) => Err(io::Error::new(io::ErrorKind::Other,
                                             "invalid UTF-8")),
            }
        } else {
            Ok(None)
        }
    }
}

impl Encoder for LineCodec {
    type Item = String;
    type Error = io::Error;

    fn encode(&mut self, msg: String, buf: &mut BytesMut) -> io::Result<()> {
        buf.extend(msg.as_bytes());
        buf.extend(b"\n");
        Ok(())
    }
}

pub struct PingPong<T> {
    upstream: T,
    server: Option<String>,
}

impl<T> PingPong<T> {
    fn new(t: T) -> Self {
        PingPong {
            upstream: t,
            server: None,
        }
    }
}

impl<T> Stream for PingPong<T>
    where T: Stream<Item = String, Error = io::Error>,
          T: Sink<SinkItem = String, SinkError = io::Error>,
{
    type Item = String;
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Option<String>, io::Error> {
        loop {
            // Poll the upstream transport. `try_ready!` will bubble up errors
            // and Async::NotReady.
            match try_ready!(self.upstream.poll()) {
                Some(ref msg) if msg.starts_with("PING ") => {
                    // Intercept [ping] messages
                    self.server = Some(msg[4..].to_owned());
                    // Try flushing the pong, only bubble up errors
                    try!(self.poll_complete());
                }
                m => return Ok(Async::Ready(m)),
            }
        }
    }
}

impl<T> Sink for PingPong<T>
    where T: Sink<SinkItem = String, SinkError = io::Error>,
{
    type SinkItem = String;
    type SinkError = io::Error;

    fn start_send(&mut self, item: String) -> StartSend<String, io::Error> {
        // Only accept the write if there are no pending pongs
        //if self.pongs_remaining > 0 {
        if self.server.is_some() {
            return Ok(AsyncSink::NotReady(item));
        }

        // If there are no pending pongs, then send the item upstream
        self.upstream.start_send(item)
    }

    fn poll_complete(&mut self) -> Poll<(), io::Error> {
        //while self.pongs_remaining > 0 {
		//while let Some(s) = self.server {
		while self.server.is_some() {
			let s = self.server.take();
            // Try to send the pong upstream
            let res = try!(self.upstream.start_send("[pong]".to_string()));

            if !res.is_ready() {
                // The upstream is not ready to accept new items
                break;
            }
        }

        // Call poll_complete on the upstream
        //
        // If there are remaining pongs to send, this call may create additional
        // capacity. One option could be to attempt to send the pongs again.
        // However, if a `start_send` returned NotReady, and this poll_complete
        // *did* create additional capacity in the upstream, then *our*
        // `poll_complete` will get called again shortly.
        //
        // Hopefully this makes sense... it probably doesn't, so please ask
        // questions in the Gitter channel and help me explain this better :)
        self.upstream.poll_complete()
    }
}

pub struct LineProto;

impl<T: AsyncRead + AsyncWrite + 'static> ClientProto<T> for LineProto {
    /// For this protocol style, `Request` matches the `Item` type of the codec's `Encoder`
    type Request = String;

    /// For this protocol style, `Response` matches the `Item` type of the codec's `Decoder`
    type Response = String;

    /// A bit of boilerplate to hook in the codec:
    type Transport = PingPong<Framed<T, LineCodec>>;
    type BindTransport = Result<Self::Transport, io::Error>;
    fn bind_transport(&self, io: T) -> Self::BindTransport {
        Ok(PingPong::new(io.framed(LineCodec)))
    }
}

pub struct Echo;

impl Service for Echo {
    // These types must match the corresponding protocol types:
    type Request = String;
    type Response = String;

    // For non-streaming protocols, service errors are always io::Error
    type Error = io::Error;

    // The future for computing the response; box it for simplicity.
    type Future = BoxFuture<Self::Response, Self::Error>;

    // Produce a future for computing a response from a request.
    fn call(&self, req: Self::Request) -> Self::Future {
        // In this case, the response is immediate.
        future::ok(req).boxed()
    }
}

fn main() {
    //let addr = "0.0.0.0:12345".parse().unwrap();
    //TcpServer::new(IntProto, addr)
    //    .serve(|| Ok(Doubler));
    let cm = "USER qj 0 * qjkx\r\nNICK qjk\r\nJOIN #test\r\n".to_string();

    //let addr = "irc.freenode.org:6697".to_socket_addrs().unwrap().next().unwrap();
    let addr = "0.0.0.0:12345".to_socket_addrs().unwrap().next().unwrap();

    //LineCodec.poll();
    //let lc = LineCodec;
    //lc.AsyncReadframed(LineCodec);

    let mut core = Core::new().unwrap();
    let handle = core.handle();
    //let cli = TcpClient::new(IntProto);
    let cli = TcpClient::new(LineProto);
    let socket = cli.connect(&addr, &handle);
    //let response = socket.and_then(|x| x.call(420));
    let response = socket.and_then(|x| x.call(cm));

    let data = core.run(response);//.unwrap();
    println!("{:?}", data);
    //let (_socket, data) = core.run(response).unwrap();
    //println!("{}", String::from_utf8_lossy(&data));
}

/*
extern crate futures;
extern crate native_tls;
extern crate tokio_core;
extern crate tokio_io;
extern crate tokio_tls;

use std::io;
use std::net::ToSocketAddrs;

use futures::Future;
use native_tls::TlsConnector;
use tokio_core::net::TcpStream;
use tokio_core::reactor::Core;
use tokio_tls::TlsConnectorExt;

fn main() {
    let cm = "USER qj 0 * qjkx\r\nNICK qjk\r\nJOIN #test\r\n".as_bytes();

    let mut core = Core::new().unwrap();
    let handle = core.handle();
    let addr = "irc.freenode.org:6697".to_socket_addrs().unwrap().next().unwrap();

    let cx = TlsConnector::builder().unwrap().build().unwrap();
    let socket = TcpStream::connect(&addr, &handle);

    let tls_handshake = socket.and_then(|socket| {
        println!("Shaking");
        let tls = cx.connect_async("irc.freenode.org", socket);
        tls.map_err(|e| {
            io::Error::new(io::ErrorKind::Other, e)
        })
    });
    let request = tls_handshake.and_then(|socket| {
        println!("Writing connect message");
        tokio_io::io::write_all(socket, cm)
    });
    let response = request.and_then(|(socket, _request)| {
        tokio_io::io::read_to_end(socket, vec![])
    });

    let (_socket, data) = core.run(response).unwrap();
    println!("{}", String::from_utf8_lossy(&data));
}

*/

/*

extern crate futures;
extern crate tokio_proto;
extern crate tokio_service;
extern crate tokio_io;
extern crate bytes;

use std::str;
use std::io::{self, ErrorKind, Write};

use futures::{future, Future, BoxFuture};
use tokio_proto::TcpServer;
use tokio_proto::pipeline::ServerProto;
use tokio_service::Service;
use tokio_io::codec::{Encoder, Decoder};
use tokio_io::codec::Framed;
use tokio_io::{AsyncRead, AsyncWrite};
use bytes::{BytesMut, BufMut};

// First, we implement a *codec*, which provides a way of encoding and
// decoding messages for the protocol. See the documentation for `Codec` in
// `tokio-core` for more details on how that works.

#[derive(Default)]
pub struct IntCodec;

fn parse_u64(from: &[u8]) -> Result<u64, io::Error> {
    Ok(str::from_utf8(from)
       .map_err(|e| io::Error::new(ErrorKind::InvalidData, e))?
       .parse()
       .map_err(|e| io::Error::new(ErrorKind::InvalidData, e))?)
}

impl Decoder for IntCodec {
    type Item = u64;
    type Error = io::Error;

    // Attempt to decode a message from the given buffer if a complete
    // message is available; returns `Ok(None)` if the buffer does not yet
    // hold a complete message.
    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<u64>, io::Error> {
        if let Some(mut i) = buf.iter().position(|&b| b == b'\n') {
            // remove the line, including the '\n', from the buffer
            let full_line = buf.split_to(i + 1);
            // strip the `\n' (and `\r' if present)
            if full_line.ends_with(&[b'\r', b'\n']) {
                i -= 1;
            }
            let slice = &full_line[..i];
            Ok(Some(parse_u64(slice)?))
        } else {
            Ok(None)
        }
    }

    // Attempt to decode a message assuming that the given buffer contains
    // *all* remaining input data.
    fn decode_eof(&mut self, buf: &mut BytesMut) -> Result<Option<u64>, io::Error> {
        let amt = buf.len();
        Ok(Some(parse_u64(&buf.split_to(amt)[..])?))
    }
}

impl Encoder for IntCodec {
    type Item = u64;
    type Error = io::Error;

    fn encode(&mut self, item: u64, into: &mut BytesMut) -> io::Result<()> {
        writeln!(into.writer(), "{}", item)?;
        Ok(())
    }
}

// Next, we implement the server protocol, which just hooks up the codec above.

pub struct IntProto;

impl<T: AsyncRead + AsyncWrite + 'static> ServerProto<T> for IntProto {
    type Request = u64;
    type Response = u64;
    type Transport = Framed<T, IntCodec>;
    type BindTransport = Result<Self::Transport, io::Error>;

    fn bind_transport(&self, io: T) -> Self::BindTransport {
        Ok(io.framed(IntCodec))
    }
}

// Now we implement a service we'd like to run on top of this protocol

pub struct Doubler;

impl Service for Doubler {
    type Request = u64;
    type Response = u64;
    type Error = io::Error;
    type Future = BoxFuture<u64, io::Error>;

    fn call(&self, req: u64) -> Self::Future {
        // Just return the request, doubled
        future::finished(req * 2).boxed()
    }
}

// Finally, we can actually host this service locally!
fn main() {
    let addr = "0.0.0.0:12345".parse().unwrap();
    TcpServer::new(IntProto, addr)
        .serve(|| Ok(Doubler));
}
*/
<commit_msg>pong works, cleaned up a bit<commit_after>#[macro_use]
extern crate futures;
extern crate tokio_proto;
extern crate tokio_service;
extern crate tokio_core;
extern crate tokio_io;
extern crate bytes;

use std::str;
use std::io;
//use std::io::{self, ErrorKind, Write};
use std::net::ToSocketAddrs;

use futures::{future, Future, BoxFuture};
use futures::{Stream, Poll, Async};
use futures::{Sink, AsyncSink, StartSend};
use tokio_proto::TcpClient;
use tokio_proto::pipeline::ClientProto;
use tokio_service::Service;
use tokio_core::reactor::Core;
use tokio_io::codec::{Encoder, Decoder};
use tokio_io::codec::Framed;
use tokio_io::{AsyncRead, AsyncWrite};
//use bytes::{BytesMut, BufMut};
use bytes::BytesMut;

#[derive(Default)]
pub struct LineCodec;

impl Decoder for LineCodec {
    type Item = String;
    type Error = io::Error;

    fn decode(&mut self, buf: &mut BytesMut) -> io::Result<Option<String>> {
        if let Some(i) = buf.iter().position(|&b| b == b'\n') {
            let line = buf.split_to(i);
            buf.split_to(1);
            match str::from_utf8(&line) {
                Ok(s) => Ok(Some(s.to_string())),
                Err(_) => Err(io::Error::new(io::ErrorKind::Other,
                                             "invalid UTF-8")),
            }
        } else {
            Ok(None)
        }
    }
}

impl Encoder for LineCodec {
    type Item = String;
    type Error = io::Error;

    fn encode(&mut self, msg: String, buf: &mut BytesMut) -> io::Result<()> {
        buf.extend(msg.as_bytes());
        buf.extend(b"\n");
        Ok(())
    }
}

pub struct PingPong<T> {
    upstream: T,
    server: Option<String>,
}

impl<T> PingPong<T> {
    fn new(t: T) -> Self {
        PingPong {
            upstream: t,
            server: None,
        }
    }
}

impl<T> Stream for PingPong<T>
    where T: Stream<Item = String, Error = io::Error>,
          T: Sink<SinkItem = String, SinkError = io::Error>,
{
    type Item = String;
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Option<String>, io::Error> {
        // Poll the upstream transport
        match try_ready!(self.upstream.poll()) {
            Some(ref msg) if msg.starts_with("PING ") => {
                // Intercept pings
                self.server = Some(msg[5..].to_owned());
                // what's this ↓ ?
                // Try flushing the pong, only bubble up errors
                Ok(Async::Ready(try_ready!(self.upstream.poll())))
            }
            m => Ok(Async::Ready(m))
        }
    }
}

impl<T> Sink for PingPong<T>
    where T: Sink<SinkItem = String, SinkError = io::Error>,
{
    type SinkItem = String;
    type SinkError = io::Error;

    fn start_send(&mut self, item: String) -> StartSend<String, io::Error> {
        // Only accept the write if there are no pending pongs
        if self.server.is_some() {
            return Ok(AsyncSink::NotReady(item));
        }
        self.upstream.start_send(item)
    }

    fn poll_complete(&mut self) -> Poll<(), io::Error> {
        if let Some(s) = self.server.take() {
            let pong = format!("PONG {}", s);
            self.upstream.start_send(pong)?;
        }
        self.upstream.poll_complete()
    }
}

pub struct LineProto;

impl<T: AsyncRead + AsyncWrite + 'static> ClientProto<T> for LineProto {
    type Request = String;
    type Response = String;
    type Transport = PingPong<Framed<T, LineCodec>>;
    type BindTransport = Result<Self::Transport, io::Error>;
    fn bind_transport(&self, io: T) -> Self::BindTransport {
        Ok(PingPong::new(io.framed(LineCodec)))
    }
}

pub struct Echo;

impl Service for Echo {
    type Request = String;
    type Response = String;
    type Error = io::Error;
    type Future = BoxFuture<Self::Response, Self::Error>;

    fn call(&self, req: Self::Request) -> Self::Future {
        future::ok(req).boxed()
    }
}

fn main() {
    //let addr = "0.0.0.0:12345".parse().unwrap();
    //TcpServer::new(IntProto, addr)
    //    .serve(|| Ok(Doubler));
    let cm = "USER qj 0 * qjkx\r\nNICK qjk\r\nJOIN #test\r\n".to_string();

    //let addr = "irc.freenode.org:6697".to_socket_addrs().unwrap().next().unwrap();
    let addr = "0.0.0.0:12345".to_socket_addrs().unwrap().next().unwrap();

    //LineCodec.poll();
    //let lc = LineCodec;
    //lc.AsyncReadframed(LineCodec);

    let mut core = Core::new().unwrap();
    let handle = core.handle();
    //let cli = TcpClient::new(IntProto);
    let cli = TcpClient::new(LineProto);
    let socket = cli.connect(&addr, &handle);
    //let response = socket.and_then(|x| x.call(420));
    let response = socket.and_then(|x| x.call(cm));

    let data = core.run(response);//.unwrap();
    println!("{:?}", data);
    //let (_socket, data) = core.run(response).unwrap();
    //println!("{}", String::from_utf8_lossy(&data));
}

<|endoftext|>


<commit_before>extern crate argparse;
extern crate glob;
extern crate regex;

use argparse::{ArgumentParser, StoreTrue, Store, StoreOption, Print};
use std::io::BufReader;
use std::io::BufRead;
use std::fs::File;
use std::io::Write;
use glob::glob;
use regex::Regex;

fn identity<T>(x: T) -> T {x}

macro_rules! printerr(
    ($($arg:tt)*) => { {
        let r = writeln!(&mut ::std::io::stderr(), $($arg)*);
        r.expect("failed printing to stderr");
    } }
);

struct Options {
    path: String,
    file: Option<String>
}

fn get_possible_files_from_glob() -> Result<Vec<std::path::PathBuf>, glob::PatternError> {
    glob("**/*").map(|paths| paths.flat_map(identity).filter(|path| path.is_file()).collect())
}

fn get_filename_minus_extension(path_str: &String) -> String {
    std::path::Path::new(path_str).file_stem().unwrap().to_str().unwrap().to_string()
}

fn is_test_file(path: &String) -> bool {
    let re = Regex::new(r"^(features/|test/|spec/|tests/)").unwrap();
    re.is_match(path.as_str())
}

fn strip_test_words(filename: &String) -> String {
    let re = Regex::new(r"(test_)?(?P<p>\w+?)(_rake_spec|_spec|_test|_steps)?(\.rb|\.exs|\.ex|\.js|\.py)?$").unwrap();
    re.replace_all(filename.as_str(), "$p")
}

fn cleanse_path(path: &str) -> String {
    let s = path.to_string();
    if s[0..2].to_string() == "./" {
        s[2..].to_string()
    } else {
        s
    }
}

fn find_longest_common_substring_length(s1: &String, s2: &String) -> i32 {
    // Currently this is implemented using a dynamic programming solution similar
    // to http://www.geeksforgeeks.org/longest-common-substring/. This is O(N*M)
    // where N is the length of one string and M is the length of the other
    // string.
    //
    // Another option would of course be to explore using something like a
    // suffix tree to solve this problem, something like, the following.
    // http://www.geeksforgeeks.org/suffix-tree-application-5-longest-common-substring-2/
    // This is O(M+N) to build a Generalized Suffix Tree and O(M+N) to find the
    // the longest common substring via depth first search.
    //
    // Beyond that we would have to explore not caring about longest substring
    // and moving to a similarity ranking algorithm that maybe cares about
    // subsequences rather that substrings, etc.
    if s1.is_empty() || s2.is_empty() {
        return 0
    }

    let mut m: Vec<Vec<i32>> = Vec::with_capacity(s1.len());
    for _ in 0..s1.len() {
        let v: Vec<i32> = vec![0; s2.len()];
        m.push(v);
    }

    let mut longest_length = 0;

    let s1_bytes = s1.as_bytes();
    let s2_bytes = s2.as_bytes();

    for i in 0..s1.len() {
        for j in 0..s2.len() {
            if s1_bytes[i] == s2_bytes[j] {
                m[i][j] = 1;
                if i > 0 && j > 0 {
                    m[i][j] += m[i-1][j-1];
                }
                if m[i][j] > longest_length {
                    longest_length = m[i][j];
                }
            }
        }
    }

    longest_length
}

fn score(s1: &String, s2: &String) -> f32 {
    let longest_match_length: f32 = find_longest_common_substring_length(s1, s2) as f32;
    (longest_match_length/s2.len() as f32) * (longest_match_length/s1.len() as f32)
}

fn find_alt(filename: &String, cleansed_path: &String, paths: Vec<String>, test_file: bool) -> String {
    let (_, alternate) = paths.iter()
        .map(|path| cleanse_path(&path))
        .filter(|path| path.contains(filename.as_str()))  // filter to paths that contain the filename
        .filter(|path| is_test_file(&path) != test_file)
        .fold((0 as f32, "".to_string()), |result, path| {
            let (highest_score, best_match) = result;
            let s = score(&path, &cleansed_path);
            if s > highest_score {
                (s, path)
            } else {
                (highest_score, best_match)
            }
        });
    alternate
}

fn main() {
    let mut options = Options {
        path: "".to_string(),
        file: None
    };

    { // block limits of borrows by refer() method calls
        let mut ap = ArgumentParser::new();
        ap.add_option(&["-v", "--version"], Print(env!("CARGO_PKG_VERSION").to_string()), "show version");
        ap.refer(&mut options.file).add_option(&["-f", "--file"], StoreOption, "possible alternates file, - for stdin");
        ap.refer(&mut options.path).add_argument("PATH", Store, "path to find alternate for").required();
        ap.parse_args_or_exit();
    }

    let cleansed_path = cleanse_path(&options.path);
    let mut filename = get_filename_minus_extension(&options.path);
    if is_test_file(&cleansed_path) {
        filename = strip_test_words(&filename);
    }

    let best_match = if let Some(unwrapped_file) = options.file {
        if unwrapped_file == "-" {
            let stdin = std::io::stdin();
            let paths: Vec<String> = stdin.lock().lines().map(|path| path.unwrap()).collect();
            find_alt(&filename, &cleansed_path, paths, is_test_file(&cleansed_path))
        } else {
            let f = match File::open(&unwrapped_file) {
                Ok(file) => file,
                Err(e) => {
                    printerr!("Failure occurred opening file {}, {}", &unwrapped_file, e);
                    std::process::exit(1)
                }
            };
            let file = BufReader::new(&f);
            let paths: Vec<String> = file.lines().map(|path| path.unwrap()).collect();
            find_alt(&filename, &cleansed_path, paths, is_test_file(&cleansed_path))
        }
    } else {
        match get_possible_files_from_glob() {
            Ok(paths) => {
                let unwrapped_paths:Vec<String> = paths.iter().map(|path| { path.to_str().unwrap().to_string() }).collect();
                find_alt(&filename, &cleansed_path, unwrapped_paths, is_test_file(&cleansed_path))
            },
            Err(e) => {
                printerr!("Error reading paths {}", e);
                std::process::exit(1)
            }
        }
    };
    print!("{}", best_match);
}
<commit_msg>Removed unused import StoreTrue<commit_after>extern crate argparse;
extern crate glob;
extern crate regex;

use argparse::{ArgumentParser, Store, StoreOption, Print};
use std::io::BufReader;
use std::io::BufRead;
use std::fs::File;
use std::io::Write;
use glob::glob;
use regex::Regex;

fn identity<T>(x: T) -> T {x}

macro_rules! printerr(
    ($($arg:tt)*) => { {
        let r = writeln!(&mut ::std::io::stderr(), $($arg)*);
        r.expect("failed printing to stderr");
    } }
);

struct Options {
    path: String,
    file: Option<String>
}

fn get_possible_files_from_glob() -> Result<Vec<std::path::PathBuf>, glob::PatternError> {
    glob("**/*").map(|paths| paths.flat_map(identity).filter(|path| path.is_file()).collect())
}

fn get_filename_minus_extension(path_str: &String) -> String {
    std::path::Path::new(path_str).file_stem().unwrap().to_str().unwrap().to_string()
}

fn is_test_file(path: &String) -> bool {
    let re = Regex::new(r"^(features/|test/|spec/|tests/)").unwrap();
    re.is_match(path.as_str())
}

fn strip_test_words(filename: &String) -> String {
    let re = Regex::new(r"(test_)?(?P<p>\w+?)(_rake_spec|_spec|_test|_steps)?(\.rb|\.exs|\.ex|\.js|\.py)?$").unwrap();
    re.replace_all(filename.as_str(), "$p")
}

fn cleanse_path(path: &str) -> String {
    let s = path.to_string();
    if s[0..2].to_string() == "./" {
        s[2..].to_string()
    } else {
        s
    }
}

fn find_longest_common_substring_length(s1: &String, s2: &String) -> i32 {
    // Currently this is implemented using a dynamic programming solution similar
    // to http://www.geeksforgeeks.org/longest-common-substring/. This is O(N*M)
    // where N is the length of one string and M is the length of the other
    // string.
    //
    // Another option would of course be to explore using something like a
    // suffix tree to solve this problem, something like, the following.
    // http://www.geeksforgeeks.org/suffix-tree-application-5-longest-common-substring-2/
    // This is O(M+N) to build a Generalized Suffix Tree and O(M+N) to find the
    // the longest common substring via depth first search.
    //
    // Beyond that we would have to explore not caring about longest substring
    // and moving to a similarity ranking algorithm that maybe cares about
    // subsequences rather that substrings, etc.
    if s1.is_empty() || s2.is_empty() {
        return 0
    }

    let mut m: Vec<Vec<i32>> = Vec::with_capacity(s1.len());
    for _ in 0..s1.len() {
        let v: Vec<i32> = vec![0; s2.len()];
        m.push(v);
    }

    let mut longest_length = 0;

    let s1_bytes = s1.as_bytes();
    let s2_bytes = s2.as_bytes();

    for i in 0..s1.len() {
        for j in 0..s2.len() {
            if s1_bytes[i] == s2_bytes[j] {
                m[i][j] = 1;
                if i > 0 && j > 0 {
                    m[i][j] += m[i-1][j-1];
                }
                if m[i][j] > longest_length {
                    longest_length = m[i][j];
                }
            }
        }
    }

    longest_length
}

fn score(s1: &String, s2: &String) -> f32 {
    let longest_match_length: f32 = find_longest_common_substring_length(s1, s2) as f32;
    (longest_match_length/s2.len() as f32) * (longest_match_length/s1.len() as f32)
}

fn find_alt(filename: &String, cleansed_path: &String, paths: Vec<String>, test_file: bool) -> String {
    let (_, alternate) = paths.iter()
        .map(|path| cleanse_path(&path))
        .filter(|path| path.contains(filename.as_str()))  // filter to paths that contain the filename
        .filter(|path| is_test_file(&path) != test_file)
        .fold((0 as f32, "".to_string()), |result, path| {
            let (highest_score, best_match) = result;
            let s = score(&path, &cleansed_path);
            if s > highest_score {
                (s, path)
            } else {
                (highest_score, best_match)
            }
        });
    alternate
}

fn main() {
    let mut options = Options {
        path: "".to_string(),
        file: None
    };

    { // block limits of borrows by refer() method calls
        let mut ap = ArgumentParser::new();
        ap.add_option(&["-v", "--version"], Print(env!("CARGO_PKG_VERSION").to_string()), "show version");
        ap.refer(&mut options.file).add_option(&["-f", "--file"], StoreOption, "possible alternates file, - for stdin");
        ap.refer(&mut options.path).add_argument("PATH", Store, "path to find alternate for").required();
        ap.parse_args_or_exit();
    }

    let cleansed_path = cleanse_path(&options.path);
    let mut filename = get_filename_minus_extension(&options.path);
    if is_test_file(&cleansed_path) {
        filename = strip_test_words(&filename);
    }

    let best_match = if let Some(unwrapped_file) = options.file {
        if unwrapped_file == "-" {
            let stdin = std::io::stdin();
            let paths: Vec<String> = stdin.lock().lines().map(|path| path.unwrap()).collect();
            find_alt(&filename, &cleansed_path, paths, is_test_file(&cleansed_path))
        } else {
            let f = match File::open(&unwrapped_file) {
                Ok(file) => file,
                Err(e) => {
                    printerr!("Failure occurred opening file {}, {}", &unwrapped_file, e);
                    std::process::exit(1)
                }
            };
            let file = BufReader::new(&f);
            let paths: Vec<String> = file.lines().map(|path| path.unwrap()).collect();
            find_alt(&filename, &cleansed_path, paths, is_test_file(&cleansed_path))
        }
    } else {
        match get_possible_files_from_glob() {
            Ok(paths) => {
                let unwrapped_paths:Vec<String> = paths.iter().map(|path| { path.to_str().unwrap().to_string() }).collect();
                find_alt(&filename, &cleansed_path, unwrapped_paths, is_test_file(&cleansed_path))
            },
            Err(e) => {
                printerr!("Error reading paths {}", e);
                std::process::exit(1)
            }
        }
    };
    print!("{}", best_match);
}
<|endoftext|>