<commit_before>//! Main module of way-cooler

#[macro_use]
extern crate lazy_static;
#[macro_use]
extern crate bitflags;
#[macro_use]
extern crate dbus_macros;
#[cfg(not(test))]
extern crate rustwlc;
extern crate getopts;
#[cfg(test)]
extern crate dummy_rustwlc as rustwlc;
#[macro_use]
extern crate log;
extern crate env_logger;
#[macro_use]
extern crate hlua;
extern crate rustc_serialize;
#[macro_use]
extern crate json_macro;
extern crate nix;
extern crate petgraph;
extern crate uuid;
extern crate dbus;
extern crate cairo;
#[macro_use]
extern crate wayland_sys;
#[macro_use]
extern crate wayland_server;

#[macro_use]
mod macros;
mod convert;

mod callbacks;
mod keys;

mod lua;
mod registry;
mod commands;
mod ipc;

mod layout;

mod render;

mod wayland;

use std::env;
use std::fs::File;
use std::io::{BufRead, BufReader};
use std::path::Path;

use getopts::Options;
use log::LogLevel;
use nix::sys::signal::{self, SigHandler, SigSet, SigAction, SaFlags};

use rustwlc::types::LogType;

const VERSION: &'static str = env!("CARGO_PKG_VERSION");
const GIT_VERSION: &'static str = include_str!(concat!(env!("OUT_DIR"), "/git-version.txt"));
const DRIVER_MOD_PATH: &'static str = "/proc/modules";

/// Callback to route wlc logs into env_logger
fn log_handler(level: LogType, message: &str) {
    match level {
        LogType::Info => info!("wlc: {}", message),
        LogType::Warn => warn!("wlc: {}", message),
        LogType::Error => error!("wlc: {}", message),
        LogType::Wayland => info!("wayland: {}", message)
    }
}

/// Formats the log strings properly
fn log_format(record: &log::LogRecord) -> String {
    let color = match record.level() {
        LogLevel::Info => "",
        LogLevel::Trace => "\x1B[37m",
        LogLevel::Debug => "\x1B[37m",
        LogLevel::Warn =>  "\x1B[33m",
        LogLevel::Error => "\x1B[31m",
    };
    let location = record.location();
    let file = location.file();
    let line = location.line();
    let mut module_path = location.module_path();
    if let Some(index) = module_path.find("way_cooler::") {
        let index = index + "way_cooler::".len();
        module_path = &module_path[index..];
    }
    format!("{} {} [{}] \x1B[37m{}:{}\x1B[0m{0} {} \x1B[0m",
            color, record.level(), module_path, file, line, record.args())
}

/// Checks the loaded modules, and reports any problematic proprietary ones
fn detect_proprietary() {
    // If DISPLAY is present, we are running embedded
    if env::var("DISPLAY").is_ok() {
        return
    }
    match File::open(Path::new(DRIVER_MOD_PATH)) {
        Ok(file) => {
            let reader = BufReader::new(&file);
            for line in reader.lines() {
                if let Ok(line) = line {
                    if line.contains("nvidia") {
                        error!("Warning: Proprietary nvidia graphics drivers are installed, \
                               but they are not compatible with Wayland. \
                               Consider using nouveau drivers for Wayland.");
                        break;
                    } else if line.contains("fglrx") {
                        error!("Warning: Proprietary AMD graphics drivers are installed, \
                                but they are not compatible with Wayland. \
                               Consider using radeon drivers for Waylad.");
                        break;
                    }
                }
            }
        },
        Err(err) => {
            warn!("Could not read proprietary modules at \"{}\", because: {:#?}",
                  DRIVER_MOD_PATH, err);
            warn!("If you are running proprietary AMD or Nvidia graphics drivers, \
                   Way Cooler may not work for you");
        }
    }
}

#[inline(always)]
/// Determines if we should build with debug symbols.
pub fn debug_enabled() -> bool {
    cfg!(not(disable_debug))
}

/// Initializes the logging system.
pub fn init_logs() {
    let mut builder = env_logger::LogBuilder::new();
    builder.format(log_format);
    builder.filter(None, log::LogLevelFilter::Trace);
    if env::var("WAY_COOLER_LOG").is_ok() {
        builder.parse(&env::var("WAY_COOLER_LOG")
                      .expect("WAY_COOLER_LOG not defined"));
    }
    builder.init().expect("Unable to initialize logging!");
    info!("Logger initialized, setting wlc handlers.");
}

fn log_environment() {
    for (key, value) in env::vars() {
        debug!("{}: {}", key, value);
    }
}

/// Handler for signals
extern "C" fn sig_handle(_: nix::libc::c_int) {
    rustwlc::terminate();
}

fn main() {
    let args: Vec<String> = env::args().collect();

    let mut opts = Options::new();
    opts.optflag("", "version", "show version information");
    let matches = match opts.parse(&args[1..]) {
        Ok(m) => m,
        Err(f) => {
            println!("{}", f.to_string());
            std::process::exit(1);
        }
    };
    if matches.opt_present("version") {
        if !GIT_VERSION.is_empty() {
            println!("Way Cooler {} @ {}", VERSION, GIT_VERSION);
        } else {
            println!("Way Cooler {}", VERSION);
        }
        println!("Way Cooler IPC version {}", ipc::VERSION);
        return
    }
    println!("Launching way-cooler...");

    let sig_action = SigAction::new(SigHandler::Handler(sig_handle), SaFlags::empty(),
                                    SigSet::empty());
    unsafe {signal::sigaction(signal::SIGINT, &sig_action).unwrap() };

    init_logs();
    log_environment();
    detect_proprietary();
    // This prepares wlc, doesn't run main loop until run_wlc is called
    let run_wlc = rustwlc::init()
        .expect("Unable to initialize wlc!");
    wayland::init_wayland_protocols();
    rustwlc::log_set_rust_handler(log_handler);
    callbacks::init();
    commands::init();
    registry::init();
    ipc::init();

    info!("Running wlc...");
    run_wlc();
}
<commit_msg>Fixed warnings in main module for latest Rust<commit_after>//! Main module of way-cooler

#[macro_use]
extern crate lazy_static;
extern crate bitflags;
extern crate dbus_macros;
#[cfg(not(test))]
extern crate rustwlc;
extern crate getopts;
#[cfg(test)]
extern crate dummy_rustwlc as rustwlc;
#[macro_use]
extern crate log;
extern crate env_logger;
extern crate hlua;
extern crate rustc_serialize;
extern crate json_macro;
extern crate nix;
extern crate petgraph;
extern crate uuid;
extern crate dbus;
extern crate cairo;
#[macro_use]
extern crate wayland_sys;
extern crate wayland_server;

#[macro_use]
mod macros;
mod convert;

mod callbacks;
mod keys;

mod lua;
mod registry;
mod commands;
mod ipc;

mod layout;

mod render;

mod wayland;

use std::env;
use std::fs::File;
use std::io::{BufRead, BufReader};
use std::path::Path;

use getopts::Options;
use log::LogLevel;
use nix::sys::signal::{self, SigHandler, SigSet, SigAction, SaFlags};

use rustwlc::types::LogType;

const VERSION: &'static str = env!("CARGO_PKG_VERSION");
const GIT_VERSION: &'static str = include_str!(concat!(env!("OUT_DIR"), "/git-version.txt"));
const DRIVER_MOD_PATH: &'static str = "/proc/modules";

/// Callback to route wlc logs into env_logger
fn log_handler(level: LogType, message: &str) {
    match level {
        LogType::Info => info!("wlc: {}", message),
        LogType::Warn => warn!("wlc: {}", message),
        LogType::Error => error!("wlc: {}", message),
        LogType::Wayland => info!("wayland: {}", message)
    }
}

/// Formats the log strings properly
fn log_format(record: &log::LogRecord) -> String {
    let color = match record.level() {
        LogLevel::Info => "",
        LogLevel::Trace => "\x1B[37m",
        LogLevel::Debug => "\x1B[37m",
        LogLevel::Warn =>  "\x1B[33m",
        LogLevel::Error => "\x1B[31m",
    };
    let location = record.location();
    let file = location.file();
    let line = location.line();
    let mut module_path = location.module_path();
    if let Some(index) = module_path.find("way_cooler::") {
        let index = index + "way_cooler::".len();
        module_path = &module_path[index..];
    }
    format!("{} {} [{}] \x1B[37m{}:{}\x1B[0m{0} {} \x1B[0m",
            color, record.level(), module_path, file, line, record.args())
}

/// Checks the loaded modules, and reports any problematic proprietary ones
fn detect_proprietary() {
    // If DISPLAY is present, we are running embedded
    if env::var("DISPLAY").is_ok() {
        return
    }
    match File::open(Path::new(DRIVER_MOD_PATH)) {
        Ok(file) => {
            let reader = BufReader::new(&file);
            for line in reader.lines() {
                if let Ok(line) = line {
                    if line.contains("nvidia") {
                        error!("Warning: Proprietary nvidia graphics drivers are installed, \
                               but they are not compatible with Wayland. \
                               Consider using nouveau drivers for Wayland.");
                        break;
                    } else if line.contains("fglrx") {
                        error!("Warning: Proprietary AMD graphics drivers are installed, \
                                but they are not compatible with Wayland. \
                               Consider using radeon drivers for Waylad.");
                        break;
                    }
                }
            }
        },
        Err(err) => {
            warn!("Could not read proprietary modules at \"{}\", because: {:#?}",
                  DRIVER_MOD_PATH, err);
            warn!("If you are running proprietary AMD or Nvidia graphics drivers, \
                   Way Cooler may not work for you");
        }
    }
}

#[inline(always)]
/// Determines if we should build with debug symbols.
pub fn debug_enabled() -> bool {
    cfg!(not(disable_debug))
}

/// Initializes the logging system.
pub fn init_logs() {
    let mut builder = env_logger::LogBuilder::new();
    builder.format(log_format);
    builder.filter(None, log::LogLevelFilter::Trace);
    if env::var("WAY_COOLER_LOG").is_ok() {
        builder.parse(&env::var("WAY_COOLER_LOG")
                      .expect("WAY_COOLER_LOG not defined"));
    }
    builder.init().expect("Unable to initialize logging!");
    info!("Logger initialized, setting wlc handlers.");
}

fn log_environment() {
    for (key, value) in env::vars() {
        debug!("{}: {}", key, value);
    }
}

/// Handler for signals
extern "C" fn sig_handle(_: nix::libc::c_int) {
    rustwlc::terminate();
}

fn main() {
    let args: Vec<String> = env::args().collect();

    let mut opts = Options::new();
    opts.optflag("", "version", "show version information");
    let matches = match opts.parse(&args[1..]) {
        Ok(m) => m,
        Err(f) => {
            println!("{}", f.to_string());
            std::process::exit(1);
        }
    };
    if matches.opt_present("version") {
        if !GIT_VERSION.is_empty() {
            println!("Way Cooler {} @ {}", VERSION, GIT_VERSION);
        } else {
            println!("Way Cooler {}", VERSION);
        }
        println!("Way Cooler IPC version {}", ipc::VERSION);
        return
    }
    println!("Launching way-cooler...");

    let sig_action = SigAction::new(SigHandler::Handler(sig_handle), SaFlags::empty(),
                                    SigSet::empty());
    unsafe {signal::sigaction(signal::SIGINT, &sig_action).unwrap() };

    init_logs();
    log_environment();
    detect_proprietary();
    // This prepares wlc, doesn't run main loop until run_wlc is called
    let run_wlc = rustwlc::init()
        .expect("Unable to initialize wlc!");
    wayland::init_wayland_protocols();
    rustwlc::log_set_rust_handler(log_handler);
    callbacks::init();
    commands::init();
    registry::init();
    ipc::init();

    info!("Running wlc...");
    run_wlc();
}
<|endoftext|>


<commit_before>extern crate piston_meta;

use piston_meta::*;

fn main() {

}

#[cfg(test)]
mod tests {
    use piston_meta::*;

    #[test]
    fn extract_is_json() {
        let _ = load_syntax_data("assets/json/syntax.txt",
            "assets/extract/test.txt");
    }

    #[test]
    fn extract() {
        let _data = load_syntax_data("assets/extract/syntax.txt",
            "assets/extract/test.txt");
        // json::print(&_data);
    }

    #[test]
    fn dependencies_is_json() {
        let _ = load_syntax_data("assets/json/syntax.txt",
            "assets/dependencies/test.txt");
    }

    #[test]
    fn dependencies() {
        let _data = load_syntax_data("assets/dependencies/syntax.txt",
            "assets/dependencies/test.txt");
        // json::print(&_data);
    }
}
<commit_msg>Added `Extract` and `Dependency`<commit_after>#![deny(missing_docs)]

//! A tool for reasoning about breaking changes in Rust ecosystems

extern crate piston_meta;

use std::rc::Rc;

/// Stores extract information.
pub struct Extract {
    /// The url of the Cargo.toml.
    pub url: String,
    /// Whether to override the library version to simulate breaking change.
    pub override_version: Option<String>,
}

/// Stores dependency information.
pub struct Dependency {
    /// The package name.
    pub name: Rc<String>,
    /// The semver version of the library.
    pub version: String,
}

fn main() {

}

#[cfg(test)]
mod tests {
    use piston_meta::*;

    #[test]
    fn extract_is_json() {
        let _ = load_syntax_data("assets/json/syntax.txt",
            "assets/extract/test.txt");
    }

    #[test]
    fn extract() {
        let _data = load_syntax_data("assets/extract/syntax.txt",
            "assets/extract/test.txt");
        // json::print(&_data);
    }

    #[test]
    fn dependencies_is_json() {
        let _ = load_syntax_data("assets/json/syntax.txt",
            "assets/dependencies/test.txt");
    }

    #[test]
    fn dependencies() {
        let _data = load_syntax_data("assets/dependencies/syntax.txt",
            "assets/dependencies/test.txt");
        // json::print(&_data);
    }
}
<|endoftext|>


<commit_before>extern crate iron;
extern crate router;
extern crate rustc_serialize;

use std::io::Read;
use std::sync::{Arc, Mutex};
use std::collections::HashMap;

use iron::prelude::*;
use iron::status;
use router::Router;
use rustc_serialize::json::Json;

#[derive(Debug)]
struct Mapping;

impl Mapping {
    fn new() -> Mapping {
        Mapping
    }
}

#[derive(Debug)]
struct Index {
    pub mappings: HashMap<&'static str, Mapping>,
}

impl Index {
    fn new() -> Index {
        Index{
            mappings: HashMap::new(),
        }
    }
}

fn index_not_found_response() -> Response {
    let mut response = Response::with((status::NotFound, "{\"message\": \"Index not found\"}"));
    response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
    return response;
}

fn main() {
    let mut indices = Arc::new(Mutex::new(HashMap::new()));
    let mut wagtail_index = Index::new();
    wagtail_index.mappings.insert("wagtaildocs_document", Mapping::new());
    indices.lock().unwrap().insert("wagtail", wagtail_index);

    let mut router = Router::new();

    router.get("/", |_: &mut Request| -> IronResult<Response> {
        Ok(Response::with((status::Ok, "Hello World!")))
    });

    {
        let indices = indices.clone();

        router.get("/:index/_count", move |req: &mut Request| -> IronResult<Response> {
            // URL parameters
            let ref index_name = req.extensions.get::<Router>().unwrap().find("index").unwrap_or("");

            // Lock index array
            let mut indices = indices.lock().unwrap();

            // Find index
            let mut index = match indices.get(index_name) {
                Some(index) => index,
                None => {
                    return Ok(index_not_found_response());
                }
            };

            // Load query from body
            let mut payload = String::new();
            req.body.read_to_string(&mut payload).unwrap();

            // TODO: Run query

            let mut response = Response::with((status::Ok, "{\"count\": 0}"));
            response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
            Ok(response)
        });
    }

    {
        let indices = indices.clone();

        router.get("/:index/_search", move |req: &mut Request| -> IronResult<Response> {
            // URL parameters
            let ref index_name = req.extensions.get::<Router>().unwrap().find("index").unwrap_or("");

            // Lock index array
            let mut indices = indices.lock().unwrap();

            // Find index
            let mut index = match indices.get(index_name) {
                Some(index) => index,
                None => {
                    return Ok(index_not_found_response());
                }
            };

            // Load query from body
            let mut payload = String::new();
            req.body.read_to_string(&mut payload).unwrap();

            // TODO: Run query

            let mut response = Response::with((status::Ok, "{\"hits\": {\"total\": 0, \"hits\": []}}"));
            response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
            Ok(response)
        });
    }

    {
        let indices = indices.clone();

        router.put("/:index/:mapping/:doc", move |req: &mut Request| -> IronResult<Response> {
            // URL parameters
            let ref index_name = req.extensions.get::<Router>().unwrap().find("index").unwrap_or("");
            let ref mapping_name = req.extensions.get::<Router>().unwrap().find("mapping").unwrap_or("");
            let ref doc_id = req.extensions.get::<Router>().unwrap().find("doc").unwrap_or("");

            // Lock index array
            let mut indices = indices.lock().unwrap();

            // Find index
            let mut index = match indices.get(index_name) {
                Some(index) => index,
                None => {
                    return Ok(index_not_found_response());
                }
            };

            // Find mapping
            let mut mapping = match index.mappings.get(mapping_name) {
                Some(mapping) => mapping,
                None => {
                    let mut response = Response::with((status::NotFound, "{\"message\": \"Mapping not found\"}"));
                    response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
                    return Ok(response);
                }
            };

            // Load data from body
            let mut payload = String::new();
            req.body.read_to_string(&mut payload).unwrap();

            let data = match Json::from_str(&payload) {
                Ok(data) => data,
                Err(error) => {
                    // TODO: What specifically is bad about the JSON?
                    let mut response = Response::with((status::BadRequest, "{\"message\": \"Couldn't parse JSON\"}"));
                    response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
                    return Ok(response);
                }
            };

            // TODO: Validate and insert document

            let mut response = Response::with((status::Ok, "{}"));
            response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
            Ok(response)
        });
    }

    Iron::new(router).http("localhost:9200").unwrap();
}
<commit_msg>Parse query JSON in count and search views<commit_after>extern crate iron;
extern crate router;
extern crate rustc_serialize;

use std::io::Read;
use std::sync::{Arc, Mutex};
use std::collections::HashMap;

use iron::prelude::*;
use iron::status;
use router::Router;
use rustc_serialize::json::Json;

#[derive(Debug)]
struct Mapping;

impl Mapping {
    fn new() -> Mapping {
        Mapping
    }
}

#[derive(Debug)]
struct Index {
    pub mappings: HashMap<&'static str, Mapping>,
}

impl Index {
    fn new() -> Index {
        Index{
            mappings: HashMap::new(),
        }
    }
}

fn index_not_found_response() -> Response {
    let mut response = Response::with((status::NotFound, "{\"message\": \"Index not found\"}"));
    response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
    return response;
}

fn main() {
    let mut indices = Arc::new(Mutex::new(HashMap::new()));
    let mut wagtail_index = Index::new();
    wagtail_index.mappings.insert("wagtaildocs_document", Mapping::new());
    indices.lock().unwrap().insert("wagtail", wagtail_index);

    let mut router = Router::new();

    router.get("/", |_: &mut Request| -> IronResult<Response> {
        Ok(Response::with((status::Ok, "Hello World!")))
    });

    {
        let indices = indices.clone();

        router.get("/:index/_count", move |req: &mut Request| -> IronResult<Response> {
            // URL parameters
            let ref index_name = req.extensions.get::<Router>().unwrap().find("index").unwrap_or("");

            // Lock index array
            let mut indices = indices.lock().unwrap();

            // Find index
            let mut index = match indices.get(index_name) {
                Some(index) => index,
                None => {
                    return Ok(index_not_found_response());
                }
            };

            // Load query from body
            let mut payload = String::new();
            req.body.read_to_string(&mut payload).unwrap();

            let data = match Json::from_str(&payload) {
                Ok(data) => data,
                Err(error) => {
                    // TODO: What specifically is bad about the JSON?
                    let mut response = Response::with((status::BadRequest, "{\"message\": \"Couldn't parse JSON\"}"));
                    response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
                    return Ok(response);
                }
            };

            // TODO: Run query

            let mut response = Response::with((status::Ok, "{\"count\": 0}"));
            response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
            Ok(response)
        });
    }

    {
        let indices = indices.clone();

        router.get("/:index/_search", move |req: &mut Request| -> IronResult<Response> {
            // URL parameters
            let ref index_name = req.extensions.get::<Router>().unwrap().find("index").unwrap_or("");

            // Lock index array
            let mut indices = indices.lock().unwrap();

            // Find index
            let mut index = match indices.get(index_name) {
                Some(index) => index,
                None => {
                    return Ok(index_not_found_response());
                }
            };

            // Load query from body
            let mut payload = String::new();
            req.body.read_to_string(&mut payload).unwrap();

            let data = match Json::from_str(&payload) {
                Ok(data) => data,
                Err(error) => {
                    // TODO: What specifically is bad about the JSON?
                    let mut response = Response::with((status::BadRequest, "{\"message\": \"Couldn't parse JSON\"}"));
                    response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
                    return Ok(response);
                }
            };

            // TODO: Run query

            let mut response = Response::with((status::Ok, "{\"hits\": {\"total\": 0, \"hits\": []}}"));
            response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
            Ok(response)
        });
    }

    {
        let indices = indices.clone();

        router.put("/:index/:mapping/:doc", move |req: &mut Request| -> IronResult<Response> {
            // URL parameters
            let ref index_name = req.extensions.get::<Router>().unwrap().find("index").unwrap_or("");
            let ref mapping_name = req.extensions.get::<Router>().unwrap().find("mapping").unwrap_or("");
            let ref doc_id = req.extensions.get::<Router>().unwrap().find("doc").unwrap_or("");

            // Lock index array
            let mut indices = indices.lock().unwrap();

            // Find index
            let mut index = match indices.get(index_name) {
                Some(index) => index,
                None => {
                    return Ok(index_not_found_response());
                }
            };

            // Find mapping
            let mut mapping = match index.mappings.get(mapping_name) {
                Some(mapping) => mapping,
                None => {
                    let mut response = Response::with((status::NotFound, "{\"message\": \"Mapping not found\"}"));
                    response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
                    return Ok(response);
                }
            };

            // Load data from body
            let mut payload = String::new();
            req.body.read_to_string(&mut payload).unwrap();

            let data = match Json::from_str(&payload) {
                Ok(data) => data,
                Err(error) => {
                    // TODO: What specifically is bad about the JSON?
                    let mut response = Response::with((status::BadRequest, "{\"message\": \"Couldn't parse JSON\"}"));
                    response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
                    return Ok(response);
                }
            };

            // TODO: Validate and insert document

            let mut response = Response::with((status::Ok, "{}"));
            response.headers.set_raw("Content-Type", vec![b"application/json".to_vec()]);
            Ok(response)
        });
    }

    Iron::new(router).http("localhost:9200").unwrap();
}
<|endoftext|>


<commit_before>extern crate time;

use std::io::{Read, Write, Seek};
use std::process::exit;

const HELP_TEXT: &'static str = "netcp (send,receive) ipaddress[:port] filename";
const CALL_SIGN: &'static str = "netcp v0.1";
const MSG_AGREE: &'static str = "AGREE   ";
const MSG_DISAGREE: &'static str = "DISAGREE";
const MSG_FILE: &'static str = "FILE";
const MSG_END: &'static str = "END ";
const TIMEOUT: i64 = 800;

fn main() {
    //get commandline arguments
    let args: Vec<String> = std::env::args().collect();

    //if there is no argument (first arg is always application name), error
    if args.len() == 1 {
        error("No arguments");
    }

    //execute option
    match args[1].as_ref() {
        "help" => {
            println!("{}", HELP_TEXT);
        }

        "send" => {
            //sending needs min. 4 args (appl. name, "send", address, filenames...)
            if args.len() < 4 {
                error("Too few arguments");
            }

            send(&args[2], &args[3..]);
        }

        "receive" => {
            //receiving needs exactly 3 args (appl. name, "receive", address)
            if args.len() != 3 {
                error("Too few or many arguments");
            }

            receive(&args[2]);
        }

        //Unknown parameter
        _ => {
            error("Unknown parameter");
        }
    }
}

fn error(msg: &str) {
    println!("Error: {}.", msg);
    std::process::exit(1);
}

fn receive_data(stream: &mut std::net::TcpStream, buf: &mut [u8]) {
    let mut begin = time::get_time();
    let mut received_bytes: usize = 0;

    while received_bytes < buf.len() && (begin - time::get_time()).num_milliseconds() < TIMEOUT {
        let res = stream.read(&mut buf[received_bytes..]);

        if res.is_err() {
            error("Connection lost");
        }

        let bytes = res.unwrap();
        received_bytes += bytes;

        if bytes != 0 {
            begin = time::get_time();
        }
    }

    if received_bytes < buf.len() {
        error("Connection lost (timeout)");
    }
}

fn send_data(stream: &mut std::net::TcpStream, data: &[u8]) {
    let mut begin = time::get_time();
    let mut sended_bytes: usize = 0;

    while sended_bytes < data.len() && (begin - time::get_time()).num_milliseconds() < TIMEOUT {
        let res = stream.write(&data[sended_bytes..]);

        if res.is_err() {
            error("Connection lost");
        }

        let bytes = res.unwrap();
        sended_bytes += bytes;

        if bytes != 0 {
            begin = time::get_time();
        }
    }

    if sended_bytes < data.len() {
        error("Connection lost (timeout)");
    }
}

fn send_u64(stream: &mut std::net::TcpStream, num: u64) {
    let data = num.to_le();
    let buf: &[u8] = unsafe { std::mem::transmute::<&u64, &[u8; 8]>(&data) };
    send_data(stream, &buf);
}

fn receive_u64(stream: &mut std::net::TcpStream) -> u64 {
    let mut num = 0u64;
    receive_data(stream, unsafe { std::mem::transmute::<&mut u64, &mut [u8; 8]>(&mut num) });
    return u64::from_le(num);
}

fn send_string(stream: &mut std::net::TcpStream, string: &str) {
    let size = string.len() as u64;
    send_u64(stream, size);
    send_data(stream, string.as_bytes());
}

fn receive_string(stream: &mut std::net::TcpStream) -> String {
    let size = receive_u64(stream);
    let mut vec = vec![0u8; size as usize];
    receive_data(stream, &mut vec[..] );
    let string_res = String::from_utf8(vec);
    match string_res {
        Err(_) => { error("Couldn't convert bytes to string"); }
        _ => {}
    }
    return string_res.unwrap();
}

fn check_agreement(stream: &mut std::net::TcpStream) -> bool {
    let mut vec = vec![0u8; MSG_AGREE.len()];
    receive_data(stream, &mut vec[..]);

    if compare_byte_array(&vec[..], MSG_AGREE.as_bytes()) {
        return true;
    }
    else if compare_byte_array(&vec[..], MSG_DISAGREE.as_bytes()) {
        return false;
    }
    else {
        error("Invalid protocol");
    }
    return false;
}

fn compare_byte_array(a: &[u8], b: &[u8]) -> bool {
    if a.len() != b.len() {
        return false;
    }

    for i in 0..a.len() {
        if a[i] != b[i] {
            return false;
        }
    }

    return true;
}

fn get_filesize(file: &mut std::fs::File) -> u64 {
    let current = file.seek(std::io::SeekFrom::Current(0)).ok().expect("Error: file seeking failed.");
    let size = file.seek(std::io::SeekFrom::End(0)).ok().expect("Error: file seeking failed.");
    file.seek(std::io::SeekFrom::Start(current)).ok().expect("Error: file seeking failed.");
    return size;
}

fn send(address: &String, file_names: &[String]) {
    //1. check file accessibilities
    //2. check address
    //3. connect with client
    //4. send files sequential

    let mut buf = [0u8; 512];

    //1. check file existences
    //get working directory to find correct files
    let work_dir_res = std::env::current_dir();
    match work_dir_res {
        Err(_) => { error("Couldn't find working directory"); },
        _ => {}
    }
    let work_dir = work_dir_res.unwrap();

    //join working directory and file_names to get absolute file addresses in a vector
    let mut file_addrs: Vec<std::path::PathBuf> = Vec::with_capacity(file_names.len());
    for file_name in file_names {
        file_addrs.push(work_dir.join(file_name));
    }

    //open and close files to check accessibility
    for file_addr in &file_addrs {
        if let Err(_) = std::fs::File::open(&file_addr) {
            println!("Error: File doesn't exist or is not accessible ({}).", file_addr.display());
            exit(1);
        }
    }

    //2. check address and bind listener
    let listener_res = std::net::TcpListener::bind(&address[..]);
    match listener_res {
        Err(ref e) => { println!("Error: {}.", e); exit(1); },
        _ => {}
    }
    let listener = listener_res.unwrap();

    //3. connect with client
    print!("Waiting for client..."); let _ = std::io::stdout().flush();
    let stream_res = listener.accept();
    match stream_res {
        Err(ref e) => { println!("Error: {}.", e); exit(1); },
        _ => {}
    }
    let (mut stream, client_addr) = stream_res.unwrap();

    //3. check if client send correct CALL_SIGN
    receive_data(&mut stream, &mut buf[..CALL_SIGN.len()]);

    if compare_byte_array(CALL_SIGN.as_bytes(), &buf[..CALL_SIGN.len()]) == false {
        error("Invalid protocol");
    }

    send_data(&mut stream, MSG_AGREE.as_bytes());

    println!("connected with {}.", client_addr); let _ = std::io::stdout().flush();

    //send files
    for file_addr in &file_addrs {

        let file_res = std::fs::File::open(&file_addr);
        match file_res {
            Err(_) => { error("Couldn't open file") },
            _ => {}
        }
        let mut file = file_res.unwrap();
        let filesize = get_filesize(&mut file);

        send_data(&mut stream, MSG_FILE.as_bytes()); //send file is ready to send
        send_u64(&mut stream, filesize); //send file size as u64

        //send filename
        let filename_os_opt = file_addr.file_name();
        if filename_os_opt.is_none() {
            error("Couldn't convert filename to utf8");
        }
        let filename_opt = filename_os_opt.unwrap().to_str();
        if filename_opt.is_none() {
            error("Couldn't convert filename to utf8");
        }
        send_string(&mut stream, &filename_opt.unwrap());

        print!("Send {}...", &filename_opt.unwrap()); let _ = std::io::stdout().flush();

        //if client send MSG_DISAGREE, continue else send file
        if check_agreement(&mut stream) == false {
            println!("cancelled by client.");
            continue;
        }

        //send file
        let mut i = 0u64;
        while i < (filesize-1) {
            if (filesize - i) >= buf.len() as u64 {
                if let Err(_) = file.read(&mut buf) {
                    error("Couldn't read from file");
                }
                send_data(&mut stream, &buf);
                i += buf.len() as u64;
            }
            else {
                if let Err(_) = file.read(&mut buf[..(filesize - i) as usize]) {
                    error("Couldn't read from file");
                }
                send_data(&mut stream, &buf[..(filesize - i) as usize]);
                i = filesize-1;
            }
        }

        println!("done."); let _ = std::io::stdout().flush();
    }

    //send end
    send_data(&mut stream, MSG_END.as_bytes());
}

fn receive(address: &String) {
    let mut buf = [0u8; 512];

    let work_dir_res = std::env::current_dir();
    match work_dir_res {
        Err(_) => { error("Couldn't find working directory"); },
        _ => {}
    }
    let work_dir = work_dir_res.unwrap();

    let stream_res = std::net::TcpStream::connect(&address[..]);
    match stream_res {
        Err(ref e) => { println!("Error: {}.", e); exit(1); },
        _ => {}
    }
    let mut stream = stream_res.unwrap();

    send_data(&mut stream, CALL_SIGN.as_bytes());

    if check_agreement(&mut stream) == false {
        error("No server found");
    }

    let mut msg_file = vec![0u8; MSG_FILE.len()];
    loop {
        receive_data(&mut stream, &mut msg_file[..]);

        if compare_byte_array(&msg_file[..], MSG_END.as_bytes()) {
            break;
        }
        else if compare_byte_array(&msg_file[..], MSG_FILE.as_bytes()) == false {
            error("Invalid protocol");
        }

        let filesize = receive_u64(&mut stream);
        let filename = receive_string(&mut stream);

        let file_res = std::fs::File::create(work_dir.join(&filename));
        match file_res {
            Err(_) => {
                println!("Couldn't create file {}. Skip file transmition.", filename);
                send_data(&mut stream, MSG_DISAGREE.as_bytes());
                continue;
            },
            _ => {}
        }
        let mut file = file_res.unwrap();

        send_data(&mut stream, MSG_AGREE.as_bytes());

        print!("Receive {}...", filename); let _ = std::io::stdout().flush();

        //get file data
        let mut i = 0u64;
        while i < (filesize-1) {
            if (filesize - i) >= buf.len() as u64 {
                receive_data(&mut stream, &mut buf);
                if let Err(_) = file.write(&buf) {
                    error("Couldn't write to file");
                }
                i += buf.len() as u64;
            }
            else {
                receive_data(&mut stream, &mut buf[..(filesize - i) as usize]);
                if let Err(_) = file.write(&buf[..(filesize - i) as usize]) {
                    error("Couldn't write to file");
                }
                i = filesize-1;
            }
        }

        println!("done."); let _ = std::io::stdout().flush();
    }
}
<commit_msg>removed unnecessary error code<commit_after>extern crate time;

use std::io::{Read, Write, Seek};

const HELP_TEXT: &'static str = "netcp (send,receive) ipaddress[:port] filename";
const CALL_SIGN: &'static str = "netcp v0.1";
const MSG_AGREE: &'static str = "AGREE   ";
const MSG_DISAGREE: &'static str = "DISAGREE";
const MSG_FILE: &'static str = "FILE";
const MSG_END: &'static str = "END ";
const TIMEOUT: i64 = 800;

macro_rules! error {
    ($($tt:tt)*) => {{
        use std::io::Write;
        write!(&mut ::std::io::stderr(), "Error: ").unwrap();
        write!(&mut ::std::io::stderr(), $($tt)*).unwrap();
        writeln!(&mut ::std::io::stderr(), ".").unwrap();
        ::std::process::exit(1)
    }}
}

fn main() {
    //get commandline arguments
    let args: Vec<String> = std::env::args().collect();

    //if there is no argument (first arg is always application name), error
    if args.len() == 1 {
        error!("No arguments");
    }

    //execute option
    match args[1].as_ref() {
        "help" => {
            println!("{}", HELP_TEXT);
        }

        "send" => {
            //sending needs min. 4 args (appl. name, "send", address, filenames...)
            if args.len() < 4 {
                error!("Too few arguments");
            }

            send(&args[2], &args[3..]);
        }

        "receive" => {
            //receiving needs exactly 3 args (appl. name, "receive", address)
            if args.len() != 3 {
                error!("Too few or many arguments");
            }

            receive(&args[2]);
        }

        //Unknown parameter
        _ => {
            error!("Unknown parameter");
        }
    }
}

fn receive_data(stream: &mut std::net::TcpStream, buf: &mut [u8]) {
    let mut begin = time::get_time();
    let mut received_bytes: usize = 0;

    while received_bytes < buf.len() && (begin - time::get_time()).num_milliseconds() < TIMEOUT {
        let bytes = match stream.read(&mut buf[received_bytes..]) {
            Err(_) => error!("Connection lost"),
            Ok(bytes) => bytes
        };

        received_bytes += bytes;

        if bytes != 0 {
            begin = time::get_time();
        }
    }

    if received_bytes < buf.len() {
        error!("Connection lost (timeout)");
    }
}

fn send_data(stream: &mut std::net::TcpStream, data: &[u8]) {
    let mut begin = time::get_time();
    let mut sended_bytes: usize = 0;

    while sended_bytes < data.len() && (begin - time::get_time()).num_milliseconds() < TIMEOUT {
        let bytes = match stream.write(&data[sended_bytes..]) {
            Err(_) => error!("Connection lost"),
            Ok(bytes) => bytes
        };

        sended_bytes += bytes;

        if bytes != 0 {
            begin = time::get_time();
        }
    }

    if sended_bytes < data.len() {
        error!("Connection lost (timeout)");
    }
}

fn send_u64(stream: &mut std::net::TcpStream, num: u64) {
    let data = num.to_le();
    let buf: &[u8] = unsafe { std::mem::transmute::<&u64, &[u8; 8]>(&data) };
    send_data(stream, &buf);
}

fn receive_u64(stream: &mut std::net::TcpStream) -> u64 {
    let mut num = 0u64;
    receive_data(stream, unsafe { std::mem::transmute::<&mut u64, &mut [u8; 8]>(&mut num) });
    return u64::from_le(num);
}

fn send_string(stream: &mut std::net::TcpStream, string: &str) {
    let size = string.len() as u64;
    send_u64(stream, size);
    send_data(stream, string.as_bytes());
}

fn receive_string(stream: &mut std::net::TcpStream) -> String {
    let size = receive_u64(stream);
    let mut vec = vec![0u8; size as usize];
    receive_data(stream, &mut vec[..] );
    let string = match String::from_utf8(vec) {
        Err(_) => error!("Couldn't convert bytes to string"),
        Ok(string) => string
    };
    return string;
}

fn check_agreement(stream: &mut std::net::TcpStream) -> bool {
    let mut vec = vec![0u8; MSG_AGREE.len()];
    receive_data(stream, &mut vec[..]);

    if compare_byte_array(&vec[..], MSG_AGREE.as_bytes()) {
        return true;
    }
    else if compare_byte_array(&vec[..], MSG_DISAGREE.as_bytes()) {
        return false;
    }
    else {
        error!("Invalid protocol");
    }
}

fn compare_byte_array(a: &[u8], b: &[u8]) -> bool {
    if a.len() != b.len() {
        return false;
    }

    for i in 0..a.len() {
        if a[i] != b[i] {
            return false;
        }
    }

    return true;
}

fn get_filesize(file: &mut std::fs::File) -> u64 {
    let current = match file.seek(std::io::SeekFrom::Current(0)) {
        Err(_) => error!("File seeking failed"),
        Ok(current) => current
    };

    let size = match file.seek(std::io::SeekFrom::End(0)) {
        Err(_) => error!("File seeking failed"),
        Ok(size) => size
    };

    if let Err(_) = file.seek(std::io::SeekFrom::Start(current)) {
        error!("File seeking failed");
    }

    return size;
}

fn send(address: &String, file_names: &[String]) {
    //1. check file accessibilities
    //2. check address
    //3. connect with client
    //4. send files sequential

    let mut buf = [0u8; 512];

    //1. check file existences
    //get working directory to find correct files
    let work_dir = match std::env::current_dir() {
        Err(_) => error!("Couldn't find working directory"),
        Ok(work_dir) => work_dir
    };

    //join working directory and file_names to get absolute file addresses in a vector
    let mut file_addrs: Vec<std::path::PathBuf> = Vec::with_capacity(file_names.len());
    for file_name in file_names {
        file_addrs.push(work_dir.join(file_name));
    }

    //open and close files to check accessibility
    for file_addr in &file_addrs {
        if let Err(_) = std::fs::File::open(&file_addr) {
            error!("File doesn't exist or is not accessible ({})", file_addr.display());
        }
    }

    //2. check address and bind listener
    let listener = match std::net::TcpListener::bind(&address[..]) {
        Err(e) => error!("{}", e),
        Ok(listener) => listener
    };

    //3. connect with client
    print!("Waiting for client..."); let _ = std::io::stdout().flush();
    let (mut stream, client_addr) = match listener.accept() {
        Err(e) => error!("{}", e),
        Ok(stream) => stream
    };

    //3. check if client send correct CALL_SIGN
    receive_data(&mut stream, &mut buf[..CALL_SIGN.len()]);

    if compare_byte_array(CALL_SIGN.as_bytes(), &buf[..CALL_SIGN.len()]) == false {
        error!("Invalid protocol");
    }

    send_data(&mut stream, MSG_AGREE.as_bytes());

    println!("connected with {}.", client_addr); let _ = std::io::stdout().flush();

    //send files
    for file_addr in &file_addrs {

        let mut file = match std::fs::File::open(&file_addr) {
            Err(_) => error!("Couldn't open file"),
            Ok(file) => file
        };

        let filesize = get_filesize(&mut file);

        send_data(&mut stream, MSG_FILE.as_bytes()); //send file is ready to send
        send_u64(&mut stream, filesize); //send file size as u64

        //send filename
        let filename_os = match file_addr.file_name() {
            None => error!("Couldn't convert filename to utf8"),
            Some(filename_os) => filename_os
        };

        let filename = match filename_os.to_str() {
            None => error!("Couldn't convert filename to utf8"),
            Some(filename) => filename
        };

        send_string(&mut stream, &filename);

        print!("Send {}...", &filename); let _ = std::io::stdout().flush();

        //if client send MSG_DISAGREE, continue else send file
        if check_agreement(&mut stream) == false {
            println!("cancelled by client.");
            continue;
        }

        //send file
        let mut i = 0u64;
        while i < (filesize-1) {
            if (filesize - i) >= buf.len() as u64 {
                if let Err(_) = file.read(&mut buf) {
                    error!("Couldn't read from file");
                }
                send_data(&mut stream, &buf);
                i += buf.len() as u64;
            }
            else {
                if let Err(_) = file.read(&mut buf[..(filesize - i) as usize]) {
                    error!("Couldn't read from file");
                }
                send_data(&mut stream, &buf[..(filesize - i) as usize]);
                i = filesize-1;
            }
        }

        println!("done."); let _ = std::io::stdout().flush();
    }

    //send end
    send_data(&mut stream, MSG_END.as_bytes());
}

fn receive(address: &String) {
    let mut buf = [0u8; 512];

    let work_dir = match std::env::current_dir() {
        Err(_) => error!("Couldn't find working directory"),
        Ok(work_dir) => work_dir
    };

    let mut stream = match std::net::TcpStream::connect(&address[..]) {
        Err(e) => error!("{}", e),
        Ok(stream) => stream
    };

    send_data(&mut stream, CALL_SIGN.as_bytes());

    if check_agreement(&mut stream) == false {
        error!("No server found");
    }

    let mut msg_file = vec![0u8; MSG_FILE.len()];
    loop {
        receive_data(&mut stream, &mut msg_file[..]);

        if compare_byte_array(&msg_file[..], MSG_END.as_bytes()) {
            break;
        }
        else if compare_byte_array(&msg_file[..], MSG_FILE.as_bytes()) == false {
            error!("Invalid protocol");
        }

        let filesize = receive_u64(&mut stream);
        let filename = receive_string(&mut stream);

        let file_res = std::fs::File::create(work_dir.join(&filename));
        match file_res {
            Err(_) => {
                println!("Couldn't create file {}. Skip file transmition.", filename);
                send_data(&mut stream, MSG_DISAGREE.as_bytes());
                continue;
            },
            _ => {}
        }
        let mut file = file_res.unwrap();

        send_data(&mut stream, MSG_AGREE.as_bytes());

        print!("Receive {}...", filename); let _ = std::io::stdout().flush();

        //get file data
        let mut i = 0u64;
        while i < (filesize-1) {
            if (filesize - i) >= buf.len() as u64 {
                receive_data(&mut stream, &mut buf);
                if let Err(_) = file.write(&buf) {
                    error!("Couldn't write to file");
                }
                i += buf.len() as u64;
            }
            else {
                receive_data(&mut stream, &mut buf[..(filesize - i) as usize]);
                if let Err(_) = file.write(&buf[..(filesize - i) as usize]) {
                    error!("Couldn't write to file");
                }
                i = filesize-1;
            }
        }

        println!("done."); let _ = std::io::stdout().flush();
    }
}
<|endoftext|>


<commit_before>extern crate ftp;
extern crate regex;
extern crate openssl;

use std::env;
use std::fs;
use std::fs::File;
use std::io::prelude::*;
use std::fs::OpenOptions;
use std::io::BufWriter;
use ftp::FtpStream;
// use openssl::ssl::*;

struct DownloadableFile {
    client_dir: String,
    server_dir: String,
    filename: String,
    clientfile_namescheme: String,
}

struct DownRequest {
    host: String,
    user: String,
    pass: String,
    remote_files: Vec<DownloadableFile>,
}

/// Makes a directory specified
fn mk_dir(d: &str) -> std::io::Result<()> {
    try!(fs::create_dir_all(d));
    Ok(())
}

/// Breaks down the config given a String of the config file's contents
fn break_conf(buff: &mut String) -> Option<DownRequest> {

    /// Find line and get the data from it, if there isn't data to get, None is returned
    fn break_line(input: &mut String, line: &str) -> Option<String> {
        if input.contains(line) {
            let loc = input.find(line).unwrap().checked_add(line.len()).unwrap();
            let temp = &input[loc..];

            let loc = temp.find(";").unwrap();
            Some(temp[..loc].trim().to_string())
        } else {
            None
        }
    }
    // Make sure we won't get an unexpected None
    if buff.contains("host:") && buff.contains("user:") && buff.contains("password:") &&
       buff.contains("remoteFiles:") {

        /// Breaks down the remoteFile: line specificly since it's kinda special in layout
        fn break_remotefile(input: &mut String) -> Option<Vec<DownloadableFile>> {

            // Trim string down to what we need
            let loc =
                input.find("remoteFiles:").unwrap().checked_add("remoteFiles:".len()).unwrap();
            let remfile = &input[loc..];
            let remfile = &remfile[..remfile.rfind(";").unwrap().checked_sub(";".len()).unwrap()];
            let remfile = &remfile[remfile.find("[").unwrap().checked_add("[".len()).unwrap()..];

            // Break it up
            let files: Vec<&str> = remfile.split(',').collect();
            let mut dlable_f: Vec<DownloadableFile> = Vec::new();

            // [LOG FUNCTION]
            println!("There are {} entries from this config to download",
                     files.len());

            for file in files {

                if file.contains("remoteDir:") && file.contains("localDir:") &&
                   file.contains("name:") {
                    let a: &mut String = &mut file.to_string();
                    dlable_f.push(DownloadableFile {
                        client_dir: break_line(a, "localDir:").unwrap(),
                        server_dir: break_line(a, "remoteDir:").unwrap(),
                        filename: break_line(a, "name:").unwrap(),
                        clientfile_namescheme: break_line(a, "nameToSaveAs:")
                            .unwrap_or(break_line(a, "name:").unwrap()),
                    });
                } else {}

            }

            if dlable_f.len() > 0 {
                Some(dlable_f)
            } else {
                // [LOG FUNCTION]
                println!("No files found to download or check on in this config. Check your \
                          config.");
                None
            }

        }

        let temp = DownRequest {
            host: break_line(buff, "host:").unwrap(),
            user: break_line(buff, "user:").unwrap(),
            pass: break_line(buff, "password:").unwrap(),
            remote_files: break_remotefile(buff).unwrap(),
        };

        Some(temp)
    } else {
        // If we would there's no point
        None
    }

}

/// Loads all configs into the folder
fn load_configs() -> std::io::Result<Vec<DownRequest>> {

    if mk_dir("./configs").is_ok() {
        println!("Made the dir needed");
    } // make directory

    let config_dir = fs::read_dir("./configs/").unwrap(); //read all files in dir, and expand them
    let mut config_files: Vec<std::fs::DirEntry> = Vec::new();

    for file in config_dir {
        config_files.push(file.unwrap());
    }

    let mut configs_fmted = Vec::new();

    for file in &config_files {
        // for all files in the directory open them to the f var., and save the contained string

        let mut f = try!(File::open(file.path()));
        let mut buff = String::new();

        try!(f.read_to_string(&mut buff));
        let temp_config = break_conf(&mut buff); // send buff off to be broken down

        if temp_config.is_some() {
            // Check to make sure we didn't get nothing.
            configs_fmted.push(temp_config.unwrap());
            println!("Config \"{}\" loaded successfully!",
                     &file.file_name().to_str().unwrap());
        } else {
            println!("Config \"{}\" couldn't be loaded.",
                     &file.file_name().to_str().unwrap());
        }

    }
    println!("There were {} configs to load!", config_files.len());
    Ok(configs_fmted)
}

/// Downloads files from a DownRequest
fn download_from_site(c: DownRequest, dr: bool) -> std::io::Result<()> {
    // TODO ftps support
    let host: String;
    let port: u16;
    if c.host.contains(":") {
        let sp: Vec<&str> = c.host.split(':').collect();
        host = sp[0].trim().to_string();
        port = sp[1].trim().parse().unwrap();

    } else {
        host = c.host;
        port = 21
    }

    println!("Downloading files from \"{}:{}\"", host, port);

    let mut fstream = FtpStream::connect((host.as_str(), port)).unwrap();

    if fstream.login(&c.user.as_str(), &c.pass.as_str()).is_ok() {

        println!("Connected to server, and logged in successfully");

        for f in c.remote_files {
            /// prints the print_progress from a usize from 0 - 100
            fn print_progress(per_done: usize) {
                let outa = per_done.checked_mul(4).unwrap().checked_div(10).unwrap();
                let mut t = 0;
                let mut pbar: String = "[".to_string();
                while t < 40 {
                    if outa < t {
                        pbar.push_str(" ")
                    } else {
                        pbar.push_str("=")
                    }

                    t = t + 1;
                }
                pbar.push_str("]");

                print!("\r{} : {}%", pbar, per_done);
            }

            let mut to_download: Vec<String> = Vec::new();
            let mut cl_names: Vec<String>;
            let _cd = fstream.cwd(&f.server_dir.as_str());

            if f.filename.contains("(*)") {
                let all_in_dir: Vec<String> = fstream.nlst(None).unwrap();
                let spieces_to_check: Vec<&str> = f.filename.split("(*)").collect();
                let cpieces_to_check: Vec<&str> = f.clientfile_namescheme.split("(*)").collect();

                for sf in all_in_dir {

                    let mut check_match = true;

                    while check_match {
                        for p in &spieces_to_check {
                            if !sf.contains(p) {
                                check_match = false;
                            }
                        }
                        break;
                    }
                    if check_match {
                        to_download.push(sf);
                    }
                }

                if f.filename.contains(f.clientfile_namescheme.as_str()) &&
                   f.filename.len() == f.clientfile_namescheme.len() {
                    cl_names = to_download.clone();
                } else {
                    cl_names = Vec::new();

                    if spieces_to_check.len() == cpieces_to_check.len() {
                        for n in &to_download {

                            let mut inv_pieces: Vec<&str> = Vec::new();
                            let mut cfn: String = "".to_string();
                            inv_pieces.push(n.as_str());
                            for piece in &spieces_to_check {
                                let mut temp: Vec<&str> = Vec::new();
                                for a in inv_pieces {
                                    if a.contains(piece) {
                                        let t: Vec<&str> = a.split(piece).collect();
                                        for tt in t {
                                            temp.push(tt);
                                        }
                                    } else {
                                        temp.push(a);
                                    }
                                }
                                inv_pieces = temp;
                            }

                            for x in 0..inv_pieces.len() {
                                if x > 0 {
                                    cfn.push_str(cpieces_to_check[x - 1])
                                }
                                cfn.push_str(inv_pieces[x]);
                            }

                            cl_names.push(cfn);
                        }
                    } else {
                        println!("Can't download files. The client and server name must have the \
                                  same number of (*)!");
                    }
                }

                println!("{} files matching \"{}\"", &to_download.len(), &f.filename);

            } else {
                cl_names = Vec::new();
                cl_names.push(f.clientfile_namescheme);
                to_download.push(f.filename);
            }

            for fz in 0..cl_names.len() {
                let dir = &f.client_dir.as_str().clone();
                let _md = mk_dir(dir); // Make Client dir

                /// Checks to see if the file is downloaded
                /// Takes a local file with it's path and the remote filee's byte size in usize
                fn is_downloaded(fileandpath: &String, rsz: &usize) -> Option<bool> {
                    let mdata = fs::metadata(fileandpath);
                    if mdata.is_ok() {

                        let lsz = mdata.unwrap().len() as usize;
                        if &lsz == rsz { Some(true) } else { Some(false) }
                    } else {
                        None
                    }
                }
                if dr == false {
                    let rsize = fstream.size(&to_download[fz].as_str()).unwrap();

                    let mut fnm = "".to_string();
                    fnm.push_str(dir);
                    fnm.push_str("/");
                    fnm.push_str(&cl_names[fz].as_str());

                    println!("Downloading \"{}\"", to_download[fz]);

                    if !is_downloaded(&fnm, &rsize.unwrap()).unwrap_or(false) {

                        // Make local file
                        let _tt = fstream.transfer_type(::ftp::types::FileType::Binary);

                        let _result = fstream.retr(&to_download[fz].as_str(), |stream| {
                            let mut buf = [0; 4096];
                            let mut flnm = "".to_string();
                            flnm.push_str(dir);
                            flnm.push_str("/");
                            flnm.push_str(&cl_names[fz].as_str());

                            let mut file = File::create(flnm).unwrap();
                            let mut lsize: usize = 0;

                            loop {
                                match stream.read(&mut buf) {
                                    Ok(0) => break,
                                    Ok(n) => {
                                        file.write_all(&buf[0..n]).unwrap();
                                        lsize += n;
                                    }
                                    Err(_err) => break,
                                };
                                print_progress(((lsize.checked_mul(100).unwrap())
                                    .checked_div(rsize.unwrap())
                                    .unwrap()));
                            }
                            print!("\r\n");
                            Ok(())
                        });
                    } else {
                        println!("File is already downloaded");
                    }
                } else {
                    println!("Would be downloading {} and saving as {}",
                             &to_download[fz],
                             &cl_names[fz]);
                }
            }
        }

        let _f = fstream.quit();

        Ok(())
    } else {
        let _f = fstream.quit();
        println!("User login credentals were incorrect");
        Ok(())
    }
}

// Adds a file to downlaod to a config
fn add_to_file(mut file: Vec<String>) -> std::io::Result<()> {
    // TODO Catch errors and check inputs
    if file.len() == 4 {
        let g = file[3].clone();
        file.push(g);
    }
    let mut name = "./configs/".to_string();
    let t = file[0].clone();
    name.push_str(t.as_str());
    let mut b = "".to_string();
    let mut _a = File::open(&name).unwrap();
    let __ = _a.read_to_string(&mut b);
    let mut _f = OpenOptions::new().write(true).truncate(true).open(name).unwrap();
    let mut _w = BufWriter::new(_f);
    let _s1: &str;
    let _s2: &str;

    let (_s1, _s2) = b.split_at(b.rfind("}").unwrap().checked_add("}".len()).unwrap());

    println!("Writing to config");

    let _ =
        _w.write_all(format!("{},\r\n{{\r\n  remoteDir: {};\r\n  localDir: {};\r\n  name: \
                              {};\r\n  nameToSaveAs: {};\r\n}}{}",
                             _s1,
                             file[1],
                             file[2],
                             file[3],
                             file[4],
                             _s2)
            .as_bytes());
    Ok(())
}

fn main() {

    let mut _should_ls = false;
    let mut dry_run = false;
    let mut should_add = false;
    let mut add_to_config: Vec<String> = Vec::new();
    let ag: Vec<String> = env::args().collect();

    for x in 0..ag.len() {
        if ag[x].contains("-dr") {
            dry_run = true;
        }
        if ag[x].contains("-af") {
            if (x + 4) < ag.len() {
                should_add = true;
                add_to_config.push(ag[x + 1].clone());
                add_to_config.push(ag[x + 2].clone());
                add_to_config.push(ag[x + 3].clone());
                add_to_config.push(ag[x + 4].clone());
                if x + 5 < ag.len() {
                    add_to_config.push(ag[x + 5].clone());
                }
            } else {
                println!("Config name, client and server directory, and server filename must be \
                          supplied");
            }
        }
    }

    if should_add {
        let _ = add_to_file(add_to_config);
    }

    // TODO ls command
    println!("\rLoading configs...");

    let configs = load_configs().unwrap();

    for con in configs {
        let _d = download_from_site(con, dry_run);
    }
    std::process::exit(0);

}
<commit_msg>Fix for compile<commit_after>extern crate ftp;
extern crate regex;
extern crate openssl;

use std::env;
use std::fs;
use std::fs::File;
use std::io::prelude::*;
use std::fs::OpenOptions;
use std::io::BufWriter;
use ftp::FtpStream;
// use openssl::ssl::*;

struct DownloadableFile {
    client_dir: String,
    server_dir: String,
    filename: String,
    clientfile_namescheme: String,
}

struct DownRequest {
    host: String,
    user: String,
    pass: String,
    remote_files: Vec<DownloadableFile>,
}

/// Makes a directory specified
fn mk_dir(d: &str) -> std::io::Result<()> {
    try!(fs::create_dir_all(d));
    Ok(())
}

/// Breaks down the config given a String of the config file's contents
fn break_conf(buff: &mut String) -> Option<DownRequest> {

    /// Find line and get the data from it, if there isn't data to get, None is returned
    fn break_line(input: &mut String, line: &str) -> Option<String> {
        if input.contains(line) {
            let loc = input.find(line).unwrap().checked_add(line.len()).unwrap();
            let temp = &input[loc..];

            let loc = temp.find(";").unwrap();
            Some(temp[..loc].trim().to_string())
        } else {
            None
        }
    }
    // Make sure we won't get an unexpected None
    if buff.contains("host:") && buff.contains("user:") && buff.contains("password:") &&
       buff.contains("remoteFiles:") {

        /// Breaks down the remoteFile: line specificly since it's kinda special in layout
        fn break_remotefile(input: &mut String) -> Option<Vec<DownloadableFile>> {

            // Trim string down to what we need
            let loc =
                input.find("remoteFiles:").unwrap().checked_add("remoteFiles:".len()).unwrap();
            let remfile = &input[loc..];
            let remfile = &remfile[..remfile.rfind(";").unwrap().checked_sub(";".len()).unwrap()];
            let remfile = &remfile[remfile.find("[").unwrap().checked_add("[".len()).unwrap()..];

            // Break it up
            let files: Vec<&str> = remfile.split(',').collect();
            let mut dlable_f: Vec<DownloadableFile> = Vec::new();

            // [LOG FUNCTION]
            println!("There are {} entries from this config to download",
                     files.len());

            for file in files {

                if file.contains("remoteDir:") && file.contains("localDir:") &&
                   file.contains("name:") {
                    let a: &mut String = &mut file.to_string();
                    dlable_f.push(DownloadableFile {
                        client_dir: break_line(a, "localDir:").unwrap(),
                        server_dir: break_line(a, "remoteDir:").unwrap(),
                        filename: break_line(a, "name:").unwrap(),
                        clientfile_namescheme: break_line(a, "nameToSaveAs:")
                            .unwrap_or(break_line(a, "name:").unwrap()),
                    });
                } else {}

            }

            if dlable_f.len() > 0 {
                Some(dlable_f)
            } else {
                // [LOG FUNCTION]
                println!("No files found to download or check on in this config. Check your \
                          config.");
                None
            }

        }

        let temp = DownRequest {
            host: break_line(buff, "host:").unwrap(),
            user: break_line(buff, "user:").unwrap(),
            pass: break_line(buff, "password:").unwrap(),
            remote_files: break_remotefile(buff).unwrap(),
        };

        Some(temp)
    } else {
        // If we would there's no point
        None
    }

}

/// Loads all configs into the folder
fn load_configs() -> std::io::Result<Vec<DownRequest>> {

    if mk_dir("./configs").is_ok() {
        println!("Made the dir needed");
    } // make directory

    let config_dir = fs::read_dir("./configs/").unwrap(); //read all files in dir, and expand them
    let mut config_files: Vec<std::fs::DirEntry> = Vec::new();

    for file in config_dir {
        config_files.push(file.unwrap());
    }

    let mut configs_fmted = Vec::new();

    for file in &config_files {
        // for all files in the directory open them to the f var., and save the contained string

        let mut f = try!(File::open(file.path()));
        let mut buff = String::new();

        try!(f.read_to_string(&mut buff));
        let temp_config = break_conf(&mut buff); // send buff off to be broken down

        if temp_config.is_some() {
            // Check to make sure we didn't get nothing.
            configs_fmted.push(temp_config.unwrap());
            println!("Config \"{}\" loaded successfully!",
                     &file.file_name().to_str().unwrap());
        } else {
            println!("Config \"{}\" couldn't be loaded.",
                     &file.file_name().to_str().unwrap());
        }

    }
    println!("There were {} configs to load!", config_files.len());
    Ok(configs_fmted)
}

/// Downloads files from a DownRequest
fn download_from_site(c: DownRequest, dr: bool) -> std::io::Result<()> {
    // TODO ftps support
    let host: String;
    let port: u16;
    if c.host.contains(":") {
        let sp: Vec<&str> = c.host.split(':').collect();
        host = sp[0].trim().to_string();
        port = sp[1].trim().parse().unwrap();

    } else {
        host = c.host;
        port = 21
    }

    println!("Downloading files from \"{}:{}\"", host, port);

    let mut fstream = FtpStream::connect((host.as_str(), port)).unwrap();

    if fstream.login(&c.user.as_str(), &c.pass.as_str()).is_ok() {

        println!("Connected to server, and logged in successfully");

        for f in c.remote_files {
            /// prints the print_progress from a usize from 0 - 100
            fn print_progress(per_done: usize) {
                let outa = per_done.checked_mul(4).unwrap().checked_div(10).unwrap();
                let mut t = 0;
                let mut pbar: String = "[".to_string();
                while t < 40 {
                    if outa < t {
                        pbar.push_str(" ")
                    } else {
                        pbar.push_str("=")
                    }

                    t = t + 1;
                }
                pbar.push_str("]");

                print!("\r{} : {}%", pbar, per_done);
            }

            let mut to_download: Vec<String> = Vec::new();
            let mut cl_names: Vec<String>;
            let _cd = fstream.cwd(&f.server_dir.as_str());

            if f.filename.contains("(*)") {
                let all_in_dir: Vec<String> = fstream.nlst(None).unwrap();
                let spieces_to_check: Vec<&str> = f.filename.split("(*)").collect();
                let cpieces_to_check: Vec<&str> = f.clientfile_namescheme.split("(*)").collect();

                for sf in all_in_dir {

                    let mut check_match = true;

                    while check_match {
                        for p in &spieces_to_check {
                            if !sf.contains(p) {
                                check_match = false;
                            }
                        }
                        break;
                    }
                    if check_match {
                        to_download.push(sf);
                    }
                }

                if f.filename.contains(f.clientfile_namescheme.as_str()) &&
                   f.filename.len() == f.clientfile_namescheme.len() {
                    cl_names = to_download.clone();
                } else {
                    cl_names = Vec::new();

                    if spieces_to_check.len() == cpieces_to_check.len() {
                        for n in &to_download {

                            let mut inv_pieces: Vec<&str> = Vec::new();
                            let mut cfn: String = "".to_string();
                            inv_pieces.push(n.as_str());
                            for piece in &spieces_to_check {
                                let mut temp: Vec<&str> = Vec::new();
                                for a in inv_pieces {
                                    if a.contains(piece) {
                                        let t: Vec<&str> = a.split(piece).collect();
                                        for tt in t {
                                            temp.push(tt);
                                        }
                                    } else {
                                        temp.push(a);
                                    }
                                }
                                inv_pieces = temp;
                            }

                            for x in 0..inv_pieces.len() {
                                if x > 0 {
                                    cfn.push_str(cpieces_to_check[x - 1])
                                }
                                cfn.push_str(inv_pieces[x]);
                            }

                            cl_names.push(cfn);
                        }
                    } else {
                        println!("Can't download files. The client and server name must have the \
                                  same number of (*)!");
                    }
                }

                println!("{} files matching \"{}\"", &to_download.len(), &f.filename);

            } else {
                cl_names = Vec::new();
                cl_names.push(f.clientfile_namescheme);
                to_download.push(f.filename);
            }

            for fz in 0..cl_names.len() {
                let dir = &f.client_dir.as_str().clone();
                let _md = mk_dir(dir); // Make Client dir

                /// Checks to see if the file is downloaded
                /// Takes a local file with it's path and the remote filee's byte size in usize
                fn is_downloaded(fileandpath: &String, rsz: &usize) -> Option<bool> {
                    let mdata = fs::metadata(fileandpath);
                    if mdata.is_ok() {

                        let lsz = mdata.unwrap().len() as usize;
                        if &lsz == rsz { Some(true) } else { Some(false) }
                    } else {
                        None
                    }
                }
                if dr == false {
                    let rsize = fstream.size(&to_download[fz].as_str()).unwrap();

                    let mut fnm = "".to_string();
                    fnm.push_str(dir);
                    fnm.push_str("/");
                    fnm.push_str(&cl_names[fz].as_str());

                    println!("Downloading \"{}\"", to_download[fz]);

                    if !is_downloaded(&fnm, &rsize.unwrap()).unwrap_or(false) {

                        // Make local file
                        let _tt = fstream.transfer_type(::ftp::types::FileType::Binary);

                        let _result = fstream.retr(&to_download[fz].as_str(), |stream| {
                            let mut buf = [0; 4096];
                            let mut flnm = "".to_string();
                            flnm.push_str(dir);
                            flnm.push_str("/");
                            flnm.push_str(&cl_names[fz].as_str());

                            let mut file = File::create(flnm).unwrap();
                            let mut lsize: usize = 0;

                            loop {
                                match stream.read(&mut buf) {
                                    Ok(0) => break,
                                    Ok(n) => {
                                        file.write_all(&buf[0..n]).unwrap();
                                        lsize += n;
                                    }
                                    Err(_err) => break,
                                };
                                print_progress(((lsize.checked_mul(100).unwrap())
                                    .checked_div(rsize.unwrap())
                                    .unwrap()));
                            }
                            print!("\r\n");
                            Ok(())
                        });
                    } else {
                        println!("File is already downloaded");
                    }
                } else {
                    println!("Would be downloading {} and saving as {}",
                             &to_download[fz],
                             &cl_names[fz]);
                }
            }
        }

        let _f = fstream.quit();

        Ok(())
    } else {
        let _f = fstream.quit();
        println!("User login credentals were incorrect");
        Ok(())
    }
}

// Adds a file to downlaod to a config
fn add_to_file(mut file: Vec<String>) -> std::io::Result<()> {
    // TODO Catch errors and check inputs
    if file.len() == 4 {
        let g = file[3].clone();
        file.push(g);
    }
    let mut name = env::current_exe().unwrap();
    let t = file[0].clone();
    name.push(t.as_str());
    let mut b = "".to_string();
    let mut _a = File::open(&name).unwrap();
    let __ = _a.read_to_string(&mut b);
    let mut _f = OpenOptions::new().write(true).truncate(true).open(name).unwrap();
    let mut _w = BufWriter::new(_f);
    let _s1: &str;
    let _s2: &str;

    let (_s1, _s2) = b.split_at(b.rfind("}").unwrap().checked_add("}".len()).unwrap());

    println!("Writing to config");

    let _ =
        _w.write_all(format!("{},\r\n{{\r\n  remoteDir: {};\r\n  localDir: {};\r\n  name: \
                              {};\r\n  nameToSaveAs: {};\r\n}}{}",
                             _s1,
                             file[1],
                             file[2],
                             file[3],
                             file[4],
                             _s2)
            .as_bytes());
    Ok(())
}

fn main() {

    let mut _should_ls = false;
    let mut dry_run = false;
    let mut should_add = false;
    let mut add_to_config: Vec<String> = Vec::new();
    let ag: Vec<String> = env::args().collect();

    for x in 0..ag.len() {
        if ag[x].contains("-dr") {
            dry_run = true;
        }
        if ag[x].contains("-af") {
            if (x + 4) < ag.len() {
                should_add = true;
                add_to_config.push(ag[x + 1].clone());
                add_to_config.push(ag[x + 2].clone());
                add_to_config.push(ag[x + 3].clone());
                add_to_config.push(ag[x + 4].clone());
                if x + 5 < ag.len() {
                    add_to_config.push(ag[x + 5].clone());
                }
            } else {
                println!("Config name, client and server directory, and server filename must be \
                          supplied");
            }
        }
    }

    if should_add {
        let _ = add_to_file(add_to_config);
    }

    // TODO ls command
    println!("\rLoading configs...");

    let configs = load_configs().unwrap();

    for con in configs {
        let _d = download_from_site(con, dry_run);
    }
    std::process::exit(0);

}
<|endoftext|>


<commit_before>#![recursion_limit = "1024"] // error_chain

extern crate websocket;
extern crate env_logger;
#[macro_use]
extern crate log;
#[macro_use]
extern crate error_chain;
extern crate url;
extern crate clap;

use std::thread;
use std::io::{stdin,stdout};

use websocket::{Message, Sender, Receiver, DataFrame};
use websocket::message::Type;
use websocket::client::request::Url;
use websocket::Client;

use std::borrow::Borrow;
use std::io::{Error as IoError, ErrorKind as IoErrorKind, Write, Read};

error_chain! {
    foreign_links {
        Io(::std::io::Error);
        Log(log::SetLoggerError);
        Url(::url::ParseError);
        Ws(::websocket::result::WebSocketError);
        VarError(::std::env::VarError);
        RE(std::sync::mpsc::RecvError);
    }
    errors {
        InvalidSpecifier(t : String) {
            description("invalid specifier")
            display("Invalid client or server specifier `{}`", t)
        }
    }
}

// Initialize logger with default "info" log level:
fn init_logger() -> Result<()> {
    let mut builder = env_logger::LogBuilder::new();
    builder.filter(None, log::LogLevelFilter::Info);
    if ::std::env::var("RUST_LOG").is_ok() {
       builder.parse(&::std::env::var("RUST_LOG")?);
    }
    builder.init()?;
    Ok(())
}

struct SenderWrapper<T: Sender> (T);

impl<T: Sender> ::std::io::Write for SenderWrapper<T> {
    fn write(&mut self, buf: &[u8]) -> ::std::io::Result<usize> {
        let message = Message::binary(buf);
        let ret;
        let len = buf.len();
        if len > 0 {
            debug!("Sending message of {} bytes", len);
            ret = self.0.send_message(&message);
        } else {
            // Interpret zero length buffer is request
            // to close communication

            debug!("Sending the closing message");
            ret = self.0.send_message(&Message::close());
        }
        ret.map_err(|e|IoError::new(IoErrorKind::BrokenPipe, e))?;
        Ok(len)
    }
    fn flush(&mut self) -> ::std::io::Result<()> {
        Ok(())
    }
}

struct ReceiverWrapper<T: Receiver<DataFrame>> (T);

impl<T:Receiver<DataFrame>> ::std::io::Read for ReceiverWrapper<T> {
    fn read(&mut self, buf: &mut [u8]) -> ::std::io::Result<usize> {
        let ret = self.0.recv_message();
        let msg : Message = ret.map_err(|e|IoError::new(IoErrorKind::BrokenPipe, e))?;

        match msg.opcode {
            Type::Close => {
                Ok(0)
            }
            Type::Ping => {
                // Sender used to be in a separate thread with a channel
                // now there's no channel, so trickier to combine ping replies
                // and usual data exchange
                unimplemented!();
            }
            _ => {
                let msgpayload : &[u8] = msg.payload.borrow();
                let len = msgpayload.len();
                debug!("Received message of {} bytes", len);

                assert!(buf.len() >= len);

                buf[0..len].copy_from_slice(msgpayload);

                Ok(len)
            }
        }
    }
}

struct Peer<R, W>
    where R : Read + Send + 'static, W: Write + Send + 'static
{
    reader: R,
    writer: W,
}

type IPeer = Peer<Box<Read+Send>, Box<Write+Send>>;

struct DataExchangeSession<R1, R2, W1, W2>
    where R1 : Read  + Send + 'static,
          R2 : Read  + Send + 'static,
          W1 : Write + Send + 'static,
          W2 : Write + Send + 'static,
{
    peer1: Peer<R1, W1>,
    peer2: Peer<R2, W2>,
}

impl<R1,R2,W1,W2> DataExchangeSession<R1,R2,W1,W2>
    where R1 : Read  + Send + 'static,
          R2 : Read  + Send + 'static,
          W1 : Write + Send + 'static,
          W2 : Write + Send + 'static,
{
    fn data_exchange(self) -> Result<()> {

        let mut reader1 = self.peer1.reader;
        let mut writer1 = self.peer1.writer;
        let mut reader2 = self.peer2.reader;
        let mut writer2 = self.peer2.writer;

        let receive_loop = thread::spawn(move || -> Result<()> {
            // Actual data transfer happens here
            ::std::io::copy(&mut reader1, &mut writer2)?;
            writer2.write(b"")?; // signal close
            Ok(())
        });

        // Actual data transfer happens here
        ::std::io::copy(&mut reader2, &mut writer1)?;
        writer1.write(b"")?; // Signal close

        debug!("Waiting for receiver side to exit");

        receive_loop.join().map_err(|x|format!("{:?}",x))?
    }
}

fn get_websocket_peer(urlstr: &str) -> Result<
        Peer<
            ReceiverWrapper<websocket::client::Receiver<websocket::WebSocketStream>>,
            SenderWrapper<websocket::client::Sender<websocket::WebSocketStream>>>
        > {
    let url = Url::parse(urlstr)?;

    info!("Connecting to {}", url);

    let request = Client::connect(url)?;

    let response = request.send()?; // Send the request and retrieve a response

    info!("Validating response...");

    response.validate()?; // Validate the response

    info!("Successfully connected");

    let (sender, receiver) = response.begin().split();

    let peer = Peer {
        reader : ReceiverWrapper(receiver),
        writer : SenderWrapper(sender),
    };
    Ok(peer)
}

fn get_stdio_peer() -> Result<Peer<std::io::Stdin, std::io::Stdout>> {
    Ok(
        Peer {
            reader : stdin(),
            writer : stdout(),
        }
    )
}

impl<R,W> Peer<R,W>
    where R : Read + Send + 'static, W: Write + Send + 'static
{
    fn upcast(self) -> IPeer  {
        Peer {
            reader: Box::new(self.reader) as Box<Read +Send>,
            writer: Box::new(self.writer) as Box<Write+Send>,
        }
    }
}

fn get_peer_by_spec(specifier: &str, server: bool) -> Result<IPeer> {
    let _ = server;
    match specifier {
        x if x == "-"             => Ok(get_stdio_peer()?.upcast()),
        x if x.starts_with("ws:") => Ok(get_websocket_peer(x)?.upcast()),
        x => Err(ErrorKind::InvalidSpecifier(x.to_string()).into()),
    }
}

fn try_main() -> Result<()> {
    //env_logger::init()?;
    init_logger()?;

    // setup command line arguments
    let matches = ::clap::App::new("WS Command Line Client")
        .version("0.1")
        .author("Vitaly \"_Vi\" Shukela <vi0oss@gmail.com>")
        .about("Send binary data from stdin to a WebSocket and back to stdout.")
        .arg(::clap::Arg::with_name("listener_spec")
             .help("Listener specifier.")
             .required(true)
             .index(1))
        .arg(::clap::Arg::with_name("connector_spec")
             .help("Connector specifier.")
             .required(true)
             .index(2))
        .get_matches();

    let listener_spec  = matches.value_of("listener_spec") .ok_or("no listener_spec" )?;
    let connector_spec = matches.value_of("connector_spec").ok_or("no connector_spec")?;

    let des = DataExchangeSession {
        peer1 : get_peer_by_spec(listener_spec,  true )?,
        peer2 : get_peer_by_spec(connector_spec, false)?,
    };

    des.data_exchange()?;

    debug!("Exited");
    Ok(())
}

fn main() {
    if let Err(x) = try_main() {
        let _ = writeln!(::std::io::stderr(), "{:?}", x);
    }
}

<commit_msg>Help message, fix for wss://<commit_after>#![recursion_limit = "1024"] // error_chain

extern crate websocket;
extern crate env_logger;
#[macro_use]
extern crate log;
#[macro_use]
extern crate error_chain;
extern crate url;
extern crate clap;

use std::thread;
use std::io::{stdin,stdout};

use websocket::{Message, Sender, Receiver, DataFrame};
use websocket::message::Type;
use websocket::client::request::Url;
use websocket::Client;

use std::borrow::Borrow;
use std::io::{Error as IoError, ErrorKind as IoErrorKind, Write, Read};

error_chain! {
    foreign_links {
        Io(::std::io::Error);
        Log(log::SetLoggerError);
        Url(::url::ParseError);
        Ws(::websocket::result::WebSocketError);
        VarError(::std::env::VarError);
        RE(std::sync::mpsc::RecvError);
    }
    errors {
        InvalidSpecifier(t : String) {
            description("invalid specifier")
            display("Invalid client or server specifier `{}`", t)
        }
    }
}

// Initialize logger with default "info" log level:
fn init_logger() -> Result<()> {
    let mut builder = env_logger::LogBuilder::new();
    builder.filter(None, log::LogLevelFilter::Info);
    if ::std::env::var("RUST_LOG").is_ok() {
       builder.parse(&::std::env::var("RUST_LOG")?);
    }
    builder.init()?;
    Ok(())
}

struct SenderWrapper<T: Sender> (T);

impl<T: Sender> ::std::io::Write for SenderWrapper<T> {
    fn write(&mut self, buf: &[u8]) -> ::std::io::Result<usize> {
        let message = Message::binary(buf);
        let ret;
        let len = buf.len();
        if len > 0 {
            debug!("Sending message of {} bytes", len);
            ret = self.0.send_message(&message);
        } else {
            // Interpret zero length buffer is request
            // to close communication

            debug!("Sending the closing message");
            ret = self.0.send_message(&Message::close());
        }
        ret.map_err(|e|IoError::new(IoErrorKind::BrokenPipe, e))?;
        Ok(len)
    }
    fn flush(&mut self) -> ::std::io::Result<()> {
        Ok(())
    }
}

struct ReceiverWrapper<T: Receiver<DataFrame>> (T);

impl<T:Receiver<DataFrame>> ::std::io::Read for ReceiverWrapper<T> {
    fn read(&mut self, buf: &mut [u8]) -> ::std::io::Result<usize> {
        let ret = self.0.recv_message();
        let msg : Message = ret.map_err(|e|IoError::new(IoErrorKind::BrokenPipe, e))?;

        match msg.opcode {
            Type::Close => {
                Ok(0)
            }
            Type::Ping => {
                // Sender used to be in a separate thread with a channel
                // now there's no channel, so trickier to combine ping replies
                // and usual data exchange
                unimplemented!();
            }
            _ => {
                let msgpayload : &[u8] = msg.payload.borrow();
                let len = msgpayload.len();
                debug!("Received message of {} bytes", len);

                assert!(buf.len() >= len);

                buf[0..len].copy_from_slice(msgpayload);

                Ok(len)
            }
        }
    }
}

struct Peer<R, W>
    where R : Read + Send + 'static, W: Write + Send + 'static
{
    reader: R,
    writer: W,
}

type IPeer = Peer<Box<Read+Send>, Box<Write+Send>>;

struct DataExchangeSession<R1, R2, W1, W2>
    where R1 : Read  + Send + 'static,
          R2 : Read  + Send + 'static,
          W1 : Write + Send + 'static,
          W2 : Write + Send + 'static,
{
    peer1: Peer<R1, W1>,
    peer2: Peer<R2, W2>,
}

impl<R1,R2,W1,W2> DataExchangeSession<R1,R2,W1,W2>
    where R1 : Read  + Send + 'static,
          R2 : Read  + Send + 'static,
          W1 : Write + Send + 'static,
          W2 : Write + Send + 'static,
{
    fn data_exchange(self) -> Result<()> {

        let mut reader1 = self.peer1.reader;
        let mut writer1 = self.peer1.writer;
        let mut reader2 = self.peer2.reader;
        let mut writer2 = self.peer2.writer;

        let receive_loop = thread::spawn(move || -> Result<()> {
            // Actual data transfer happens here
            ::std::io::copy(&mut reader1, &mut writer2)?;
            writer2.write(b"")?; // signal close
            Ok(())
        });

        // Actual data transfer happens here
        ::std::io::copy(&mut reader2, &mut writer1)?;
        writer1.write(b"")?; // Signal close

        debug!("Waiting for receiver side to exit");

        receive_loop.join().map_err(|x|format!("{:?}",x))?
    }
}

fn get_websocket_peer(urlstr: &str) -> Result<
        Peer<
            ReceiverWrapper<websocket::client::Receiver<websocket::WebSocketStream>>,
            SenderWrapper<websocket::client::Sender<websocket::WebSocketStream>>>
        > {
    let url = Url::parse(urlstr)?;

    info!("Connecting to {}", url);

    let request = Client::connect(url)?;

    let response = request.send()?; // Send the request and retrieve a response

    info!("Validating response...");

    response.validate()?; // Validate the response

    info!("Successfully connected");

    let (sender, receiver) = response.begin().split();

    let peer = Peer {
        reader : ReceiverWrapper(receiver),
        writer : SenderWrapper(sender),
    };
    Ok(peer)
}

fn get_stdio_peer() -> Result<Peer<std::io::Stdin, std::io::Stdout>> {
    Ok(
        Peer {
            reader : stdin(),
            writer : stdout(),
        }
    )
}

impl<R,W> Peer<R,W>
    where R : Read + Send + 'static, W: Write + Send + 'static
{
    fn upcast(self) -> IPeer  {
        Peer {
            reader: Box::new(self.reader) as Box<Read +Send>,
            writer: Box::new(self.writer) as Box<Write+Send>,
        }
    }
}

fn get_peer_by_spec(specifier: &str, server: bool) -> Result<IPeer> {
    let _ = server;
    match specifier {
        x if x == "-"               => Ok(get_stdio_peer()?.upcast()),
        x if x.starts_with("ws:")   => Ok(get_websocket_peer(x)?.upcast()),
        x if x.starts_with("wss:")  => Ok(get_websocket_peer(x)?.upcast()),
        x => Err(ErrorKind::InvalidSpecifier(x.to_string()).into()),
    }
}

fn try_main() -> Result<()> {
    //env_logger::init()?;
    init_logger()?;

    // setup command line arguments
    let matches = ::clap::App::new("websocat")
        .version("0.1")
        .author("Vitaly \"_Vi\" Shukela <vi0oss@gmail.com>")
        .about("Exchange binary data between websocket and something.\nSocat analogue with websockets.")
        .arg(::clap::Arg::with_name("listener_spec")
             .help("Listener specifier.")
             .required(true)
             .index(1))
        .arg(::clap::Arg::with_name("connector_spec")
             .help("Connector specifier.")
             .required(true)
             .index(2))
        .after_help(r#"
Specifiers are:
  ws[s]://<rest of websocket URL>    websockets
  -                                  stdin/stdout
  (more to be implemented)

Examples:
  websocat - wss://myserver/mysocket
    Connect stdin/stdout to secure web socket once.
    Currently it is the only working example.
"#)
        .get_matches();

    let listener_spec  = matches.value_of("listener_spec") .ok_or("no listener_spec" )?;
    let connector_spec = matches.value_of("connector_spec").ok_or("no connector_spec")?;

    let des = DataExchangeSession {
        peer1 : get_peer_by_spec(listener_spec,  true )?,
        peer2 : get_peer_by_spec(connector_spec, false)?,
    };

    des.data_exchange()?;

    debug!("Exited");
    Ok(())
}

fn main() {
    if let Err(x) = try_main() {
        let _ = writeln!(::std::io::stderr(), "{:?}", x);
    }
}

<|endoftext|>