<commit_before>/// The underlying OsString/OsStr implementation on Windows is a
/// wrapper around the "WTF-8" encoding; see the `wtf8` module for more.
use crate::borrow::Cow;
use crate::fmt;
use crate::mem;
use crate::rc::Rc;
use crate::sync::Arc;
use crate::sys_common::wtf8::{Wtf8, Wtf8Buf};
use crate::sys_common::{AsInner, FromInner, IntoInner};

#[derive(Clone, Hash)]
pub struct Buf {
    pub inner: Wtf8Buf,
}

impl IntoInner<Wtf8Buf> for Buf {
    fn into_inner(self) -> Wtf8Buf {
        self.inner
    }
}

impl FromInner<Wtf8Buf> for Buf {
    fn from_inner(inner: Wtf8Buf) -> Self {
        Buf { inner }
    }
}

impl AsInner<Wtf8> for Buf {
    fn as_inner(&self) -> &Wtf8 {
        &self.inner
    }
}

impl fmt::Debug for Buf {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Debug::fmt(self.as_slice(), formatter)
    }
}

impl fmt::Display for Buf {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Display::fmt(self.as_slice(), formatter)
    }
}

pub struct Slice {
    pub inner: Wtf8,
}

impl fmt::Debug for Slice {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Debug::fmt(&self.inner, formatter)
    }
}

impl fmt::Display for Slice {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Display::fmt(&self.inner, formatter)
    }
}

impl Buf {
    pub fn with_capacity(capacity: usize) -> Buf {
        Buf { inner: Wtf8Buf::with_capacity(capacity) }
    }

    pub fn clear(&mut self) {
        self.inner.clear()
    }

    pub fn capacity(&self) -> usize {
        self.inner.capacity()
    }

    pub fn from_string(s: String) -> Buf {
        Buf { inner: Wtf8Buf::from_string(s) }
    }

    pub fn as_slice(&self) -> &Slice {
        // Safety: Slice is just a wrapper for Wtf8,
        // and as_slice returns &Wtf8. Therefore,
        // transmute &Wtf8 to &Slice is safe.
        unsafe { mem::transmute(self.inner.as_slice()) }
    }

    #[inline]
    pub fn as_mut_slice(&mut self) -> &mut Slice {
        // Safety: Slice is just a wrapper for Wtf8,
        // and as_slice returns &Wtf8. Therefore,
        // transmute &mut Wtf8 to &mut Slice is safe.
        unsafe { mem::transmute(self.inner.as_mut_slice()) }
    }

    pub fn into_string(self) -> Result<String, Buf> {
        self.inner.into_string().map_err(|buf| Buf { inner: buf })
    }

    pub fn push_slice(&mut self, s: &Slice) {
        self.inner.push_wtf8(&s.inner)
    }

    pub fn reserve(&mut self, additional: usize) {
        self.inner.reserve(additional)
    }

    pub fn reserve_exact(&mut self, additional: usize) {
        self.inner.reserve_exact(additional)
    }

    pub fn shrink_to_fit(&mut self) {
        self.inner.shrink_to_fit()
    }

    #[inline]
    pub fn shrink_to(&mut self, min_capacity: usize) {
        self.inner.shrink_to(min_capacity)
    }

    #[inline]
    pub fn into_box(self) -> Box<Slice> {
        unsafe { mem::transmute(self.inner.into_box()) }
    }

    #[inline]
    pub fn from_box(boxed: Box<Slice>) -> Buf {
        let inner: Box<Wtf8> = unsafe { mem::transmute(boxed) };
        Buf { inner: Wtf8Buf::from_box(inner) }
    }

    #[inline]
    pub fn into_arc(&self) -> Arc<Slice> {
        self.as_slice().into_arc()
    }

    #[inline]
    pub fn into_rc(&self) -> Rc<Slice> {
        self.as_slice().into_rc()
    }
}

impl Slice {
    #[inline]
    pub fn from_str(s: &str) -> &Slice {
        unsafe { mem::transmute(Wtf8::from_str(s)) }
    }

    pub fn to_str(&self) -> Option<&str> {
        self.inner.as_str()
    }

    pub fn to_string_lossy(&self) -> Cow<'_, str> {
        self.inner.to_string_lossy()
    }

    pub fn to_owned(&self) -> Buf {
        let mut buf = Wtf8Buf::with_capacity(self.inner.len());
        buf.push_wtf8(&self.inner);
        Buf { inner: buf }
    }

    #[inline]
    pub fn into_box(&self) -> Box<Slice> {
        unsafe { mem::transmute(self.inner.into_box()) }
    }

    pub fn empty_box() -> Box<Slice> {
        unsafe { mem::transmute(Wtf8::empty_box()) }
    }

    #[inline]
    pub fn into_arc(&self) -> Arc<Slice> {
        let arc = self.inner.into_arc();
        unsafe { Arc::from_raw(Arc::into_raw(arc) as *const Slice) }
    }

    #[inline]
    pub fn into_rc(&self) -> Rc<Slice> {
        let rc = self.inner.into_rc();
        unsafe { Rc::from_raw(Rc::into_raw(rc) as *const Slice) }
    }
}
<commit_msg>remove #[inline] for consistency in windows/os_str<commit_after>/// The underlying OsString/OsStr implementation on Windows is a
/// wrapper around the "WTF-8" encoding; see the `wtf8` module for more.
use crate::borrow::Cow;
use crate::fmt;
use crate::mem;
use crate::rc::Rc;
use crate::sync::Arc;
use crate::sys_common::wtf8::{Wtf8, Wtf8Buf};
use crate::sys_common::{AsInner, FromInner, IntoInner};

#[derive(Clone, Hash)]
pub struct Buf {
    pub inner: Wtf8Buf,
}

impl IntoInner<Wtf8Buf> for Buf {
    fn into_inner(self) -> Wtf8Buf {
        self.inner
    }
}

impl FromInner<Wtf8Buf> for Buf {
    fn from_inner(inner: Wtf8Buf) -> Self {
        Buf { inner }
    }
}

impl AsInner<Wtf8> for Buf {
    fn as_inner(&self) -> &Wtf8 {
        &self.inner
    }
}

impl fmt::Debug for Buf {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Debug::fmt(self.as_slice(), formatter)
    }
}

impl fmt::Display for Buf {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Display::fmt(self.as_slice(), formatter)
    }
}

pub struct Slice {
    pub inner: Wtf8,
}

impl fmt::Debug for Slice {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Debug::fmt(&self.inner, formatter)
    }
}

impl fmt::Display for Slice {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Display::fmt(&self.inner, formatter)
    }
}

impl Buf {
    pub fn with_capacity(capacity: usize) -> Buf {
        Buf { inner: Wtf8Buf::with_capacity(capacity) }
    }

    pub fn clear(&mut self) {
        self.inner.clear()
    }

    pub fn capacity(&self) -> usize {
        self.inner.capacity()
    }

    pub fn from_string(s: String) -> Buf {
        Buf { inner: Wtf8Buf::from_string(s) }
    }

    pub fn as_slice(&self) -> &Slice {
        // Safety: Slice is just a wrapper for Wtf8,
        // and as_slice returns &Wtf8. Therefore,
        // transmute &Wtf8 to &Slice is safe.
        unsafe { mem::transmute(self.inner.as_slice()) }
    }

    pub fn as_mut_slice(&mut self) -> &mut Slice {
        // Safety: Slice is just a wrapper for Wtf8,
        // and as_slice returns &Wtf8. Therefore,
        // transmute &mut Wtf8 to &mut Slice is safe.
        unsafe { mem::transmute(self.inner.as_mut_slice()) }
    }

    pub fn into_string(self) -> Result<String, Buf> {
        self.inner.into_string().map_err(|buf| Buf { inner: buf })
    }

    pub fn push_slice(&mut self, s: &Slice) {
        self.inner.push_wtf8(&s.inner)
    }

    pub fn reserve(&mut self, additional: usize) {
        self.inner.reserve(additional)
    }

    pub fn reserve_exact(&mut self, additional: usize) {
        self.inner.reserve_exact(additional)
    }

    pub fn shrink_to_fit(&mut self) {
        self.inner.shrink_to_fit()
    }

    #[inline]
    pub fn shrink_to(&mut self, min_capacity: usize) {
        self.inner.shrink_to(min_capacity)
    }

    #[inline]
    pub fn into_box(self) -> Box<Slice> {
        unsafe { mem::transmute(self.inner.into_box()) }
    }

    #[inline]
    pub fn from_box(boxed: Box<Slice>) -> Buf {
        let inner: Box<Wtf8> = unsafe { mem::transmute(boxed) };
        Buf { inner: Wtf8Buf::from_box(inner) }
    }

    #[inline]
    pub fn into_arc(&self) -> Arc<Slice> {
        self.as_slice().into_arc()
    }

    #[inline]
    pub fn into_rc(&self) -> Rc<Slice> {
        self.as_slice().into_rc()
    }
}

impl Slice {
    #[inline]
    pub fn from_str(s: &str) -> &Slice {
        unsafe { mem::transmute(Wtf8::from_str(s)) }
    }

    pub fn to_str(&self) -> Option<&str> {
        self.inner.as_str()
    }

    pub fn to_string_lossy(&self) -> Cow<'_, str> {
        self.inner.to_string_lossy()
    }

    pub fn to_owned(&self) -> Buf {
        let mut buf = Wtf8Buf::with_capacity(self.inner.len());
        buf.push_wtf8(&self.inner);
        Buf { inner: buf }
    }

    #[inline]
    pub fn into_box(&self) -> Box<Slice> {
        unsafe { mem::transmute(self.inner.into_box()) }
    }

    pub fn empty_box() -> Box<Slice> {
        unsafe { mem::transmute(Wtf8::empty_box()) }
    }

    #[inline]
    pub fn into_arc(&self) -> Arc<Slice> {
        let arc = self.inner.into_arc();
        unsafe { Arc::from_raw(Arc::into_raw(arc) as *const Slice) }
    }

    #[inline]
    pub fn into_rc(&self) -> Rc<Slice> {
        let rc = self.inner.into_rc();
        unsafe { Rc::from_raw(Rc::into_raw(rc) as *const Slice) }
    }
}
<|endoftext|>


<commit_before>use encoding_rs::Encoding;
use html5ever::rcdom::RcDom;
use reqwest::blocking::Client;
use reqwest::header::{HeaderMap, HeaderValue, USER_AGENT};
use std::collections::HashMap;
use std::fs;
use std::io::{self, prelude::*, Error, Write};
use std::path::Path;
use std::process;
use std::time::Duration;
use url::Url;

use monolith::html::{
    add_favicon, create_metadata_tag, get_base_url, get_charset, has_favicon, html_to_dom,
    serialize_document, set_base_url, set_charset, walk_and_embed_assets,
};
use monolith::opts::Options;
use monolith::url::{create_data_url, resolve_url};
use monolith::utils::retrieve_asset;

enum Output {
    Stdout(io::Stdout),
    File(fs::File),
}

impl Output {
    fn new(file_path: &str) -> Result<Output, Error> {
        if file_path.is_empty() || file_path.eq("-") {
            Ok(Output::Stdout(io::stdout()))
        } else {
            Ok(Output::File(fs::File::create(file_path)?))
        }
    }

    fn write(&mut self, bytes: &Vec<u8>) -> Result<(), Error> {
        match self {
            Output::Stdout(stdout) => {
                stdout.write_all(bytes)?;
                // Ensure newline at end of output
                if bytes.last() != Some(&b"\n"[0]) {
                    stdout.write(b"\n")?;
                }
                stdout.flush()
            }
            Output::File(file) => {
                file.write_all(bytes)?;
                // Ensure newline at end of output
                if bytes.last() != Some(&b"\n"[0]) {
                    file.write(b"\n")?;
                }
                file.flush()
            }
        }
    }
}

pub fn read_stdin() -> Vec<u8> {
    let mut buffer: Vec<u8> = vec![];

    match io::stdin().lock().read_to_end(&mut buffer) {
        Ok(_) => buffer,
        Err(_) => buffer,
    }
}

fn main() {
    let options = Options::from_args();

    // Check if target was provided
    if options.target.len() == 0 {
        if !options.silent {
            eprintln!("No target specified");
        }
        process::exit(1);
    }

    // Check if custom charset is valid
    if let Some(custom_charset) = options.charset.clone() {
        if !Encoding::for_label_no_replacement(custom_charset.as_bytes()).is_some() {
            eprintln!("Unknown encoding: {}", &custom_charset);
            process::exit(1);
        }
    }

    let mut use_stdin: bool = false;

    let target_url = match options.target.as_str() {
        "-" => {
            // Read from pipe (stdin)
            use_stdin = true;
            // Set default target URL to an empty data URL; the user can set it via --base-url
            Url::parse("data:text/html,").unwrap()
        }
        target => match Url::parse(&target) {
            Ok(url) => match url.scheme() {
                "data" | "file" | "http" | "https" => url,
                unsupported_scheme => {
                    if !options.silent {
                        eprintln!("Unsupported target URL type: {}", unsupported_scheme);
                    }
                    process::exit(1)
                }
            },
            Err(_) => {
                // Failed to parse given base URL (perhaps it's a filesystem path?)
                let path: &Path = Path::new(&target);
                match path.exists() {
                    true => match path.is_file() {
                        true => {
                            let canonical_path = fs::canonicalize(&path).unwrap();
                            match Url::from_file_path(canonical_path) {
                                Ok(url) => url,
                                Err(_) => {
                                    if !options.silent {
                                        eprintln!(
                                            "Could not generate file URL out of given path: {}",
                                            &target
                                        );
                                    }
                                    process::exit(1);
                                }
                            }
                        }
                        false => {
                            if !options.silent {
                                eprintln!("Local target is not a file: {}", &target);
                            }
                            process::exit(1);
                        }
                    },
                    false => {
                        // It is not a FS path, now we do what browsers do:
                        // prepend "http://" and hope it points to a website
                        Url::parse(&format!("http://{hopefully_url}", hopefully_url = &target))
                            .unwrap()
                    }
                }
            }
        },
    };

    // Initialize client
    let mut cache = HashMap::new();
    let mut header_map = HeaderMap::new();
    if let Some(user_agent) = &options.user_agent {
        header_map.insert(
            USER_AGENT,
            HeaderValue::from_str(&user_agent).expect("Invalid User-Agent header specified"),
        );
    }
    let client = if options.timeout > 0 {
        Client::builder().timeout(Duration::from_secs(options.timeout))
    } else {
        // No timeout is default
        Client::builder()
    }
    .danger_accept_invalid_certs(options.insecure)
    .default_headers(header_map)
    .build()
    .expect("Failed to initialize HTTP client");

    // At first we assume that base URL is the same as target URL
    let mut base_url: Url = target_url.clone();

    let data: Vec<u8>;
    let mut document_encoding: String = "".to_string();
    let mut dom: RcDom;

    // Retrieve target document
    if use_stdin {
        data = read_stdin();
    } else if target_url.scheme() == "file"
        || (target_url.scheme() == "http" || target_url.scheme() == "https")
        || target_url.scheme() == "data"
    {
        match retrieve_asset(&mut cache, &client, &target_url, &target_url, &options, 0) {
            Ok((retrieved_data, final_url, media_type, charset)) => {
                // Make sure the media type is text/html
                if !media_type.eq_ignore_ascii_case("text/html") {
                    if !options.silent {
                        eprintln!("Unsupported document media type");
                    }
                    process::exit(1);
                }

                if options
                    .base_url
                    .clone()
                    .unwrap_or("".to_string())
                    .is_empty()
                {
                    base_url = final_url;
                }

                data = retrieved_data;
                document_encoding = charset;
            }
            Err(_) => {
                if !options.silent {
                    eprintln!("Could not retrieve target document");
                }
                process::exit(1);
            }
        }
    } else {
        process::exit(1);
    }

    // Initial parse
    dom = html_to_dom(&data, document_encoding.clone());

    // TODO: investigate if charset from filesystem/data URL/HTTP headers
    //       has say over what's specified in HTML

    // Attempt to determine document's charset
    if let Some(html_charset) = get_charset(&dom.document) {
        if !html_charset.is_empty() {
            // Check if the charset specified inside HTML is valid
            if let Some(encoding) = Encoding::for_label_no_replacement(html_charset.as_bytes()) {
                document_encoding = html_charset;
                dom = html_to_dom(&data, encoding.name().to_string());
            }
        }
    }

    // Use custom base URL if specified, read and use what's in the DOM otherwise
    let custom_base_url: String = options.base_url.clone().unwrap_or("".to_string());
    if custom_base_url.is_empty() {
        // No custom base URL is specified
        // Try to see if document has BASE element
        if let Some(existing_base_url) = get_base_url(&dom.document) {
            base_url = resolve_url(&target_url, &existing_base_url);
        }
    } else {
        // Custom base URL provided
        match Url::parse(&custom_base_url) {
            Ok(parsed_url) => {
                if parsed_url.scheme() == "file" {
                    // File base URLs can only work with
                    // documents saved from filesystem
                    if target_url.scheme() == "file" {
                        base_url = parsed_url;
                    }
                } else {
                    base_url = parsed_url;
                }
            }
            Err(_) => {
                // Failed to parse given base URL, perhaps it's a filesystem path?
                if target_url.scheme() == "file" {
                    // Relative paths could work for documents saved from filesystem
                    let path: &Path = Path::new(&custom_base_url);
                    if path.exists() {
                        match Url::from_file_path(fs::canonicalize(&path).unwrap()) {
                            Ok(file_url) => {
                                base_url = file_url;
                            }
                            Err(_) => {
                                if !options.silent {
                                    eprintln!(
                                        "Could not map given path to base URL: {}",
                                        custom_base_url
                                    );
                                }
                                process::exit(1);
                            }
                        }
                    }
                }
            }
        }
    }

    // Traverse through the document and embed remote assets
    walk_and_embed_assets(&mut cache, &client, &base_url, &dom.document, &options, 0);

    // Update or add new BASE element to reroute network requests and hash-links
    if let Some(new_base_url) = options.base_url.clone() {
        dom = set_base_url(&dom.document, new_base_url);
    }

    // Request and embed /favicon.ico (unless it's already linked in the document)
    if !options.no_images
        && (target_url.scheme() == "http" || target_url.scheme() == "https")
        && !has_favicon(&dom.document)
    {
        let favicon_ico_url: Url = resolve_url(&base_url, "/favicon.ico");

        match retrieve_asset(
            &mut cache,
            &client,
            &target_url,
            &favicon_ico_url,
            &options,
            0,
        ) {
            Ok((data, final_url, media_type, charset)) => {
                let favicon_data_url: Url =
                    create_data_url(&media_type, &charset, &data, &final_url);
                dom = add_favicon(&dom.document, favicon_data_url.to_string());
            }
            Err(_) => {
                // Failed to retrieve /favicon.ico
            }
        }
    }

    // Save using specified charset, if given
    if let Some(custom_charset) = options.charset.clone() {
        document_encoding = custom_charset;
        dom = set_charset(dom, document_encoding.clone());
    }

    // Serialize DOM tree
    let mut result: Vec<u8> = serialize_document(dom, document_encoding, &options);

    // Prepend metadata comment tag
    if !options.no_metadata {
        let mut metadata_comment: String = create_metadata_tag(&target_url);
        metadata_comment += "\n";
        result.splice(0..0, metadata_comment.as_bytes().to_vec());
    }

    // Define output
    let mut output = Output::new(&options.output).expect("Could not prepare output");

    // Write result into stdout or file
    output.write(&result).expect("Could not write HTML output");
}
<commit_msg>parse XML documents, save non-HTML and non-XML targets unparsed<commit_after>use encoding_rs::Encoding;
use html5ever::rcdom::RcDom;
use reqwest::blocking::Client;
use reqwest::header::{HeaderMap, HeaderValue, USER_AGENT};
use std::collections::HashMap;
use std::fs;
use std::io::{self, prelude::*, Error, Write};
use std::path::Path;
use std::process;
use std::time::Duration;
use url::Url;

use monolith::html::{
    add_favicon, create_metadata_tag, get_base_url, get_charset, has_favicon, html_to_dom,
    serialize_document, set_base_url, set_charset, walk_and_embed_assets,
};
use monolith::opts::Options;
use monolith::url::{create_data_url, resolve_url};
use monolith::utils::retrieve_asset;

enum Output {
    Stdout(io::Stdout),
    File(fs::File),
}

impl Output {
    fn new(file_path: &str) -> Result<Output, Error> {
        if file_path.is_empty() || file_path.eq("-") {
            Ok(Output::Stdout(io::stdout()))
        } else {
            Ok(Output::File(fs::File::create(file_path)?))
        }
    }

    fn write(&mut self, bytes: &Vec<u8>) -> Result<(), Error> {
        match self {
            Output::Stdout(stdout) => {
                stdout.write_all(bytes)?;
                // Ensure newline at end of output
                if bytes.last() != Some(&b"\n"[0]) {
                    stdout.write(b"\n")?;
                }
                stdout.flush()
            }
            Output::File(file) => {
                file.write_all(bytes)?;
                // Ensure newline at end of output
                if bytes.last() != Some(&b"\n"[0]) {
                    file.write(b"\n")?;
                }
                file.flush()
            }
        }
    }
}

pub fn read_stdin() -> Vec<u8> {
    let mut buffer: Vec<u8> = vec![];

    match io::stdin().lock().read_to_end(&mut buffer) {
        Ok(_) => buffer,
        Err(_) => buffer,
    }
}

fn main() {
    let options = Options::from_args();

    // Check if target was provided
    if options.target.len() == 0 {
        if !options.silent {
            eprintln!("No target specified");
        }
        process::exit(1);
    }

    // Check if custom charset is valid
    if let Some(custom_charset) = options.charset.clone() {
        if !Encoding::for_label_no_replacement(custom_charset.as_bytes()).is_some() {
            eprintln!("Unknown encoding: {}", &custom_charset);
            process::exit(1);
        }
    }

    let mut use_stdin: bool = false;

    let target_url = match options.target.as_str() {
        "-" => {
            // Read from pipe (stdin)
            use_stdin = true;
            // Set default target URL to an empty data URL; the user can set it via --base-url
            Url::parse("data:text/html,").unwrap()
        }
        target => match Url::parse(&target) {
            Ok(url) => match url.scheme() {
                "data" | "file" | "http" | "https" => url,
                unsupported_scheme => {
                    if !options.silent {
                        eprintln!("Unsupported target URL type: {}", unsupported_scheme);
                    }
                    process::exit(1)
                }
            },
            Err(_) => {
                // Failed to parse given base URL (perhaps it's a filesystem path?)
                let path: &Path = Path::new(&target);
                match path.exists() {
                    true => match path.is_file() {
                        true => {
                            let canonical_path = fs::canonicalize(&path).unwrap();
                            match Url::from_file_path(canonical_path) {
                                Ok(url) => url,
                                Err(_) => {
                                    if !options.silent {
                                        eprintln!(
                                            "Could not generate file URL out of given path: {}",
                                            &target
                                        );
                                    }
                                    process::exit(1);
                                }
                            }
                        }
                        false => {
                            if !options.silent {
                                eprintln!("Local target is not a file: {}", &target);
                            }
                            process::exit(1);
                        }
                    },
                    false => {
                        // It is not a FS path, now we do what browsers do:
                        // prepend "http://" and hope it points to a website
                        Url::parse(&format!("http://{hopefully_url}", hopefully_url = &target))
                            .unwrap()
                    }
                }
            }
        },
    };

    // Initialize client
    let mut cache = HashMap::new();
    let mut header_map = HeaderMap::new();
    if let Some(user_agent) = &options.user_agent {
        header_map.insert(
            USER_AGENT,
            HeaderValue::from_str(&user_agent).expect("Invalid User-Agent header specified"),
        );
    }
    let client = if options.timeout > 0 {
        Client::builder().timeout(Duration::from_secs(options.timeout))
    } else {
        // No timeout is default
        Client::builder()
    }
    .danger_accept_invalid_certs(options.insecure)
    .default_headers(header_map)
    .build()
    .expect("Failed to initialize HTTP client");

    // At first we assume that base URL is the same as target URL
    let mut base_url: Url = target_url.clone();

    let data: Vec<u8>;
    let mut document_encoding: String = "".to_string();
    let mut dom: RcDom;

    // Retrieve target document
    if use_stdin {
        data = read_stdin();
    } else if target_url.scheme() == "file"
        || (target_url.scheme() == "http" || target_url.scheme() == "https")
        || target_url.scheme() == "data"
    {
        match retrieve_asset(&mut cache, &client, &target_url, &target_url, &options, 0) {
            Ok((retrieved_data, final_url, media_type, charset)) => {
                // Provide output as text without processing it, the way browsers do
                if !media_type.eq_ignore_ascii_case("text/html")
                    && !media_type.eq_ignore_ascii_case("application/xhtml+xml")
                {
                    // Define output
                    let mut output =
                        Output::new(&options.output).expect("Could not prepare output");

                    // Write retrieved data into STDOUT or file
                    output
                        .write(&retrieved_data)
                        .expect("Could not write output");

                    // Nothing else to do past this point
                    process::exit(0);
                }

                if options
                    .base_url
                    .clone()
                    .unwrap_or("".to_string())
                    .is_empty()
                {
                    base_url = final_url;
                }

                data = retrieved_data;
                document_encoding = charset;
            }
            Err(_) => {
                if !options.silent {
                    eprintln!("Could not retrieve target document");
                }
                process::exit(1);
            }
        }
    } else {
        process::exit(1);
    }

    // Initial parse
    dom = html_to_dom(&data, document_encoding.clone());

    // TODO: investigate if charset from filesystem/data URL/HTTP headers
    //       has say over what's specified in HTML

    // Attempt to determine document's charset
    if let Some(html_charset) = get_charset(&dom.document) {
        if !html_charset.is_empty() {
            // Check if the charset specified inside HTML is valid
            if let Some(encoding) = Encoding::for_label_no_replacement(html_charset.as_bytes()) {
                document_encoding = html_charset;
                dom = html_to_dom(&data, encoding.name().to_string());
            }
        }
    }

    // Use custom base URL if specified, read and use what's in the DOM otherwise
    let custom_base_url: String = options.base_url.clone().unwrap_or("".to_string());
    if custom_base_url.is_empty() {
        // No custom base URL is specified
        // Try to see if document has BASE element
        if let Some(existing_base_url) = get_base_url(&dom.document) {
            base_url = resolve_url(&target_url, &existing_base_url);
        }
    } else {
        // Custom base URL provided
        match Url::parse(&custom_base_url) {
            Ok(parsed_url) => {
                if parsed_url.scheme() == "file" {
                    // File base URLs can only work with
                    // documents saved from filesystem
                    if target_url.scheme() == "file" {
                        base_url = parsed_url;
                    }
                } else {
                    base_url = parsed_url;
                }
            }
            Err(_) => {
                // Failed to parse given base URL, perhaps it's a filesystem path?
                if target_url.scheme() == "file" {
                    // Relative paths could work for documents saved from filesystem
                    let path: &Path = Path::new(&custom_base_url);
                    if path.exists() {
                        match Url::from_file_path(fs::canonicalize(&path).unwrap()) {
                            Ok(file_url) => {
                                base_url = file_url;
                            }
                            Err(_) => {
                                if !options.silent {
                                    eprintln!(
                                        "Could not map given path to base URL: {}",
                                        custom_base_url
                                    );
                                }
                                process::exit(1);
                            }
                        }
                    }
                }
            }
        }
    }

    // Traverse through the document and embed remote assets
    walk_and_embed_assets(&mut cache, &client, &base_url, &dom.document, &options, 0);

    // Update or add new BASE element to reroute network requests and hash-links
    if let Some(new_base_url) = options.base_url.clone() {
        dom = set_base_url(&dom.document, new_base_url);
    }

    // Request and embed /favicon.ico (unless it's already linked in the document)
    if !options.no_images
        && (target_url.scheme() == "http" || target_url.scheme() == "https")
        && !has_favicon(&dom.document)
    {
        let favicon_ico_url: Url = resolve_url(&base_url, "/favicon.ico");

        match retrieve_asset(
            &mut cache,
            &client,
            &target_url,
            &favicon_ico_url,
            &options,
            0,
        ) {
            Ok((data, final_url, media_type, charset)) => {
                let favicon_data_url: Url =
                    create_data_url(&media_type, &charset, &data, &final_url);
                dom = add_favicon(&dom.document, favicon_data_url.to_string());
            }
            Err(_) => {
                // Failed to retrieve /favicon.ico
            }
        }
    }

    // Save using specified charset, if given
    if let Some(custom_charset) = options.charset.clone() {
        document_encoding = custom_charset;
        dom = set_charset(dom, document_encoding.clone());
    }

    // Serialize DOM tree
    let mut result: Vec<u8> = serialize_document(dom, document_encoding, &options);

    // Prepend metadata comment tag
    if !options.no_metadata {
        let mut metadata_comment: String = create_metadata_tag(&target_url);
        metadata_comment += "\n";
        result.splice(0..0, metadata_comment.as_bytes().to_vec());
    }

    // Define output
    let mut output = Output::new(&options.output).expect("Could not prepare output");

    // Write result into STDOUT or file
    output.write(&result).expect("Could not write output");
}
<|endoftext|>


<commit_before>use std::io::{self, BufReader};
use std::io::BufRead;
use std::fs::File;
extern crate time;
use time::Tm;
use time::strptime;
use time::Duration;
extern crate stats;
use stats::median;
extern crate clap;
use clap::{Arg, App};

mod percentile;
use percentile::percentile;

mod http_status;
use http_status::HttpStatus;

#[derive(Eq, PartialEq, Clone)]
#[derive(Debug)]
pub struct Request {
    id: i32,
    time: Tm,
    url: String,
}

impl Request {
    fn new_from_log_line(log_line: &String) -> Result<Request, io::Error> {
        let parts: Vec<&str> = log_line.split(" ").collect();

        let id = parts[2];
        let url = parts[5];

        Ok(Request {
            id: id[1..id.len()-1].parse().unwrap(),
            time: strptime(parts[0], "%d/%b/%Y:%H:%M:%S").unwrap(),
            url: url.to_string()
        })
    }

    fn get_matching_response<'a>(&'a self, responses: &'a Vec<Response>) -> Option<&Response> {
        match responses.binary_search_by_key(&self.id, |r| r.id) {
            Ok(index) => Some(&responses[index]),
            Err(_) => None
        }
    }
}

#[derive(Eq, PartialEq, Clone)]
#[derive(Debug)]
pub struct Response {
    id: i32,
    time: Tm,
    mime_type: String,
    response_time: Duration,
    http_status: HttpStatus,
}

impl Response {
    fn new_from_log_line(log_line: &String) -> Result<Response, io::Error> {
        let parts: Vec<&str> = log_line.split(" ").collect();

        let id = parts[2];
        let response_time = parts[parts.len()-1];
        let mime_type = parts[5];
        let status_code = parts[4];

        Ok(Response {
            id: id[1..id.len()-1].parse().unwrap(),
            time: strptime(parts[0], "%d/%b/%Y:%H:%M:%S").unwrap(),
            response_time: Duration::milliseconds(response_time[0..response_time.len()-2].parse().unwrap()),
            mime_type: mime_type.to_string(),
            http_status: HttpStatus::from_code(status_code.parse().unwrap()),
        })
    }
}

pub fn open_logfile(path: &str) -> Result<(Vec<Request>,Vec<Response>), io::Error> {
    let f = try!(File::open(path));

    let f = BufReader::new(f);

    let mut requests: Vec<Request> = Vec::new();
    let mut responses: Vec<Response> = Vec::new();

    for line in f.lines() {
        let line_value = &line.unwrap();

        if line_value.contains("->") {
            let r = try!(Request::new_from_log_line(&line_value));
            requests.push(r)
        }

        if line_value.contains("<-") {
            let r = try!(Response::new_from_log_line(&line_value));
            responses.push(r)
        }

    }

    responses.sort_by_key(|r| r.id);

    Ok((requests, responses))
}

pub struct RequestResponsePair {
    request: Request,
    response: Response
}

pub fn pair_requests_responses(requests:Vec<Request>, responses: Vec<Response>) -> Vec<RequestResponsePair> {
    let mut request_response_pairs: Vec<RequestResponsePair> = Vec::new();

    for request in requests  {
        if let Some(response) = request.get_matching_response(&responses) {
            request_response_pairs.push(RequestResponsePair{
                request: request.clone(),
                response: response.clone()
            })
        }
    }

    request_response_pairs
}

#[derive(Eq, PartialEq)]
#[derive(Debug)]
pub struct RequestLogAnalyzerResult {
    count: usize,
    max: usize,
    min: usize,
    avg: usize,
    median: usize,
    percentile90: usize,
}

pub fn analyze(request_response_pairs: &Vec<RequestResponsePair>) -> RequestLogAnalyzerResult {
    let times: Vec<i64> = request_response_pairs.iter()
        .map(|rr: &RequestResponsePair| -> i64 {rr.response.response_time.num_milliseconds() })
        .collect();

    let sum: usize = times.iter().sum::<i64>() as usize;
    let avg: usize = sum / times.len();

    let max: usize = *times.iter().max().unwrap() as usize;
    let min: usize = *times.iter().min().unwrap() as usize;

    let percentile90: usize = percentile(&times, 0.9) as usize;

    let median = median(times.into_iter()).unwrap() as usize;

    RequestLogAnalyzerResult {
        count: request_response_pairs.len().into(),
        max: max,
        min: min,
        avg: avg,
        median: median,
        percentile90: percentile90,
    }
}

fn render_terminal(result: RequestLogAnalyzerResult) {
    println!("count:\t{}", result.count);
    println!("time.avg:\t{}", result.avg);
    println!("time.min:\t{}", result.min);
    println!("time.median:\t{}", result.median);
    println!("time.90percent:\t{}", result.percentile90);
    println!("time.max:\t{}", result.max);
}

fn main() {
    let matches = App::new("Request.log Analyzer")
        .arg(Arg::with_name("filename")
            .index(1)
            .value_name("FILE")
            .help("Log file to analyze")
            .takes_value(true))
            .get_matches();

    let filename = matches.value_of("filename").unwrap();

    let lines = open_logfile(filename);
    let (requests, responses) = lines.unwrap();

    let pairs: Vec<RequestResponsePair> = pair_requests_responses(requests, responses);

    let result: RequestLogAnalyzerResult = analyze(&pairs);

    render_terminal(result);
}

#[cfg(test)]
mod tests {
	use super::*;
    extern crate time;
    use time::strptime;
    use::time::Duration;
    use http_status::HttpStatus;

    #[test]
    fn test_parse_request_line() {
        let line = "08/Apr/2016:09:58:47 +0200 [02] -> GET /content/some/other.html HTTP/1.1".to_string();

        let expected = Request {
            id: 2,
            time: strptime("08/Apr/2016:09:58:47 +0200", "%d/%b/%Y:%H:%M:%S").unwrap(),
            url: "/content/some/other.html".to_string()
        };

        let result = Request::new_from_log_line(&line);

        assert_eq!(result.unwrap(), expected)
    }

    #[test]
    fn test_parse_response_line() {
        let line = "08/Apr/2016:09:58:48 +0200 [02] <- 200 text/html 10ms".to_string();

        let expected = Response {
            id: 2,
            time: strptime("08/Apr/2016:09:58:48 +0200", "%d/%b/%Y:%H:%M:%S").unwrap(),
            mime_type: "text/html".to_string(),
            response_time: Duration::milliseconds(10),
            http_status: HttpStatus::OK,
        };

        let result = Response::new_from_log_line(&line);

        assert_eq!(result.unwrap(), expected)
    }

    #[test]
    fn test_open_logfile() {
        let lines = open_logfile("src/test/simple-1.log");
        let (requests, responses) = lines.unwrap();

        assert_eq!(requests.len(), 2);
        assert_eq!(responses.len(), 2);
    }

    #[test]
    fn test_get_matching_response() {
        let lines = open_logfile("src/test/simple-1.log");
        let (requests, responses) = lines.unwrap();

        let result = requests[0].get_matching_response(&responses);

        let expected = Response {
            id: 1,
            time: strptime("08/Apr/2016:09:57:47 +0200", "%d/%b/%Y:%H:%M:%S").unwrap(),
            mime_type: "text/html".to_string(),
            response_time: Duration::milliseconds(7),
            http_status: HttpStatus::OK,
        };

        assert_eq!(*result.unwrap(), expected);
    }

    #[test]
    fn test_get_matching_response_none_found() {
        let lines = open_logfile("src/test/simple-1.log");
        let (_, responses) = lines.unwrap();

        let request_without_matching = Request {
            id: 999,
            time: strptime("08/Apr/2016:09:58:47 +0200", "%d/%b/%Y:%H:%M:%S").unwrap(),
            url: "/content/some/other.html".to_string()
        };

        let result = request_without_matching.get_matching_response(&responses);

        assert!(result.is_none());
    }

    #[test]
    fn test_pair_requests_resonses() {
        let lines = open_logfile("src/test/simple-1.log");
        let (requests, responses) = lines.unwrap();

        let result = pair_requests_responses(requests, responses);

        assert_eq!(result.len(), 2);

        assert_eq!(result[0].request.id, result[0].response.id);
        assert_eq!(result[1].request.id, result[1].response.id);
    }

    #[test]
    fn test_request_log_analyzer_result() {
        let lines = open_logfile("src/test/response-time-calculations.log");
        let (requests, responses) = lines.unwrap();

        let request_response_pairs = pair_requests_responses(requests, responses);

        let result: RequestLogAnalyzerResult = analyze(&request_response_pairs);

        let expected = RequestLogAnalyzerResult {
            count: 3,
            max: 100,
            min: 1,
            avg: 37,
            median: 10,
            percentile90: 100,
        };

        assert_eq!(result, expected);
    }

    #[test]
    fn test_90_percentile_calculation() {
        let lines = open_logfile("src/test/percentile.log");
        let (requests, responses) = lines.unwrap();

        let request_response_pairs = pair_requests_responses(requests, responses);

        let result: RequestLogAnalyzerResult = analyze(&request_response_pairs);

        assert_eq!(result.percentile90, 9);
    }
}
<commit_msg>Handle special case with space within MIME type field<commit_after>use std::io::{self, BufReader};
use std::io::BufRead;
use std::fs::File;
extern crate time;
use time::Tm;
use time::strptime;
use time::Duration;
extern crate stats;
use stats::median;
extern crate clap;
use clap::{Arg, App};

mod percentile;
use percentile::percentile;

mod http_status;
use http_status::HttpStatus;

#[derive(Eq, PartialEq, Clone)]
#[derive(Debug)]
pub struct Request {
    id: i32,
    time: Tm,
    url: String,
}

impl Request {
    fn new_from_log_line(log_line: &String) -> Result<Request, io::Error> {
        let parts: Vec<&str> = log_line.split(" ").collect();

        let id = parts[2];
        let url = parts[5];

        Ok(Request {
            id: id[1..id.len()-1].parse().unwrap(),
            time: strptime(parts[0], "%d/%b/%Y:%H:%M:%S").unwrap(),
            url: url.to_string()
        })
    }

    fn get_matching_response<'a>(&'a self, responses: &'a Vec<Response>) -> Option<&Response> {
        match responses.binary_search_by_key(&self.id, |r| r.id) {
            Ok(index) => Some(&responses[index]),
            Err(_) => None
        }
    }
}

#[derive(Eq, PartialEq, Clone)]
#[derive(Debug)]
pub struct Response {
    id: i32,
    time: Tm,
    mime_type: String,
    response_time: Duration,
    http_status: HttpStatus,
}

impl Response {
    fn new_from_log_line(log_line: &String) -> Result<Response, io::Error> {
        let parts: Vec<&str> = log_line.split(" ").collect();

        let id = parts[2];
        let response_time = parts[parts.len()-1];

        // Handle special case where the mime type sometimes contains
        // a space, so we need to re-assemble it
        let mime_type = match parts.len() {
            8 => format!("{} {}", parts[5], parts[6]),
            _ => parts[5].to_string()
        };

        let status_code = parts[4];

        Ok(Response {
            id: id[1..id.len()-1].parse().unwrap(),
            time: strptime(parts[0], "%d/%b/%Y:%H:%M:%S").unwrap(),
            response_time: Duration::milliseconds(response_time[0..response_time.len()-2].parse().unwrap()),
            mime_type: mime_type,
            http_status: HttpStatus::from_code(status_code.parse().unwrap()),
        })
    }
}

pub fn open_logfile(path: &str) -> Result<(Vec<Request>,Vec<Response>), io::Error> {
    let f = try!(File::open(path));

    let f = BufReader::new(f);

    let mut requests: Vec<Request> = Vec::new();
    let mut responses: Vec<Response> = Vec::new();

    for line in f.lines() {
        let line_value = &line.unwrap();

        if line_value.contains("->") {
            let r = try!(Request::new_from_log_line(&line_value));
            requests.push(r)
        }

        if line_value.contains("<-") {
            let r = try!(Response::new_from_log_line(&line_value));
            responses.push(r)
        }

    }

    responses.sort_by_key(|r| r.id);

    Ok((requests, responses))
}

pub struct RequestResponsePair {
    request: Request,
    response: Response
}

pub fn pair_requests_responses(requests:Vec<Request>, responses: Vec<Response>) -> Vec<RequestResponsePair> {
    let mut request_response_pairs: Vec<RequestResponsePair> = Vec::new();

    for request in requests  {
        if let Some(response) = request.get_matching_response(&responses) {
            request_response_pairs.push(RequestResponsePair{
                request: request.clone(),
                response: response.clone()
            })
        }
    }

    request_response_pairs
}

#[derive(Eq, PartialEq)]
#[derive(Debug)]
pub struct RequestLogAnalyzerResult {
    count: usize,
    max: usize,
    min: usize,
    avg: usize,
    median: usize,
    percentile90: usize,
}

pub fn analyze(request_response_pairs: &Vec<RequestResponsePair>) -> RequestLogAnalyzerResult {
    let times: Vec<i64> = request_response_pairs.iter()
        .map(|rr: &RequestResponsePair| -> i64 {rr.response.response_time.num_milliseconds() })
        .collect();

    let sum: usize = times.iter().sum::<i64>() as usize;
    let avg: usize = sum / times.len();

    let max: usize = *times.iter().max().unwrap() as usize;
    let min: usize = *times.iter().min().unwrap() as usize;

    let percentile90: usize = percentile(&times, 0.9) as usize;

    let median = median(times.into_iter()).unwrap() as usize;

    RequestLogAnalyzerResult {
        count: request_response_pairs.len().into(),
        max: max,
        min: min,
        avg: avg,
        median: median,
        percentile90: percentile90,
    }
}

fn render_terminal(result: RequestLogAnalyzerResult) {
    println!("count:\t{}", result.count);
    println!("time.avg:\t{}", result.avg);
    println!("time.min:\t{}", result.min);
    println!("time.median:\t{}", result.median);
    println!("time.90percent:\t{}", result.percentile90);
    println!("time.max:\t{}", result.max);
}

fn main() {
    let matches = App::new("Request.log Analyzer")
        .arg(Arg::with_name("filename")
            .index(1)
            .value_name("FILE")
            .help("Log file to analyze")
            .takes_value(true))
            .get_matches();

    let filename = matches.value_of("filename").unwrap();

    let lines = open_logfile(filename);
    let (requests, responses) = lines.unwrap();

    let pairs: Vec<RequestResponsePair> = pair_requests_responses(requests, responses);

    let result: RequestLogAnalyzerResult = analyze(&pairs);

    render_terminal(result);
}

#[cfg(test)]
mod tests {
	use super::*;
    extern crate time;
    use time::strptime;
    use::time::Duration;
    use http_status::HttpStatus;

    #[test]
    fn test_parse_request_line() {
        let line = "08/Apr/2016:09:58:47 +0200 [02] -> GET /content/some/other.html HTTP/1.1".to_string();

        let expected = Request {
            id: 2,
            time: strptime("08/Apr/2016:09:58:47 +0200", "%d/%b/%Y:%H:%M:%S").unwrap(),
            url: "/content/some/other.html".to_string()
        };

        let result = Request::new_from_log_line(&line);

        assert_eq!(result.unwrap(), expected)
    }

    #[test]
    fn test_parse_response_line() {
        let line = "08/Apr/2016:09:58:48 +0200 [02] <- 200 text/html 10ms".to_string();

        let expected = Response {
            id: 2,
            time: strptime("08/Apr/2016:09:58:48 +0200", "%d/%b/%Y:%H:%M:%S").unwrap(),
            mime_type: "text/html".to_string(),
            response_time: Duration::milliseconds(10),
            http_status: HttpStatus::OK,
        };

        let result = Response::new_from_log_line(&line);

        assert_eq!(result.unwrap(), expected)
    }

    #[test]
    fn test_parse_response_line_inconsistent_space() {
        let line = "06/Apr/2016:14:54:16 +0200 [200] <- 200 text/html; charset=utf-8 250ms".to_string();

        let expected = Response {
            id: 200,
            time: strptime("06/Apr/2016:14:54:16 +0200", "%d/%b/%Y:%H:%M:%S").unwrap(),
            mime_type: "text/html; charset=utf-8".to_string(),
            response_time: Duration::milliseconds(250),
            http_status: HttpStatus::OK,
        };

        let result = Response::new_from_log_line(&line);

        assert_eq!(result.unwrap(), expected)
    }

    #[test]
    fn test_open_logfile() {
        let lines = open_logfile("src/test/simple-1.log");
        let (requests, responses) = lines.unwrap();

        assert_eq!(requests.len(), 2);
        assert_eq!(responses.len(), 2);
    }

    #[test]
    fn test_get_matching_response() {
        let lines = open_logfile("src/test/simple-1.log");
        let (requests, responses) = lines.unwrap();

        let result = requests[0].get_matching_response(&responses);

        let expected = Response {
            id: 1,
            time: strptime("08/Apr/2016:09:57:47 +0200", "%d/%b/%Y:%H:%M:%S").unwrap(),
            mime_type: "text/html".to_string(),
            response_time: Duration::milliseconds(7),
            http_status: HttpStatus::OK,
        };

        assert_eq!(*result.unwrap(), expected);
    }

    #[test]
    fn test_get_matching_response_none_found() {
        let lines = open_logfile("src/test/simple-1.log");
        let (_, responses) = lines.unwrap();

        let request_without_matching = Request {
            id: 999,
            time: strptime("08/Apr/2016:09:58:47 +0200", "%d/%b/%Y:%H:%M:%S").unwrap(),
            url: "/content/some/other.html".to_string()
        };

        let result = request_without_matching.get_matching_response(&responses);

        assert!(result.is_none());
    }

    #[test]
    fn test_pair_requests_resonses() {
        let lines = open_logfile("src/test/simple-1.log");
        let (requests, responses) = lines.unwrap();

        let result = pair_requests_responses(requests, responses);

        assert_eq!(result.len(), 2);

        assert_eq!(result[0].request.id, result[0].response.id);
        assert_eq!(result[1].request.id, result[1].response.id);
    }

    #[test]
    fn test_request_log_analyzer_result() {
        let lines = open_logfile("src/test/response-time-calculations.log");
        let (requests, responses) = lines.unwrap();

        let request_response_pairs = pair_requests_responses(requests, responses);

        let result: RequestLogAnalyzerResult = analyze(&request_response_pairs);

        let expected = RequestLogAnalyzerResult {
            count: 3,
            max: 100,
            min: 1,
            avg: 37,
            median: 10,
            percentile90: 100,
        };

        assert_eq!(result, expected);
    }

    #[test]
    fn test_90_percentile_calculation() {
        let lines = open_logfile("src/test/percentile.log");
        let (requests, responses) = lines.unwrap();

        let request_response_pairs = pair_requests_responses(requests, responses);

        let result: RequestLogAnalyzerResult = analyze(&request_response_pairs);

        assert_eq!(result.percentile90, 9);
    }
}
<|endoftext|>


<commit_before>// main.rs
#![feature(plugin)]
#![plugin(docopt_macros)]

extern crate rustc_serialize;
extern crate docopt;
extern crate regex;

use std::fmt::{ Display, Formatter };
use std::error::Error;
use std::fs::File;
use std::io::prelude::*;
//use std::io::{BufReader,BufRead};
use std::io::BufReader;
use std::fs::OpenOptions;
use std::io::BufWriter;
use std::path::Path;
// https://github.com/docopt/docopt.rs
use docopt::Docopt;

// https://doc.rust-lang.org/regex/regex/index.html
use regex::Regex;

docopt!(Args derive Debug, "
Alexey Mekhanoshin

Usage:
	rus_fritz -e <engfile>  -r <rusfile> ( -o <outfile> | --stdout )
	rus_fritz (-h | --help)
	rus_fritz --version
Options:
  -h --help		Show this screen.
  --version		Show version.
");

struct EnglishName {
	name: String,
	opis: String,
}

struct RussianName {
	name: String,
	opis: String,
	using: bool,
}

impl EnglishName {
	fn new<S: Into<String>>(name: S, opis: S) -> EnglishName {
		EnglishName {
			name: name.into(),
			opis: opis.into(),
		}
	}
}

impl RussianName {
	fn new<S: Into<String>>(name: S, opis: S) -> RussianName {
		RussianName {
			name: name.into(),
			opis: opis.into(),
			using: false,
		}
	}
}

impl std::fmt::Display for EnglishName {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> Result<(), std::fmt::Error> {
        //Ok(())
        write!(f, "Display Eng: {},\t\t\t{}", self.name, self.opis)
    }
}

impl std::fmt::Display for RussianName {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> Result<(), std::fmt::Error> {
        	write!(f, "Display Rus: {},\t\t\t{}", self.name, self.opis)
    }
}

fn open_eng_file(engvec: &mut Vec<EnglishName>, file_name: &String) {
	let path = Path::new(file_name);
	let display = path.display();
	let file = match File::open(path) {
		Ok(f) => f,
		Err(why) => panic!("Not open file {}: {}", display, Error::description(&why) ),
	};
	let re = Regex::new(r"(.*)(,)(.*)(\x22.*\x22)").unwrap();
	for line in BufReader::new(file).lines() {
		let s = line.unwrap();
		//println!("{}", s)
		for cap in re.captures_iter(&s) {
			let ubs = EnglishName::new( cap.at(1).unwrap(), cap.at(4).unwrap() );
			engvec.push(ubs);
		}
	}
}

fn open_rus_file(engvec: &mut Vec<RussianName>, file_name: &String) {
	let path = Path::new(file_name);
	let display = path.display();
	let file = match File::open(path) {
		Ok(f) => f,
		Err(why) => panic!("Not open file: {}: {}", display, Error::description(&why) ),
	};
	let re = Regex::new(r"(.*)(,)(.*)(\x22.*\x22)").unwrap();
	for line in BufReader::new(file).lines() {
		let s = line.unwrap();
		//println!("{}", s)
		for cap in re.captures_iter(&s) {
			let ubs = RussianName::new( cap.at(1).unwrap(), cap.at(4).unwrap() );
			engvec.push(ubs);
		}
	}
}

fn create_out_file(file_name: &String) -> BufWriter<File> {
	let path = Path::new(file_name);
	let display = path.display();
	let mut options = OpenOptions::new();
	let file = match options.create(true).write(true).open(path) {
		Ok(file) => file,
		Err(why) => panic!("Not create/open file {}: {}", display, Error::description(&why)),
	};
	let writer = BufWriter::new(file);
	return writer;
	// return true;
}

fn main() {

	let args: Args = Args::docopt().decode().unwrap_or_else(|e| e.exit());
	if args.flag_stdout {
		println!("{:?}", args);
		if args.flag_h { println!("Print Help coming soon ..." ); return; }
		if args.flag_version { println!("Print Version coming soon ..." ); return; }
	}

	let mut engvec: Vec<EnglishName> = Vec::new();
	let mut rusvec: Vec<RussianName> = Vec::new();
	//let wr: BufWriter<File>;

	open_eng_file(&mut engvec, &args.arg_engfile.to_string());
	assert!( !engvec.is_empty() );
	open_rus_file(&mut rusvec, &args.arg_rusfile.to_string());
	assert!( !rusvec.is_empty() );

	//  Подготовка закончилась, начинаем работу.

	if args.flag_stdout {
		println!("/*\n\tRecreated by RusFritz project\n*/");
	} else {
		// Write to file
	}

	let mut i=0;
	let mut found = false;
	let mut outstr = String::new();

    for e in &mut engvec {
    	i = i+1;
    	for r in &mut rusvec {
    		if e.name == r.name {
    			e.opis = r.opis.clone();
    			found = true;
    			r.using = true;
    			if args.flag_stdout {
    				println!("{},\t\t{};", r.name, r.opis);
    			} else {
    				let s = String::from( format!("{},\t\t{};\n", r.name, r.opis) );
    				outstr.push_str(&s);
    			}
    		}
    	}
    	if !found {
    		//println!("ERR: {} - не найдено соответствие", e.name );
    		if args.flag_stdout {
    			println!("{},\t\t{};",e.name, e.opis);
    		} else {
    			// Wite output to file
    		}

    	}
    	found = false;
    }
    // Print if not using from Russian Names
    println!("------ Not using Found from Russian File --------");
    for r in rusvec {
    	if r.using == false {
    		println!("{},\t\t{}",r.name,r.opis);

    	}
    }

    // File write, if needed
    if args.flag_stdout == false {
		let mut wr = create_out_file(&args.arg_outfile.to_string());
		if outstr.len() > 0 {
			match wr.write_all(&outstr.as_bytes()) {
				Err(why) => {
					panic!("couldn't write to{}: {}",
						&args.arg_outfile.to_string(),
						Error::description(&why))
				},
				Ok(_) => println!("Good write file"),
			};
			wr.flush();
		}
	}

    //println!("Обработано {} строк из {}",i, engvec.len());
    //let writer = std::io::file_writer(&Path(std::os::args()[2]), [io::Append, io::Create]).unwrap();

}
<commit_msg>Reformat text<commit_after>// main.rs
#![feature(plugin)]
#![plugin(docopt_macros)]

extern crate rustc_serialize;
extern crate docopt;
extern crate regex;

use std::fmt::{ Display, Formatter };
use std::error::Error;
use std::fs::File;
use std::io::prelude::*;
use std::io::BufReader;
use std::fs::OpenOptions;
use std::io::BufWriter;
use std::path::Path;
// https://github.com/docopt/docopt.rs
use docopt::Docopt;

// https://doc.rust-lang.org/regex/regex/index.html
use regex::Regex;

docopt!(Args derive Debug, "
Alexey Mekhanoshin

Usage:
	rus_fritz -e <engfile>  -r <rusfile> ( -o <outfile> | --stdout )
	rus_fritz (-h | --help)
	rus_fritz --version
Options:
  -h --help		Show this screen.
  --version		Show version.
");

struct EnglishName {
    name: String,
    opis: String,
}

struct RussianName {
    name: String,
    opis: String,
    using: bool,
}

impl EnglishName {
    fn new<S: Into<String>>(name: S, opis: S) -> EnglishName {
        EnglishName {
            name: name.into(),
            opis: opis.into(),
        }
    }
}

impl RussianName {
    fn new<S: Into<String>>(name: S, opis: S) -> RussianName {
        RussianName {
            name: name.into(),
            opis: opis.into(),
            using: false,
        }
    }
}

impl std::fmt::Display for EnglishName {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> Result<(), std::fmt::Error> {
        //Ok(())
        write!(f, "Display Eng: {},\t\t\t{}", self.name, self.opis)
    }
}

impl std::fmt::Display for RussianName {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> Result<(), std::fmt::Error> {
        write!(f, "Display Rus: {},\t\t\t{}", self.name, self.opis)
        }
}

fn open_eng_file(engvec: &mut Vec<EnglishName>, file_name: &String) {
    let path = Path::new(file_name);
    let display = path.display();
    let file = match File::open(path) {
        Ok(f) => f,
        Err(why) => panic!("Not open file {}: {}", display, Error::description(&why) ),
    };
    let re = Regex::new(r"(.*)(,)(.*)(\x22.*\x22)").unwrap();
    for line in BufReader::new(file).lines() {
        let s = line.unwrap();
        for cap in re.captures_iter(&s) {
            let ubs = EnglishName::new( cap.at(1).unwrap(), cap.at(4).unwrap() );
            engvec.push(ubs);
    }
    }
}

fn open_rus_file(engvec: &mut Vec<RussianName>, file_name: &String) {
    let path = Path::new(file_name);
    let display = path.display();
    let file = match File::open(path) {
        Ok(f) => f,
        Err(why) => panic!("Not open file: {}: {}", display, Error::description(&why) ),
    };
    let re = Regex::new(r"(.*)(,)(.*)(\x22.*\x22)").unwrap();
    for line in BufReader::new(file).lines() {
        let s = line.unwrap();
        for cap in re.captures_iter(&s) {
            let ubs = RussianName::new( cap.at(1).unwrap(), cap.at(4).unwrap() );
            engvec.push(ubs);
        }
    }
}

fn create_out_file(file_name: &String) -> BufWriter<File> {
    let path = Path::new(file_name);
    let display = path.display();
    let mut options = OpenOptions::new();
    let file = match options.create(true).write(true).open(path) {
	    Ok(file) => file,
	    Err(why) => panic!("Not create/open file {}: {}", display, Error::description(&why)),
	};
    let writer = BufWriter::new(file);
    return writer;
}

fn main() {

    let args: Args = Args::docopt().decode().unwrap_or_else(|e| e.exit());
    if args.flag_stdout {
        println!("{:?}", args);
        if args.flag_h { println!("Print Help coming soon ..." ); return; }
        if args.flag_version { println!("Print Version coming soon ..." ); return; }
    }

    let mut engvec: Vec<EnglishName> = Vec::new();
    let mut rusvec: Vec<RussianName> = Vec::new();

    open_eng_file(&mut engvec, &args.arg_engfile.to_string());
    assert!( !engvec.is_empty() );
    open_rus_file(&mut rusvec, &args.arg_rusfile.to_string());
    assert!( !rusvec.is_empty() );

	//  Подготовка закончилась, начинаем работу.

    if args.flag_stdout {
        println!("/*\n\tRecreated by RusFritz project\n*/");
    } else {
		// Write to file
    }

    let mut i=0;
    let mut found = false;
    let mut outstr = String::new();

    for e in &mut engvec {
        i = i+1;
        for r in &mut rusvec {
            if e.name == r.name {
                e.opis = r.opis.clone();
                found = true;
                r.using = true;
                if args.flag_stdout {
                    println!("{},\t\t{};", r.name, r.opis);
                } else {
                    let s = String::from( format!("{},\t\t{};\n", r.name, r.opis) );
                    outstr.push_str(&s);
                }
            }
        }
        if !found {
            //println!("ERR: {} - не найдено соответствие", e.name );
            if args.flag_stdout {
                println!("{},\t\t{};",e.name, e.opis);
            } else {
                // Wite output to file
            }
        }
        found = false;
    }
    // Print if not using from Russian Names
    println!("------ Not using Found from Russian File --------");
    for r in rusvec {
        if r.using == false {
            println!("{},\t\t{}",r.name,r.opis);
        }
    }

    // File write, if needed
    if args.flag_stdout == false {
        let mut wr = create_out_file(&args.arg_outfile.to_string());
        if outstr.len() > 0 {
            match wr.write_all(&outstr.as_bytes()) {
                Err(why) => {
                    panic!("couldn't write to{}: {}",
                        &args.arg_outfile.to_string(),
                        Error::description(&why))
                },
                Ok(_) => println!("Good write file"),
            };
            wr.flush();
        }
    }
    //println!("Обработано {} строк из {}",i, engvec.len());
    //let writer = std::io::file_writer(&Path(std::os::args()[2]), [io::Append, io::Create]).unwrap();
}
<|endoftext|>


<commit_before>use std::fs::File;
use std::io::{Read,Write};

enum Direction {
    Left,
    Right,
    Up,
    Down,
}

struct State {
    xpos: i32,
    ypos: i32,
    dir: Direction,
    stack: Vec<i32>,
}

impl State {
    fn push(&mut self, val: i32) {
        self.push(val);
    }

    fn pop(&mut self) -> i32 {
        self.pop()
    }

    fn add(&mut self) {
        let a1 = self.pop();
        let a2 = self.pop();
        self.push(a1+a2);
    }

    fn sub(&mut self) {
        let a1 = self.pop();
        let a2 = self.pop();
        self.push(a1-a2);
    }

    fn mult(&mut self) {
        let a1 = self.pop();
        let a2 = self.pop();
        self.push(a1*a2);
    }

    fn div(&mut self) {
        let a1 = self.pop();
        let a2 = self.pop();
        self.push(a1/a2);
    }

    fn dup(&mut self) {
        let top = self.pop();
        self.push(top);
        self.push(top);
    }

}

/* lines converts a Vector containing file data into vectors of lines
 * each element is a vector of chars and the line length for easy access */
fn lines(v : Vec<u8>) -> Vec<(Vec<char>, i32)> {
    let mut ret = Vec::new();
    let mut new_line = Vec::new();
    for ch in v {
        if ch as char == '\n' {
            let length = new_line.iter().count();
            ret.push((new_line, length as i32));
            new_line = Vec::new();
        }
        else {
            new_line.push(ch as char);
        }
    }
    ret
}

fn main() {
    let mut f = File::open("tempy.txt").unwrap();
    let mut read_in = Vec::new();
    f.read_to_end(&mut read_in).unwrap();
    let lines_in = lines(read_in);
    let mut st = State { xpos: 0, ypos: 0, dir: Direction::Left, stack: Vec::new()};

    while true {
        let target_char = lines_in[st.ypos as usize].0[st.xpos as usize];
        match target_char {
            '>' => println!("Hit >"),
            '<' => println!("Hit <"),
             _  => {
                 std::io::stderr()
                    .write(format!("Unknown Character in program: '{}' ", target_char)
                    .as_bytes())
                    .unwrap();
                    std::process::exit(1);
                },
        }
    }
}
<commit_msg>save point<commit_after>use std::fs::File;
use std::io::{Read,Write};

enum Direction {
    Left,
    Right,
    Up,
    Down,
}

struct State {
    xpos: i32,
    ypos: i32,
    dir: Direction,
    stack: Vec<i32>,
}

impl State {
    fn push(&mut self, val: i32) {
        self.stack.push(val);
    }

    fn pop(&mut self) -> i32 {
        let res = self.stack.pop();
        if std::option::Option::is_none(&res) {
            std::io::stderr()
               .write("Stack is empty!"
               .as_bytes())
               .unwrap();
               std::process::exit(1);
        }
        res.unwrap()
    }

    fn add(&mut self) {
        let a1 = self.pop();
        let a2 = self.pop();
        self.push(a1+a2);
    }

    fn sub(&mut self) {
        let a1 = self.pop();
        let a2 = self.pop();
        self.push(a1-a2);
    }

    fn mult(&mut self) {
        let a1 = self.pop();
        let a2 = self.pop();
        self.push(a1*a2);
    }

    fn div(&mut self) {
        let a1 = self.pop();
        let a2 = self.pop();
        self.push(a1/a2);
    }

    fn dup(&mut self) {
        let top = self.pop();
        self.push(top);
        self.push(top);
    }

    fn is_zero(&mut self) -> bool {
        let top = self.pop();
        top == 0
    }

    fn is_gt(&mut self) -> bool {
        let top = self.pop();
        top > 0
    }

    fn is_lt(&mut self) -> bool {
        let top = self.pop();
        top < 0
    }

    fn negate(&mut self) {
        let top = self.pop();
        self.push(-1 * top);
    }

    fn next_step(&mut self) -> (i32, i32) {
        match self.dir {
            Direction::Up => {
                self.ypos -= 1;
            },
            Direction::Down => {
                self.ypos += 1;
            },
            Direction::Left => {
                self.xpos -= 1;
            },
            Direction::Right => {
                self.xpos += 1;
            }
        }
        (self.ypos, self.xpos)
    }

    fn logical_neg(&mut self) {
        let top = self.pop();
        if top != 0 {
            self.push(0);
        }
        else {
            self.push(1);
        }
    }

    fn is_empty(&self) -> bool {
        self.stack.is_empty()
    }
}

/* lines converts a Vector containing file data into vectors of lines
 * each element is a vector of chars and the line length for easy access */
fn lines(v : Vec<u8>) -> Vec<(Vec<char>, i32)> {
    let mut ret = Vec::new();
    let mut new_line = Vec::new();
    for ch in v {
        if ch as char == '\n' {
            let length = new_line.iter().count();
            ret.push((new_line, length as i32));
            new_line = Vec::new();
        }
        else {
            new_line.push(ch as char);
        }
    }
    ret
}

fn print_char(val : i32) {
    let mut buf = [val as u8; 1];
    std::io::stdout.write(&buf).unwrap();
}

fn read_char() -> i32 {
    0
}

fn main() {
    let mut f = File::open("tempy.txt").unwrap();
    let mut read_in = Vec::new();
    f.read_to_end(&mut read_in).unwrap();
    let lines_in = lines(read_in);
    let mut st = State { xpos: 0, ypos: 0, dir: Direction::Right, stack: Vec::new()};

    while true {
        let target_char = lines_in[st.ypos as usize].0[st.xpos as usize];
        match target_char {
            '>' => {
                st.dir = Direction::Right;
            },
            '<' => {
                st.dir = Direction::Left;
            },
            '^' => {
                st.dir = Direction.Up;
            },
            'v' => {
                st.dir = Direction.Down;
            },
            '0' => st.push(0),
            '1' => st.push(1),
            '2' => st.push(2),
            '3' => st.push(3),
            '4' => st.push(4),
            '5' => st.push(5),
            '6' => st.push(6),
            '7' => st.push(7),
            '8' => st.push(8),
            '9' => st.push(9),
            '@' => st.dup(),
            '+' => st.add(),
            '-' => st.sub(),
            '*' => st.mult(),
            '/' => st.div(),
            ';' => std::process:exit(0),

             _  => {
                 std::io::stderr()
                    .write(format!("Unknown Character in program: '{}' ", target_char)
                    .as_bytes())
                    .unwrap();
                    std::process::exit(1);
                },
        }
    }
}
<|endoftext|>


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

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

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

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

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

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

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

  let mut app = App { print_ascii: false, print_binary: false, sql_logging: false };
  let binary_matches = [
    ("SSL3.0 Record Protocol", r"^\x16\x03\x00..\x01"),
    ("TLS1.0 Record Protocol", r"^\x16\x03\x01..\x01"),
    ("TLS1.1 Record Protocol", r"^\x16\x03\x02..\x01"),
    ("TLS1.2 Record Protocol", r"^\x16\x03\x03..\x01"),
    ("SSL3.0 CLIENT_HELLO", r"^\x16....\x01...\x03\x00"),
    ("TLS1.0 CLIENT_HELLO", r"^\x16....\x01...\x03\x01"),
    ("TLS1.1 CLIENT_HELLO", r"^\x16....\x01...\x03\x02"),
    ("TLS1.2 CLIENT_HELLO", r"^\x16....\x01...\x03\x03"),
    ("SMB1 COMMAND NEGOTIATE", r"^....\xffSMB\x72"),
    ("SMB1 NT_STATUS Success", r"^....\xffSMB.[\x00-\x0f]"),
    ("SMB1 NT_STATUS Information", r"^....\xffSMB.[\x40-\x4f]"),
    ("SMB1 NT_STATUS Warning", r"^....\xffSMB.[\x80-\x8f]"),
    ("SMB1 NT_STATUS Error", r"^....\xffSMB.[\xc0-\xcf]"),
    ("SMB2 COMMAND NEGOTIATE", r"^\x00...\xfeSMB........\x00\x00"),
    ("SMB2 NT_STATUS Success", r"^\x00...\xfeSMB....[\x00-\x0f]"),
    ("SMB2 NT_STATUS Information", r"^\x00...\xfeSMB....[\x40-\x4f]"),
    ("SMB2 NT_STATUS Warning", r"^\x00...\xfeSMB....[\x80-\x8f]"),
    ("SMB2 NT_STATUS Error", r"^\x00...\xfeSMB....[\xc0-\xcf]"),
    ("MS-TDS PRELOGIN Request", r"^\x12\x01\x00.\x00\x00"),
    ("MS-TDS LOGIN Request", r"^\x10\x01\x00.\x00\x00"),
    ("SOCKS4 NOAUTH Request", r"^\x04\x01\x00\x50"),
    ("SOCKS5 NOAUTH Request", r"^\x05\x01\x00$"), // Tested ok-ish
    ("SOCKS5 USER/PASS Request", r"^\x05\x02\x00\x02$") // possibly broken
  ];
  let io_timeout = Duration::new(300, 0); // 5 minutes

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

  let docs = YamlLoader::load_from_str(&config_str).unwrap();
  let config = &docs[0];
  //println!("{:?}", config);
  if config["general"].is_badvalue() {
    println!("No 'general' section found in configuration file");
    exit(-1);
  }
  if config["ports"].is_badvalue() {
    println!("No 'ports' section found in configuration file");
    exit(-1);
  }
  if !config["general"]["print_ascii"].is_badvalue() {
    if config["general"]["print_ascii"].as_bool().unwrap() {
      app.print_ascii = true;
      println!("Printing ASCII");
    }
  }
  if !config["general"]["print_binary"].is_badvalue() {
    if config["general"]["print_binary"].as_bool().unwrap() {
      app.print_binary = true;
      println!("Printing binary in hexadecimal");
    }
  }
  if !config["general"]["sql_logging"].is_badvalue() {
    if config["general"]["sql_logging"].as_bool().unwrap() {
      app.sql_logging = true;
      println!("Logging events to portlurker.sqlite");
      let conn = Connection::open("portlurker.sqlite").expect("Failed to open or create database! Check your permissions or disk space.");
      conn.execute("CREATE TABLE IF NOT EXISTS connections (
        id         INTEGER PRIMARY KEY,
        time       INTEGER NOT NULL DEFAULT (strftime('%s', 'now')),
        remoteip   TEXT NOT NULL,
        remoteport INTEGER NOT NULL,
        localport  INTEGER NOT NULL
      )", &[]).expect("Failed to create table inside database!");
    }
  }
  let app = app; // Revert to immutable

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

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

  for port in config["ports"].as_vec().unwrap() {
    if !port["tcp"].is_badvalue() {
      let portno = port["tcp"].as_i64().unwrap();
      println!("TCP port {}", portno);
      let mut banner = Arc::new(String::new());
      if let Some(x) = port["banner"].as_str() {
        Arc::get_mut(& mut banner).unwrap().push_str(x);
        println!("  with banner: {}", *banner);
      }
      let regexset = regexset.clone();
      thread::spawn(move || {
        let server = TcpListener::bind(("0.0.0.0", portno as u16)).unwrap();
        for res in server.incoming() {
          let mut stream = match res {
            Ok(stream) => stream,
            Err(e) => { println!("ACCEPT ERROR TCP {}: {}", portno, e.to_string()); continue; }
          };
          stream.set_read_timeout(Some(io_timeout)).expect("Failed to set read timeout on TcpStream");
          stream.set_write_timeout(Some(io_timeout)).expect("Failed to set write timeout on TcpStream");
          let addr = stream.peer_addr().unwrap();
          println!("CONNECT TCP {} from {}", portno, addr);

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

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

            let conn = Connection::open("./portlurker.sqlite").unwrap();
            conn.execute("INSERT INTO connections (remoteip, remoteport, localport) VALUES (?1, ?2, ?3)", &[&newdbentry.remoteip, &newdbentry.remoteport, &newdbentry.localport]).expect("Can't write new row into table!");
          }

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

  loop { thread::sleep(Duration::new(60, 0)); }
}
<commit_msg>Config file error handling, replaced panic for message + exit<commit_after>extern crate yaml_rust;
extern crate regex;
extern crate rusqlite;

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

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

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

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

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

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

  let mut app = App { print_ascii: false, print_binary: false, sql_logging: false };
  let binary_matches = [
    ("SSL3.0 Record Protocol", r"^\x16\x03\x00..\x01"),
    ("TLS1.0 Record Protocol", r"^\x16\x03\x01..\x01"),
    ("TLS1.1 Record Protocol", r"^\x16\x03\x02..\x01"),
    ("TLS1.2 Record Protocol", r"^\x16\x03\x03..\x01"),
    ("SSL3.0 CLIENT_HELLO", r"^\x16....\x01...\x03\x00"),
    ("TLS1.0 CLIENT_HELLO", r"^\x16....\x01...\x03\x01"),
    ("TLS1.1 CLIENT_HELLO", r"^\x16....\x01...\x03\x02"),
    ("TLS1.2 CLIENT_HELLO", r"^\x16....\x01...\x03\x03"),
    ("SMB1 COMMAND NEGOTIATE", r"^....\xffSMB\x72"),
    ("SMB1 NT_STATUS Success", r"^....\xffSMB.[\x00-\x0f]"),
    ("SMB1 NT_STATUS Information", r"^....\xffSMB.[\x40-\x4f]"),
    ("SMB1 NT_STATUS Warning", r"^....\xffSMB.[\x80-\x8f]"),
    ("SMB1 NT_STATUS Error", r"^....\xffSMB.[\xc0-\xcf]"),
    ("SMB2 COMMAND NEGOTIATE", r"^\x00...\xfeSMB........\x00\x00"),
    ("SMB2 NT_STATUS Success", r"^\x00...\xfeSMB....[\x00-\x0f]"),
    ("SMB2 NT_STATUS Information", r"^\x00...\xfeSMB....[\x40-\x4f]"),
    ("SMB2 NT_STATUS Warning", r"^\x00...\xfeSMB....[\x80-\x8f]"),
    ("SMB2 NT_STATUS Error", r"^\x00...\xfeSMB....[\xc0-\xcf]"),
    ("MS-TDS PRELOGIN Request", r"^\x12\x01\x00.\x00\x00"),
    ("MS-TDS LOGIN Request", r"^\x10\x01\x00.\x00\x00"),
    ("SOCKS4 NOAUTH Request", r"^\x04\x01\x00\x50"),
    ("SOCKS5 NOAUTH Request", r"^\x05\x01\x00$"), // Tested ok-ish
    ("SOCKS5 USER/PASS Request", r"^\x05\x02\x00\x02$") // possibly broken
  ];
  let io_timeout = Duration::new(300, 0); // 5 minutes

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

  let docs = YamlLoader::load_from_str(&config_str).unwrap();
  let config = &docs[0];
  //println!("{:?}", config);
  if config["general"].is_badvalue() {
    println!("No 'general' section found in configuration file");
    exit(-1);
  }
  if config["ports"].is_badvalue() {
    println!("No 'ports' section found in configuration file");
    exit(-1);
  }
  if !config["general"]["print_ascii"].is_badvalue() {
    if config["general"]["print_ascii"].as_bool().unwrap() {
      app.print_ascii = true;
      println!("Printing ASCII");
    }
  }
  if !config["general"]["print_binary"].is_badvalue() {
    if config["general"]["print_binary"].as_bool().unwrap() {
      app.print_binary = true;
      println!("Printing binary in hexadecimal");
    }
  }
  if !config["general"]["sql_logging"].is_badvalue() {
    if config["general"]["sql_logging"].as_bool().unwrap() {
      app.sql_logging = true;
      println!("Logging events to portlurker.sqlite");
      let conn = Connection::open("portlurker.sqlite").expect("Failed to open or create database! Check your permissions or disk space.");
      conn.execute("CREATE TABLE IF NOT EXISTS connections (
        id         INTEGER PRIMARY KEY,
        time       INTEGER NOT NULL DEFAULT (strftime('%s', 'now')),
        remoteip   TEXT NOT NULL,
        remoteport INTEGER NOT NULL,
        localport  INTEGER NOT NULL
      )", &[]).expect("Failed to create table inside database!");
    }
  }
  let app = app; // Revert to immutable

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

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

  for port in config["ports"].as_vec().unwrap() {
    if !port["tcp"].is_badvalue() {
      let portno = port["tcp"].as_i64().unwrap();
      println!("TCP port {}", portno);
      let mut banner = Arc::new(String::new());
      if let Some(x) = port["banner"].as_str() {
        Arc::get_mut(& mut banner).unwrap().push_str(x);
        println!("  with banner: {}", *banner);
      }
      let regexset = regexset.clone();
      thread::spawn(move || {
        let server = TcpListener::bind(("0.0.0.0", portno as u16)).unwrap();
        for res in server.incoming() {
          let mut stream = match res {
            Ok(stream) => stream,
            Err(e) => { println!("ACCEPT ERROR TCP {}: {}", portno, e.to_string()); continue; }
          };
          stream.set_read_timeout(Some(io_timeout)).expect("Failed to set read timeout on TcpStream");
          stream.set_write_timeout(Some(io_timeout)).expect("Failed to set write timeout on TcpStream");
          let addr = stream.peer_addr().unwrap();
          println!("CONNECT TCP {} from {}", portno, addr);

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

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

            let conn = Connection::open("./portlurker.sqlite").unwrap();
            conn.execute("INSERT INTO connections (remoteip, remoteport, localport) VALUES (?1, ?2, ?3)", &[&newdbentry.remoteip, &newdbentry.remoteport, &newdbentry.localport]).expect("Can't write new row into table!");
          }

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

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