<commit_before>//! Low level and internal http and https implementation.

use std::collections::HashMap;
use std::time::Duration;

use async_std::io::{self, stderr, Read, Result as IoResult, Write};
use async_std::net::{SocketAddr, TcpStream};
use async_std::prelude::*;
use log::Level::Info;

use super::asynctls::TLSStream;
use super::constants;
use super::dns::Resolver;
use super::request::Request;
use super::results::{CabotError, CabotResult};

#[derive(Debug, PartialEq)]
enum TransferEncoding {
    ContentLength(usize),
    Chunked,
    Unkown,
    None,
}

#[derive(Debug, PartialEq)]
enum TransferEncodingStatus {
    ReadingHeader,
    ChunkHeader,
    ReadingBody(usize),
}

impl From<&[u8]> for TransferEncoding {
    fn from(hdr: &[u8]) -> Self {
        let hdr = String::from_utf8_lossy(hdr);
        let hdrup = hdr.to_ascii_uppercase();
        match hdrup.as_str() {
            "CHUNKED" => TransferEncoding::Chunked,
            _ => TransferEncoding::Unkown,
        }
    }
}

struct HttpDecoder<'a> {
    reader: &'a mut (dyn Read + Unpin), // tls require Write
    writer: &'a mut (dyn Write + Unpin),
    buffer: Vec<u8>,
    transfer_encoding: TransferEncoding,
    transfer_encoding_status: TransferEncodingStatus,
    read_timeout: Duration,
}

impl<'a> HttpDecoder<'a> {
    fn new(
        writer: &'a mut (dyn Write + Unpin),
        reader: &'a mut (dyn Read + Unpin),
        read_timeout: u64,
    ) -> Self {
        HttpDecoder {
            writer,
            reader,
            buffer: Vec::with_capacity(constants::BUFFER_PAGE_SIZE),
            transfer_encoding: TransferEncoding::None,
            transfer_encoding_status: TransferEncodingStatus::ReadingHeader,
            read_timeout: Duration::from_millis(read_timeout),
        }
    }

    async fn chunk_read(&mut self) -> IoResult<usize> {
        let ret = io::timeout(self.read_timeout, async {
            let mut buf = [0; constants::BUFFER_PAGE_SIZE];
            let ret = self.reader.read(&mut buf[..]).await;
            if let Ok(count) = ret {
                if count > 0 {
                    self.buffer.extend_from_slice(&buf[0..count]);
                }
            }
            ret
        });
        ret.await
    }
    async fn read_write_headers(&mut self) -> IoResult<()> {
        info!("Reading headers");
        loop {
            let _count = self.chunk_read().await?;
            if let Some(_) = self.process_headers().await {
                break;
            }
        }
        Ok(())
    }

    async fn read_write_no_transfer_encoding(&mut self) -> IoResult<()> {
        loop {
            self.writer.write(self.buffer.as_slice()).await.unwrap();
            self.buffer.clear();
            let cnt = self.chunk_read().await?;
            if cnt == 0 {
                break;
            }
        }

        Ok(())
    }

    async fn read_content_length(&mut self, size: usize) -> IoResult<()> {
        let mut read_count = self.buffer.len();
        loop {
            self.writer.write(self.buffer.as_slice()).await.unwrap();
            self.buffer.clear();

            if read_count >= size {
                break;
            }
            read_count = read_count + self.chunk_read().await?;
            warn!("< {}", read_count);
        }
        Ok(())
    }

    async fn read_write_chunk(&mut self) -> IoResult<()> {
        loop {
            let done = self.process_chunk().await?;
            if done {
                break;
            }
            let cnt = self.chunk_read().await?;
            if cnt == 0 {
                break;
            }
        }

        Ok(())
    }

    async fn read_write_body(&mut self) -> IoResult<()> {
        info!("Reading body");
        match self.transfer_encoding {
            TransferEncoding::ContentLength(size) => {
                self.read_content_length(size).await?;
            }
            TransferEncoding::Chunked => {
                self.read_write_chunk().await?;
            }
            _ => {
                warn!("Neither Content-Length, not chunk is response header");
                self.read_write_no_transfer_encoding().await?;
            }
        }

        if self.buffer.len() > 0 {
            let b = String::from_utf8_lossy(self.buffer.as_slice());
            error!("Buffer not clear: {}", b);
        }

        Ok(())
    }

    async fn read_write(&mut self) -> IoResult<()> {
        self.read_write_headers().await?;
        self.read_write_body().await?;
        self.writer.flush().await
    }

    async fn process_headers(&mut self) -> Option<usize> {
        let ret = {
            let resp_header: Vec<&[u8]> = constants::SPLIT_HEADERS_RE
                .splitn(self.buffer.as_slice(), 2)
                .collect();
            if resp_header.len() == 2 {
                // We have headers
                let headers = resp_header.get(0).unwrap();
                if let Some(header) = constants::TRANSFER_ENCODING.captures(headers) {
                    if let Some(tenc) = header.get(1) {
                        self.transfer_encoding = TransferEncoding::from(tenc.as_bytes());
                    }
                } else if let Some(header) = constants::CONTENT_LENGTH.captures(headers) {
                    if let Some(clength) = header.get(1) {
                        let clength = String::from_utf8_lossy(clength.as_bytes()).into_owned();
                        let clength = usize::from_str_radix(clength.as_str(), 10).unwrap();
                        self.transfer_encoding = TransferEncoding::ContentLength(clength);
                    }
                }
                let resp = headers.len();
                self.writer.write(headers).await.unwrap();
                Some(resp + 4) // + CRLF CRLF
            } else {
                None
            }
        };
        if let Some(to_drain) = ret {
            let buffer = self.buffer.drain(to_drain..).collect();
            self.buffer = buffer;
            info!("End of Headers readched");
            debug!("Transfer encoding: {:?}", self.transfer_encoding);
            debug!("{:?}", String::from_utf8_lossy(self.buffer.as_slice()));
        }
        ret
    }
    async fn process_chunk(&mut self) -> IoResult<bool> {
        debug!(
            "transfer_encoding_status: {:?}",
            self.transfer_encoding_status
        );
        let mut can_process_buffer = self.buffer.len() > 0;
        let mut body_chunk_size = 0;
        let mut header_len: usize;
        while can_process_buffer {
            header_len = 0;
            if self.transfer_encoding_status == TransferEncodingStatus::ReadingHeader {
                // we read the chunk size to drain
                let header: Vec<&[u8]> = constants::SPLIT_HEADER_BRE
                    .splitn(self.buffer.as_slice(), 2)
                    .collect();
                if header.len() == 2 {
                    if let Some(header) = constants::GET_CHUNK_SIZE.captures(header[0]) {
                        if let Some(size) = header.get(1) {
                            let size = size.as_bytes();
                            header_len = size.len() + 2;
                            let size = String::from_utf8_lossy(size).into_owned();
                            body_chunk_size = usize::from_str_radix(size.as_str(), 16).unwrap();
                            self.transfer_encoding_status = TransferEncodingStatus::ChunkHeader;
                        }
                    } else {
                        // else return Error
                        // break;
                    }
                } else {
                    debug!(
                        "Chunked not complete: {}",
                        String::from_utf8_lossy(self.buffer.as_slice())
                    );
                }
            }

            if self.transfer_encoding_status == TransferEncodingStatus::ChunkHeader {
                if body_chunk_size == 0 {
                    debug!(
                        "0 chunked size received: {}",
                        String::from_utf8_lossy(self.buffer.as_slice())
                    );
                    self.buffer.clear(); // should we check that it is '0\r\n' ?
                    return Ok(true);
                }

                debug!(
                    "Before header cleanup: {}",
                    String::from_utf8_lossy(self.buffer.as_slice())
                );
                let buffer: Vec<u8> = self.buffer.drain(header_len..).collect();
                self.buffer = buffer;
                debug!(
                    "After header cleanup: {}",
                    String::from_utf8_lossy(self.buffer.as_slice())
                );
                self.transfer_encoding_status =
                    TransferEncodingStatus::ReadingBody(body_chunk_size);
            }

            if let TransferEncodingStatus::ReadingBody(body_chunk_size) =
                self.transfer_encoding_status
            {
                error!("!! {} > {}", self.buffer.len(), body_chunk_size);
                if self.buffer.len() > body_chunk_size {
                    let mut buffer: Vec<u8> = self.buffer.drain(body_chunk_size..).collect();
                    self.writer.write(self.buffer.as_slice()).await?;

                    let buffer2 = buffer.drain(2..).collect(); // CRLF
                    self.buffer = buffer2;
                    self.transfer_encoding_status = TransferEncodingStatus::ReadingHeader;
                } else {
                    can_process_buffer = false;
                }
            } else {
                can_process_buffer = false;
            }
        }

        return Ok(false);
    }
}

async fn log_request(request: &[u8], verbose: bool) {
    if !log_enabled!(Info) && !verbose {
        return;
    }
    let request: Vec<&[u8]> = constants::SPLIT_HEADERS_RE.splitn(request, 2).collect();
    let headers = String::from_utf8_lossy(&request[0]);
    let headers: Vec<&str> = constants::SPLIT_HEADER_RE.split(&headers).collect();
    let bodylen = if request.len() == 2 {
        let body = &request[1];
        body.len()
    } else {
        0
    };
    if log_enabled!(Info) {
        for header in headers {
            info!("> {}", header);
        }
        if bodylen > 0 {
            info!("> [{} bytes]", bodylen);
        }
        info!(">");
    } else if verbose {
        for header in headers {
            writeln!(&mut stderr(), "> {}", header).await.unwrap();
        }
        if bodylen > 0 {
            writeln!(&mut stderr(), "> [{} bytes]", bodylen)
                .await
                .unwrap();
        }
        writeln!(&mut stderr(), ">").await.unwrap();
    }
}

async fn from_http(
    request: &Request,
    client: &mut TcpStream,
    out: &mut (dyn Write + Unpin),
    verbose: bool,
    read_timeout: u64,
) -> CabotResult<()> {
    let request_bytes = request.to_bytes();
    let raw_request = request_bytes.as_slice();
    log_request(&raw_request, verbose).await;

    debug!("Sending request...");
    client.write_all(&raw_request).await?;

    debug!("Reading response headers...");

    let mut http_decoder = HttpDecoder::new(out, client, read_timeout);
    http_decoder.read_write().await?;
    Ok(())
}

async fn from_https(
    request: &Request,
    client: &mut TcpStream,
    out: &mut (dyn Write + Unpin),
    verbose: bool,
    read_timeout: u64,
) -> CabotResult<()> {
    let request_bytes = request.to_bytes();
    let raw_request = request_bytes.as_slice();
    log_request(&raw_request, verbose).await;

    let mut tls_client = TLSStream::new(client, request.host())?;
    tls_client.starttls().await?;

    debug!("Sending request...");
    tls_client.write_all(&raw_request).await?;
    debug!("Request sent");

    debug!("Decoding response...");
    let mut http_decoder = HttpDecoder::new(out, &mut tls_client, read_timeout);
    http_decoder.read_write().await?;
    Ok(())
}

pub async fn http_query(
    request: &Request,
    mut out: &mut (dyn Write + Unpin),
    authorities: &HashMap<String, SocketAddr>,
    verbose: bool,
    ipv4: bool,
    ipv6: bool,
    dns_timeout: u64,
    connect_timeout: u64,
    read_timeout: u64,
) -> CabotResult<()> {
    debug!(
        "HTTP Query {} {}",
        request.http_method(),
        request.request_uri()
    );

    let authority = request.authority();

    let addr = match authorities.get(authority) {
        Some(val) => {
            info!("Fetch authority {} using autorities map", authority);
            *val
        }
        None => {
            info!("Fetch authority {} using resolver", authority);
            let resolver = Resolver::new(verbose);
            resolver
                .get_addr(authority, ipv4, ipv6, dns_timeout)
                .await?
        }
    };

    info!("Connecting to {}", addr);
    let mut client = io::timeout(Duration::from_millis(connect_timeout), async {
        TcpStream::connect(addr).await
    })
    .await?;

    // client.set_read_timeout(Some(Duration::new(5, 0))).unwrap();

    match request.scheme() {
        "http" => from_http(request, &mut client, &mut out, verbose, read_timeout).await?,
        "https" => from_https(request, &mut client, &mut out, verbose, read_timeout).await?,
        _ => {
            return Err(CabotError::SchemeError(format!(
                "Unrecognized scheme {}",
                request.scheme()
            )))
        }
    };

    out.flush().await.unwrap();

    Ok(())
}
<commit_msg>Update message for connection timeout<commit_after>//! Low level and internal http and https implementation.

use std::collections::HashMap;
use std::time::Duration;

use async_std::io::{self, stderr, Read, Result as IoResult, Write};
use async_std::net::{SocketAddr, TcpStream};
use async_std::prelude::*;
use log::Level::Info;

use super::asynctls::TLSStream;
use super::constants;
use super::dns::Resolver;
use super::request::Request;
use super::results::{CabotError, CabotResult};

#[derive(Debug, PartialEq)]
enum TransferEncoding {
    ContentLength(usize),
    Chunked,
    Unkown,
    None,
}

#[derive(Debug, PartialEq)]
enum TransferEncodingStatus {
    ReadingHeader,
    ChunkHeader,
    ReadingBody(usize),
}

impl From<&[u8]> for TransferEncoding {
    fn from(hdr: &[u8]) -> Self {
        let hdr = String::from_utf8_lossy(hdr);
        let hdrup = hdr.to_ascii_uppercase();
        match hdrup.as_str() {
            "CHUNKED" => TransferEncoding::Chunked,
            _ => TransferEncoding::Unkown,
        }
    }
}

struct HttpDecoder<'a> {
    reader: &'a mut (dyn Read + Unpin), // tls require Write
    writer: &'a mut (dyn Write + Unpin),
    buffer: Vec<u8>,
    transfer_encoding: TransferEncoding,
    transfer_encoding_status: TransferEncodingStatus,
    read_timeout: Duration,
}

impl<'a> HttpDecoder<'a> {
    fn new(
        writer: &'a mut (dyn Write + Unpin),
        reader: &'a mut (dyn Read + Unpin),
        read_timeout: u64,
    ) -> Self {
        HttpDecoder {
            writer,
            reader,
            buffer: Vec::with_capacity(constants::BUFFER_PAGE_SIZE),
            transfer_encoding: TransferEncoding::None,
            transfer_encoding_status: TransferEncodingStatus::ReadingHeader,
            read_timeout: Duration::from_millis(read_timeout),
        }
    }

    async fn chunk_read(&mut self) -> IoResult<usize> {
        let ret = io::timeout(self.read_timeout, async {
            let mut buf = [0; constants::BUFFER_PAGE_SIZE];
            let ret = self.reader.read(&mut buf[..]).await;
            if let Ok(count) = ret {
                if count > 0 {
                    self.buffer.extend_from_slice(&buf[0..count]);
                }
            }
            ret
        });
        ret.await
    }
    async fn read_write_headers(&mut self) -> IoResult<()> {
        info!("Reading headers");
        loop {
            let _count = self.chunk_read().await?;
            if let Some(_) = self.process_headers().await {
                break;
            }
        }
        Ok(())
    }

    async fn read_write_no_transfer_encoding(&mut self) -> IoResult<()> {
        loop {
            self.writer.write(self.buffer.as_slice()).await.unwrap();
            self.buffer.clear();
            let cnt = self.chunk_read().await?;
            if cnt == 0 {
                break;
            }
        }

        Ok(())
    }

    async fn read_content_length(&mut self, size: usize) -> IoResult<()> {
        let mut read_count = self.buffer.len();
        loop {
            self.writer.write(self.buffer.as_slice()).await.unwrap();
            self.buffer.clear();

            if read_count >= size {
                break;
            }
            read_count = read_count + self.chunk_read().await?;
            warn!("< {}", read_count);
        }
        Ok(())
    }

    async fn read_write_chunk(&mut self) -> IoResult<()> {
        loop {
            let done = self.process_chunk().await?;
            if done {
                break;
            }
            let cnt = self.chunk_read().await?;
            if cnt == 0 {
                break;
            }
        }

        Ok(())
    }

    async fn read_write_body(&mut self) -> IoResult<()> {
        info!("Reading body");
        match self.transfer_encoding {
            TransferEncoding::ContentLength(size) => {
                self.read_content_length(size).await?;
            }
            TransferEncoding::Chunked => {
                self.read_write_chunk().await?;
            }
            _ => {
                warn!("Neither Content-Length, not chunk is response header");
                self.read_write_no_transfer_encoding().await?;
            }
        }

        if self.buffer.len() > 0 {
            let b = String::from_utf8_lossy(self.buffer.as_slice());
            error!("Buffer not clear: {}", b);
        }

        Ok(())
    }

    async fn read_write(&mut self) -> IoResult<()> {
        self.read_write_headers().await?;
        self.read_write_body().await?;
        self.writer.flush().await
    }

    async fn process_headers(&mut self) -> Option<usize> {
        let ret = {
            let resp_header: Vec<&[u8]> = constants::SPLIT_HEADERS_RE
                .splitn(self.buffer.as_slice(), 2)
                .collect();
            if resp_header.len() == 2 {
                // We have headers
                let headers = resp_header.get(0).unwrap();
                if let Some(header) = constants::TRANSFER_ENCODING.captures(headers) {
                    if let Some(tenc) = header.get(1) {
                        self.transfer_encoding = TransferEncoding::from(tenc.as_bytes());
                    }
                } else if let Some(header) = constants::CONTENT_LENGTH.captures(headers) {
                    if let Some(clength) = header.get(1) {
                        let clength = String::from_utf8_lossy(clength.as_bytes()).into_owned();
                        let clength = usize::from_str_radix(clength.as_str(), 10).unwrap();
                        self.transfer_encoding = TransferEncoding::ContentLength(clength);
                    }
                }
                let resp = headers.len();
                self.writer.write(headers).await.unwrap();
                Some(resp + 4) // + CRLF CRLF
            } else {
                None
            }
        };
        if let Some(to_drain) = ret {
            let buffer = self.buffer.drain(to_drain..).collect();
            self.buffer = buffer;
            info!("End of Headers readched");
            debug!("Transfer encoding: {:?}", self.transfer_encoding);
            debug!("{:?}", String::from_utf8_lossy(self.buffer.as_slice()));
        }
        ret
    }
    async fn process_chunk(&mut self) -> IoResult<bool> {
        debug!(
            "transfer_encoding_status: {:?}",
            self.transfer_encoding_status
        );
        let mut can_process_buffer = self.buffer.len() > 0;
        let mut body_chunk_size = 0;
        let mut header_len: usize;
        while can_process_buffer {
            header_len = 0;
            if self.transfer_encoding_status == TransferEncodingStatus::ReadingHeader {
                // we read the chunk size to drain
                let header: Vec<&[u8]> = constants::SPLIT_HEADER_BRE
                    .splitn(self.buffer.as_slice(), 2)
                    .collect();
                if header.len() == 2 {
                    if let Some(header) = constants::GET_CHUNK_SIZE.captures(header[0]) {
                        if let Some(size) = header.get(1) {
                            let size = size.as_bytes();
                            header_len = size.len() + 2;
                            let size = String::from_utf8_lossy(size).into_owned();
                            body_chunk_size = usize::from_str_radix(size.as_str(), 16).unwrap();
                            self.transfer_encoding_status = TransferEncodingStatus::ChunkHeader;
                        }
                    } else {
                        // else return Error
                        // break;
                    }
                } else {
                    debug!(
                        "Chunked not complete: {}",
                        String::from_utf8_lossy(self.buffer.as_slice())
                    );
                }
            }

            if self.transfer_encoding_status == TransferEncodingStatus::ChunkHeader {
                if body_chunk_size == 0 {
                    debug!(
                        "0 chunked size received: {}",
                        String::from_utf8_lossy(self.buffer.as_slice())
                    );
                    self.buffer.clear(); // should we check that it is '0\r\n' ?
                    return Ok(true);
                }

                debug!(
                    "Before header cleanup: {}",
                    String::from_utf8_lossy(self.buffer.as_slice())
                );
                let buffer: Vec<u8> = self.buffer.drain(header_len..).collect();
                self.buffer = buffer;
                debug!(
                    "After header cleanup: {}",
                    String::from_utf8_lossy(self.buffer.as_slice())
                );
                self.transfer_encoding_status =
                    TransferEncodingStatus::ReadingBody(body_chunk_size);
            }

            if let TransferEncodingStatus::ReadingBody(body_chunk_size) =
                self.transfer_encoding_status
            {
                error!("!! {} > {}", self.buffer.len(), body_chunk_size);
                if self.buffer.len() > body_chunk_size {
                    let mut buffer: Vec<u8> = self.buffer.drain(body_chunk_size..).collect();
                    self.writer.write(self.buffer.as_slice()).await?;

                    let buffer2 = buffer.drain(2..).collect(); // CRLF
                    self.buffer = buffer2;
                    self.transfer_encoding_status = TransferEncodingStatus::ReadingHeader;
                } else {
                    can_process_buffer = false;
                }
            } else {
                can_process_buffer = false;
            }
        }

        return Ok(false);
    }
}

async fn log_request(request: &[u8], verbose: bool) {
    if !log_enabled!(Info) && !verbose {
        return;
    }
    let request: Vec<&[u8]> = constants::SPLIT_HEADERS_RE.splitn(request, 2).collect();
    let headers = String::from_utf8_lossy(&request[0]);
    let headers: Vec<&str> = constants::SPLIT_HEADER_RE.split(&headers).collect();
    let bodylen = if request.len() == 2 {
        let body = &request[1];
        body.len()
    } else {
        0
    };
    if log_enabled!(Info) {
        for header in headers {
            info!("> {}", header);
        }
        if bodylen > 0 {
            info!("> [{} bytes]", bodylen);
        }
        info!(">");
    } else if verbose {
        for header in headers {
            writeln!(&mut stderr(), "> {}", header).await.unwrap();
        }
        if bodylen > 0 {
            writeln!(&mut stderr(), "> [{} bytes]", bodylen)
                .await
                .unwrap();
        }
        writeln!(&mut stderr(), ">").await.unwrap();
    }
}

async fn from_http(
    request: &Request,
    client: &mut TcpStream,
    out: &mut (dyn Write + Unpin),
    verbose: bool,
    read_timeout: u64,
) -> CabotResult<()> {
    let request_bytes = request.to_bytes();
    let raw_request = request_bytes.as_slice();
    log_request(&raw_request, verbose).await;

    debug!("Sending request...");
    client.write_all(&raw_request).await?;

    debug!("Reading response headers...");

    let mut http_decoder = HttpDecoder::new(out, client, read_timeout);
    http_decoder.read_write().await?;
    Ok(())
}

async fn from_https(
    request: &Request,
    client: &mut TcpStream,
    out: &mut (dyn Write + Unpin),
    verbose: bool,
    read_timeout: u64,
) -> CabotResult<()> {
    let request_bytes = request.to_bytes();
    let raw_request = request_bytes.as_slice();
    log_request(&raw_request, verbose).await;

    let mut tls_client = TLSStream::new(client, request.host())?;
    tls_client.starttls().await?;

    debug!("Sending request...");
    tls_client.write_all(&raw_request).await?;
    debug!("Request sent");

    debug!("Decoding response...");
    let mut http_decoder = HttpDecoder::new(out, &mut tls_client, read_timeout);
    http_decoder.read_write().await?;
    Ok(())
}

pub async fn http_query(
    request: &Request,
    mut out: &mut (dyn Write + Unpin),
    authorities: &HashMap<String, SocketAddr>,
    verbose: bool,
    ipv4: bool,
    ipv6: bool,
    dns_timeout: u64,
    connect_timeout: u64,
    read_timeout: u64,
) -> CabotResult<()> {
    debug!(
        "HTTP Query {} {}",
        request.http_method(),
        request.request_uri()
    );

    let authority = request.authority();

    let addr = match authorities.get(authority) {
        Some(val) => {
            info!("Fetch authority {} using autorities map", authority);
            *val
        }
        None => {
            info!("Fetch authority {} using resolver", authority);
            let resolver = Resolver::new(verbose);
            resolver
                .get_addr(authority, ipv4, ipv6, dns_timeout)
                .await?
        }
    };

    info!("Connecting to {}", addr);
    let mut client = io::timeout(Duration::from_millis(connect_timeout), async {
        TcpStream::connect(addr).await
    })
    .await
    .map_err(|err| match err.kind() {
        io::ErrorKind::TimedOut => io::Error::new(err.kind(), "Connection Timeout".to_owned()),
        _ => err,
    })?;

    // client.set_read_timeout(Some(Duration::new(5, 0))).unwrap();

    match request.scheme() {
        "http" => from_http(request, &mut client, &mut out, verbose, read_timeout).await?,
        "https" => from_https(request, &mut client, &mut out, verbose, read_timeout).await?,
        _ => {
            return Err(CabotError::SchemeError(format!(
                "Unrecognized scheme {}",
                request.scheme()
            )))
        }
    };

    out.flush().await.unwrap();

    Ok(())
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate log;
extern crate loggerv;

#[macro_use]
extern crate clap;

extern crate serde_json;

extern crate time;

extern crate regex;

extern crate colored;

use clap::App;
use colored::*;
use log::LogLevel;
use regex::Regex;
use serde_json::Value;
use std::collections::BTreeMap as DataMap;
use std::fmt;
use std::fs::File;
use std::io::{self, Read};
use std::io::Write;
use std::process;
use time::get_time;

#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Default)]
struct MinionResult {
    command: Option<String>,
    retcode: Retcode,
    output: Option<String>,
    result: Option<String>,
    host: String,
}

type MinionResults = Vec<MinionResult>;

#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
enum Retcode {
    Success,
    Failure,
}

impl Retcode {
    fn is_success(&self) -> bool {
        self == &Retcode::Success
    }
}

impl Default for Retcode {
    fn default() -> Retcode {
        Retcode::Failure
    }
}

impl From<u64> for Retcode {
    fn from(input: u64) -> Self {
        if input == 0 {
            Retcode::Success
        } else {
            Retcode::Failure
        }
    }
}

fn main() {
    let yaml = load_yaml!("cli.yml");
    let matches = App::from_yaml(yaml)
        .version(crate_version!())
        .get_matches();
    trace!("matches: {:?}", matches);

    {
        let loglevel: LogLevel = value_t!(matches, "loglevel", LogLevel)
            .expect("can not parse loglevel from args");
        loggerv::init_with_level(loglevel).expect("can not initialize logger with parsed loglevel");
    }

    let changed = matches.is_present("changed");
    let no_save_file = matches.is_present("no_save_file");

    let host_data = {
        let input = matches.value_of("input").expect("can not get input file from args");

        let input = if input == "-" {
            let mut buffer = String::new();
            io::stdin().read_to_string(&mut buffer).expect("can not read from stdin");
            buffer
        } else {
            let mut file = File::open(input).expect("can not open input file");
            let mut input = String::new();
            file.read_to_string(&mut input).expect("can not read input file to string");
            input
        };

        // match all hosts that have not returned as they are not in the json data
        // format is normally like "Minion pricesearch did not respond. No job will be sent."
        let catch_not_returned_minions =
            Regex::new(r"(?m)^Minion (\S*) did not respond\. No job will be sent\.$")
                .expect("regex for catching not returned minions is not valid");

        for host in catch_not_returned_minions.captures_iter(input.as_str()) {
            warn!("host {} did not return", &host[1]);
        }

        // clean up hosts that have not returned from the json data
        catch_not_returned_minions.replace_all(input.as_str(), "").into_owned()
    };

    let value: Value = serde_json::from_str(host_data.as_str())
        .expect("can not convert input data to value. have you run the salt command with \
                 --static?");

    let results = match get_results(&value) {
        Ok(r) => r,
        Err(e) => {
            error!("can not get results from serde value: {}", e);
            if !no_save_file {
                write_save_file(host_data);
            }
            process::exit(1)
        }
    };

    trace!("results: {:#?}", results);

    let compressed = get_compressed(results);
    trace!("compressed: {:#?}", compressed);

    print_compressed(compressed, changed);
}

#[derive(Debug)]
enum ResultError {
    ConvertDiffToString,
    ConvertValueToString,
    ReturnCodeNotNumber,
    RetValueIsNone,
    RetValueIsNull,
    RetValueIsNumber,
    ValueNotAnObject,
}

impl fmt::Display for ResultError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            ResultError::ConvertDiffToString => write!(f, "can not convert diff to string"),
            ResultError::ConvertValueToString => write!(f, "can not convert value to string"),
            ResultError::ReturnCodeNotNumber => write!(f, "returncode is not a number"),
            ResultError::RetValueIsNone => write!(f, "ret value is none"),
            ResultError::RetValueIsNull => write!(f, "ret value is null"),
            ResultError::RetValueIsNumber => write!(f, "ret value is number"),
            ResultError::ValueNotAnObject => write!(f, "value it not an object"),
        }
    }
}

fn get_results(value: &Value) -> Result<MinionResults, ResultError> {
    if !value.is_object() {
        return Err(ResultError::ValueNotAnObject);
    }

    let mut results: MinionResults = Vec::new();

    for (host, values) in value.as_object().unwrap().iter() {
        trace!("host: {:#?}", host);
        trace!("values: {:#?}", values);

        let retcode: Retcode = match values.get("retcode") {
            Some(o) => {
                match o.as_u64() {
                    Some(v) => v.into(),
                    None => return Err(ResultError::ReturnCodeNotNumber),
                }
            }
            None => {
                warn!("host {} does not have a return code", host);
                Retcode::Failure
            }
        };

        let ret = match values.get("ret") {
            None => return Err(ResultError::RetValueIsNone),
            Some(r) => r,
        };

        match *ret {
            Value::Null => return Err(ResultError::RetValueIsNull),
            Value::Bool(r) => {
                let (command, result, output) = if r {
                    (None, Some("true".to_string()), None)
                } else {
                    (None, Some("false".to_string()), None)
                };

                results.push(MinionResult {
                    command: command,
                    host: host.clone(),
                    output: output,
                    result: result,
                    retcode: retcode,
                });
            }
            Value::Number(_) => return Err(ResultError::RetValueIsNumber),
            Value::String(ref r) => {
                results.push(MinionResult {
                    host: host.clone(),
                    result: Some(r.clone()),
                    retcode: retcode,
                    ..MinionResult::default()
                });
            }

            Value::Array(ref r) => {
                let values: Vec<_> = r.iter()
                    .map(|v| v.as_str().expect("can not convert the array value to a string"))
                    .collect();

                results.push(MinionResult {
                    host: host.clone(),
                    result: Some(values.join("\n").to_string()),
                    retcode: retcode.clone(),
                    ..MinionResult::default()
                });
            }
            Value::Object(ref r) => {
                if r.is_empty() {
                    results.push(MinionResult {
                        host: host.clone(),
                        retcode: retcode.clone(),
                        ..MinionResult::default()
                    });
                }

                for (command, command_result) in r.iter() {
                    trace!("command: {:#?}", command);
                    trace!("command_result: {:#?}", command_result);

                    let result = match command_result.get("comment") {
                        Some(r) => {
                            match r.as_str() {
                                Some(s) => Some(s.to_string()),
                                None => return Err(ResultError::ConvertValueToString),
                            }
                        }
                        None => None,
                    };

                    let output = match command_result.get("changes") {
                        Some(r) => {
                            match r.get("diff") {
                                Some(d) => {
                                    match d.as_str() {
                                        Some(i) => Some(i.to_string()),
                                        None => return Err(ResultError::ConvertDiffToString),
                                    }
                                }
                                None => None,
                            }
                        }
                        None => None,
                    };

                    results.push(MinionResult {
                        command: Some(command.to_string()),
                        host: host.clone(),
                        output: output,
                        result: result,
                        retcode: retcode.clone(),
                    });
                }
            }
        };
    }

    Ok(results)
}

fn get_compressed(results: MinionResults) -> DataMap<MinionResult, Vec<String>> {
    // compress output by changeing the hostname to the same value for all results and then just
    // adding all hosts with that value to the map.
    let mut compressed: DataMap<MinionResult, Vec<String>> = DataMap::new();
    for result in results {
        let mut result_no_host = result.clone();
        result_no_host.host = String::new();

        compressed.entry(result_no_host).or_insert_with(Vec::new).push(result.host);
    }

    compressed
}

fn print_compressed(compressed: DataMap<MinionResult, Vec<String>>, changed: bool) {
    let mut unchanged = 0;
    for (result, hosts) in compressed {
        // continue if we only want to print out changes and there are none and the command was a
        // success
        // TODO: make this a filter of the map
        if changed && !result.output.is_some() && result.retcode.is_success() {
            unchanged = unchanged + 1;
            continue;
        }

        println!("");
        println!("{}", "----------".bold());
        println!("");

        // state, command info
        {
            if result.command.is_some() {
                println!("{}", "------".purple());

                if result.command.is_some() {
                    println!("{}",
                             format!("COMMAND: {}", result.command.clone().unwrap()).purple());
                }

                println!("{}\n", "------".purple());
            }
        }

        // hosts
        {
            println!("{}", "------".cyan());
            println!("{}{}", "HOSTS: ".cyan(), hosts.join(", ").as_str());
            println!("{}\n", "------".cyan());
        }

        // output
        {
            println!("{}", "------".yellow());

            match result.retcode {
                Retcode::Success => println!("{}{}", "RETURN CODE: ".yellow(), "Success".green()),
                Retcode::Failure => println!("{}{}", "RETURN CODE: ".yellow(), "Failure".red()),
            }

            if result.result.is_some() {
                println!("{}", "RESULT:".yellow());
                println!("{}\n", result.result.unwrap());
            }

            println!("{}", "OUTPUT:".yellow());
            if result.output.is_some() {
                for line in result.output.unwrap().lines() {
                    if line.starts_with('-') {
                        println!("{}", line.red());
                        continue;
                    }

                    if line.starts_with('+') {
                        println!("{}", line.green());
                        continue;
                    }

                    println!("{}", line);
                }
            } else {
                println!("No changes");
            }
            println!("{}", "------".yellow());
        }
    }

    if changed {
        println!("");
        info!("{} unchanged states", unchanged);
    }
}

fn write_save_file(host_data: String) {
    let save_filename = format!("/tmp/salt-compressor_{}.json", get_time().sec);
    let mut save_file = File::create(save_filename.clone()).expect("can not create save_file");
    save_file.write_all(host_data.as_bytes()).expect("can not write host data to save_file");
    info!("please send me the save file under {} which contains the json \
           data from salt",
          save_filename);
}
<commit_msg>Now intigrating minions that did not return into results.<commit_after>#[macro_use]
extern crate log;
extern crate loggerv;

#[macro_use]
extern crate clap;

extern crate serde_json;

extern crate time;

extern crate regex;

extern crate colored;

use clap::App;
use colored::*;
use log::LogLevel;
use regex::Regex;
use serde_json::Value;
use std::collections::BTreeMap as DataMap;
use std::fmt;
use std::fs::File;
use std::io::{self, Read};
use std::io::Write;
use std::process;
use time::get_time;

#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Default)]
struct MinionResult {
    command: Option<String>,
    retcode: Retcode,
    output: Option<String>,
    result: Option<String>,
    host: String,
}

type MinionResults = Vec<MinionResult>;

#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
enum Retcode {
    Success,
    Failure,
}

impl Retcode {
    fn is_success(&self) -> bool {
        self == &Retcode::Success
    }
}

impl Default for Retcode {
    fn default() -> Retcode {
        Retcode::Failure
    }
}

impl From<u64> for Retcode {
    fn from(input: u64) -> Self {
        if input == 0 {
            Retcode::Success
        } else {
            Retcode::Failure
        }
    }
}

fn main() {
    let yaml = load_yaml!("cli.yml");
    let matches = App::from_yaml(yaml)
        .version(crate_version!())
        .get_matches();
    trace!("matches: {:?}", matches);

    {
        let loglevel: LogLevel = value_t!(matches, "loglevel", LogLevel)
            .expect("can not parse loglevel from args");
        loggerv::init_with_level(loglevel).expect("can not initialize logger with parsed loglevel");
    }

    let changed = matches.is_present("changed");
    let no_save_file = matches.is_present("no_save_file");

    let (host_data, no_return) = {
        let input = matches.value_of("input").expect("can not get input file from args");

        let input = if input == "-" {
            let mut buffer = String::new();
            io::stdin().read_to_string(&mut buffer).expect("can not read from stdin");
            buffer
        } else {
            let mut file = File::open(input).expect("can not open input file");
            let mut input = String::new();
            file.read_to_string(&mut input).expect("can not read input file to string");
            input
        };

        // match all hosts that have not returned as they are not in the json data
        // format is normally like "Minion pricesearch did not respond. No job will be sent."
        let catch_not_returned_minions =
            Regex::new(r"(?m)^Minion (\S*) did not respond\. No job will be sent\.$")
                .expect("regex for catching not returned minions is not valid");

        let mut no_return = Vec::new();
        for host in catch_not_returned_minions.captures_iter(input.as_str()) {
            no_return.push(host[1].to_string());
        }

        // clean up hosts that have not returned from the json data
        (catch_not_returned_minions.replace_all(input.as_str(), "").into_owned(), no_return)
    };

    let value: Value = serde_json::from_str(host_data.as_str())
        .expect("can not convert input data to value. have you run the salt command with \
                 --static?");

    let results = match get_results(&value, no_return) {
        Ok(r) => r,
        Err(e) => {
            error!("can not get results from serde value: {}", e);
            if !no_save_file {
                write_save_file(host_data);
            }
            process::exit(1)
        }
    };

    trace!("results: {:#?}", results);

    let compressed = get_compressed(results);
    trace!("compressed: {:#?}", compressed);

    print_compressed(compressed, changed);
}

#[derive(Debug)]
enum ResultError {
    ConvertDiffToString,
    ConvertValueToString,
    ReturnCodeNotNumber,
    RetValueIsNone,
    RetValueIsNull,
    RetValueIsNumber,
    ValueNotAnObject,
}

impl fmt::Display for ResultError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            ResultError::ConvertDiffToString => write!(f, "can not convert diff to string"),
            ResultError::ConvertValueToString => write!(f, "can not convert value to string"),
            ResultError::ReturnCodeNotNumber => write!(f, "returncode is not a number"),
            ResultError::RetValueIsNone => write!(f, "ret value is none"),
            ResultError::RetValueIsNull => write!(f, "ret value is null"),
            ResultError::RetValueIsNumber => write!(f, "ret value is number"),
            ResultError::ValueNotAnObject => write!(f, "value it not an object"),
        }
    }
}

fn get_results(value: &Value, no_return: Vec<String>) -> Result<MinionResults, ResultError> {
    if !value.is_object() {
        return Err(ResultError::ValueNotAnObject);
    }

    let mut results: MinionResults = Vec::new();

    for (host, values) in value.as_object().unwrap().iter() {
        trace!("host: {:#?}", host);
        trace!("values: {:#?}", values);

        let retcode: Retcode = match values.get("retcode") {
            Some(o) => {
                match o.as_u64() {
                    Some(v) => v.into(),
                    None => return Err(ResultError::ReturnCodeNotNumber),
                }
            }
            None => {
                warn!("host {} does not have a return code", host);
                Retcode::Failure
            }
        };

        let ret = match values.get("ret") {
            None => return Err(ResultError::RetValueIsNone),
            Some(r) => r,
        };

        match *ret {
            Value::Null => return Err(ResultError::RetValueIsNull),
            Value::Bool(r) => {
                let (command, result, output) = if r {
                    (None, Some("true".to_string()), None)
                } else {
                    (None, Some("false".to_string()), None)
                };

                results.push(MinionResult {
                    command: command,
                    host: host.clone(),
                    output: output,
                    result: result,
                    retcode: retcode,
                });
            }
            Value::Number(_) => return Err(ResultError::RetValueIsNumber),
            Value::String(ref r) => {
                results.push(MinionResult {
                    host: host.clone(),
                    result: Some(r.clone()),
                    retcode: retcode,
                    ..MinionResult::default()
                });
            }

            Value::Array(ref r) => {
                let values: Vec<_> = r.iter()
                    .map(|v| v.as_str().expect("can not convert the array value to a string"))
                    .collect();

                results.push(MinionResult {
                    host: host.clone(),
                    result: Some(values.join("\n").to_string()),
                    retcode: retcode.clone(),
                    ..MinionResult::default()
                });
            }
            Value::Object(ref r) => {
                if r.is_empty() {
                    results.push(MinionResult {
                        host: host.clone(),
                        retcode: retcode.clone(),
                        ..MinionResult::default()
                    });
                }

                for (command, command_result) in r.iter() {
                    trace!("command: {:#?}", command);
                    trace!("command_result: {:#?}", command_result);

                    let result = match command_result.get("comment") {
                        Some(r) => {
                            match r.as_str() {
                                Some(s) => Some(s.to_string()),
                                None => return Err(ResultError::ConvertValueToString),
                            }
                        }
                        None => None,
                    };

                    let output = match command_result.get("changes") {
                        Some(r) => {
                            match r.get("diff") {
                                Some(d) => {
                                    match d.as_str() {
                                        Some(i) => Some(i.to_string()),
                                        None => return Err(ResultError::ConvertDiffToString),
                                    }
                                }
                                None => None,
                            }
                        }
                        None => None,
                    };

                    results.push(MinionResult {
                        command: Some(command.to_string()),
                        host: host.clone(),
                        output: output,
                        result: result,
                        retcode: retcode.clone(),
                    });
                }
            }
        };
    }

    for host in no_return {
        results.push(MinionResult {
            host: host,
            retcode: Retcode::Failure,
            output: Some("Minion did not respond. No job will be sent.".to_string()),
            ..MinionResult::default()
        });
    }

    Ok(results)
}

fn get_compressed(results: MinionResults) -> DataMap<MinionResult, Vec<String>> {
    // compress output by changeing the hostname to the same value for all results and then just
    // adding all hosts with that value to the map.
    let mut compressed: DataMap<MinionResult, Vec<String>> = DataMap::new();
    for result in results {
        let mut result_no_host = result.clone();
        result_no_host.host = String::new();

        compressed.entry(result_no_host).or_insert_with(Vec::new).push(result.host);
    }

    compressed
}

fn print_compressed(compressed: DataMap<MinionResult, Vec<String>>, changed: bool) {
    let mut unchanged = 0;
    for (result, hosts) in compressed {
        // continue if we only want to print out changes and there are none and the command was a
        // success
        // TODO: make this a filter of the map
        if changed && !result.output.is_some() && result.retcode.is_success() {
            unchanged = unchanged + 1;
            continue;
        }

        println!("");
        println!("{}", "----------".bold());
        println!("");

        // state, command info
        {
            if result.command.is_some() {
                println!("{}", "------".purple());

                if result.command.is_some() {
                    println!("{}",
                             format!("COMMAND: {}", result.command.clone().unwrap()).purple());
                }

                println!("{}\n", "------".purple());
            }
        }

        // hosts
        {
            println!("{}", "------".cyan());
            println!("{}{}", "HOSTS: ".cyan(), hosts.join(", ").as_str());
            println!("{}\n", "------".cyan());
        }

        // output
        {
            println!("{}", "------".yellow());

            match result.retcode {
                Retcode::Success => println!("{}{}", "RETURN CODE: ".yellow(), "Success".green()),
                Retcode::Failure => println!("{}{}", "RETURN CODE: ".yellow(), "Failure".red()),
            }

            if result.result.is_some() {
                println!("{}", "RESULT:".yellow());
                println!("{}\n", result.result.unwrap());
            }

            println!("{}", "OUTPUT:".yellow());
            if result.output.is_some() {
                for line in result.output.unwrap().lines() {
                    if line.starts_with('-') {
                        println!("{}", line.red());
                        continue;
                    }

                    if line.starts_with('+') {
                        println!("{}", line.green());
                        continue;
                    }

                    println!("{}", line);
                }
            } else {
                println!("No changes");
            }
            println!("{}", "------".yellow());
        }
    }

    if changed {
        println!("");
        info!("{} unchanged states", unchanged);
    }
}

fn write_save_file(host_data: String) {
    let save_filename = format!("/tmp/salt-compressor_{}.json", get_time().sec);
    let mut save_file = File::create(save_filename.clone()).expect("can not create save_file");
    save_file.write_all(host_data.as_bytes()).expect("can not write host data to save_file");
    info!("please send me the save file under {} which contains the json \
           data from salt",
          save_filename);
}
<|endoftext|>


<commit_before>use std::env;
use std::fs;
use std::io::{self, ErrorKind, Read, Write, BufWriter};
use std::path::Path;
use std::process::exit;
use std::time::{Duration, Instant};

use getopts::Options;

use ppbert::prelude::*;
use ppbert::parsers::*;

const PROG_NAME: &str = "ppbert";

#[derive(Clone, Copy)]
enum ParserChoice {
    ByExtension,
    ForceBert1,
    ForceBert2,
    ForceDiskLog,
}

fn opt_usize(m: &getopts::Matches, opt: &str, default: usize) -> usize {
    match m.opt_get_default(opt, default) {
        Ok(n) => n,
        Err(_) => {
            eprintln!("'{}' must be a number", opt);
            exit(1);
        }
    }
}

fn main() {
    let mut opts = Options::new();
    opts.optflag("V", "version", "display version");
    opts.optflag("h", "help", "display this help");
    opts.optopt("i", "indent", "indent with NUM spaces", "NUM");
    opts.optopt("m", "per-line", "print at most NUM basic terms per line", "NUM");
    opts.optflag("p", "parse", "parse only, not pretty print");
    opts.optflag("1", "bert1", "force ppbert to use regular BERT parser");
    opts.optflag("2", "bert2", "force ppbert to use BERT2 parser");
    opts.optflag("d", "disk-log", "force ppbert to use DiskLog parser");
    opts.optflag("v", "verbose", "show diagnostics on stderr");
    opts.optflag("j", "json", "print as JSON");
    opts.optflag("t", "transform-proplists", "convert proplists to JSON objects");

    let mut matches = match opts.parse(env::args().skip(1)) {
        Ok(m) => m,
        Err(e) => {
            eprintln!("{}: {}", PROG_NAME, e);
            eprintln!("{}", opts.usage(&format!("{} {}", PROG_NAME, VERSION)));
            exit(1);
        }
    };

    if matches.opt_present("help") {
        println!("{}", opts.usage(&format!("{} {}", PROG_NAME, VERSION)));
        exit(0);
    }

    if matches.opt_present("version") {
        println!("{} {}", PROG_NAME, VERSION);
        exit(0);
    }

    // If no files to process, use stdin.
    if matches.free.is_empty() {
        matches.free.push("-".to_owned());
    }

    let indent_width = opt_usize(&matches, "indent", 2);
    let max_per_line = opt_usize(&matches, "per-line", 6);
    let parse_only = matches.opt_present("parse");
    let json = matches.opt_present("json");
    let transform_proplists = matches.opt_present("transform-proplists");
    let verbose = matches.opt_present("verbose");

    let parser_choice =
        if matches.opt_present("bert1") {
            ParserChoice::ForceBert1
        } else if matches.opt_present("bert2") {
            ParserChoice::ForceBert2
        } else if matches.opt_present("disk-log") {
            ParserChoice::ForceDiskLog
        } else {
            ParserChoice::ByExtension
        };

    let output_fn = match (json, transform_proplists) {
        (true, false)  => pp_json,
        (true, true)   => pp_json_proplist,
        (false, false) => pp_bert,
        (false, true)  => {
            eprintln!("{}: --transform-proplists is only valid with the --json flag", PROG_NAME);
            pp_bert
        }
    };

    let mut return_code = 0;
    for file in &matches.free {
        let res = handle_file(file, parse_only, verbose,
                              indent_width, max_per_line,
                              parser_choice, output_fn);
        match res {
            Ok(()) => (),
            Err(ref e) => {
                if broken_pipe(e) {
                    break;
                }
                return_code = 1;
                eprintln!("{}: {}: {}", PROG_NAME, file, e);
            }
        }
    }
    exit(return_code);
}

fn broken_pipe(err: &BertError) -> bool {
    match *err {
        BertError::IoError(ref ioerr) =>
            ioerr.kind() == ErrorKind::BrokenPipe,
        _ => false
    }
}

fn read_bytes(filename: &str) -> Result<Vec<u8>> {
    if filename == "-" {
        let stdin = io::stdin();
        let mut stdin = stdin.lock();
        let mut buf: Vec<u8> = Vec::with_capacity(4096);
        stdin.read_to_end(&mut buf)?;
        return Ok(buf);
    } else {
        let buf = fs::read(filename)?;
        return Ok(buf);
    }
}

fn parser_from_ext(filename: &str, bytes: Vec<u8>) -> Box<Parser> {
    let ext: Option<&str> =
        Path::new(filename)
        .extension()
        .and_then(|x| x.to_str());
    match ext {
        Some("bert") | Some("bert1") => Box::new(Bert1Parser::new(bytes)),
        Some("bert2") => Box::new(Bert2Parser::new(bytes)),
        Some("log") => Box::new(DiskLogParser::new(bytes)),
        _ => {
            eprintln!("{}: cannot find an appropriate parser for {}; using BERT",
                      PROG_NAME, filename);
            Box::new(Bert1Parser::new(bytes))
        },
    }
}

fn handle_file(
    filename: &str,
    parse_only: bool,
    verbose: bool,
    indent: usize,
    terms_per_line: usize,
    parser_choice: ParserChoice,
    pp_fn: fn(BertTerm, usize, usize) -> Result<()>
) -> Result<()> {

    // Read file or stdin into buffer
    let now = Instant::now();
    let bytes = read_bytes(filename)?;
    let read_dur = now.elapsed();

    let mut parser: Box<Parser> = match parser_choice {
        ParserChoice::ForceBert1 => Box::new(Bert1Parser::new(bytes)),
        ParserChoice::ForceBert2 => Box::new(Bert2Parser::new(bytes)),
        ParserChoice::ForceDiskLog => Box::new(DiskLogParser::new(bytes)),
        ParserChoice::ByExtension => parser_from_ext(filename, bytes),
    };

    let mut parse_dur = Duration::new(0, 0);
    let mut pp_dur = Duration::new(0, 0);

    loop {
        let now = Instant::now();
        let term = match parser.next() {
            Some(term) => term?,
            None => break,
        };
        parse_dur += now.elapsed();
        if !parse_only {
            let now = Instant::now();
            pp_fn(term, indent, terms_per_line)?;
            pp_dur += now.elapsed();
        }
    }

    if verbose {
        eprintln!("{}: {} read time: {}.{:06} seconds", PROG_NAME, filename, read_dur.as_secs(), read_dur.subsec_micros());
        eprintln!("{}: {} parse time: {}.{:06} seconds", PROG_NAME, filename, parse_dur.as_secs(), parse_dur.subsec_micros());
        if !parse_only {
            eprintln!("{}: {} print time: {}.{:06} seconds", PROG_NAME, filename, pp_dur.as_secs(), pp_dur.subsec_micros());
        }
    }

    return Ok(());
}

/// Outputs a BertTerm to stdout.
fn pp_bert(term: BertTerm, indent_width: usize, terms_per_line: usize) -> Result<()> {
    let stdout = io::stdout();
    let stdout = stdout.lock();
    let mut stdout = BufWriter::new(stdout);
    term.write_as_erlang(&mut stdout, indent_width, terms_per_line)?;
    stdout.flush()?;
    writeln!(&mut stdout, "")?;
    return Ok(());
}

/// Outputs a BertTerm as JSON to stdout.
fn pp_json(term: BertTerm, _: usize, _: usize) -> Result<()> {
    let stdout = io::stdout();
    let stdout = stdout.lock();
    let mut stdout = BufWriter::new(stdout);
    term.write_as_json(&mut stdout, false)?;
    writeln!(&mut stdout, "")?;
    stdout.flush()?;
    return Ok(());
}

/// Outputs a BertTerm as JSON to stdout;
/// Erlang proplists are converted to JSON objects.
fn pp_json_proplist(term: BertTerm, _: usize, _: usize) -> Result<()> {
    let stdout = io::stdout();
    let stdout = stdout.lock();
    let mut stdout = BufWriter::new(stdout);
    term.write_as_json(&mut stdout, true)?;
    writeln!(&mut stdout, "")?;
    stdout.flush()?;
    return Ok(());
}
<commit_msg>Use Duration's Debug impl to show times<commit_after>use std::env;
use std::fs;
use std::io::{self, ErrorKind, Read, Write, BufWriter};
use std::path::Path;
use std::process::exit;
use std::time::{Duration, Instant};

use getopts::Options;

use ppbert::prelude::*;
use ppbert::parsers::*;

const PROG_NAME: &str = "ppbert";

#[derive(Clone, Copy)]
enum ParserChoice {
    ByExtension,
    ForceBert1,
    ForceBert2,
    ForceDiskLog,
}

fn opt_usize(m: &getopts::Matches, opt: &str, default: usize) -> usize {
    match m.opt_get_default(opt, default) {
        Ok(n) => n,
        Err(_) => {
            eprintln!("'{}' must be a number", opt);
            exit(1);
        }
    }
}

fn main() {
    let mut opts = Options::new();
    opts.optflag("V", "version", "display version");
    opts.optflag("h", "help", "display this help");
    opts.optopt("i", "indent", "indent with NUM spaces", "NUM");
    opts.optopt("m", "per-line", "print at most NUM basic terms per line", "NUM");
    opts.optflag("p", "parse", "parse only, not pretty print");
    opts.optflag("1", "bert1", "force ppbert to use regular BERT parser");
    opts.optflag("2", "bert2", "force ppbert to use BERT2 parser");
    opts.optflag("d", "disk-log", "force ppbert to use DiskLog parser");
    opts.optflag("v", "verbose", "show diagnostics on stderr");
    opts.optflag("j", "json", "print as JSON");
    opts.optflag("t", "transform-proplists", "convert proplists to JSON objects");

    let mut matches = match opts.parse(env::args().skip(1)) {
        Ok(m) => m,
        Err(e) => {
            eprintln!("{}: {}", PROG_NAME, e);
            eprintln!("{}", opts.usage(&format!("{} {}", PROG_NAME, VERSION)));
            exit(1);
        }
    };

    if matches.opt_present("help") {
        println!("{}", opts.usage(&format!("{} {}", PROG_NAME, VERSION)));
        exit(0);
    }

    if matches.opt_present("version") {
        println!("{} {}", PROG_NAME, VERSION);
        exit(0);
    }

    // If no files to process, use stdin.
    if matches.free.is_empty() {
        matches.free.push("-".to_owned());
    }

    let indent_width = opt_usize(&matches, "indent", 2);
    let max_per_line = opt_usize(&matches, "per-line", 6);
    let parse_only = matches.opt_present("parse");
    let json = matches.opt_present("json");
    let transform_proplists = matches.opt_present("transform-proplists");
    let verbose = matches.opt_present("verbose");

    let parser_choice =
        if matches.opt_present("bert1") {
            ParserChoice::ForceBert1
        } else if matches.opt_present("bert2") {
            ParserChoice::ForceBert2
        } else if matches.opt_present("disk-log") {
            ParserChoice::ForceDiskLog
        } else {
            ParserChoice::ByExtension
        };

    let output_fn = match (json, transform_proplists) {
        (true, false)  => pp_json,
        (true, true)   => pp_json_proplist,
        (false, false) => pp_bert,
        (false, true)  => {
            eprintln!("{}: --transform-proplists is only valid with the --json flag", PROG_NAME);
            pp_bert
        }
    };

    let mut return_code = 0;
    for file in &matches.free {
        let res = handle_file(file, parse_only, verbose,
                              indent_width, max_per_line,
                              parser_choice, output_fn);
        match res {
            Ok(()) => (),
            Err(ref e) => {
                if broken_pipe(e) {
                    break;
                }
                return_code = 1;
                eprintln!("{}: {}: {}", PROG_NAME, file, e);
            }
        }
    }
    exit(return_code);
}

fn broken_pipe(err: &BertError) -> bool {
    match *err {
        BertError::IoError(ref ioerr) =>
            ioerr.kind() == ErrorKind::BrokenPipe,
        _ => false
    }
}

fn read_bytes(filename: &str) -> Result<Vec<u8>> {
    if filename == "-" {
        let stdin = io::stdin();
        let mut stdin = stdin.lock();
        let mut buf: Vec<u8> = Vec::with_capacity(4096);
        stdin.read_to_end(&mut buf)?;
        return Ok(buf);
    } else {
        let buf = fs::read(filename)?;
        return Ok(buf);
    }
}

fn parser_from_ext(filename: &str, bytes: Vec<u8>) -> Box<Parser> {
    let ext: Option<&str> =
        Path::new(filename)
        .extension()
        .and_then(|x| x.to_str());
    match ext {
        Some("bert") | Some("bert1") => Box::new(Bert1Parser::new(bytes)),
        Some("bert2") => Box::new(Bert2Parser::new(bytes)),
        Some("log") => Box::new(DiskLogParser::new(bytes)),
        _ => {
            eprintln!("{}: cannot find an appropriate parser for {}; using BERT",
                      PROG_NAME, filename);
            Box::new(Bert1Parser::new(bytes))
        },
    }
}

fn handle_file(
    filename: &str,
    parse_only: bool,
    verbose: bool,
    indent: usize,
    terms_per_line: usize,
    parser_choice: ParserChoice,
    pp_fn: fn(BertTerm, usize, usize) -> Result<()>
) -> Result<()> {

    // Read file or stdin into buffer
    let now = Instant::now();
    let bytes = read_bytes(filename)?;
    let read_dur = now.elapsed();

    let mut parser: Box<Parser> = match parser_choice {
        ParserChoice::ForceBert1 => Box::new(Bert1Parser::new(bytes)),
        ParserChoice::ForceBert2 => Box::new(Bert2Parser::new(bytes)),
        ParserChoice::ForceDiskLog => Box::new(DiskLogParser::new(bytes)),
        ParserChoice::ByExtension => parser_from_ext(filename, bytes),
    };

    let mut parse_dur = Duration::new(0, 0);
    let mut pp_dur = Duration::new(0, 0);

    loop {
        let now = Instant::now();
        let term = match parser.next() {
            Some(term) => term?,
            None => break,
        };
        parse_dur += now.elapsed();
        if !parse_only {
            let now = Instant::now();
            pp_fn(term, indent, terms_per_line)?;
            pp_dur += now.elapsed();
        }
    }

    if verbose {
        eprintln!("{}: {} read time: {:?}", PROG_NAME, filename, read_dur);
        eprintln!("{}: {} parse time: {:?}", PROG_NAME, filename, parse_dur);
        if !parse_only {
            eprintln!("{}: {} print time: {:?}", PROG_NAME, filename, pp_dur);
        }
    }

    return Ok(());
}

/// Outputs a BertTerm to stdout.
fn pp_bert(term: BertTerm, indent_width: usize, terms_per_line: usize) -> Result<()> {
    let stdout = io::stdout();
    let stdout = stdout.lock();
    let mut stdout = BufWriter::new(stdout);
    term.write_as_erlang(&mut stdout, indent_width, terms_per_line)?;
    stdout.flush()?;
    writeln!(&mut stdout, "")?;
    return Ok(());
}

/// Outputs a BertTerm as JSON to stdout.
fn pp_json(term: BertTerm, _: usize, _: usize) -> Result<()> {
    let stdout = io::stdout();
    let stdout = stdout.lock();
    let mut stdout = BufWriter::new(stdout);
    term.write_as_json(&mut stdout, false)?;
    writeln!(&mut stdout, "")?;
    stdout.flush()?;
    return Ok(());
}

/// Outputs a BertTerm as JSON to stdout;
/// Erlang proplists are converted to JSON objects.
fn pp_json_proplist(term: BertTerm, _: usize, _: usize) -> Result<()> {
    let stdout = io::stdout();
    let stdout = stdout.lock();
    let mut stdout = BufWriter::new(stdout);
    term.write_as_json(&mut stdout, true)?;
    writeln!(&mut stdout, "")?;
    stdout.flush()?;
    return Ok(());
}
<|endoftext|>


<commit_before>
#![feature(old_io, collections)]

extern crate lib_word_count;
use lib_word_count as word_count;

use std::old_io as io;
use std::old_io::Reader;

fn main() {

    let mut word_index = Vec::new();

    let input = match io::stdio::stdin().read_to_string() {
        Ok(result) => result,
        Err(error) => panic!("Could not read from stdin: {:?}", error)
    };

    word_count::count_words(&input, &mut word_index);

    for indexed_word in word_index {
        if indexed_word.word.len() >= 5 {
            println!("'{}':\t{}", indexed_word.word, indexed_word.appeared);
        } else {
            println!("'{}':\t\t{}", indexed_word.word, indexed_word.appeared)
        }
    }
}
<commit_msg>Add documentation to src/main.rs<commit_after>//!
//! Main executable.
//!
//! This file compiles to an executable.
//!
//! # Panics
//!
//! Panics when it can't read from stdin.
//!
//! # Examples
//!
//! ```sh
//! $ echo 'I like cookies. Mmm... Cookies.' | word_count
//! 'cookies':	  2
//! 'mmm':		  1
//! 'like':		  1
//! 'i':		  1
//! ```

#![feature(old_io, collections)]

extern crate lib_word_count;
use lib_word_count as word_count;

use std::old_io as io;
use std::old_io::Reader;

fn main() {

    let mut word_index = Vec::new();

    let input = match io::stdio::stdin().read_to_string() {
        Ok(result) => result,
        Err(error) => panic!("Could not read from stdin: {:?}", error)
    };

    word_count::count_words(&input, &mut word_index);

    for indexed_word in word_index {
        if indexed_word.word.len() >= 5 {
            println!("'{}':\t{}", indexed_word.word, indexed_word.appeared);
        } else {
            println!("'{}':\t\t{}", indexed_word.word, indexed_word.appeared)
        }
    }
}
<|endoftext|>


<commit_before>/*
** A simple program that allows basic arithmetic operations using roman numerals and conversions to decimal numbers.
*/

fn main() {
    println!("Bismillah");
}<commit_msg>Add function to convert roman numerals into decimal numbers.<commit_after>/*
** A simple program that allows basic arithmetic operations using roman numerals and conversions to decimal numbers.
*/

fn main() {

	let roman_number = "CLIII";
	let decimal_number = to_decimal(&roman_number);

	println!("The roman number {} is equal to {} in decimal numbers", roman_number, decimal_number);

	let roman_number = "XXXIV";
	let decimal_number = to_decimal(&roman_number);

	println!("The roman number {} is equal to {} in decimal numbers", roman_number, decimal_number);
}

fn to_decimal(roman_number: &str) -> i32 {

	let mut num = 0;

	for c in roman_number.chars() {
		num += roman_char_to_decimal(c);
	}

	num
}

fn roman_char_to_decimal(roman_char: char) -> i32 {

	match roman_char {
	    'I' => 1,
	    'V' => 5,
	    'X' => 10,
	    'L' => 50,
	    'C' => 100,
	    'D' => 500,
	    'M' => 1000,
	    _ => 0,
	}
}<|endoftext|>


<commit_before>/*
    Copyright © 2021 Alastair Feille

    This Source Code Form is subject to the terms of the Mozilla Public
    License, v. 2.0. If a copy of the MPL was not distributed with this
    file, You can obtain one at https://mozilla.org/MPL/2.0/.

    SPDX-License-Identifier: MPL-2.0
*/

use std::{io,
          io::Write,
          process,
          process::Command};

fn main()
{
    loop
    {
        // Print prompt
        print!("> ");
        io::stdout().flush().expect("failed to print prompt");

        // Read in line
        let mut line = String::new();
        io::stdin().read_line(&mut line)
                   .expect("failed to read from stdin");

        // Skip empty lines
        if line.trim().is_empty()
        {
            continue;
        }

        // Split line into command and arguments
        let tokens: Vec<&str> = line.split_whitespace().collect();
        let (head, args) = tokens.split_at(1);
        if let Some(cmd) = head.get(0)
        {
            if cmd.to_string() == "exit"
            {
                process::exit(0);
            }
            match Command::new(cmd).args(args).spawn()
            {
                Ok(mut child) =>
                {
                    let _ecode = child.wait().expect("failed to wait on child");
                },
                Err(e) => eprintln!("mush: could not run command {}: {}", cmd, e),
            }
        }
    }
}
<commit_msg>feat: add current working directory to prompt<commit_after>/*
    Copyright © 2021 Alastair Feille

    This Source Code Form is subject to the terms of the Mozilla Public
    License, v. 2.0. If a copy of the MPL was not distributed with this
    file, You can obtain one at https://mozilla.org/MPL/2.0/.

    SPDX-License-Identifier: MPL-2.0
*/

use std::{env,
          fs,
          io,
          io::Write,
          process,
          process::Command};

fn main()
{
    loop
    {
        let cwd: String = env::current_dir().map(|p| {
                                                // Attempt to canonicalize the path else return "???"
                                                fs::canonicalize(p).unwrap_or("???".into())
                                                                   .to_string_lossy()
                                                                   .into()
                                            })
                                            // Attempt to get the current path else return "???"
                                            .unwrap_or("???".into());
        // Print prompt
        print!("{}> ", cwd);
        io::stdout().flush().expect("failed to print prompt");

        // Read in line
        let mut line = String::new();
        io::stdin().read_line(&mut line)
                   .expect("failed to read from stdin");

        // Skip empty lines
        if line.trim().is_empty()
        {
            continue;
        }

        // Split line into command and arguments
        let tokens: Vec<&str> = line.split_whitespace().collect();
        let (head, args) = tokens.split_at(1);
        if let Some(cmd) = head.get(0)
        {
            if cmd.to_string() == "exit"
            {
                process::exit(0);
            }
            match Command::new(cmd).args(args).spawn()
            {
                Ok(mut child) =>
                {
                    let _ecode = child.wait().expect("failed to wait on child");
                },
                Err(e) => eprintln!("mush: could not run command {}: {}", cmd, e),
            }
        }
    }
}
<|endoftext|>