<commit_before>extern crate libc;
extern crate clap;

mod iopoll;

#[macro_use]
mod util;

use clap::{Arg, App};
use iopoll::{Token, EventLoop, EventSet, EventHandler};

use std::io::prelude::*;
use std::io::{stdin, stdout};
use std::net::{TcpStream, TcpListener, UdpSocket, SocketAddr};
use std::process::exit;
use std::time::Duration;
use std::str::FromStr;

#[derive(Default,Debug)]
struct ProgramOptions {
    ipv4_only: bool,
    ipv6_only: bool,
    use_listen: bool,
    detach_stdin: bool,
    use_unix: bool,
    use_udp: bool,
    interval_secs: u32,
    zero_io_mode: bool,
    hostname: String,
    source_port: u16,
    target_port: u16,
    verbosity: u8,
    wait_time: Option<u32>,
}

enum NetworkConnection {
    TcpClient(TcpStream),
    UdpClient(UdpSocket),
}

fn parse_commandline() -> ProgramOptions {
    let matches = App::new("netcat")
        .version("0.1.0")
        .arg(Arg::with_name("ipv4").short("4").help("Forces use of IPv4 addresses only."))
        .arg(Arg::with_name("ipv6").short("6").help("Forces use of IPv6 addresses only."))
        .arg(Arg::with_name("udp").short("u").long("udp").help("UDP mode"))
        .arg(Arg::with_name("hostname").required(false))
        .arg(Arg::with_name("target-port").required(false))
        .arg(Arg::with_name("verbose")
            .short("v")
            .multiple(true)
            .help("Set verbosity level (can be used several times)"))
        .arg(Arg::with_name("no-dns")
            .short("n")
            .long("nodns")
            .help("Suppress name/port resolutions"))
        .arg(Arg::with_name("listen")
            .short("l")
            .long("listen")
            .help("Listen mode, for inbound connects"))
        .arg(Arg::with_name("zero-io")
            .short("z")
            .help("Zero-I/O mode [used for scanning]"))
        .arg(Arg::with_name("wait-time")
            .short("w")
            .value_name("secs")
            .help("Timeout for connects and final net reads"))
        .arg(Arg::with_name("detach-stdin").short("d").help("Detach from stdin"))
        .arg(Arg::with_name("source-port").short("p").long("source-port").value_name("port"))
        .get_matches();

    let mut options: ProgramOptions = Default::default();
    options.hostname = matches.value_of("hostname").unwrap_or("localhost").to_string();
    options.source_port =
        matches.value_of("source-port").map_or(Ok(0), |v| u16::from_str(v)).unwrap_or(0);
    options.target_port =
        matches.value_of("target-port").map_or(Ok(0), |v| u16::from_str(v)).unwrap_or(0);
    options.use_udp = matches.is_present("udp");
    options.ipv4_only = matches.is_present("ipv4");
    options.ipv6_only = matches.is_present("ipv6");
    options.verbosity = matches.occurrences_of("verbose") as u8;
    if let Some(wait_time) = matches.value_of("wait-time") {
        options.wait_time = Some(u32::from_str(wait_time).expect("Invalid wait time"));
    }
    options
}

pub struct NetcatClientEventHandler {
    stdin_open: bool,
    network_open: bool,
    network_client: NetworkConnection,
}

impl NetcatClientEventHandler {
    fn new(network_connection: NetworkConnection) -> NetcatClientEventHandler {
        NetcatClientEventHandler {
            stdin_open: true,
            network_open: true,
            network_client: network_connection,
        }
    }
}

impl EventHandler for NetcatClientEventHandler {
    fn ready_for_io(&mut self, event_loop: &mut EventLoop, token: Token, eventset: EventSet) {
        let mut shutdown_loop = false;
        if eventset.is_readable() {
            if token == Token(0) {
                if self.stdin_open {
                    let mut stdin = stdin();
                    let mut buf = [0; 1024];
                    match stdin.read(&mut buf) {
                        Ok(n) => {
                            //println!("stdin: {} bytes read", n);
                            match self.network_client {
                                NetworkConnection::TcpClient(ref mut tcpstream) => {
                                    tcpstream.write(&buf[0..n]).unwrap();
                                }
                                NetworkConnection::UdpClient(ref mut udpsocket) => {
                                    match udpsocket.send(&buf[..n]) {
                                        Ok(sent) => {
                                            shutdown_loop = sent != n;
                                        }
                                        Err(_) => shutdown_loop = true,
                                    }
                                }
                            }
                            if n == 0 {
                                self.stdin_open = false;
                                event_loop.unregister_stdin();
                                if let NetworkConnection::TcpClient(ref mut tcpstream) =
                                       self.network_client {
                                    tcpstream.shutdown(std::net::Shutdown::Write);
                                }
                            }
                        }
                        Err(_) => {
                            shutdown_loop = true;
                        }
                    }
                }
            } else {
                let mut stdout = stdout();
                let mut buf = [0; 1024];
                if let Ok(n) = match self.network_client {
                    NetworkConnection::TcpClient(ref mut tcpstream) => tcpstream.read(&mut buf),
                    NetworkConnection::UdpClient(ref mut udpsocket) => udpsocket.recv(&mut buf),
                } {
                    stdout.write(&buf[0..n]).unwrap();
                    if n == 0 {
                        self.network_open = false;
                        if let NetworkConnection::TcpClient(ref mut tcpstream) =
                               self.network_client {
                            tcpstream.shutdown(std::net::Shutdown::Read);
                            shutdown_loop = true;
                        }

                    }
                }
            }
        }

        if shutdown_loop {
            event_loop.shutdown();
        }
    }

    fn timeout(&mut self, eventloop: &mut EventLoop ) {
        eventloop.shutdown();
    }
}

fn main() {
    let options = parse_commandline();
    if options.verbosity > 1 {
        println_stderr!("options: {:?}", options);
    }

    let mut exit_code = 0;

    let stdin = stdin();
    let connection;
    let mut eventloop = EventLoop::new(options.wait_time);
    if !options.detach_stdin {
        eventloop.register_stdin(&stdin);
    }

    if options.use_udp {
        let sock = UdpSocket::bind(("127.0.0.1", options.source_port)).unwrap();
        sock.connect((options.hostname.as_str(), options.target_port)).unwrap();
        if let Some(timeout) = options.wait_time {
            sock.set_read_timeout(Some(Duration::new(timeout as u64, 0)));
        }
        eventloop.register_read(&sock);
        connection = NetworkConnection::UdpClient(sock);
    } else {
        let tcpstream = TcpStream::connect((options.hostname.as_str(), options.target_port))
            .unwrap();
        if let Some(timeout) = options.wait_time {
            tcpstream.set_read_timeout(Some(Duration::new(timeout as u64, 0)));
        }
        eventloop.register_read(&tcpstream);
        connection = NetworkConnection::TcpClient(tcpstream);
    }

    let mut eh = NetcatClientEventHandler::new(connection);
    if let Err(err) = eventloop.run(&mut eh) {
        exit_code = 1;
        if options.verbosity > 0 {
            println_stderr!("{}", err);
        }
    }
    exit(exit_code);
}
<commit_msg>fix unhandled result compiler warnings<commit_after>extern crate libc;
extern crate clap;

mod iopoll;

#[macro_use]
mod util;

use clap::{Arg, App};
use iopoll::{Token, EventLoop, EventSet, EventHandler};

use std::io::prelude::*;
use std::io::{stdin, stdout};
use std::net::{TcpStream, TcpListener, UdpSocket, SocketAddr};
use std::process::exit;
use std::time::Duration;
use std::str::FromStr;

#[derive(Default,Debug)]
struct ProgramOptions {
    ipv4_only: bool,
    ipv6_only: bool,
    use_listen: bool,
    detach_stdin: bool,
    use_unix: bool,
    use_udp: bool,
    interval_secs: u32,
    zero_io_mode: bool,
    hostname: String,
    source_port: u16,
    target_port: u16,
    verbosity: u8,
    wait_time_ms: Option<u32>,
}

enum NetworkConnection {
    TcpClient(TcpStream),
    UdpClient(UdpSocket),
}

fn parse_commandline() -> ProgramOptions {
    let matches = App::new("netcat")
        .version("0.1.0")
        .arg(Arg::with_name("ipv4").short("4").help("Forces use of IPv4 addresses only."))
        .arg(Arg::with_name("ipv6").short("6").help("Forces use of IPv6 addresses only."))
        .arg(Arg::with_name("udp").short("u").long("udp").help("UDP mode"))
        .arg(Arg::with_name("hostname").required(false))
        .arg(Arg::with_name("target-port").required(false))
        .arg(Arg::with_name("verbose")
            .short("v")
            .multiple(true)
            .help("Set verbosity level (can be used several times)"))
        .arg(Arg::with_name("no-dns")
            .short("n")
            .long("nodns")
            .help("Suppress name/port resolutions"))
        .arg(Arg::with_name("listen")
            .short("l")
            .long("listen")
            .help("Listen mode, for inbound connects"))
        .arg(Arg::with_name("zero-io")
            .short("z")
            .help("Zero-I/O mode [used for scanning]"))
        .arg(Arg::with_name("wait-time")
            .short("w")
            .value_name("secs")
            .help("Timeout for connects and final net reads"))
        .arg(Arg::with_name("detach-stdin").short("d").help("Detach from stdin"))
        .arg(Arg::with_name("source-port").short("p").long("source-port").value_name("port"))
        .get_matches();

    let mut options: ProgramOptions = Default::default();
    options.hostname = matches.value_of("hostname").unwrap_or("localhost").to_string();
    options.source_port =
        matches.value_of("source-port").map_or(Ok(0), |v| u16::from_str(v)).unwrap_or(0);
    options.target_port =
        matches.value_of("target-port").map_or(Ok(0), |v| u16::from_str(v)).unwrap_or(0);
    options.use_udp = matches.is_present("udp");
    options.ipv4_only = matches.is_present("ipv4");
    options.ipv6_only = matches.is_present("ipv6");
    options.verbosity = matches.occurrences_of("verbose") as u8;
    if let Some(wait_time) = matches.value_of("wait-time") {
        options.wait_time_ms = Some(u32::from_str(wait_time).expect("Invalid wait time") * 1000);
    }
    options
}

pub struct NetcatClientEventHandler {
    stdin_open: bool,
    network_open: bool,
    network_client: NetworkConnection,
}

impl NetcatClientEventHandler {
    fn new(network_connection: NetworkConnection) -> NetcatClientEventHandler {
        NetcatClientEventHandler {
            stdin_open: true,
            network_open: true,
            network_client: network_connection,
        }
    }
}

impl EventHandler for NetcatClientEventHandler {
    fn ready_for_io(&mut self, event_loop: &mut EventLoop, token: Token, eventset: EventSet) {
        let mut shutdown_loop = false;
        if eventset.is_readable() {
            if token == Token(0) {
                if self.stdin_open {
                    let mut stdin = stdin();
                    let mut buf = [0; 1024];
                    match stdin.read(&mut buf) {
                        Ok(n) => {
                            // println!("stdin: {} bytes read", n);
                            match self.network_client {
                                NetworkConnection::TcpClient(ref mut tcpstream) => {
                                    tcpstream.write(&buf[0..n]).unwrap();
                                }
                                NetworkConnection::UdpClient(ref mut udpsocket) => {
                                    match udpsocket.send(&buf[..n]) {
                                        Ok(sent) => {
                                            shutdown_loop = sent != n;
                                        }
                                        Err(_) => shutdown_loop = true,
                                    }
                                }
                            }
                            if n == 0 {
                                self.stdin_open = false;
                                event_loop.unregister_stdin();
                                if let NetworkConnection::TcpClient(ref mut tcpstream) =
                                       self.network_client {
                                    tcpstream.shutdown(std::net::Shutdown::Write).unwrap(); // TODO
                                }
                            }
                        }
                        Err(_) => {
                            shutdown_loop = true;
                        }
                    }
                }
            } else {
                let mut stdout = stdout();
                let mut buf = [0; 1024];
                if let Ok(n) = match self.network_client {
                    NetworkConnection::TcpClient(ref mut tcpstream) => tcpstream.read(&mut buf),
                    NetworkConnection::UdpClient(ref mut udpsocket) => udpsocket.recv(&mut buf),
                } {
                    stdout.write(&buf[0..n]).unwrap();
                    if n == 0 {
                        shutdown_loop = true;
                        self.network_open = false;
                    }
                }
            }
        }

        if shutdown_loop {
            event_loop.shutdown();
        }
    }

    fn timeout(&mut self, eventloop: &mut EventLoop) {
        eventloop.shutdown();
    }
}

fn main() {
    let options = parse_commandline();
    if options.verbosity > 1 {
        println_stderr!("options: {:?}", options);
    }

    let mut exit_code = 0;

    let stdin = stdin();
    let connection;
    let mut eventloop = EventLoop::new(options.wait_time_ms);
    if !options.detach_stdin {
        eventloop.register_stdin(&stdin);
    }

    if options.use_udp {
        let sock = UdpSocket::bind(("127.0.0.1", options.source_port)).unwrap();
        sock.connect((options.hostname.as_str(), options.target_port)).unwrap();
        if let Some(timeout) = options.wait_time_ms {
            sock.set_read_timeout(Some(Duration::new(timeout as u64, 0))).unwrap(); // TODO
        }
        eventloop.register_read(&sock);
        connection = NetworkConnection::UdpClient(sock);
    } else {
        let tcpstream = TcpStream::connect((options.hostname.as_str(), options.target_port))
            .unwrap();
        if let Some(timeout) = options.wait_time_ms {
            tcpstream.set_read_timeout(Some(Duration::new(timeout as u64, 0))).unwrap(); // TODO
        }
        eventloop.register_read(&tcpstream);
        connection = NetworkConnection::TcpClient(tcpstream);
    }

    let mut eh = NetcatClientEventHandler::new(connection);
    if let Err(err) = eventloop.run(&mut eh) {
        exit_code = 1;
        if options.verbosity > 0 {
            println_stderr!("{}", err);
        }
    }
    exit(exit_code);
}
<|endoftext|>


<commit_before>extern crate hyper;

use std::io::{TcpListener, TcpStream};
use std::io::{Acceptor, Listener};

fn handle_client(mut stream: TcpStream) {
    stream.set_timeout(Some(10));
    match stream.read_to_string() {
        Err(e) => { println!("Error reading request: {}", e); }
        Ok(req) => { println!("The reqeust was: {}", req); }
    }

    let resp = b"HTTP/1.1 200 OK\r\nContent-Type: text/html; Connection: close;\r\n\r\nHello!\n";
    match stream.write(resp) {
        Err(e) => { println!("Error writing response: {}", e); }
        Ok(_) => {}
    }
}

fn main() {
    let listener = TcpListener::bind("127.0.0.1:42007");

    let mut acceptor = listener.listen();

    for stream in acceptor.incoming() {
        match stream {
            Err(e) => { println!("Error in stream: {}", e); }
            Ok(stream) => { handle_client(stream) }
        }
    }

    drop(acceptor);
}
<commit_msg>working on reading the request<commit_after>extern crate hyper;

use std::io::{TcpListener, TcpStream};
use std::io::{Acceptor, Listener};

fn handle_client(mut stream: TcpStream) {
    const SIZE: usize = 1024;
    let mut buf: Vec<u8> = Vec::with_capacity(SIZE);
    match stream.read(buf.as_mut_slice()) {
        Err(e) => { println!("Error reading request: {}", e); }
        Ok(_) => {
            match std::str::from_utf8(buf.as_mut_slice()) {
                Err(e) => { println!("Error making buffer: {}", e); }
                Ok(req) => { println!("Here is what I got: {}", req); }
            }
        }
    }

    let resp = b"HTTP/1.1 200 OK\r\nContent-Type: text/html; Connection: close;\r\n\r\nHello!\n";
    match stream.write(resp) {
        Err(e) => { println!("Error writing response: {}", e); }
        Ok(_) => {}
    }
}

fn main() {
    let listener = TcpListener::bind("127.0.0.1:42007");

    let mut acceptor = listener.listen();

    for stream in acceptor.incoming() {
        match stream {
            Err(e) => { println!("Error in stream: {}", e); }
            Ok(stream) => { handle_client(stream) }
        }
    }

    drop(acceptor);
}
<|endoftext|>


<commit_before>#![feature(decode_utf8)]
#![feature(io)]

extern crate argparse;
extern crate bit_set;
extern crate byteorder;
extern crate compress;
extern crate libc;

use std::fs;
use std::io;
use std::mem;
use std::path;
use std::slice;

use argparse::{Store, StoreTrue};

use bit_set::BitSet;

use byteorder::{BigEndian, ReadBytesExt};

use compress::lz4;

use libc::c_void;

//magic:
use std::io::Read;
use std::io::Seek;
use std::os::unix::io::AsRawFd;

type CharResult = Result<char, io::CharsError>;

static TRI_MAX: usize = 64 * 64 * 64;

fn simplify(wut: char) -> u8 {
    let c = match wut {
        'a' ... 'z' => (wut as u8 - 'a' as u8 + 'A' as u8) as char,
        _ => wut,
    };

    if c > 128 as char {
        if c.is_whitespace() {
            return 2;
        }
        if c.is_control() {
            return 0;
        }
        return 63;
    }

    let sym_end = 33u8;
    let letters = 21u8;

    match c {
        '\r' | '\n' => 1,
        '\t' | '\x0c' | '\x0b' | ' ' => 2,
        '!' => 3,
        '"' | '\'' | '`' => 4,
        '$' => 5,
        '%' => 6,
        '&' => 7,
        '(' ... '@' => 8 + (c as u8 - '(' as u8),
        'A' ... 'I' => sym_end + (c as u8 - 'A' as u8),
        'J' | 'K' => sym_end + letters,
        'L' ... 'P' => sym_end + (c as u8 - 'A' as u8) - 2,
        'Q' => sym_end + letters,
        'R' ... 'W' => sym_end + (c as u8 - 'A' as u8) - 3,
        'X' | 'Z' => sym_end + letters,
        'Y' => sym_end + letters - 1,
        '[' => 55,
        '\\' => 56,
        ']'=> 57,
        '^' | '~' | '#' => 58,
        '_' => 59,
        '{' ... '}' => 60 + (c as u8 - '{' as u8),
        _ => 0,
    }
}

fn explain(wut: u8) -> char {
    match wut {
        0 => 'X',
        1 => 'N',
        2 => ' ',
        3 => '!',
        4 => '"',
        5 => '$',
        6 => '%',
        7 => '&',
        8 ... 32 => (wut - 8 + '(' as u8) as char,
        33 ... 41 => (wut - 33 + 'a' as u8) as char,
        42 ... 46 => (wut - 42 + 'l' as u8) as char,
        47 ... 52 => (wut - 47 + 'r' as u8) as char,
        53 => 'y',
        54 => 'X',
        55 => '[',
        56 => '\\',
        57 => ']',
        58 => '#',
        59 => '_',
        60 ... 62 => (wut - 60 + '{' as u8) as char,
        63 => 'U',
        _ => 'D',
    }
}

fn unpack(wut: usize) -> String {
    let mut ret = String::with_capacity(3);
    ret.push(explain((wut / 64 / 64 % 64) as u8));
    ret.push(explain((wut / 64 % 64) as u8));
    ret.push(explain((wut % 64) as u8));
    return ret;
}

fn trigrams_for<T: Iterator<Item=CharResult>>(input: T) -> Result<BitSet, String> {
    let mut line: u64 = 1;
    let mut prev: [u8; 3] = [0; 3];
    let mut ret: BitSet = BitSet::with_capacity(64 * 64 * 64);

    for (off, maybe_char) in input.enumerate() {
        let c = try!(maybe_char.map_err(|e| {
            format!("line {}: file char {}: failed: {}", line, off, e)
        }));
        if '\n' == c {
            line += 1;
        }
        if '\0' == c {
            return Err(format!("line {}: null found: not a text file", line));
        }
        prev[0] = prev[1];
        prev[1] = prev[2];
        prev[2] = simplify(c);
        let tri: usize = 64 * 64 * prev[0] as usize + 64 * prev[1] as usize + prev[2] as usize;
        ret.insert(tri);
    }
    return Ok(ret);
}

struct Mapped<'a, T: 'a> {
    file: fs::File,
    map: *mut c_void,
    data: &'a mut [T],
}

impl <'a, T: 'a> Mapped<'a, T> {
    fn fixed_len<P>(path: P, len: usize) -> io::Result<Mapped<'a, T>>
        where P: AsRef<path::Path> {
        let file = fs::OpenOptions::new().read(true).write(true).create(true).open(path)?;
        file.set_len(mem::size_of::<T>() as u64 * len as u64)?;
        let map: *mut c_void = unsafe {
            libc::mmap(0 as *mut c_void,
                       len * mem::size_of::<T>(),
                       libc::PROT_READ | libc::PROT_WRITE,
                       libc::MAP_SHARED,
                       file.as_raw_fd(),
                       0)
        };

        if libc::MAP_FAILED == map {
            return Err(io::Error::last_os_error());
        }

        let data = unsafe { slice::from_raw_parts_mut(map as *mut T, len) };
        Ok(Mapped { file, map, data })
    }

    fn remap(&mut self, len: usize) -> io::Result<()> {
        self.file.set_len(mem::size_of::<T>() as u64 * len as u64)?;
        let new_map = unsafe {
            libc::mremap(self.map,
                         self.data.len() * mem::size_of::<T>(),
                         len * mem::size_of::<T>(),
                         libc::MREMAP_MAYMOVE)
        };

        if libc::MAP_FAILED == new_map {
            return Err(io::Error::last_os_error());
        }

        self.data = unsafe { slice::from_raw_parts_mut(new_map as *mut T, len) };
        self.map = new_map;
        Ok(())
    }
}

impl <'a, T: 'a> Drop for Mapped<'a, T> {
    fn drop(&mut self) {
        unsafe {
            assert_eq!(0, libc::munmap(self.map, self.data.len() * mem::size_of::<T>()));
        }
    }
}

struct Index<'a> {
    idx: Mapped<'a, u32>,
    pages: Mapped<'a, u64>,
    free_page: usize,
    page_size: usize,
}

impl<'a> Index<'a> {
    fn new() -> io::Result<Index<'a>> {
        let idx: Mapped<u32> = Mapped::fixed_len("idx", TRI_MAX).unwrap();

        let page_size: usize = 1024;

        let pages_len: usize = match fs::metadata("pages") {
            Ok(m) => {
                let proposed: u64 = m.len() / mem::size_of::<u64>() as u64;
                assert!(proposed < usize::max_value() as u64);
                proposed as usize
            },
            Err(e) => if e.kind() == io::ErrorKind::NotFound {
                2 * page_size
            } else {
                panic!("couldn't get info on pages file: {}", e)
            }
        };

        let pages: Mapped<u64> = Mapped::fixed_len("pages", pages_len)?;
        let mut free_page: usize = pages.data.len() / page_size;

        loop {
            if 0 != pages.data[(free_page - 1) * page_size] {
                break;
            }
            if 1 == free_page {
                break;
            }
            free_page -= 1;
        }

        Ok(Index { idx, pages, free_page, page_size })
    }

    fn page_for(&mut self, trigram: u32) -> io::Result<usize> {
        assert!(trigram < TRI_MAX as u32);

        let page = self.idx.data[trigram as usize] as usize;
        if 0 != page {
            return Ok(page);
        }

        let found_page = self.next_page()?;
        self.idx.data[trigram as usize] = found_page as u32;
        Ok(found_page)
    }

    fn next_page(&mut self) -> io::Result<usize> {
        let ret = self.free_page;
        self.free_page += 1;
        if self.free_page >= self.pages.data.len() / self.page_size {
            let old_len = self.pages.data.len();
            self.pages.remap(old_len + 100 * self.page_size)?;
        }

        Ok(ret)
    }

    fn append(&mut self, trigram: u32, document: u64) -> io::Result<()> {
        let mut page = self.page_for(trigram)?;
        let mut header_loc;
        let mut header;
        loop {
            header_loc = page * self.page_size;
            header = self.pages.data[header_loc];
            if header == self.page_size as u64 {
                page = self.next_page()?;
                self.pages.data[header_loc] = page as u64 + self.page_size as u64;
                header = 0;
                header_loc = page * self.page_size;
                break;
            } else if header > self.page_size as u64 {
                page = header as usize - self.page_size;
            } else {
                break;
            }
        }
        self.pages.data[header_loc] += 1;
        self.pages.data[page * self.page_size + 1 + header as usize] = document;
        Ok(())
    }

    fn append_trigrams(&mut self, trigrams: BitSet, document: u64) -> io::Result<()> {
        for found in trigrams.iter() {
            self.append(found as u32, document)?;
        }
        Ok(())
    }
}

fn round_up(x: u64) -> u64 {
    let mut ret = x;
    loop {
        if ret % 16 == 0 {
            return ret;
        }
        ret += 1;
    }
}

fn eat_chunk(mut fh: &mut fs::File) -> io::Result<BitSet> {
    let end = fh.read_u64::<BigEndian>()?;
    let extra_len = fh.read_u64::<BigEndian>()?;
    let start = fh.seek(io::SeekFrom::Current(extra_len as i64))?;
    let ret = {
        let decoder = lz4::Decoder::new(&mut fh);
        let range = decoder.bytes();
        let exploding = range.map(|x| x.unwrap());
        let decoder = std::char::decode_utf8(exploding);
        let errors = decoder.map(|x| x.map_err(|_| io::CharsError::NotUtf8));
        trigrams_for(errors).map_err(|msg| io::Error::new(io::ErrorKind::Other, msg))
    };
    let next = round_up(start + end - extra_len - 16);
    fh.seek(io::SeekFrom::Start(next))?;
    ret
}

fn main() {
    let mut from: String = "".to_string();
    let mut simple = false;
    let mut addendum: u64 = 0;
    {
        let mut ap = argparse::ArgumentParser::new();
        ap.set_description("totally not a load of tools glued together");
        ap.refer(&mut from)
                .required()
                .add_option(&["-f", "--input-file"], Store,
                            "pack file to read");
        ap.refer(&mut addendum)
                .add_option(&["-i", "--addendum"], Store,
                            "number to add to file offset");
        ap.refer(&mut simple)
                .add_option(&["--simple"], StoreTrue,
                            "not a pack, just a normal decompressed file");
        ap.parse_args_or_exit();
    }

    let mut idx = Index::new().unwrap();

    let mut fh = fs::File::open(from).expect("input file must exist and be readable");
    if simple {
        let trigrams = trigrams_for(fh.chars()).expect("trigramming must work");
        idx.append_trigrams(trigrams, addendum).unwrap();
        return;
    }

    fh.seek(io::SeekFrom::Start(16)).unwrap();
    loop {

        let document = fh.seek(io::SeekFrom::Current(0)).unwrap() + addendum;
        let trigrams = eat_chunk(&mut fh).unwrap();
        idx.append_trigrams(trigrams, document).unwrap();
    }
}
<commit_msg>dumb off-by-one<commit_after>#![feature(decode_utf8)]
#![feature(io)]

extern crate argparse;
extern crate bit_set;
extern crate byteorder;
extern crate compress;
extern crate libc;

use std::fs;
use std::io;
use std::mem;
use std::path;
use std::slice;

use argparse::{Store, StoreTrue};

use bit_set::BitSet;

use byteorder::{BigEndian, ReadBytesExt};

use compress::lz4;

use libc::c_void;

//magic:
use std::io::Read;
use std::io::Seek;
use std::os::unix::io::AsRawFd;

type CharResult = Result<char, io::CharsError>;

static TRI_MAX: usize = 64 * 64 * 64;

fn simplify(wut: char) -> u8 {
    let c = match wut {
        'a' ... 'z' => (wut as u8 - 'a' as u8 + 'A' as u8) as char,
        _ => wut,
    };

    if c > 128 as char {
        if c.is_whitespace() {
            return 2;
        }
        if c.is_control() {
            return 0;
        }
        return 63;
    }

    let sym_end = 33u8;
    let letters = 21u8;

    match c {
        '\r' | '\n' => 1,
        '\t' | '\x0c' | '\x0b' | ' ' => 2,
        '!' => 3,
        '"' | '\'' | '`' => 4,
        '$' => 5,
        '%' => 6,
        '&' => 7,
        '(' ... '@' => 8 + (c as u8 - '(' as u8),
        'A' ... 'I' => sym_end + (c as u8 - 'A' as u8),
        'J' | 'K' => sym_end + letters,
        'L' ... 'P' => sym_end + (c as u8 - 'A' as u8) - 2,
        'Q' => sym_end + letters,
        'R' ... 'W' => sym_end + (c as u8 - 'A' as u8) - 3,
        'X' | 'Z' => sym_end + letters,
        'Y' => sym_end + letters - 1,
        '[' => 55,
        '\\' => 56,
        ']'=> 57,
        '^' | '~' | '#' => 58,
        '_' => 59,
        '{' ... '}' => 60 + (c as u8 - '{' as u8),
        _ => 0,
    }
}

fn explain(wut: u8) -> char {
    match wut {
        0 => 'X',
        1 => 'N',
        2 => ' ',
        3 => '!',
        4 => '"',
        5 => '$',
        6 => '%',
        7 => '&',
        8 ... 32 => (wut - 8 + '(' as u8) as char,
        33 ... 41 => (wut - 33 + 'a' as u8) as char,
        42 ... 46 => (wut - 42 + 'l' as u8) as char,
        47 ... 52 => (wut - 47 + 'r' as u8) as char,
        53 => 'y',
        54 => 'X',
        55 => '[',
        56 => '\\',
        57 => ']',
        58 => '#',
        59 => '_',
        60 ... 62 => (wut - 60 + '{' as u8) as char,
        63 => 'U',
        _ => 'D',
    }
}

fn unpack(wut: usize) -> String {
    let mut ret = String::with_capacity(3);
    ret.push(explain((wut / 64 / 64 % 64) as u8));
    ret.push(explain((wut / 64 % 64) as u8));
    ret.push(explain((wut % 64) as u8));
    return ret;
}

fn trigrams_for<T: Iterator<Item=CharResult>>(input: T) -> Result<BitSet, String> {
    let mut line: u64 = 1;
    let mut prev: [u8; 3] = [0; 3];
    let mut ret: BitSet = BitSet::with_capacity(64 * 64 * 64);

    for (off, maybe_char) in input.enumerate() {
        let c = try!(maybe_char.map_err(|e| {
            format!("line {}: file char {}: failed: {}", line, off, e)
        }));
        if '\n' == c {
            line += 1;
        }
        if '\0' == c {
            return Err(format!("line {}: null found: not a text file", line));
        }
        prev[0] = prev[1];
        prev[1] = prev[2];
        prev[2] = simplify(c);
        let tri: usize = 64 * 64 * prev[0] as usize + 64 * prev[1] as usize + prev[2] as usize;
        ret.insert(tri);
    }
    return Ok(ret);
}

struct Mapped<'a, T: 'a> {
    file: fs::File,
    map: *mut c_void,
    data: &'a mut [T],
}

impl <'a, T: 'a> Mapped<'a, T> {
    fn fixed_len<P>(path: P, len: usize) -> io::Result<Mapped<'a, T>>
        where P: AsRef<path::Path> {
        let file = fs::OpenOptions::new().read(true).write(true).create(true).open(path)?;
        file.set_len(mem::size_of::<T>() as u64 * len as u64)?;
        let map: *mut c_void = unsafe {
            libc::mmap(0 as *mut c_void,
                       len * mem::size_of::<T>(),
                       libc::PROT_READ | libc::PROT_WRITE,
                       libc::MAP_SHARED,
                       file.as_raw_fd(),
                       0)
        };

        if libc::MAP_FAILED == map {
            return Err(io::Error::last_os_error());
        }

        let data = unsafe { slice::from_raw_parts_mut(map as *mut T, len) };
        Ok(Mapped { file, map, data })
    }

    fn remap(&mut self, len: usize) -> io::Result<()> {
        self.file.set_len(mem::size_of::<T>() as u64 * len as u64)?;
        let new_map = unsafe {
            libc::mremap(self.map,
                         self.data.len() * mem::size_of::<T>(),
                         len * mem::size_of::<T>(),
                         libc::MREMAP_MAYMOVE)
        };

        if libc::MAP_FAILED == new_map {
            return Err(io::Error::last_os_error());
        }

        self.data = unsafe { slice::from_raw_parts_mut(new_map as *mut T, len) };
        self.map = new_map;
        Ok(())
    }
}

impl <'a, T: 'a> Drop for Mapped<'a, T> {
    fn drop(&mut self) {
        unsafe {
            assert_eq!(0, libc::munmap(self.map, self.data.len() * mem::size_of::<T>()));
        }
    }
}

struct Index<'a> {
    idx: Mapped<'a, u32>,
    pages: Mapped<'a, u64>,
    free_page: usize,
    page_size: usize,
}

impl<'a> Index<'a> {
    fn new() -> io::Result<Index<'a>> {
        let idx: Mapped<u32> = Mapped::fixed_len("idx", TRI_MAX).unwrap();

        let page_size: usize = 1024;

        let pages_len: usize = match fs::metadata("pages") {
            Ok(m) => {
                let proposed: u64 = m.len() / mem::size_of::<u64>() as u64;
                assert!(proposed < usize::max_value() as u64);
                proposed as usize
            },
            Err(e) => if e.kind() == io::ErrorKind::NotFound {
                2 * page_size
            } else {
                panic!("couldn't get info on pages file: {}", e)
            }
        };

        let pages: Mapped<u64> = Mapped::fixed_len("pages", pages_len)?;
        let mut free_page: usize = pages.data.len() / page_size;

        loop {
            if 0 != pages.data[(free_page - 1) * page_size] {
                break;
            }
            if 1 == free_page {
                break;
            }
            free_page -= 1;
        }

        Ok(Index { idx, pages, free_page, page_size })
    }

    fn page_for(&mut self, trigram: u32) -> io::Result<usize> {
        assert!(trigram < TRI_MAX as u32);

        let page = self.idx.data[trigram as usize] as usize;
        if 0 != page {
            return Ok(page);
        }

        let found_page = self.next_page()?;
        self.idx.data[trigram as usize] = found_page as u32;
        Ok(found_page)
    }

    fn next_page(&mut self) -> io::Result<usize> {
        let ret = self.free_page;
        self.free_page += 1;
        if self.free_page >= self.pages.data.len() / self.page_size {
            let old_len = self.pages.data.len();
            self.pages.remap(old_len + 100 * self.page_size)?;
        }

        Ok(ret)
    }

    fn append(&mut self, trigram: u32, document: u64) -> io::Result<()> {
        let mut page = self.page_for(trigram)?;
        let mut header_loc;
        let mut header;
        loop {
            header_loc = page * self.page_size;
            header = self.pages.data[header_loc];
            if header == (self.page_size - 1) as u64 {
                page = self.next_page()?;
                self.pages.data[header_loc] = page as u64 + self.page_size as u64;
                header = 0;
                header_loc = page * self.page_size;
                break;
            } else if header >= self.page_size as u64 {
                page = header as usize - self.page_size;
            } else {
                break;
            }
        }
        self.pages.data[header_loc] += 1;
        self.pages.data[page * self.page_size + 1 + header as usize] = document;
        Ok(())
    }

    fn append_trigrams(&mut self, trigrams: BitSet, document: u64) -> io::Result<()> {
        for found in trigrams.iter() {
            self.append(found as u32, document)?;
        }
        Ok(())
    }
}

fn round_up(x: u64) -> u64 {
    let mut ret = x;
    loop {
        if ret % 16 == 0 {
            return ret;
        }
        ret += 1;
    }
}

fn eat_chunk(mut fh: &mut fs::File) -> io::Result<BitSet> {
    let end = fh.read_u64::<BigEndian>()?;
    let extra_len = fh.read_u64::<BigEndian>()?;
    let start = fh.seek(io::SeekFrom::Current(extra_len as i64))?;
    let ret = {
        let decoder = lz4::Decoder::new(&mut fh);
        let range = decoder.bytes();
        let exploding = range.map(|x| x.unwrap());
        let decoder = std::char::decode_utf8(exploding);
        let errors = decoder.map(|x| x.map_err(|_| io::CharsError::NotUtf8));
        trigrams_for(errors).map_err(|msg| io::Error::new(io::ErrorKind::Other, msg))
    };
    let next = round_up(start + end - extra_len - 16);
    fh.seek(io::SeekFrom::Start(next))?;
    ret
}

fn main() {
    let mut from: String = "".to_string();
    let mut simple = false;
    let mut addendum: u64 = 0;
    {
        let mut ap = argparse::ArgumentParser::new();
        ap.set_description("totally not a load of tools glued together");
        ap.refer(&mut from)
                .required()
                .add_option(&["-f", "--input-file"], Store,
                            "pack file to read");
        ap.refer(&mut addendum)
                .add_option(&["-i", "--addendum"], Store,
                            "number to add to file offset");
        ap.refer(&mut simple)
                .add_option(&["--simple"], StoreTrue,
                            "not a pack, just a normal decompressed file");
        ap.parse_args_or_exit();
    }

    let mut idx = Index::new().unwrap();

    let mut fh = fs::File::open(from).expect("input file must exist and be readable");
    if simple {
        let trigrams = trigrams_for(fh.chars()).expect("trigramming must work");
        idx.append_trigrams(trigrams, addendum).unwrap();
        return;
    }

    fh.seek(io::SeekFrom::Start(16)).unwrap();
    loop {

        let document = fh.seek(io::SeekFrom::Current(0)).unwrap() + addendum;
        let trigrams = eat_chunk(&mut fh).unwrap();
        idx.append_trigrams(trigrams, document).unwrap();
    }
}
<|endoftext|>


<commit_before>// Copyright © 2016 Bart Massey
// [This program is licensed under the "MIT License"]
// Please see the file COPYING in the source
// distribution of this software for license terms.

// Rocket pastebin demo

#![feature(plugin)]
#![feature(plugin, custom_derive)]
#![feature(proc_macro)]
#![plugin(rocket_codegen)]

static UPLOAD: &'static str = "upload";
static NID: usize = 8;

extern crate rocket;
extern crate rocket_contrib;
extern crate rand;
extern crate serde;
#[macro_use] extern crate serde_derive;

use std::io;
use std::io::{Read, BufRead, Write};
use std::path::Path;
use std::collections::HashMap;
use std::fs::{File, create_dir, metadata, remove_file};

use rocket::Data;
use rocket::response::NamedFile;
use rocket_contrib::Template;
use rocket::request::Form;
use rocket::response::Redirect;

mod paste_id;
use paste_id::PasteID;

#[derive(Serialize)]
struct PasteInfo {
    paste_id: String,
    snippet: String,
    ellipsis: bool
}

#[get("/")]
fn index() -> io::Result<Template> {
    let dirents = std::fs::read_dir(UPLOAD)
                  .expect("could not open upload directory");
    let pastes = dirents.map(|d| {
        let paste_id = d.unwrap().file_name().into_string().unwrap();
        let paste_filename = Path::new(UPLOAD).join(&paste_id);
        let paste_file = NamedFile::open(&paste_filename)
                         .expect("could not open paste file");
        let mut paste_reader = io::BufReader::new(paste_file);
        let mut paste_line = String::new();
        let paste_len = paste_reader.read_line(&mut paste_line)
                        .expect("could not read first line of paste");
        let (ellipsis, snippet) =
            if paste_len < 24 {
                (false, paste_line)
            } else {
                let result = (&paste_line)[0..24].to_string();
                (true, result)
            };
        PasteInfo { paste_id: paste_id, snippet: snippet, ellipsis: ellipsis }
    }).collect::<Vec<PasteInfo>>();
    let mut context: HashMap<&str, _> = HashMap::new();
    context.insert("pastes", &pastes);
    Ok(Template::render("index", &context))
}

#[post("/", data = "<paste>")]
fn upload(paste: Data) -> io::Result<String> {
    let paste_id = PasteID::new(NID);
    let path = Path::new(UPLOAD).join(&paste_id.to_string());
    paste.stream_to_file(path)?;
    Ok(format!("/{}\n", paste_id))
}

fn open_paste(id: &PasteID) -> io::Result<NamedFile> {
    let filename = Path::new(UPLOAD).join(id.to_string());
    NamedFile::open(&filename)
}

fn create_paste(id: &PasteID) -> io::Result<File> {
    let filename = Path::new(UPLOAD).join(id.to_string());
    File::create(&filename)
}

#[get("/<id>")]
fn retrieve(id: PasteID) -> Option<NamedFile> {
    open_paste(&id).ok()
}

#[get("/new")]
fn edit_new() -> io::Result<Redirect> {
    let paste_id = PasteID::new(NID);
    let f = try!(create_paste(&paste_id));
    drop(f);
    Ok(Redirect::to(&format!("/edit/{}", paste_id.to_string())))
}

#[get("/edit/<id>")]
fn make_edit(id: PasteID) -> io::Result<Template> {
    let paste = open_paste(&id)?;
    let mut paste_reader = io::BufReader::new(paste);
    let mut paste_contents = String::new();
    paste_reader.read_to_string(&mut paste_contents)?;
    let paste_id = id.to_string();
    let paste_rows = format!("{}", 100);
    let mut context: HashMap<&str, _> = HashMap::new();
    context.insert("contents", &paste_contents);
    context.insert("paste_id", &paste_id);
    context.insert("paste_rows", &paste_rows);
    Ok(Template::render("edit", &context))
}

#[derive(FromForm)]
struct EditForm {
    paste: String
}

#[post("/edit/<id>", data = "<form>")]
fn accept_edit(id: PasteID, form: Form<EditForm>) -> io::Result<Redirect> {
    let edit = form.get();
    let paste = create_paste(&id)?;
    let mut paste_writer = io::BufWriter::new(paste);
    let paste_contents =
        edit.paste.chars().map(|c|{c as u8}).collect::<Vec<u8>>();
    paste_writer.write_all(paste_contents.as_slice())?;
    paste_writer.flush()?;
    Ok(Redirect::to("/"))
}

#[get("/delete/<id>")]
fn delete(id: PasteID) -> io::Result<Redirect> {
    let filename = Path::new(UPLOAD).join(id.to_string());
    try!(remove_file(&filename));
    Ok(Redirect::to("/"))
}

fn ok_dir<P: AsRef<Path>>(path: P) -> bool {
    match metadata(path.as_ref()) {
        Ok(md) => md.is_dir(),
        _ => false
    }
}

fn check_upload_dir() {
    if !ok_dir(UPLOAD) {
        create_dir(UPLOAD).expect("could not create upload directory");
    }
}

fn main() {
    check_upload_dir();
    let rocket = rocket::ignite();
    let rocket = rocket.mount("/", routes![index, upload, edit_new, retrieve,
                                           make_edit, accept_edit, delete]);
    rocket.launch()
}

#[cfg(test)]
mod test {
    extern crate scraper;

    use super::rocket;
    use rocket::testing::MockRequest;
    use rocket::http::{Status, Method};
    use self::scraper::{Html, Selector};

    #[test]
    fn index() {
        let rocket = rocket::ignite().mount("/", routes![super::index]);

        let mut req = MockRequest::new(Method::Get, "/");
        let mut response = req.dispatch_with(&rocket);

        assert_eq!(response.status(), Status::Ok);

        let body_str = response.body().and_then(|b| b.into_string());
        assert_has_text(body_str.unwrap_or(String::from("")), String::from("Pastebin!"));
    }

    fn assert_has_text(body: String, expected: String) {
      let body_html = Html::parse_fragment(&body.to_string());
      let selector = Selector::parse("*").unwrap();
      let elements = body_html.select(&selector);
      let texts = elements.map(|e| e.text().collect::<Vec<_>>().concat());
      let actual = texts.fold("".to_string(), (|acc, str| format!("{}:{}", acc, str)));
      assert!(&actual.contains(&expected));
    }
}
<commit_msg>add some more assertions.<commit_after>// Copyright © 2016 Bart Massey
// [This program is licensed under the "MIT License"]
// Please see the file COPYING in the source
// distribution of this software for license terms.

// Rocket pastebin demo

#![feature(plugin)]
#![feature(plugin, custom_derive)]
#![feature(proc_macro)]
#![plugin(rocket_codegen)]

static UPLOAD: &'static str = "upload";
static NID: usize = 8;

extern crate rocket;
extern crate rocket_contrib;
extern crate rand;
extern crate serde;
#[macro_use] extern crate serde_derive;

use std::io;
use std::io::{Read, BufRead, Write};
use std::path::Path;
use std::collections::HashMap;
use std::fs::{File, create_dir, metadata, remove_file};

use rocket::Data;
use rocket::response::NamedFile;
use rocket_contrib::Template;
use rocket::request::Form;
use rocket::response::Redirect;

mod paste_id;
use paste_id::PasteID;

#[derive(Serialize)]
struct PasteInfo {
    paste_id: String,
    snippet: String,
    ellipsis: bool
}

#[get("/")]
fn index() -> io::Result<Template> {
    let dirents = std::fs::read_dir(UPLOAD)
                  .expect("could not open upload directory");
    let pastes = dirents.map(|d| {
        let paste_id = d.unwrap().file_name().into_string().unwrap();
        let paste_filename = Path::new(UPLOAD).join(&paste_id);
        let paste_file = NamedFile::open(&paste_filename)
                         .expect("could not open paste file");
        let mut paste_reader = io::BufReader::new(paste_file);
        let mut paste_line = String::new();
        let paste_len = paste_reader.read_line(&mut paste_line)
                        .expect("could not read first line of paste");
        let (ellipsis, snippet) =
            if paste_len < 24 {
                (false, paste_line)
            } else {
                let result = (&paste_line)[0..24].to_string();
                (true, result)
            };
        PasteInfo { paste_id: paste_id, snippet: snippet, ellipsis: ellipsis }
    }).collect::<Vec<PasteInfo>>();
    let mut context: HashMap<&str, _> = HashMap::new();
    context.insert("pastes", &pastes);
    Ok(Template::render("index", &context))
}

#[post("/", data = "<paste>")]
fn upload(paste: Data) -> io::Result<String> {
    let paste_id = PasteID::new(NID);
    let path = Path::new(UPLOAD).join(&paste_id.to_string());
    paste.stream_to_file(path)?;
    Ok(format!("/{}\n", paste_id))
}

fn open_paste(id: &PasteID) -> io::Result<NamedFile> {
    let filename = Path::new(UPLOAD).join(id.to_string());
    NamedFile::open(&filename)
}

fn create_paste(id: &PasteID) -> io::Result<File> {
    let filename = Path::new(UPLOAD).join(id.to_string());
    File::create(&filename)
}

#[get("/<id>")]
fn retrieve(id: PasteID) -> Option<NamedFile> {
    open_paste(&id).ok()
}

#[get("/new")]
fn edit_new() -> io::Result<Redirect> {
    let paste_id = PasteID::new(NID);
    let f = try!(create_paste(&paste_id));
    drop(f);
    Ok(Redirect::to(&format!("/edit/{}", paste_id.to_string())))
}

#[get("/edit/<id>")]
fn make_edit(id: PasteID) -> io::Result<Template> {
    let paste = open_paste(&id)?;
    let mut paste_reader = io::BufReader::new(paste);
    let mut paste_contents = String::new();
    paste_reader.read_to_string(&mut paste_contents)?;
    let paste_id = id.to_string();
    let paste_rows = format!("{}", 100);
    let mut context: HashMap<&str, _> = HashMap::new();
    context.insert("contents", &paste_contents);
    context.insert("paste_id", &paste_id);
    context.insert("paste_rows", &paste_rows);
    Ok(Template::render("edit", &context))
}

#[derive(FromForm)]
struct EditForm {
    paste: String
}

#[post("/edit/<id>", data = "<form>")]
fn accept_edit(id: PasteID, form: Form<EditForm>) -> io::Result<Redirect> {
    let edit = form.get();
    let paste = create_paste(&id)?;
    let mut paste_writer = io::BufWriter::new(paste);
    let paste_contents =
        edit.paste.chars().map(|c|{c as u8}).collect::<Vec<u8>>();
    paste_writer.write_all(paste_contents.as_slice())?;
    paste_writer.flush()?;
    Ok(Redirect::to("/"))
}

#[get("/delete/<id>")]
fn delete(id: PasteID) -> io::Result<Redirect> {
    let filename = Path::new(UPLOAD).join(id.to_string());
    try!(remove_file(&filename));
    Ok(Redirect::to("/"))
}

fn ok_dir<P: AsRef<Path>>(path: P) -> bool {
    match metadata(path.as_ref()) {
        Ok(md) => md.is_dir(),
        _ => false
    }
}

fn check_upload_dir() {
    if !ok_dir(UPLOAD) {
        create_dir(UPLOAD).expect("could not create upload directory");
    }
}

fn main() {
    check_upload_dir();
    let rocket = rocket::ignite();
    let rocket = rocket.mount("/", routes![index, upload, edit_new, retrieve,
                                           make_edit, accept_edit, delete]);
    rocket.launch()
}

#[cfg(test)]
mod test {
    extern crate scraper;

    use super::rocket;
    use rocket::testing::MockRequest;
    use rocket::http::{Status, Method};
    use self::scraper::{Html, Selector};

    #[test]
    fn index() {
        let rocket = rocket::ignite().mount("/", routes![super::index]);

        let mut req = MockRequest::new(Method::Get, "/");
        let mut response = req.dispatch_with(&rocket);

        assert_eq!(response.status(), Status::Ok);

        let body_str = response.body().and_then(|b| b.into_string()).unwrap_or(String::from(""));
        assert_has_text(body_str.clone(), String::from("Pastebin!"));
        assert_has_selector(body_str.clone(), String::from("ul"));
        assert_no_selector(body_str.clone(), String::from("ul li"));
    }

    fn assert_has_text(body: String, expected: String) {
      let body_html = Html::parse_fragment(&body.to_string());
      let selector = Selector::parse("*").unwrap();
      let elements = body_html.select(&selector);
      let texts = elements.map(|e| e.text().collect::<Vec<_>>().concat());
      let actual = texts.fold("".to_string(), (|acc, str| format!("{}:{}", acc, str)));
      assert!(&actual.contains(&expected));
    }

    fn assert_has_selector(body: String, selector_text: String) {
      let body_html = Html::parse_fragment(&body.to_string());
      let selector = Selector::parse(&selector_text).unwrap();
      let count = body_html.select(&selector).count();
      assert!(count > 0);
    }

    fn assert_no_selector(body: String, selector_text: String) {
      let body_html = Html::parse_fragment(&body.to_string());
      let selector = Selector::parse(&selector_text).unwrap();
      let count = body_html.select(&selector).count();
      assert_eq!(0, count);
    }
}
<|endoftext|>


<commit_before>use clap::{App, Arg, SubCommand, crate_version};
use console::{style, Emoji};
use indicatif::ProgressBar;
use syntect::easy::HighlightFile;
use syntect::parsing::SyntaxSet;
use syntect::highlighting::{ThemeSet, Style};
use syntect::util::{as_24_bit_terminal_escaped};
use std::fs::remove_file;
use std::io::BufRead;
use std::process::Command;
use std::sync::mpsc::channel;
use std::time::Duration;
use notify::DebouncedEvent;
use notify::{RecommendedWatcher, Watcher, RecursiveMode};

fn main() {
    let matches = App::new("rustlings")
        .version(crate_version!())
        .author("Olivia Hugger")
        .about("Test")
        .subcommand(SubCommand::with_name("verify").alias("v"))
        .subcommand(SubCommand::with_name("watch").alias("w"))
        .subcommand(
            SubCommand::with_name("run")
                .alias("r")
                .arg(Arg::with_name("file").required(true).index(1)),
        ).get_matches();

    let ss = SyntaxSet::load_defaults_newlines();
    let ts = ThemeSet::load_defaults();

    println!(r#"                 _   _ _                  "#);
    println!(r#"  _ __ _   _ ___| |_| (_)_ __   __ _ ___  "#);
    println!(r#" | '__| | | / __| __| | | '_ \ / _` / __| "#);
    println!(r#" | |  | |_| \__ \ |_| | | | | | (_| \__ \ "#);
    println!(r#" |_|   \__,_|___/\__|_|_|_| |_|\__, |___/ "#);
    println!(r#"                               |___/      "#);
    println!("");

    if let Some(matches) = matches.subcommand_matches("run") {
        if let Some(filename) = matches.value_of("file") {
            let bar = ProgressBar::new_spinner();
            bar.set_message(format!("Compiling {}...", filename).as_str());
            bar.enable_steady_tick(100);
            let compilecmd = Command::new("rustc")
                .args(&[filename, "-o", "temp"])
                .output()
                .expect("fail");
            bar.set_message(format!("Running {}...", filename).as_str());
            if compilecmd.status.success() {
                let runcmd = Command::new("./temp").output().expect("fail");
                bar.finish_and_clear();

                if runcmd.status.success() {
                    println!("{}", String::from_utf8_lossy(&runcmd.stdout));
                    let formatstr =
                        format!("{} Successfully ran {}", Emoji("✅", "✓"), filename);
                    println!("{}", style(formatstr).green());
                    clean();
                } else {
                    println!("{}", String::from_utf8_lossy(&runcmd.stdout));
                    println!("{}", String::from_utf8_lossy(&runcmd.stderr));

                    let formatstr =
                        format!("{} Ran {} with errors", Emoji("⚠ ", "!"), filename);
                    println!("{}", style(formatstr).red());
                    clean();
                }
            } else {
                bar.finish_and_clear();
                let formatstr = format!(
                    "{} Compilation of {} failed! Compiler error message:\n",
                    Emoji("⚠ ", "!"),
                    filename
                );
                println!("{}", style(formatstr).red());
                println!("{}", String::from_utf8_lossy(&compilecmd.stderr));
                clean();
            }
        } else {
            panic!("Please supply a filename!");
        }
    }

    if let Some(_) = matches.subcommand_matches("verify") {
        match verify() {
            Ok(_) => {}
            Err(_) => std::process::exit(1),
        }
    }

    if let Some(_) = matches.subcommand_matches("watch") {
        watch().unwrap();
    }

    if let None = matches.subcommand_name() {
        let mut highlighter = HighlightFile::new("default_out.md", &ss, &ts.themes["base16-eighties.dark"]).unwrap();
        for maybe_line in highlighter.reader.lines() {
            let line = maybe_line.unwrap();
            let regions: Vec<(Style, &str)> = highlighter.highlight_lines.highlight(&line, &ss);
            println!("{}", as_24_bit_terminal_escaped(&regions[..], true));
        }
    }

    println!("\x1b[0m");
}

fn watch() -> notify::Result<()> {
    let (tx, rx) = channel();

    let mut watcher: RecommendedWatcher = Watcher::new(tx, Duration::from_secs(2))?;
    watcher.watch("./exercises", RecursiveMode::Recursive)?;

    let _ignored = verify();

    loop {
        match rx.recv() {
            Ok(event) => {
                match event {
                    DebouncedEvent::Chmod(_)
                    | DebouncedEvent::Write(_) => {
                        let _ignored = verify();
                    }
                    _ => {}
                }
            },
            Err(e) => println!("watch error: {:?}", e),
        }
    }
}

fn verify() -> Result<(), ()> {
    compile_only("exercises/ex1.rs")?;
    compile_only("exercises/ex2.rs")?;
    compile_only("exercises/ex3.rs")?;
    compile_only("exercises/ex4.rs")?;
    compile_only("exercises/ex5.rs")?;
    compile_only("exercises/variables/variables1.rs")?;
    compile_only("exercises/variables/variables2.rs")?;
    compile_only("exercises/variables/variables3.rs")?;
    compile_only("exercises/variables/variables4.rs")?;
    compile_only("exercises/functions/functions1.rs")?;
    compile_only("exercises/functions/functions2.rs")?;
    compile_only("exercises/functions/functions3.rs")?;
    compile_only("exercises/functions/functions4.rs")?;
    compile_only("exercises/functions/functions5.rs")?;
    compile_only("exercises/primitive_types/primitive_types1.rs")?;
    compile_only("exercises/primitive_types/primitive_types2.rs")?;
    compile_only("exercises/primitive_types/primitive_types3.rs")?;
    compile_only("exercises/primitive_types/primitive_types4.rs")?;
    compile_only("exercises/primitive_types/primitive_types5.rs")?;
    compile_only("exercises/primitive_types/primitive_types6.rs")?;
    test("exercises/tests/tests1.rs")?;
    test("exercises/tests/tests2.rs")?;
    test("exercises/tests/tests3.rs")?;
    test("exercises/tests/tests4.rs")?;
    test("exercises/if/if1.rs")?;
    compile_only("exercises/strings/strings1.rs")?;
    compile_only("exercises/strings/strings2.rs")?;
    compile_only("exercises/strings/strings3.rs")?;
    compile_only("exercises/move_semantics/move_semantics1.rs")?;
    compile_only("exercises/move_semantics/move_semantics2.rs")?;
    compile_only("exercises/move_semantics/move_semantics3.rs")?;
    compile_only("exercises/move_semantics/move_semantics4.rs")?;
    compile_only("exercises/modules/modules1.rs")?;
    compile_only("exercises/modules/modules2.rs")?;
    compile_only("exercises/macros/macros1.rs")?;
    compile_only("exercises/macros/macros2.rs")?;
    compile_only("exercises/macros/macros3.rs")?;
    compile_only("exercises/macros/macros4.rs")?;
    test("exercises/error_handling/errors1.rs")?;
    test("exercises/error_handling/errors2.rs")?;
    test("exercises/error_handling/errors3.rs")?;
    test("exercises/error_handling/errorsn.rs")?;
    compile_only("exercises/error_handling/option1.rs")?;
    test("exercises/error_handling/result1.rs")?;
    Ok(())
}

fn compile_only(filename: &str) -> Result<(), ()> {
    let bar = ProgressBar::new_spinner();
    bar.set_message(format!("Compiling {}...", filename).as_str());
    bar.enable_steady_tick(100);
    let compilecmd = Command::new("rustc")
        .args(&[filename, "-o", "temp", "--color", "always"])
        .output()
        .expect("fail");
    bar.finish_and_clear();
    if compilecmd.status.success() {
        let formatstr = format!(
            "{} Successfully compiled {}!",
            Emoji("✅", "✓"),
            filename
        );
        println!("{}", style(formatstr).green());
        clean();
        Ok(())
    } else {
        let formatstr = format!(
            "{} Compilation of {} failed! Compiler error message:\n",
            Emoji("⚠ ", "!"),
            filename
        );
        println!("{}", style(formatstr).red());
        println!("{}", String::from_utf8_lossy(&compilecmd.stderr));
        clean();
        Err(())
    }
}

fn test(filename: &str) -> Result<(), ()> {
    let bar = ProgressBar::new_spinner();
    bar.set_message(format!("Testing {}...", filename).as_str());
    bar.enable_steady_tick(100);
    let testcmd = Command::new("rustc")
        .args(&["--test", filename, "-o", "temp"])
        .output()
        .expect("fail");
    bar.finish_and_clear();
    if testcmd.status.success() {
        let formatstr = format!("{} Successfully tested {}!", Emoji("✅", "✓"), filename);
        println!("{}", style(formatstr).green());
        clean();
        Ok(())
    } else {
        let formatstr = format!(
            "{} Testing of {} failed! Please try again.",
            Emoji("⚠ ", "!"),
            filename
        );
        println!("{}", style(formatstr).red());
        clean();
        Err(())
    }
}

fn clean() {
    let _ignored = remove_file("temp");
}
<commit_msg>revert<commit_after>#[macro_use]
extern crate clap;
extern crate console;
extern crate indicatif;
extern crate syntect;
extern crate notify;

use clap::{App, Arg, SubCommand};
use console::{style, Emoji};
use indicatif::ProgressBar;
use syntect::easy::HighlightFile;
use syntect::parsing::SyntaxSet;
use syntect::highlighting::{ThemeSet, Style};
use syntect::util::{as_24_bit_terminal_escaped, LinesWithEndings};
use std::fs::remove_file;
use std::io::BufRead;
use std::process::Command;
use std::fs::read_to_string;
use std::sync::mpsc::channel;
use std::time::Duration;
use notify::DebouncedEvent;
use notify::{RecommendedWatcher, Watcher, RecursiveMode};

fn main() {
    let matches = App::new("rustlings")
        .version(crate_version!())
        .author("Olivia Hugger")
        .about("Test")
        .subcommand(SubCommand::with_name("verify").alias("v"))
        .subcommand(SubCommand::with_name("watch").alias("w"))
        .subcommand(
            SubCommand::with_name("run")
                .alias("r")
                .arg(Arg::with_name("file").required(true).index(1)),
        ).get_matches();

    let ss = SyntaxSet::load_defaults_newlines();
    let ts = ThemeSet::load_defaults();

    println!(r#"                 _   _ _                  "#);
    println!(r#"  _ __ _   _ ___| |_| (_)_ __   __ _ ___  "#);
    println!(r#" | '__| | | / __| __| | | '_ \ / _` / __| "#);
    println!(r#" | |  | |_| \__ \ |_| | | | | | (_| \__ \ "#);
    println!(r#" |_|   \__,_|___/\__|_|_|_| |_|\__, |___/ "#);
    println!(r#"                               |___/      "#);
    println!("");

    if let Some(matches) = matches.subcommand_matches("run") {
        if let Some(filename) = matches.value_of("file") {
            let bar = ProgressBar::new_spinner();
            bar.set_message(format!("Compiling {}...", filename).as_str());
            bar.enable_steady_tick(100);
            let compilecmd = Command::new("rustc")
                .args(&[filename, "-o", "temp"])
                .output()
                .expect("fail");
            bar.set_message(format!("Running {}...", filename).as_str());
            if compilecmd.status.success() {
                let runcmd = Command::new("./temp").output().expect("fail");
                bar.finish_and_clear();

                if runcmd.status.success() {
                    println!("{}", String::from_utf8_lossy(&runcmd.stdout));
                    let formatstr =
                        format!("{} Successfully ran {}", Emoji("✅", "✓"), filename);
                    println!("{}", style(formatstr).green());
                    clean();
                } else {
                    println!("{}", String::from_utf8_lossy(&runcmd.stdout));
                    println!("{}", String::from_utf8_lossy(&runcmd.stderr));

                    let formatstr =
                        format!("{} Ran {} with errors", Emoji("⚠ ", "!"), filename);
                    println!("{}", style(formatstr).red());
                    clean();
                }
            } else {
                bar.finish_and_clear();
                let formatstr = format!(
                    "{} Compilation of {} failed! Compiler error message:\n",
                    Emoji("⚠ ", "!"),
                    filename
                );
                println!("{}", style(formatstr).red());
                println!("{}", String::from_utf8_lossy(&compilecmd.stderr));
                clean();
            }
        } else {
            panic!("Please supply a filename!");
        }
    }

    if let Some(_) = matches.subcommand_matches("verify") {
        match verify() {
            Ok(_) => {}
            Err(_) => std::process::exit(1),
        }
    }

    if let Some(_) = matches.subcommand_matches("watch") {
        watch().unwrap();
    }

    if let None = matches.subcommand_name() {
        let mut highlighter = HighlightFile::new("default_out.md", &ss, &ts.themes["base16-eighties.dark"]).unwrap();
        for maybe_line in highlighter.reader.lines() {
            let line = maybe_line.unwrap();
            let regions: Vec<(Style, &str)> = highlighter.highlight_lines.highlight(&line, &ss);
            println!("{}", as_24_bit_terminal_escaped(&regions[..], true));
        }
    }

    println!("\x1b[0m");
}

fn watch() -> notify::Result<()> {
    let (tx, rx) = channel();

    let mut watcher: RecommendedWatcher = Watcher::new(tx, Duration::from_secs(2))?;
    watcher.watch("./exercises", RecursiveMode::Recursive)?;

    let _ignored = verify();

    loop {
        match rx.recv() {
            Ok(event) => {
                match event {
                    DebouncedEvent::Chmod(_)
                    | DebouncedEvent::Write(_) => {
                        let _ignored = verify();
                    }
                    _ => {}
                }
            },
            Err(e) => println!("watch error: {:?}", e),
        }
    }
}

fn verify() -> Result<(), ()> {
    compile_only("exercises/ex1.rs")?;
    compile_only("exercises/ex2.rs")?;
    compile_only("exercises/ex3.rs")?;
    compile_only("exercises/ex4.rs")?;
    compile_only("exercises/ex5.rs")?;
    compile_only("exercises/variables/variables1.rs")?;
    compile_only("exercises/variables/variables2.rs")?;
    compile_only("exercises/variables/variables3.rs")?;
    compile_only("exercises/variables/variables4.rs")?;
    compile_only("exercises/functions/functions1.rs")?;
    compile_only("exercises/functions/functions2.rs")?;
    compile_only("exercises/functions/functions3.rs")?;
    compile_only("exercises/functions/functions4.rs")?;
    compile_only("exercises/functions/functions5.rs")?;
    compile_only("exercises/primitive_types/primitive_types1.rs")?;
    compile_only("exercises/primitive_types/primitive_types2.rs")?;
    compile_only("exercises/primitive_types/primitive_types3.rs")?;
    compile_only("exercises/primitive_types/primitive_types4.rs")?;
    compile_only("exercises/primitive_types/primitive_types5.rs")?;
    compile_only("exercises/primitive_types/primitive_types6.rs")?;
    test("exercises/tests/tests1.rs")?;
    test("exercises/tests/tests2.rs")?;
    test("exercises/tests/tests3.rs")?;
    test("exercises/tests/tests4.rs")?;
    test("exercises/if/if1.rs")?;
    compile_only("exercises/strings/strings1.rs")?;
    compile_only("exercises/strings/strings2.rs")?;
    compile_only("exercises/strings/strings3.rs")?;
    compile_only("exercises/move_semantics/move_semantics1.rs")?;
    compile_only("exercises/move_semantics/move_semantics2.rs")?;
    compile_only("exercises/move_semantics/move_semantics3.rs")?;
    compile_only("exercises/move_semantics/move_semantics4.rs")?;
    compile_only("exercises/modules/modules1.rs")?;
    compile_only("exercises/modules/modules2.rs")?;
    compile_only("exercises/macros/macros1.rs")?;
    compile_only("exercises/macros/macros2.rs")?;
    compile_only("exercises/macros/macros3.rs")?;
    compile_only("exercises/macros/macros4.rs")?;
    test("exercises/error_handling/errors1.rs")?;
    test("exercises/error_handling/errors2.rs")?;
    test("exercises/error_handling/errors3.rs")?;
    test("exercises/error_handling/errorsn.rs")?;
    compile_only("exercises/error_handling/option1.rs")?;
    test("exercises/error_handling/result1.rs")?;
    Ok(())
}

fn compile_only(filename: &str) -> Result<(), ()> {
    let bar = ProgressBar::new_spinner();
    bar.set_message(format!("Compiling {}...", filename).as_str());
    bar.enable_steady_tick(100);
    let compilecmd = Command::new("rustc")
        .args(&[filename, "-o", "temp", "--color", "always"])
        .output()
        .expect("fail");
    bar.finish_and_clear();
    if compilecmd.status.success() {
        let formatstr = format!(
            "{} Successfully compiled {}!",
            Emoji("✅", "✓"),
            filename
        );
        println!("{}", style(formatstr).green());
        clean();
        Ok(())
    } else {
        let formatstr = format!(
            "{} Compilation of {} failed! Compiler error message:\n",
            Emoji("⚠ ", "!"),
            filename
        );
        println!("{}", style(formatstr).red());
        println!("{}", String::from_utf8_lossy(&compilecmd.stderr));
        clean();
        Err(())
    }
}

fn test(filename: &str) -> Result<(), ()> {
    let bar = ProgressBar::new_spinner();
    bar.set_message(format!("Testing {}...", filename).as_str());
    bar.enable_steady_tick(100);
    let testcmd = Command::new("rustc")
        .args(&["--test", filename, "-o", "temp"])
        .output()
        .expect("fail");
    bar.finish_and_clear();
    if testcmd.status.success() {
        let formatstr = format!("{} Successfully tested {}!", Emoji("✅", "✓"), filename);
        println!("{}", style(formatstr).green());
        clean();
        Ok(())
    } else {
        let formatstr = format!(
            "{} Testing of {} failed! Please try again.",
            Emoji("⚠ ", "!"),
            filename
        );
        println!("{}", style(formatstr).red());
        clean();
        Err(())
    }
}

fn clean() {
    let _ignored = remove_file("temp");
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate clap;
#[macro_use]
extern crate error_chain;
extern crate mhost;
extern crate resolv_conf;
extern crate tokio_core;
extern crate trust_dns;

use mhost::{multiple_lookup, Statistics, Query};

use clap::{App, Arg, ArgMatches, Shell};
use error_chain::ChainedError;
use std::fmt;
use std::fs::File;
use std::io::Read;
use std::net::IpAddr;
use std::str::FromStr;
use std::time::Duration;
use tokio_core::reactor::Core;
use trust_dns::rr::{RData, RecordType};

// Arbitrary list of public DNS servers
static DEFAULT_DNS_SERVERS: &'static [&str] = &[
    // "Level3",
    "209.244.0.3",
    "209.244.0.4",
    // "Verisign",
    "64.6.64.6",
    "64.6.65.6",
    // "Google",
    "8.8.8.8",
    "8.8.4.4",
    // "DNS.WATCH",
    "84.200.69.80",
    "84.200.70.40",
    // "OpenDNS Home",
    "208.67.222.222",
    "208.67.220.220",
    // "SafeDNS",
    "195.46.39.39",
    "195.46.39.40",
    // "Dyn",
    "216.146.35.35",
    "216.146.36.36",
    // "FreeDNS",
    "37.235.1.174",
    "37.235.1.177",
    // "Alternate DNS",
    "198.101.242.72",
    "23.253.163.53",
];

static DEFAULT_RECORD_TYPES: &'static [&str] = &["a", "aaaa", "mx"];

fn run() -> Result<()> {
    let args = build_cli().get_matches();

    if args.is_present("completions") {
        return generate_completion(&args);
    }

    let dns_endpoints = get_dns_servers(&args)
        .chain_err(|| ErrorKind::ServerIpAddrParsingError)?
        .into_iter()
        .map(|s| (s, 53))
        .collect();

    let timeout = value_t!(args, "timeout", u64)
        .map(Duration::from_secs)
        .unwrap();

    let mut query = match IpAddr::from_str(args.value_of("domain name").unwrap()) {
        Ok(ip) => Query::from(ip, vec![RecordType::PTR]),
        Err(_) => {
            let record_types = get_record_types(&args).chain_err(|| {
                ErrorKind::ResoureRecordTypeParsingError
            })?;
            Query::new(args.value_of("domain name").unwrap(), record_types)
        }
    };
    query = query.set_timeout(timeout);

    let mut io_loop = Core::new().unwrap();
    let lookup = multiple_lookup(&io_loop.handle(), query, dns_endpoints);
    let result = io_loop.run(lookup);

    /*
    match result {
        Ok(responses) => {
            for response in responses {
                match response {
                    Ok(ref x) => println!("{}", DnsResponse(x)),
                    Err(ref e) => print_error(e),
                }
            }
        }
        Err(_) => {
            println!("General Error");
        }
    }
    */
    let statistics = Statistics::from(result.as_ref().unwrap());

    println!("Received {} (min {}, max {} records) answers from {} servers",
             statistics.num_of_ok_samples,
             statistics.min_num_of_records,
             statistics.max_num_of_records,
             statistics.num_of_samples,
    );
    let mut records: Vec<_> = statistics
        .record_counts
        .values()
        .map(|rr| {
            let record = DnsRecord(rr.0);
            let count = rr.1;
            format!("{} ({})", record, count)
        })
        .collect();
    records.sort();
    for r in records {
        println!("* {}", r);
    }
    if !statistics.failures.is_empty() {
        println!("Failures");
        for f in statistics.failures {
            print!("* ");
            print_error(f);
        }
    }

    Ok(())
}

fn build_cli() -> App<'static, 'static> {
    let name = env!("CARGO_PKG_NAME");
    let version = env!("CARGO_PKG_VERSION");
    let about = env!("CARGO_PKG_DESCRIPTION");

    App::new(name)
        .version(version)
        .about(about)
        .arg(
            Arg::with_name("domain name")
                .index(1)
                .required(true)
                .conflicts_with("completions")
                .help("domain name or IP address to look up"),
        )
        .arg(
            Arg::with_name("timeout")
                .long("timeout")
                .default_value("5")
                .help("timeout for server responses in sec")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("record types")
                .long("type")
                .short("t")
                .takes_value(true)
                .multiple(true)
                .require_delimiter(true)
                //.number_of_values(1)
                .possible_values(
                    &[
                        "a",
                        "aaaa",
                        "any",
                        "cname",
                        "dnskey",
                        "mx",
                        "ns",
                        "opt",
                        "ptr",
                        "soa",
                        "srv",
                        "txt",
                    ],
                )
                .help("Select resource record type [default: a, aaaa, mx]"),
        )
        .arg(
            Arg::with_name("dont use local dns servers")
                .short("L")
                .help("Do not use local (/etc/resolv.conf) DNS servers"),
        )
        .arg(Arg::with_name("predefined server").short("p").help(
            "use predefined DNS server set",
        ))
        .arg(
            Arg::with_name("DNS servers")
                .long("server")
                .short("s")
                .takes_value(true)
                .multiple(true)
                .number_of_values(1)
                .help(
                    "DNS servers to use; if empty use predefined, public DNS server",
                ),
        )
        .arg(
            Arg::with_name("completions")
                .long("completions")
                .takes_value(true)
                .hidden(true)
                .possible_values(&["bash", "fish", "zsh"])
                .help("The shell to generate the script for"),
        )
}

fn generate_completion(args: &ArgMatches) -> Result<()> {
    let bin_name = env!("CARGO_PKG_NAME");
    let shell = args.value_of("completions").ok_or(
        ErrorKind::CliArgsParsingError,
    )?;
    build_cli().gen_completions_to(
        bin_name,
        shell.parse::<Shell>().unwrap(),
        &mut std::io::stdout(),
    );
    Ok(())
}

fn get_local_dns_servers() -> Result<Vec<IpAddr>> {
    let mut buf = Vec::with_capacity(4096);
    let mut f = File::open("/etc/resolv.conf").chain_err(
        || ErrorKind::ResolvConfError,
    )?;
    f.read_to_end(&mut buf).unwrap();
    let cfg = resolv_conf::Config::parse(&buf[..]).chain_err(|| {
        ErrorKind::ResolvConfError
    })?;
    Ok(cfg.nameservers)
}

fn get_dns_servers(args: &ArgMatches) -> Result<Vec<IpAddr>> {
    let mut dns_servers: Vec<IpAddr> = Vec::new();
    if let Some(servers) = args.values_of_lossy("DNS servers") {
        dns_servers.extend(servers.into_iter().map(|server| {
            IpAddr::from_str(&server).unwrap()
        }));
    }
    if args.is_present("predefined server") {
        dns_servers.extend(DEFAULT_DNS_SERVERS.iter().map(|server| {
            IpAddr::from_str(server).unwrap()
        }));
    }
    if !args.is_present("dont use local dns servers") {
        dns_servers.extend(get_local_dns_servers()?);
    }

    Ok(dns_servers)
}

fn get_record_types(args: &ArgMatches) -> Result<Vec<RecordType>> {
    let record_types = if let Some(record_types) = args.values_of_lossy("record types") {
        record_types
            .iter()
            .map(|rt| RecordType::from_str(&rt.to_uppercase()).unwrap())
            .collect()
    } else {
        DEFAULT_RECORD_TYPES
            .iter()
            .map(|rt| RecordType::from_str(&rt.to_uppercase()).unwrap())
            .collect()
    };

    Ok(record_types)
}

// Newtype pattern for Display implementation
struct DnsResponse<'a>(pub &'a mhost::Response);

// Display impl for plain, basic output
impl<'a> fmt::Display for DnsResponse<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let DnsResponse(dns_response) = *self;
        if dns_response.answers.is_empty() {
            return write!(f, "DNS server {} has no records.", dns_response.server);
        }
        let _ = write!(f, "DNS server {} responded with\n", dns_response.server);
        let mut answers: Vec<String> = dns_response
            .answers
            .iter()
            .map(|answer| format!("* {}", DnsRecord(answer)))
            .collect();
        answers.sort();
        write!(f, "{}", answers.join("\n"))
    }
}

// Newtype pattern for Display implementation
struct DnsRecord<'a>(pub &'a trust_dns::rr::Record);

impl<'a> fmt::Display for DnsRecord<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let DnsRecord(record) = *self;

        let str = match *record.rdata() {
            RData::A(ip) => format!("IPv4: {}", ip),
            RData::AAAA(ip) => format!("IPv6: {}", ip),
            RData::CNAME(ref name) => format!("CNAME: {}", name),
            RData::MX(ref mx) => {
                format!("MX: {} with preference {}", mx.exchange(), mx.preference())
            }
            RData::NS(ref name) => format!("NS: {}", name),
            RData::SOA(ref soa) => {
                format!(
                    "SOA: {} {} {} {} {} {} {}",
                    soa.mname(),
                    soa.rname(),
                    soa.serial(),
                    soa.refresh(),
                    soa.retry(),
                    soa.expire(),
                    soa.minimum()
                )
            }
            RData::TXT(ref txt) => format!("TXT: {}", txt.txt_data().join(" ")),
            RData::PTR(ref ptr) => format!("PTR: {}", ptr.to_string()),
            ref x => format!(" * unclassified answer: {:?}", x),
        };
        write!(f, "{}", str)
    }
}

fn print_error<T: ChainedError>(err: &T) {
    print!("{} ", err);
    for e in err.iter().skip(1) {
        print!("because {}", e);
    }
    println!();

    // The backtrace is only available if run with `RUST_BACKTRACE=1`.
    if let Some(backtrace) = err.backtrace() {
        println!("backtrace: {:?}", backtrace);
    }
}

error_chain! {
    errors {
        CliArgsParsingError {
            description("failed to parse CLI arguments")
            display("failed to parse CLI arguments")
        }

        ResolvConfError {
            description("failed to parse /etc/resolv.conf")
            display("failed to parse /etc/resolv.cons")
        }

        ServerIpAddrParsingError {
            description("failed to parse server IP address")
            display("failed to parse server IP address")
        }

        ResoureRecordTypeParsingError {
            description("failed to parse resource record type")
            display("failed to parse resource record type")
        }
    }

    links {
        Lookup(::mhost::lookup::Error, ::mhost::lookup::ErrorKind);
    }
}

quick_main!(run);
<commit_msg>Add output modules -- first attempt<commit_after>#[macro_use]
extern crate clap;
#[macro_use]
extern crate error_chain;
extern crate mhost;
extern crate resolv_conf;
extern crate tokio_core;
extern crate trust_dns;

use mhost::{multiple_lookup, Statistics, Query, Response};
use mhost::lookup::Result as LookupResult;

use clap::{App, Arg, ArgMatches, Shell};
use error_chain::ChainedError;
use std::fmt;
use std::fs::File;
use std::io::Read;
use std::net::IpAddr;
use std::str::FromStr;
use std::time::Duration;
use tokio_core::reactor::Core;
use trust_dns::rr::{RData, RecordType};

// Arbitrary list of public DNS servers
static DEFAULT_DNS_SERVERS: &'static [&str] = &[
    // "Level3",
    "209.244.0.3",
    "209.244.0.4",
    // "Verisign",
    "64.6.64.6",
    "64.6.65.6",
    // "Google",
    "8.8.8.8",
    "8.8.4.4",
    // "DNS.WATCH",
    "84.200.69.80",
    "84.200.70.40",
    // "OpenDNS Home",
    "208.67.222.222",
    "208.67.220.220",
    // "SafeDNS",
    "195.46.39.39",
    "195.46.39.40",
    // "Dyn",
    "216.146.35.35",
    "216.146.36.36",
    // "FreeDNS",
    "37.235.1.174",
    "37.235.1.177",
    // "Alternate DNS",
    "198.101.242.72",
    "23.253.163.53",
];

static DEFAULT_RECORD_TYPES: &'static [&str] = &["a", "aaaa", "mx"];

fn run() -> Result<()> {
    let args = build_cli().get_matches();

    if args.is_present("completions") {
        return generate_completion(&args);
    }

    let dns_endpoints = get_dns_servers(&args)
        .chain_err(|| ErrorKind::ServerIpAddrParsingError)?
        .into_iter()
        .map(|s| (s, 53))
        .collect();

    let timeout = value_t!(args, "timeout", u64)
        .map(Duration::from_secs)
        .unwrap();

    let mut query = match IpAddr::from_str(args.value_of("domain name").unwrap()) {
        Ok(ip) => Query::from(ip, vec![RecordType::PTR]),
        Err(_) => {
            let record_types = get_record_types(&args).chain_err(|| {
                ErrorKind::ResoureRecordTypeParsingError
            })?;
            Query::new(args.value_of("domain name").unwrap(), record_types)
        }
    };
    query = query.set_timeout(timeout);

    let mut io_loop = Core::new().unwrap();
    let lookup = multiple_lookup(&io_loop.handle(), query, dns_endpoints);
    let result = io_loop.run(lookup);

    let responses = result.as_ref().unwrap();

    for output in args.values_of_lossy("output module").unwrap() {
        match output.as_ref() {
            "details" => DetailsOutput::new(responses).output(),
            _ => SummaryOutput::new(responses).output(),
        }
    }

    Ok(())
}

trait OutputModule {
    fn output(self) -> ();
}

struct DetailsOutput<'a> {
    responses: &'a [LookupResult<Response>],
}

impl<'a> DetailsOutput<'a> {
    pub fn new(responses: &'a [LookupResult<Response>]) -> Self {
        DetailsOutput { responses }
    }
}

impl<'a> OutputModule for DetailsOutput<'a> {
    fn output(self) -> () {
        for response in self.responses {
            match *response {
                Ok(ref x) => println!("{}", DnsResponse(x)),
                Err(ref e) => print_error(e),
            }
        }
    }
}

struct SummaryOutput<'a> {
    statistics: Statistics<'a>,
}

impl<'a> SummaryOutput<'a> {
    pub fn new(responses: &'a [LookupResult<Response>]) -> Self {
        let statistics = Statistics::from(responses);
        SummaryOutput { statistics }
    }
}

impl<'a> OutputModule for SummaryOutput<'a> {
    fn output(self) -> () {
        println!("Received {} (min {}, max {} records) answers from {} servers",
                 self.statistics.num_of_ok_samples,
                 self.statistics.min_num_of_records,
                 self.statistics.max_num_of_records,
                 self.statistics.num_of_samples,
        );
        let mut records: Vec<_> = self.statistics
            .record_counts
            .values()
            .map(|rr| {
                let record = DnsRecord(rr.0);
                let count = rr.1;
                format!("{} ({})", record, count)
            })
            .collect();
        records.sort();
        for r in records {
            println!("* {}", r);
        }
        if !self.statistics.failures.is_empty() {
            println!("Failures");
            for f in self.statistics.failures {
                print!("* ");
                print_error(f);
            }
        }
    }
}

fn build_cli() -> App<'static, 'static> {
    let name = env!("CARGO_PKG_NAME");
    let version = env!("CARGO_PKG_VERSION");
    let about = env!("CARGO_PKG_DESCRIPTION");

    App::new(name)
        .version(version)
        .about(about)
        .arg(
            Arg::with_name("domain name")
                .index(1)
                .required(true)
                .conflicts_with("completions")
                .help("domain name or IP address to look up")
        )
        .arg(
            Arg::with_name("timeout")
                .long("timeout")
                .default_value("5")
                .help("timeout for server responses in sec")
                .takes_value(true)
        )
        .arg(
            Arg::with_name("record types")
                .long("type")
                .short("t")
                .takes_value(true)
                .multiple(true)
                .require_delimiter(true)
                .possible_values(
                    &[
                        "a",
                        "aaaa",
                        "any",
                        "cname",
                        "dnskey",
                        "mx",
                        "ns",
                        "opt",
                        "ptr",
                        "soa",
                        "srv",
                        "txt",
                    ],
                )
                .help("Select resource record type [default: a, aaaa, mx]")
        )
        .arg(
            Arg::with_name("dont use local dns servers")
                .short("L")
                .help("Do not use local (/etc/resolv.conf) DNS servers")
        )
        .arg(Arg::with_name("predefined server").short("p").help(
            "use predefined DNS server set"
        ))
        .arg(
            Arg::with_name("DNS servers")
                .long("server")
                .short("s")
                .takes_value(true)
                .multiple(true)
                .require_delimiter(true)
                .number_of_values(1)
                .help(
                    "DNS servers to use; if empty use predefined, public DNS server"
                )
        )
        .arg(
            Arg::with_name("output module")
                .long("module")
                .short("m")
                .takes_value(true)
                .multiple(true)
                .require_delimiter(true)
                .default_value("summary")
                .possible_values(
                    &[
                        "summary",
                        "details",
                    ],
                )
                .help("Select output module")
        )
        .arg(
            Arg::with_name("completions")
                .long("completions")
                .takes_value(true)
                .hidden(true)
                .possible_values(&["bash", "fish", "zsh"])
                .help("The shell to generate the script for")
        )
}

fn generate_completion(args: &ArgMatches) -> Result<()> {
    let bin_name = env!("CARGO_PKG_NAME");
    let shell = args.value_of("completions").ok_or(
        ErrorKind::CliArgsParsingError,
    )?;
    build_cli().gen_completions_to(
        bin_name,
        shell.parse::<Shell>().unwrap(),
        &mut std::io::stdout(),
    );
    Ok(())
}

fn get_local_dns_servers() -> Result<Vec<IpAddr>> {
    let mut buf = Vec::with_capacity(4096);
    let mut f = File::open("/etc/resolv.conf").chain_err(
        || ErrorKind::ResolvConfError,
    )?;
    f.read_to_end(&mut buf).unwrap();
    let cfg = resolv_conf::Config::parse(&buf[..]).chain_err(|| {
        ErrorKind::ResolvConfError
    })?;
    Ok(cfg.nameservers)
}

fn get_dns_servers(args: &ArgMatches) -> Result<Vec<IpAddr>> {
    let mut dns_servers: Vec<IpAddr> = Vec::new();
    if let Some(servers) = args.values_of_lossy("DNS servers") {
        dns_servers.extend(servers.into_iter().map(|server| {
            IpAddr::from_str(&server).unwrap()
        }));
    }
    if args.is_present("predefined server") {
        dns_servers.extend(DEFAULT_DNS_SERVERS.iter().map(|server| {
            IpAddr::from_str(server).unwrap()
        }));
    }
    if !args.is_present("dont use local dns servers") {
        dns_servers.extend(get_local_dns_servers()?);
    }

    Ok(dns_servers)
}

fn get_record_types(args: &ArgMatches) -> Result<Vec<RecordType>> {
    let record_types = if let Some(record_types) = args.values_of_lossy("record types") {
        record_types
            .iter()
            .map(|rt| RecordType::from_str(&rt.to_uppercase()).unwrap())
            .collect()
    } else {
        DEFAULT_RECORD_TYPES
            .iter()
            .map(|rt| RecordType::from_str(&rt.to_uppercase()).unwrap())
            .collect()
    };

    Ok(record_types)
}

// Newtype pattern for Display implementation
struct DnsResponse<'a>(pub &'a mhost::Response);

// Display impl for plain, basic output
impl<'a> fmt::Display for DnsResponse<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let DnsResponse(dns_response) = *self;
        if dns_response.answers.is_empty() {
            return write!(f, "DNS server {} has no records.", dns_response.server);
        }
        let _ = write!(f, "DNS server {} responded with\n", dns_response.server);
        let mut answers: Vec<String> = dns_response
            .answers
            .iter()
            .map(|answer| format!("* {}", DnsRecord(answer)))
            .collect();
        answers.sort();
        write!(f, "{}", answers.join("\n"))
    }
}

// Newtype pattern for Display implementation
struct DnsRecord<'a>(pub &'a trust_dns::rr::Record);

impl<'a> fmt::Display for DnsRecord<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let DnsRecord(record) = *self;

        let str = match *record.rdata() {
            RData::A(ip) => format!("IPv4: {}", ip),
            RData::AAAA(ip) => format!("IPv6: {}", ip),
            RData::CNAME(ref name) => format!("CNAME: {}", name),
            RData::MX(ref mx) => {
                format!("MX: {} with preference {}", mx.exchange(), mx.preference())
            }
            RData::NS(ref name) => format!("NS: {}", name),
            RData::SOA(ref soa) => {
                format!(
                    "SOA: {} {} {} {} {} {} {}",
                    soa.mname(),
                    soa.rname(),
                    soa.serial(),
                    soa.refresh(),
                    soa.retry(),
                    soa.expire(),
                    soa.minimum()
                )
            }
            RData::TXT(ref txt) => format!("TXT: {}", txt.txt_data().join(" ")),
            RData::PTR(ref ptr) => format!("PTR: {}", ptr.to_string()),
            ref x => format!(" * unclassified answer: {:?}", x),
        };
        write!(f, "{}", str)
    }
}

fn print_error<T: ChainedError>(err: &T) {
    print!("{} ", err);
    for e in err.iter().skip(1) {
        print!("because {}", e);
    }
    println!();

    // The backtrace is only available if run with `RUST_BACKTRACE=1`.
    if let Some(backtrace) = err.backtrace() {
        println!("backtrace: {:?}", backtrace);
    }
}

error_chain! {
    errors {
        CliArgsParsingError {
            description("failed to parse CLI arguments")
            display("failed to parse CLI arguments")
        }

        ResolvConfError {
            description("failed to parse /etc/resolv.conf")
            display("failed to parse /etc/resolv.cons")
        }

        ServerIpAddrParsingError {
            description("failed to parse server IP address")
            display("failed to parse server IP address")
        }

        ResoureRecordTypeParsingError {
            description("failed to parse resource record type")
            display("failed to parse resource record type")
        }
    }

    links {
        Lookup(::mhost::lookup::Error, ::mhost::lookup::ErrorKind);
    }
}

quick_main!(run);
<|endoftext|>


<commit_before>extern crate hkg;

use hkg::utility::client::download;
use hkg::utility::cache::readfile;

fn main(){

    // let body = download();
    // println!("Response: {}", body);
    let s = readfile("demo.html".to_string());
    print!("contains:\n{}", s)
}
<commit_msg>1st html parse<commit_after>extern crate hkg;

extern crate kuchiki;
use kuchiki::traits::*;

use hkg::utility::client::download;
use hkg::utility::cache::readfile;

fn main(){

    // let body = download();
    // println!("Response: {}", body);

    // let html = readfile("demo.html".to_string());
    // print!("contains:\n{}", html)

    let path = "demo.html";

    let document = kuchiki::parse_html().from_utf8().from_file(&path).unwrap();
    // println!("{:?}", document.to_string());

    let topic_list_panel = document.select(".Topic_ListPanel").unwrap().collect::<Vec<_>>();

    if topic_list_panel.len() >  0 {
        let trs = topic_list_panel[0].as_node().select("tr").unwrap().collect::<Vec<_>>();
        println!("trs len: {:?}", trs.len());

        for (index, tr) in trs.iter().enumerate() {
            let tr_node = tr.as_node();
            // println!("{:?}", tr_node.text_contents());

            let tds = tr_node.select("td").unwrap().collect::<Vec<_>>();
            println!("tds[{:?}]={:?}", index, tds.len());
        }
    }
    else
    {
        println!("not found");
    }

}
<|endoftext|>