<commit_before>extern crate nom;
extern crate device_tree_source;
#[macro_use]
extern crate clap;

mod change_tracker;

use std::process::Command;
use std::path::{Path, PathBuf};
use std::fs::remove_file;
use std::io::{self, BufRead, Write};
use std::iter::Iterator;
use std::fmt::{self, Display, Formatter};
// use std::fs::File;

use device_tree_source::parser::parse_dt;
use device_tree_source::tree::Offset;

use device_tree_source::include::{IncludeBounds, IncludeMethod, include_files};

use change_tracker::LabelStore;

const CPP_OUTPUT_NAME: &'static str = "dts_viewer_tmp.dts";

// General idea:
//  Run CPP
//  Parse file for DTS includes and replace with include contents
//  Find byte starts/ends for each file
//  Parse file to create device tree
//  create mapping of objects to byte offset starting points
//  - not perfect with offsets stored in object, but it will do
//  create mapping of full paths to objects
//  create mapping of labels to full paths
//  Parse device tree to create changes
//  Use mapping and file starts to pair changes to file byte offsets
//  Turn file byte offsets to file lines/cols
//  TODO: ???
//  TODO: Profit!
//  Oh, and somehow display the damn information
fn main() {
    let matches = clap_app!(dts_viewer =>
            (version: crate_version!())
            (author: "Gabriel s. <ga29smith@gmail.com>")
            (@arg file: +required "DTS file to parse")
            (@arg no_defaults: -n --no_defaults "Disable default includes. \
                An 'include' directory, if it exists, is automatically included")
            (@arg include: -I ... +takes_value "Additional files to pass to CPP as an include")
        )
        .get_matches();

    let file_name = matches.value_of("file").unwrap();

    let file = Path::new(file_name);
    let parent = file.parent().expect("Could not get parent directory of file");

    let mut cpp_command = Command::new("gcc");
    cpp_command.args(&["-H", "-E", "-nostdinc"])
        .args(&["-undef", "-D__DTS__", "-x", "assembler-with-cpp"])
        .args(&["-o", CPP_OUTPUT_NAME])
        .arg(&file_name);

    if !matches.is_present("no_defaults") {
        cpp_command.args(&["-I", parent.to_str().unwrap()]);
        let include_dir = parent.join("include/");
        if include_dir.is_dir() {
            cpp_command.args(&["-I", include_dir.to_str().unwrap()]);
        }
    }

    if let Some(includes) = matches.values_of("include") {
        for include in includes {
            let include_dir = parent.join(include);
            if include_dir.is_dir() {
                cpp_command.args(&["-I", include_dir.to_str().unwrap()]);
            }
        }
    }

    let include_output = cpp_command.output().expect("Failed to execute CPP");
    let cpp_stderr = String::from_utf8_lossy(&include_output.stderr);
    if !include_output.status.success() {
        println!("Failed to execute CPP. Error message is below.");
        print!("{}", cpp_stderr);
        return;
    }

    let (buffer, bounds) = match include_files(Path::new(CPP_OUTPUT_NAME), 0) {
        Ok(x) => x,
        Err(e) => {
            println!("{}", e);
            return;
        }
    };

    if let Err(e) = remove_file(Path::new(CPP_OUTPUT_NAME)) {
        println!("Failed to delete temp file: {:?}", e);
    }

    // println!("{:#?}", bounds);
    // {
    //     let mut total_dts_dump = File::create("total_dts_dump.dts").unwrap();
    //     total_dts_dump.write_all(&buffer).unwrap()
    // }

    let include_tree = IncludeTree::bounds_to_tree(&bounds);
    if let Some(tree) = include_tree {
        println!("{}", tree);
    } else {
        println!("-- Could not constuct include tree from bounds!");
    }

    let (boot_info, ammends) = match parse_dt(&buffer) {
        Ok(dt) => dt,
        Err(err) => {
            println!("{:?}", err);
            return;
        }
    };

    // TODO: perform secondary checks and jazz (only smooth)
    //
    // checkes to do:
    // duplicate node names
    // duplicate property names
    // node name format (only one @)
    // unit addr vs reg/ranges property
    // duplicate labels -- handled by creation of alias store
    // duplicate explict phandles
    // name property does not match name
    // check if X is cell only
    // #address-cells
    // #size-cells
    // #interrupt-cells
    // check if X is string only
    // device_type
    // model
    // status
    // fixup addr size cells
    // check reg property format
    // check ranges property format
    // avoid default addr size?
    // obsolete chosen iterrupt controller
    //
    // maybe:
    // fixup phandle refs
    // fixup path refs

    let mut store = LabelStore::new();
    store.fill(&boot_info, &ammends);

    loop {
        print!("Enter alias or path: ");
        io::stdout().flush().expect("Error flushing stdout");

        let mut line = String::new();
        let stdin = io::stdin();
        stdin.lock().read_line(&mut line).expect("Error reading from stdin");
        let line = line.trim();

        if line.is_empty() {
            break;
        }

        let path = if line.starts_with('/') {
            Some(PathBuf::from(line))
        } else {
            let line = if line.starts_with('&') {
                &line[1..]
            } else {
                line
            };
            // TODO:
            // if label covers multiple paths (after saving old aliases is implemented)
            //     ask user which one they want as well as why this is being asked
            // show all changes along the way
            match store.path_from_label(line) {
                Some(path) => {
                    println!("Path: {}", path.to_string_lossy());
                    Some(path.to_path_buf())
                }
                None => {
                    println!("Label points to no path");
                    None
                }
            }
        };

        if let Some(path) = path {
            match store.changes_from_path(&path) {
                Some(changes) => {
                    for change in changes {
                        let offset = change.get_offset();
                        // println!("Start offset: {}", offset);
                        match bounds.binary_search_by(|b| {
                            use std::cmp::Ordering::*;
                            match (b.global_start.cmp(&offset), b.global_end().cmp(&offset)) {
                                (Less, Greater) | (Equal, Greater) => Equal,
                                (Greater, Greater) => Greater,
                                (Equal, Less) | (Less, Less) | (Less, Equal) | (Equal, Equal)
                                    => Less,
                                _ => unreachable!(),
                            }
                        }) {
                            Ok(off) => {
                                let bound = &bounds[off];
                                print!("File: {}", bound.path.to_string_lossy());

                                match bound.file_line_from_global(&buffer, offset) {
                                    Ok((line, col)) => {
                                        println!(", Line: {}, Column: {}", line, col)
                                    }
                                    Err(err) => println!(" {}", err),
                                }
                            }
                            Err(_) => println!("Could not find file for offset {}", offset),
                        }
                        println!("{}\n", change);
                    }
                }
                None => println!("Nothing at path"),
            }
        }
    }
}

#[derive(Debug)]
struct IncludeTree {
    path: PathBuf,
    method: IncludeMethod,
    includes: Vec<IncludeTree>,
}

impl IncludeTree {
    fn bounds_to_tree(bounds: &[IncludeBounds]) -> Option<IncludeTree> {
        if let Some(first) = bounds.first() {
            let mut tree = IncludeTree {
                path: first.path.clone(),
                includes: Vec::new(),
                method: first.method.clone(),
            };

            //TODO: we don't really need the filter, benchmark speed w/wo
            for sub_bounds in bounds.split(|b| b.path == first.path)
                                    .filter(|s| !s.is_empty()) {
                if let Some(sub_tree) = Self::bounds_to_tree(sub_bounds) {
                    tree.includes.push(sub_tree);
                }
            }

            Some(tree)
        } else {
            None
        }
    }

    fn write(&self, f: &mut Formatter, prefix: &str) -> fmt::Result {
        let mut next_prefix = prefix.to_string();
        next_prefix.push_str(" |-");

        let method = match self.method {
            IncludeMethod::CPP => "CPP",
            IncludeMethod::DTS => "DTS",
        };
        writeln!(f, "{} {}: {}", prefix, self.path.to_string_lossy(), method)?;
        for t in &self.includes {
            t.write(f, &next_prefix)?;
        }

        Ok(())
    }
}

impl Display for IncludeTree {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        self.write(f, "")
    }
}
<commit_msg>Differentiate non-fatal error messages from normal output<commit_after>extern crate nom;
extern crate device_tree_source;
#[macro_use]
extern crate clap;

mod change_tracker;

use std::process::Command;
use std::path::{Path, PathBuf};
use std::fs::remove_file;
use std::io::{self, BufRead, Write};
use std::iter::Iterator;
use std::fmt::{self, Display, Formatter};
// use std::fs::File;

use device_tree_source::parser::parse_dt;
use device_tree_source::tree::Offset;

use device_tree_source::include::{IncludeBounds, IncludeMethod, include_files};

use change_tracker::LabelStore;

const CPP_OUTPUT_NAME: &'static str = "dts_viewer_tmp.dts";

// General idea:
//  Run CPP
//  Parse file for DTS includes and replace with include contents
//  Find byte starts/ends for each file
//  Parse file to create device tree
//  create mapping of objects to byte offset starting points
//  - not perfect with offsets stored in object, but it will do
//  create mapping of full paths to objects
//  create mapping of labels to full paths
//  Parse device tree to create changes
//  Use mapping and file starts to pair changes to file byte offsets
//  Turn file byte offsets to file lines/cols
//  TODO: ???
//  TODO: Profit!
//  Oh, and somehow display the damn information
fn main() {
    let matches = clap_app!(dts_viewer =>
            (version: crate_version!())
            (author: "Gabriel s. <ga29smith@gmail.com>")
            (@arg file: +required "DTS file to parse")
            (@arg no_defaults: -n --no_defaults "Disable default includes. \
                An 'include' directory, if it exists, is automatically included")
            (@arg include: -I ... +takes_value "Additional files to pass to CPP as an include")
        )
        .get_matches();

    let file_name = matches.value_of("file").unwrap();

    let file = Path::new(file_name);
    let parent = file.parent().expect("Could not get parent directory of file");

    let mut cpp_command = Command::new("gcc");
    cpp_command.args(&["-H", "-E", "-nostdinc"])
        .args(&["-undef", "-D__DTS__", "-x", "assembler-with-cpp"])
        .args(&["-o", CPP_OUTPUT_NAME])
        .arg(&file_name);

    if !matches.is_present("no_defaults") {
        cpp_command.args(&["-I", parent.to_str().unwrap()]);
        let include_dir = parent.join("include/");
        if include_dir.is_dir() {
            cpp_command.args(&["-I", include_dir.to_str().unwrap()]);
        }
    }

    if let Some(includes) = matches.values_of("include") {
        for include in includes {
            let include_dir = parent.join(include);
            if include_dir.is_dir() {
                cpp_command.args(&["-I", include_dir.to_str().unwrap()]);
            }
        }
    }

    let include_output = cpp_command.output().expect("Failed to execute CPP");
    let cpp_stderr = String::from_utf8_lossy(&include_output.stderr);
    if !include_output.status.success() {
        println!("Failed to execute CPP. Error message is below.");
        print!("{}", cpp_stderr);
        return;
    }

    let (buffer, bounds) = match include_files(Path::new(CPP_OUTPUT_NAME), 0) {
        Ok(x) => x,
        Err(e) => {
            println!("{}", e);
            return;
        }
    };

    if let Err(e) = remove_file(Path::new(CPP_OUTPUT_NAME)) {
        println!("-- Failed to delete temp file: {:?}", e);
    }

    // println!("{:#?}", bounds);
    // {
    //     let mut total_dts_dump = File::create("total_dts_dump.dts").unwrap();
    //     total_dts_dump.write_all(&buffer).unwrap()
    // }

    let include_tree = IncludeTree::bounds_to_tree(&bounds);
    if let Some(tree) = include_tree {
        println!("{}", tree);
    } else {
        println!("-- Could not constuct include tree from bounds!");
    }

    let (boot_info, ammends) = match parse_dt(&buffer) {
        Ok(dt) => dt,
        Err(err) => {
            println!("{:?}", err);
            return;
        }
    };

    // TODO: perform secondary checks and jazz (only smooth)
    //
    // checkes to do:
    // duplicate node names
    // duplicate property names
    // node name format (only one @)
    // unit addr vs reg/ranges property
    // duplicate labels -- handled by creation of alias store
    // duplicate explict phandles
    // name property does not match name
    // check if X is cell only
    // #address-cells
    // #size-cells
    // #interrupt-cells
    // check if X is string only
    // device_type
    // model
    // status
    // fixup addr size cells
    // check reg property format
    // check ranges property format
    // avoid default addr size?
    // obsolete chosen iterrupt controller
    //
    // maybe:
    // fixup phandle refs
    // fixup path refs

    let mut store = LabelStore::new();
    store.fill(&boot_info, &ammends);

    loop {
        print!("Enter alias or path: ");
        io::stdout().flush().expect("Error flushing stdout");

        let mut line = String::new();
        let stdin = io::stdin();
        stdin.lock().read_line(&mut line).expect("Error reading from stdin");
        let line = line.trim();

        if line.is_empty() {
            break;
        }

        let path = if line.starts_with('/') {
            Some(PathBuf::from(line))
        } else {
            let line = if line.starts_with('&') {
                &line[1..]
            } else {
                line
            };
            // TODO:
            // if label covers multiple paths (after saving old aliases is implemented)
            //     ask user which one they want as well as why this is being asked
            // show all changes along the way
            match store.path_from_label(line) {
                Some(path) => {
                    println!("Path: {}", path.to_string_lossy());
                    Some(path.to_path_buf())
                }
                None => {
                    println!("Label points to no path");
                    None
                }
            }
        };

        if let Some(path) = path {
            match store.changes_from_path(&path) {
                Some(changes) => {
                    for change in changes {
                        let offset = change.get_offset();
                        // println!("Start offset: {}", offset);
                        match bounds.binary_search_by(|b| {
                            use std::cmp::Ordering::*;
                            match (b.global_start.cmp(&offset), b.global_end().cmp(&offset)) {
                                (Less, Greater) | (Equal, Greater) => Equal,
                                (Greater, Greater) => Greater,
                                (Equal, Less) | (Less, Less) | (Less, Equal) | (Equal, Equal)
                                    => Less,
                                _ => unreachable!(),
                            }
                        }) {
                            Ok(off) => {
                                let bound = &bounds[off];
                                print!("File: {}", bound.path.to_string_lossy());

                                match bound.file_line_from_global(&buffer, offset) {
                                    Ok((line, col)) => {
                                        println!(", Line: {}, Column: {}", line, col)
                                    }
                                    Err(err) => println!(" {}", err),
                                }
                            }
                            Err(_) => println!("-- Could not find file for offset {}", offset),
                        }
                        println!("{}\n", change);
                    }
                }
                None => println!("Nothing at path"),
            }
        }
    }
}

#[derive(Debug)]
struct IncludeTree {
    path: PathBuf,
    method: IncludeMethod,
    includes: Vec<IncludeTree>,
}

impl IncludeTree {
    fn bounds_to_tree(bounds: &[IncludeBounds]) -> Option<IncludeTree> {
        if let Some(first) = bounds.first() {
            let mut tree = IncludeTree {
                path: first.path.clone(),
                includes: Vec::new(),
                method: first.method.clone(),
            };

            //TODO: we don't really need the filter, benchmark speed w/wo
            for sub_bounds in bounds.split(|b| b.path == first.path)
                                    .filter(|s| !s.is_empty()) {
                if let Some(sub_tree) = Self::bounds_to_tree(sub_bounds) {
                    tree.includes.push(sub_tree);
                }
            }

            Some(tree)
        } else {
            None
        }
    }

    fn write(&self, f: &mut Formatter, prefix: &str) -> fmt::Result {
        let mut next_prefix = prefix.to_string();
        next_prefix.push_str(" |-");

        let method = match self.method {
            IncludeMethod::CPP => "CPP",
            IncludeMethod::DTS => "DTS",
        };
        writeln!(f, "{} {}: {}", prefix, self.path.to_string_lossy(), method)?;
        for t in &self.includes {
            t.write(f, &next_prefix)?;
        }

        Ok(())
    }
}

impl Display for IncludeTree {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        self.write(f, "")
    }
}
<|endoftext|>


<commit_before>//! **Wa** ve **f** unction Solv **er**: a parallelized, 3D, Schrödinger equation solver.
//!
//! Wafer exploits a Wick-rotated time-dependent Schrödinger equation to solve for
//! time-independent solutions in three dimensions.
//!
//! Inspired by [quantumfdtd](http://sourceforge.net/projects/quantumfdtd/),
//! which is a proof of concept tool and falls short of a general purpose utility.
//! Wafer attempts to remedy this issue.
//!
//! If you use Wafer in your research, please reference the following articles:
//!
//! M. Strickland and D. Yager-Elorriaga, “A parallel algorithm for solving the 3d
//! Schrödinger equation”,
//! [Journal of Computational Physics __229__, 6015–6026 (2010)](http://dx.doi.org/10.1016/j.jcp.2010.04.032).

#![cfg_attr(feature="clippy", feature(plugin))]
#![cfg_attr(feature="clippy", plugin(clippy))]
#![cfg_attr(feature="clippy", warn(missing_docs_in_private_items))]
#![cfg_attr(feature="clippy", warn(single_match_else))]

extern crate ansi_term;
extern crate chrono;
#[macro_use]
extern crate clap;
extern crate csv;
extern crate indicatif;
#[macro_use]
extern crate lazy_static;
#[macro_use(s)]
extern crate ndarray;
extern crate ndarray_parallel;
extern crate num;
extern crate num_cpus;
extern crate ordinal;
extern crate rand;
extern crate rayon;
extern crate rmp_serde as rmps;
extern crate serde;
#[macro_use]
extern crate serde_derive;
extern crate serde_json;
#[macro_use]
extern crate slog;
extern crate slog_async;
extern crate slog_term;
extern crate term_size;

use clap::{App, Arg};
use slog::{Drain, Duplicate, Logger, Fuse, LevelFilter, Level};
use std::fs::OpenOptions;
use std::process;
use std::thread;
use std::time::{Duration, Instant};
use config::Config;

include!(concat!(env!("OUT_DIR"), "/version.rs"));

/// Config is a (mostly) public module which reads the configuration file `wafer.cfg`
/// and poplulates the `Config` struct with the information required to run the current
/// instance of the application.
pub mod config;
/// The meat of the caclulation is performed on a finite grid. Basically all of the computation
/// work is done within this module.
mod grid;
/// Any required file input (apart from configuration) is handled here. Plain text and binary formats.
mod input;
/// All file output is handled in this module. Plain text and binary options are both here.
mod output;
/// Handles the potential generation, binding energy offsets, callouts to files or scripts
/// if needed etc.
mod potential;

/// Exits (no error) after a short pause. Because we're using async logs, sometimes we dump before
/// the log system outputs information. We spool for a little first in these instances so we get the
/// logging info.
fn exit_with_pause() {
    thread::sleep(Duration::from_millis(10));
    process::exit(1);
}

/// System entry point
fn main() {
    //Simulation timer.
    let start_time = Instant::now();

    //CLI options. For the moment, a custom configuration file and custom debug levels.
    let matches = App::new("Wafer")
                        .version(crate_version!())
                        .author(crate_authors!())
                        .about("Exploits a Wick-rotated time-dependent Schrödinger equation to solve for time-independent solutions in three dimensions.")
                        .arg(Arg::with_name("config")
                                    .short("c")
                                    .long("config")
                                    .value_name("FILE")
                                    .help("The configuration file to use (default is \"wafer.cfg\")")
                                    .takes_value(true))
                        .arg(Arg::with_name("debug")
                                    .short("d")
                                    .multiple(true)
                                    .help("Raises screen debug level"))
                        .get_matches();

    //Load configuation parameters.
    let config_file = matches.value_of("config").unwrap_or("wafer.cfg");
    let config = match Config::load(config_file) {
        Ok(c) => c,
        Err(err) => {
            println!("Error processing configuration: {}", err);
            //This match isn't happy using the `exit_with_pause` function, but fine with it explicit like this...
            thread::sleep(Duration::from_millis(10));
            process::exit(1);
        }
    };
    if let Err(err) = output::check_output_dir(&config.project_name) {
        println!("Could not communicate with output directory: {}", err);
        exit_with_pause();
    };

    //Setup logging.
    let log_file = OpenOptions::new()
        .create(true)
        .write(true)
        .truncate(true)
        .open(output::get_project_dir(&config.project_name) + "/simulation.log")
        .expect("Cannot connect to log file");
    let syslog = slog_term::PlainDecorator::new(log_file);
    let sys_drain = slog_term::FullFormat::new(syslog).build().fuse();
    let sys_drain = slog_async::Async::new(sys_drain).build().fuse();
    let screen = slog_term::TermDecorator::new().build();
    let screen_drain = slog_term::FullFormat::new(screen).build().fuse();
    let screen_drain = slog_async::Async::new(screen_drain).build().fuse();

    let screen_level = match matches.occurrences_of("debug") {
        0 => Level::Warning,
        1 => Level::Info,
        2 | _ => Level::Debug,
    };
    let log = Logger::root(Fuse::new(Duplicate::new(LevelFilter::new(screen_drain, screen_level), sys_drain)),
                           o!());

    info!(log, "Starting Wafer solver"; "version" => crate_version!(), "build-id" => short_sha());

    info!(log, "Checking/creating directories");
    if let Err(err) = input::check_input_dir() {
        crit!(log, "Could not communicate with input directory: {}", err);
        exit_with_pause();
    };

    //Override rayon's defaults of threads (including HT cores) to physical cores
    if let Err(err) = rayon::initialize(rayon::Configuration::new()
                                            .num_threads(num_cpus::get_physical())) {
        crit!(log, "Failed to initialise thread pool: {}", err);
        exit_with_pause();
    };

    let term_width = *output::TERMWIDTH;

    let sha = if term_width <= 97 { short_sha() } else { sha() };

    output::print_banner(sha);

    info!(log, "Loading Configuation from disk");
    config.print(term_width);
    if let Err(err) = output::copy_config(&config.project_name) {
        //TODO: Input from CLI if non-default config file
        warn!(log,
              "Configuration file not copied to output directory: {}",
              err);
    };

    if let Err(err) = grid::run(&config, &log) {
        crit!(log, "{}", err);
        exit_with_pause();
    };

    let elapsed = start_time.elapsed();
    let time_taken = (elapsed.as_secs() as f64) + (elapsed.subsec_nanos() as f64 / 1000_000_000.0);
    match time_taken {
        0.0...60.0 => {
            println!("Simulation complete. Elapsed time: {:.3} seconds.",
                     time_taken)
        }
        60.0...3600.0 => {
            let minutes = (time_taken / 60.).floor();
            let seconds = time_taken - 60. * minutes;
            println!("Simulation complete. Elapsed time: {} minutes, {:.3} seconds.",
                     minutes,
                     seconds);
        }
        _ => {
            let hours = (time_taken / 3600.).floor();
            let minutes = ((time_taken - 3600. * hours) / 60.).floor();
            let seconds = time_taken - 3600. * hours - 60. * minutes;
            println!("Simulation complete. Elapsed time: {} hours, {} minutes, {:.3} seconds.",
                     hours,
                     minutes,
                     seconds);
        }
    }
    info!(log, "Simulation completed");
}

#[cfg(test)]
mod tests {
    #[test]
    fn placeholder() {
        let num = 5;
        assert_eq!(num, 5);
    }
}
<commit_msg>Handled overlooked panic point<commit_after>//! **Wa** ve **f** unction Solv **er**: a parallelized, 3D, Schrödinger equation solver.
//!
//! Wafer exploits a Wick-rotated time-dependent Schrödinger equation to solve for
//! time-independent solutions in three dimensions.
//!
//! Inspired by [quantumfdtd](http://sourceforge.net/projects/quantumfdtd/),
//! which is a proof of concept tool and falls short of a general purpose utility.
//! Wafer attempts to remedy this issue.
//!
//! If you use Wafer in your research, please reference the following articles:
//!
//! M. Strickland and D. Yager-Elorriaga, “A parallel algorithm for solving the 3d
//! Schrödinger equation”,
//! [Journal of Computational Physics __229__, 6015–6026 (2010)](http://dx.doi.org/10.1016/j.jcp.2010.04.032).

#![cfg_attr(feature="clippy", feature(plugin))]
#![cfg_attr(feature="clippy", plugin(clippy))]
#![cfg_attr(feature="clippy", warn(missing_docs_in_private_items))]
#![cfg_attr(feature="clippy", warn(single_match_else))]

extern crate ansi_term;
extern crate chrono;
#[macro_use]
extern crate clap;
extern crate csv;
extern crate indicatif;
#[macro_use]
extern crate lazy_static;
#[macro_use(s)]
extern crate ndarray;
extern crate ndarray_parallel;
extern crate num;
extern crate num_cpus;
extern crate ordinal;
extern crate rand;
extern crate rayon;
extern crate rmp_serde as rmps;
extern crate serde;
#[macro_use]
extern crate serde_derive;
extern crate serde_json;
#[macro_use]
extern crate slog;
extern crate slog_async;
extern crate slog_term;
extern crate term_size;

use clap::{App, Arg};
use slog::{Drain, Duplicate, Logger, Fuse, LevelFilter, Level};
use std::fs::OpenOptions;
use std::process;
use std::thread;
use std::time::{Duration, Instant};
use config::Config;

include!(concat!(env!("OUT_DIR"), "/version.rs"));

/// Config is a (mostly) public module which reads the configuration file `wafer.cfg`
/// and poplulates the `Config` struct with the information required to run the current
/// instance of the application.
pub mod config;
/// The meat of the caclulation is performed on a finite grid. Basically all of the computation
/// work is done within this module.
mod grid;
/// Any required file input (apart from configuration) is handled here. Plain text and binary formats.
mod input;
/// All file output is handled in this module. Plain text and binary options are both here.
mod output;
/// Handles the potential generation, binding energy offsets, callouts to files or scripts
/// if needed etc.
mod potential;

/// Exits (no error) after a short pause. Because we're using async logs, sometimes we dump before
/// the log system outputs information. We spool for a little first in these instances so we get the
/// logging info.
fn exit_with_pause() {
    thread::sleep(Duration::from_millis(10));
    process::exit(1);
}

/// System entry point
fn main() {
    //Simulation timer.
    let start_time = Instant::now();

    //CLI options. For the moment, a custom configuration file and custom debug levels.
    let matches = App::new("Wafer")
                        .version(crate_version!())
                        .author(crate_authors!())
                        .about("Exploits a Wick-rotated time-dependent Schrödinger equation to solve for time-independent solutions in three dimensions.")
                        .arg(Arg::with_name("config")
                                    .short("c")
                                    .long("config")
                                    .value_name("FILE")
                                    .help("The configuration file to use (default is \"wafer.cfg\")")
                                    .takes_value(true))
                        .arg(Arg::with_name("debug")
                                    .short("d")
                                    .multiple(true)
                                    .help("Raises screen debug level"))
                        .get_matches();

    //Load configuation parameters.
    let config_file = matches.value_of("config").unwrap_or("wafer.cfg");
    let config = match Config::load(config_file) {
        Ok(c) => c,
        Err(err) => {
            println!("Error processing configuration: {}", err);
            process::exit(1);
        }
    };
    if let Err(err) = output::check_output_dir(&config.project_name) {
        println!("Could not communicate with output directory: {}", err);
        process::exit(1);
    };

    //Setup logging.
    let log_file = match OpenOptions::new()
        .create(true)
        .write(true)
        .truncate(true)
        .open(output::get_project_dir(&config.project_name) + "/simulation.log") {
        Ok(f) => f,
        Err(err) => {
            println!("Could not connect to log file: {}", err);
            process::exit(1);
        }
    };
    let syslog = slog_term::PlainDecorator::new(log_file);
    let sys_drain = slog_term::FullFormat::new(syslog).build().fuse();
    let sys_drain = slog_async::Async::new(sys_drain).build().fuse();
    let screen = slog_term::TermDecorator::new().build();
    let screen_drain = slog_term::FullFormat::new(screen).build().fuse();
    let screen_drain = slog_async::Async::new(screen_drain).build().fuse();

    let screen_level = match matches.occurrences_of("debug") {
        0 => Level::Warning,
        1 => Level::Info,
        2 | _ => Level::Debug,
    };
    let log = Logger::root(Fuse::new(Duplicate::new(LevelFilter::new(screen_drain, screen_level), sys_drain)),
                           o!());

    info!(log, "Starting Wafer solver"; "version" => crate_version!(), "build-id" => short_sha());

    info!(log, "Checking/creating directories");
    if let Err(err) = input::check_input_dir() {
        crit!(log, "Could not communicate with input directory: {}", err);
        exit_with_pause();
    };

    //Override rayon's defaults of threads (including HT cores) to physical cores
    if let Err(err) = rayon::initialize(rayon::Configuration::new()
                                            .num_threads(num_cpus::get_physical())) {
        crit!(log, "Failed to initialise thread pool: {}", err);
        exit_with_pause();
    };

    let term_width = *output::TERMWIDTH;

    let sha = if term_width <= 97 { short_sha() } else { sha() };

    output::print_banner(sha);

    info!(log, "Loading Configuation from disk");
    config.print(term_width);
    if let Err(err) = output::copy_config(&config.project_name) {
        //TODO: Input from CLI if non-default config file
        warn!(log,
              "Configuration file not copied to output directory: {}",
              err);
    };

    if let Err(err) = grid::run(&config, &log) {
        crit!(log, "{}", err);
        exit_with_pause();
    };

    let elapsed = start_time.elapsed();
    let time_taken = (elapsed.as_secs() as f64) + (elapsed.subsec_nanos() as f64 / 1000_000_000.0);
    match time_taken {
        0.0...60.0 => {
            println!("Simulation complete. Elapsed time: {:.3} seconds.",
                     time_taken)
        }
        60.0...3600.0 => {
            let minutes = (time_taken / 60.).floor();
            let seconds = time_taken - 60. * minutes;
            println!("Simulation complete. Elapsed time: {} minutes, {:.3} seconds.",
                     minutes,
                     seconds);
        }
        _ => {
            let hours = (time_taken / 3600.).floor();
            let minutes = ((time_taken - 3600. * hours) / 60.).floor();
            let seconds = time_taken - 3600. * hours - 60. * minutes;
            println!("Simulation complete. Elapsed time: {} hours, {} minutes, {:.3} seconds.",
                     hours,
                     minutes,
                     seconds);
        }
    }
    info!(log, "Simulation completed");
}

#[cfg(test)]
mod tests {
    #[test]
    fn placeholder() {
        let num = 5;
        assert_eq!(num, 5);
    }
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate clap;
#[macro_use]
extern crate lazy_static;

extern crate libc;
extern crate scmp;
extern crate tempdir;

use std::fs;
use std::env;
use std::ffi::OsStr;
use std::ffi::CString;
use std::io::Error;
use std::os::raw::c_void;
use std::os::unix::process::CommandExt;
use std::os::unix::ffi::OsStrExt;
use std::path::Path;
use std::path::PathBuf;
use std::process::Command;
use scmp::*;

unsafe fn allow_syscall(ctx : *mut c_void, nr : u32) {
    assert!(seccomp_rule_add_array(ctx, SCMP_ACT_ALLOW, nr as i32, 0, std::ptr::null()) == 0);
}

fn setup_seccomp() -> Result<(), Error> {
    unsafe {
        let ctx0 = seccomp_init(SCMP_ACT_TRAP);
        assert!(!ctx0.is_null());
        allow_syscall(ctx0, __NR_execve);
        allow_syscall(ctx0, __NR_brk);
        allow_syscall(ctx0, __NR_access);
        allow_syscall(ctx0, __NR_mmap);
        allow_syscall(ctx0, __NR_open);
        allow_syscall(ctx0, __NR_read);
        allow_syscall(ctx0, __NR_write);
        allow_syscall(ctx0, __NR_stat);
        allow_syscall(ctx0, __NR_fstat);
        allow_syscall(ctx0, __NR_close);
        allow_syscall(ctx0, __NR_mprotect);
        allow_syscall(ctx0, __NR_arch_prctl);
        allow_syscall(ctx0, __NR_munmap);
        allow_syscall(ctx0, __NR_set_tid_address);
        allow_syscall(ctx0, __NR_set_robust_list);
        allow_syscall(ctx0, __NR_rt_sigaction);
        allow_syscall(ctx0, __NR_rt_sigprocmask);
        allow_syscall(ctx0, __NR_getrlimit);
        allow_syscall(ctx0, __NR_statfs);
        allow_syscall(ctx0, __NR_fcntl);
        allow_syscall(ctx0, __NR_geteuid);
        allow_syscall(ctx0, __NR_umask);
        allow_syscall(ctx0, __NR_pipe);
        allow_syscall(ctx0, __NR_clone);
        allow_syscall(ctx0, __NR_wait4);
        allow_syscall(ctx0, __NR_exit_group);
        allow_syscall(ctx0, __NR_dup);
        allow_syscall(ctx0, __NR_ioctl);
        allow_syscall(ctx0, __NR_getpid);
        allow_syscall(ctx0, __NR_gettid);
        allow_syscall(ctx0, __NR_tgkill);
        allow_syscall(ctx0, __NR_rt_sigreturn);
        allow_syscall(ctx0, __NR_socket);
        allow_syscall(ctx0, __NR_connect);
        allow_syscall(ctx0, __NR_lseek);
        allow_syscall(ctx0, __NR_openat);
        allow_syscall(ctx0, __NR_unlinkat);
        allow_syscall(ctx0, __NR_utimensat);
        allow_syscall(ctx0, __NR_lstat);
        allow_syscall(ctx0, __NR_fchmod);
        allow_syscall(ctx0, __NR_unlink);
        allow_syscall(ctx0, __NR_utime);
        allow_syscall(ctx0, __NR_futex);
        allow_syscall(ctx0, __NR_chmod);
        allow_syscall(ctx0, __NR_getuid);
        allow_syscall(ctx0, __NR_getgid);
        allow_syscall(ctx0, __NR_getppid);
        allow_syscall(ctx0, __NR_getcwd);
        allow_syscall(ctx0, __NR_getegid);
        allow_syscall(ctx0, __NR_getdents);
        allow_syscall(ctx0, __NR_readlink);
        assert!(seccomp_load(ctx0) == 0);
    }
    return Ok(());
}

const RENAME_NOREPLACE : u64 = 1;

fn path_to_cstring(path : &Path) -> CString {
    return CString::new(path.as_os_str().as_bytes()).unwrap();
}

fn rename_noreplace(path0 : &Path, path1 : &Path) -> Result<(), Error> {
    unsafe {
        if libc::syscall(__NR_renameat2 as i64, libc::AT_FDCWD, path_to_cstring(path0).as_ptr(), libc::AT_FDCWD, path_to_cstring(path1).as_ptr(), RENAME_NOREPLACE, 0) == 0 {
            return Ok(());
        } else {
            return Err(Error::last_os_error());
        }
    }
}

fn rename_or_suffix(path0 : &Path, path1 : &Path) -> Result<PathBuf, Error> {
    let rv = rename_noreplace(path0, path1);
    if let Err(err) = rv {
        if err.raw_os_error().expect("Error fails to wrap errno") != libc::EEXIST {
            return Err(err);
        }
        for suffix in 1.. {
            let mut builder = path1.as_os_str().to_os_string();
            builder.push(format!(".{}", suffix));
            let rv = rename_noreplace(path0, Path::new(&builder));
            if let Err(err) = rv {
                if err.raw_os_error().expect("Error fails to wrap errno") != libc::EEXIST {
                    return Err(err);
                }
                continue;
            }
            return Ok(Path::new(&builder).to_path_buf());
        }
        return Ok(PathBuf::new()); // XXX Unreachable
    } else {
        return Ok(path1.to_path_buf());
    }
}

struct ArchiveType {
    extensions: Vec<&'static str>,
    extract_cmd: Vec<&'static str>,
}

lazy_static! {
    static ref ARCHIVE_TYPES : Vec<ArchiveType> = vec![
        ArchiveType {
            extensions: vec![
                ".tgz",
                ".tar.gz",
                ".tar.Z",
                ".tbz2",
                ".tar.bz2",
                ".tlz",
                ".tar.lz",
                ".tar.lzma",
                ".tar.lzo",
                ".tar.xz",
            ],
            extract_cmd: vec!["tar", "-xf"],
        },
        ArchiveType {
            extensions: vec![
                ".zip",
            ],
            extract_cmd: vec!["unzip", "--"],
        },
        ArchiveType {
            extensions: vec![
                ".part1.rar",
                ".rar",
            ],
            extract_cmd: vec!["unrar", "x", "--"],
        },
        ArchiveType {
            extensions: vec![
                ".7z",
            ],
            extract_cmd: vec!["7z", "x", "--"],
        },
    ];
}

fn main() {
    let matches = clap_app!(myapp =>
        (@arg ARCHIVE: +required ... "Archives to extract")
    ).get_matches();

    let mut found = false;

    for arch in matches.values_of_os("ARCHIVE").unwrap() {
        let arch = Path::new(arch);
        println!("Extract {:?}", arch);
        for at in ARCHIVE_TYPES.iter() {
            for &ext in &at.extensions {
                if !arch.as_os_str().as_bytes().ends_with(ext.as_bytes()) {
                    continue;
                }
                found = true;

                let noext : &OsStr = OsStrExt::from_bytes(&arch.as_os_str().as_bytes()[0..arch.as_os_str().as_bytes().len()-ext.as_bytes().len()]);
                let stripped =
                    if let Some(noext) = Path::new(noext).file_name() {
                        noext
                    } else {
                        OsStr::new("extracted")
                    };
                let cmd = at.extract_cmd[0];
                let tmpdir = tempdir::TempDir::new_in(
                    env::current_dir().expect("Unable to get the current directory"),
                    stripped.to_string_lossy().to_mut().as_str()).expect("Unable to create temporary directory");
                let status = Command::new(cmd)
                    .args(&at.extract_cmd[1..])
                    .arg(&arch)
                    .current_dir(&tmpdir)
                    .before_exec(setup_seccomp)
                    .status()
                    .expect(format!("Failed to launch {}", cmd).as_str());
                assert!(status.success(), format!("{} wasn't successful", cmd));
                let extracted_path = tmpdir.into_path();
                let mut dir_list = fs::read_dir(&extracted_path).expect("Unable to read the extraction directory");
                if let Some(path) = dir_list.next() {
                    let path = path.expect("Unable to iterate the extraction directory");
                    if let Some(path1) = dir_list.next() {
                        path1.expect("Unable to iterate the extraction directory");
                        let newname = rename_or_suffix(&extracted_path, Path::new(stripped)).expect("Failed to rename");
                        println!("Extracted to {:?}", newname);
                    } else {
                        let newname = rename_or_suffix(&path.path(), Path::new(stripped)).expect("Failed to rename");
                        println!("Extracted to {:?}", newname);
                        fs::remove_dir(&extracted_path).expect("Failed to remove empty extraction directory");
                    }
                } else {
                    println!("No files were extracted");
                }
                break;
            }
        }
        if !found {
            println!("Can't handle {:?}", arch);
        }
    }
}

<commit_msg>Factor in some byte strings<commit_after>#[macro_use]
extern crate clap;
#[macro_use]
extern crate lazy_static;

extern crate libc;
extern crate scmp;
extern crate tempdir;

use std::fs;
use std::env;
use std::ffi::OsStr;
use std::ffi::CString;
use std::io::Error;
use std::os::raw::c_void;
use std::os::unix::process::CommandExt;
use std::os::unix::ffi::OsStrExt;
use std::path::Path;
use std::path::PathBuf;
use std::process::Command;
use scmp::*;

unsafe fn allow_syscall(ctx : *mut c_void, nr : u32) {
    assert!(seccomp_rule_add_array(ctx, SCMP_ACT_ALLOW, nr as i32, 0, std::ptr::null()) == 0);
}

fn setup_seccomp() -> Result<(), Error> {
    unsafe {
        let ctx0 = seccomp_init(SCMP_ACT_TRAP);
        assert!(!ctx0.is_null());
        allow_syscall(ctx0, __NR_execve);
        allow_syscall(ctx0, __NR_brk);
        allow_syscall(ctx0, __NR_access);
        allow_syscall(ctx0, __NR_mmap);
        allow_syscall(ctx0, __NR_open);
        allow_syscall(ctx0, __NR_read);
        allow_syscall(ctx0, __NR_write);
        allow_syscall(ctx0, __NR_stat);
        allow_syscall(ctx0, __NR_fstat);
        allow_syscall(ctx0, __NR_close);
        allow_syscall(ctx0, __NR_mprotect);
        allow_syscall(ctx0, __NR_arch_prctl);
        allow_syscall(ctx0, __NR_munmap);
        allow_syscall(ctx0, __NR_set_tid_address);
        allow_syscall(ctx0, __NR_set_robust_list);
        allow_syscall(ctx0, __NR_rt_sigaction);
        allow_syscall(ctx0, __NR_rt_sigprocmask);
        allow_syscall(ctx0, __NR_getrlimit);
        allow_syscall(ctx0, __NR_statfs);
        allow_syscall(ctx0, __NR_fcntl);
        allow_syscall(ctx0, __NR_geteuid);
        allow_syscall(ctx0, __NR_umask);
        allow_syscall(ctx0, __NR_pipe);
        allow_syscall(ctx0, __NR_clone);
        allow_syscall(ctx0, __NR_wait4);
        allow_syscall(ctx0, __NR_exit_group);
        allow_syscall(ctx0, __NR_dup);
        allow_syscall(ctx0, __NR_ioctl);
        allow_syscall(ctx0, __NR_getpid);
        allow_syscall(ctx0, __NR_gettid);
        allow_syscall(ctx0, __NR_tgkill);
        allow_syscall(ctx0, __NR_rt_sigreturn);
        allow_syscall(ctx0, __NR_socket);
        allow_syscall(ctx0, __NR_connect);
        allow_syscall(ctx0, __NR_lseek);
        allow_syscall(ctx0, __NR_openat);
        allow_syscall(ctx0, __NR_unlinkat);
        allow_syscall(ctx0, __NR_utimensat);
        allow_syscall(ctx0, __NR_lstat);
        allow_syscall(ctx0, __NR_fchmod);
        allow_syscall(ctx0, __NR_unlink);
        allow_syscall(ctx0, __NR_utime);
        allow_syscall(ctx0, __NR_futex);
        allow_syscall(ctx0, __NR_chmod);
        allow_syscall(ctx0, __NR_getuid);
        allow_syscall(ctx0, __NR_getgid);
        allow_syscall(ctx0, __NR_getppid);
        allow_syscall(ctx0, __NR_getcwd);
        allow_syscall(ctx0, __NR_getegid);
        allow_syscall(ctx0, __NR_getdents);
        allow_syscall(ctx0, __NR_readlink);
        assert!(seccomp_load(ctx0) == 0);
    }
    return Ok(());
}

const RENAME_NOREPLACE : u64 = 1;

fn path_to_cstring(path : &Path) -> CString {
    return CString::new(path.as_os_str().as_bytes()).unwrap();
}

fn rename_noreplace(path0 : &Path, path1 : &Path) -> Result<(), Error> {
    unsafe {
        if libc::syscall(__NR_renameat2 as i64, libc::AT_FDCWD, path_to_cstring(path0).as_ptr(), libc::AT_FDCWD, path_to_cstring(path1).as_ptr(), RENAME_NOREPLACE, 0) == 0 {
            return Ok(());
        } else {
            return Err(Error::last_os_error());
        }
    }
}

fn rename_or_suffix(path0 : &Path, path1 : &Path) -> Result<PathBuf, Error> {
    let rv = rename_noreplace(path0, path1);
    if let Err(err) = rv {
        if err.raw_os_error().expect("Error fails to wrap errno") != libc::EEXIST {
            return Err(err);
        }
        for suffix in 1.. {
            let mut builder = path1.as_os_str().to_os_string();
            builder.push(format!(".{}", suffix));
            let rv = rename_noreplace(path0, Path::new(&builder));
            if let Err(err) = rv {
                if err.raw_os_error().expect("Error fails to wrap errno") != libc::EEXIST {
                    return Err(err);
                }
                continue;
            }
            return Ok(Path::new(&builder).to_path_buf());
        }
        return Ok(PathBuf::new()); // XXX Unreachable
    } else {
        return Ok(path1.to_path_buf());
    }
}

struct ArchiveType {
    extensions: Vec<&'static str>,
    extract_cmd: Vec<&'static str>,
}

lazy_static! {
    static ref ARCHIVE_TYPES : Vec<ArchiveType> = vec![
        ArchiveType {
            extensions: vec![
                ".tgz",
                ".tar.gz",
                ".tar.Z",
                ".tbz2",
                ".tar.bz2",
                ".tlz",
                ".tar.lz",
                ".tar.lzma",
                ".tar.lzo",
                ".tar.xz",
            ],
            extract_cmd: vec!["tar", "-xf"],
        },
        ArchiveType {
            extensions: vec![
                ".zip",
            ],
            extract_cmd: vec!["unzip", "--"],
        },
        ArchiveType {
            extensions: vec![
                ".part1.rar",
                ".rar",
            ],
            extract_cmd: vec!["unrar", "x", "--"],
        },
        ArchiveType {
            extensions: vec![
                ".7z",
            ],
            extract_cmd: vec!["7z", "x", "--"],
        },
    ];
}

fn main() {
    let matches = clap_app!(myapp =>
        (@arg ARCHIVE: +required ... "Archives to extract")
    ).get_matches();

    let mut found = false;

    for arch in matches.values_of_os("ARCHIVE").unwrap() {
        let arch = Path::new(arch);
        let arch_bytes = arch.as_os_str().as_bytes();
        println!("Extract {:?}", arch);
        for at in ARCHIVE_TYPES.iter() {
            for &ext in &at.extensions {
                let ext_bytes = ext.as_bytes();
                if !arch_bytes.ends_with(ext_bytes) {
                    continue;
                }
                found = true;

                let noext : &OsStr = OsStrExt::from_bytes(&arch_bytes[0..arch_bytes.len()-ext_bytes.len()]);
                let stripped =
                    if let Some(noext) = Path::new(noext).file_name() {
                        noext
                    } else {
                        OsStr::new("extracted")
                    };
                let cmd = at.extract_cmd[0];
                let tmpdir = tempdir::TempDir::new_in(
                    env::current_dir().expect("Unable to get the current directory"),
                    stripped.to_string_lossy().to_mut().as_str()).expect("Unable to create temporary directory");
                let status = Command::new(cmd)
                    .args(&at.extract_cmd[1..])
                    .arg(&arch)
                    .current_dir(&tmpdir)
                    .before_exec(setup_seccomp)
                    .status()
                    .expect(format!("Failed to launch {}", cmd).as_str());
                assert!(status.success(), format!("{} wasn't successful", cmd));
                let extracted_path = tmpdir.into_path();
                let mut dir_list = fs::read_dir(&extracted_path).expect("Unable to read the extraction directory");
                if let Some(path) = dir_list.next() {
                    let path = path.expect("Unable to iterate the extraction directory");
                    if let Some(path1) = dir_list.next() {
                        path1.expect("Unable to iterate the extraction directory");
                        let newname = rename_or_suffix(&extracted_path, Path::new(stripped)).expect("Failed to rename");
                        println!("Extracted to {:?}", newname);
                    } else {
                        let newname = rename_or_suffix(&path.path(), Path::new(stripped)).expect("Failed to rename");
                        println!("Extracted to {:?}", newname);
                        fs::remove_dir(&extracted_path).expect("Failed to remove empty extraction directory");
                    }
                } else {
                    println!("No files were extracted");
                }
                break;
            }
        }
        if !found {
            println!("Can't handle {:?}", arch);
        }
    }
}

<|endoftext|>


<commit_before>#[macro_use]
extern crate quicli;
extern crate ansi_term;

use quicli::prelude::*;
use ansi_term::Color::Yellow;

#[macro_use] mod macros;
mod about_variables;

#[derive(Debug, StructOpt)]
struct Cli {
    exercise: Option<String>,
}

main!(|args: Cli| if let Some(e) = args.exercise {
    println!("selected {}", e);
} else {
    println!("Welcome to {}", Yellow.paint("rustlings"));
    about_variables::exec();
});
<commit_msg>make the initial output prettier<commit_after>#[macro_use]
extern crate quicli;
extern crate ansi_term;

use quicli::prelude::*;
use ansi_term::Color::Yellow;

#[macro_use] mod macros;
mod about_variables;

#[derive(Debug, StructOpt)]
struct Cli {
    exercise: Option<String>,
}

main!(|args: Cli| if let Some(e) = args.exercise {
    println!("selected {}", e);
} else {
    println!("Welcome to {}!\n", Yellow.paint("Rustlings"));
    about_variables::exec();
});
<|endoftext|>


<commit_before>fn main() {
	println!("platypie");
}
<commit_msg>A simple display of gratitude<commit_after>fn main() {
  let message = ["Thanks for installing Platypus!","We're trying to make the world a better place.","We're very happy you've decided to get involved."];

  for sptr in message.iter() {
    println!("{}",*sptr);
  }
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate clap;

#[macro_use]
extern crate log;
extern crate env_logger;

extern crate xmlrpc;
extern crate rustc_serialize;

extern crate regex;

use clap::{Arg, App};
use std::env;
use std::io;
use std::io::prelude::*;

use regex::Regex;

const GANDI_URL_PROD: &'static str = "https://rpc.gandi.net/xmlrpc/";

trait DNSProvider {
    fn init(&mut self, domain: &str);
    fn is_record_already_declared(&self, record_name: &str) -> Option<String>;
    fn update_record(&self, record_name: &str, ip_addr: &str) -> bool;
    fn create_record(&self, record_name: &str, ip_addr: &str);
}

struct GandiDNSProvider<'a> {
  xmlrpc_server: &'a str,
  apikey: &'a str,
  zone_id: Option<u32>,
  zone_id_version: u16,
}

impl<'a> GandiDNSProvider<'a> {

    fn new(gandi_url: &'a str, gandi_apikey: &'a str) -> GandiDNSProvider<'a> {
        GandiDNSProvider {
            xmlrpc_server: gandi_url,
            apikey: gandi_apikey,
            zone_id: None,
            zone_id_version: 0,
        }
    }
}

impl<'a> DNSProvider for GandiDNSProvider<'a> {

    fn init(&mut self, domain: &str) {

        let (client, mut request) = self.get_gandi_client("domain.info");
        request = request.argument(&domain.to_string());
        request = request.finalize();

        let response = client.remote_call(&request).unwrap();

        // TODO: fix this ugly code
        // Handle errors
        let zone_id_pos = response.body.find("zone_id").unwrap();
        let (_, body_end) = response.body.split_at(zone_id_pos);
        let first_int_markup = body_end.find("<int>").unwrap();
        let (_, body_end) = body_end.split_at(first_int_markup + "<int>".len());
        let first_end_int_markup = body_end.find("</int>").unwrap();
        let (zone_id, _) = body_end.split_at(first_end_int_markup);
        self.zone_id = zone_id.parse::<u32>().ok();

        debug!("Zone id: {}", self.zone_id.unwrap());
    }

    fn is_record_already_declared(&self, record_name: &str) -> Option<String> {

        let response = &self.get_record_list(record_name, &self.zone_id_version);

        // Extract already configured IP address
        // We are looking for something like that: <value><string>55.32.210.10</string></value>
        let regex = Regex::new(r"<value><string>([0-9.]*)</string></value>").unwrap();

        let caps = regex.captures(&*response.body);

        caps
            .map_or(None, |caps| caps.at(1))
            .map(|val| val.to_string())
    }

    fn update_record(&self, record_name: &str, ip_addr: &str) -> bool {

        // Create a new zone and get return version

        let new_zone_version = &self.create_new_zone();

        debug!("New zone version: {}", new_zone_version);

        // Extract new record id

        let new_record_id = &self.get_record_id(record_name, &new_zone_version);

        debug!("New record id: {}", new_record_id);

        // Update zone with the new record

        &self.update_zone_with_record(record_name, ip_addr, new_zone_version, new_record_id);

        // Activate the new zone
        debug!("Activate version '{}' of the zone '{}'", new_zone_version, &self.zone_id.unwrap());

        let (client, mut request) = self.get_gandi_client("domain.zone.version.set");
        request = request.argument(&self.zone_id.unwrap());
        request = request.argument(&new_zone_version);
        request = request.finalize();

        let response = client.remote_call(&request).unwrap();

        let regex = Regex::new(r"<boolean>([0-1]*)</boolean>").unwrap();

        let caps = regex.captures(&*response.body).unwrap();

        let result = caps.at(1).unwrap();

        debug!("Activate version result: {}", result);

        match result {
            "1" => true,
            "0" | _ => false,
        }
    }

    fn create_record(&self, record_name: &str, ip_addr: &str) {
        unimplemented!();
    }
}

impl<'a> GandiDNSProvider<'a> {

    fn get_gandi_client(&self, rpc_action: &str) -> (xmlrpc::Client, xmlrpc::Request) {
        let client = xmlrpc::Client::new(self.xmlrpc_server);
        let mut request = xmlrpc::Request::new(rpc_action);
        request = request.argument(&self.apikey.to_string());
        (client, request)
    }

    fn get_record_list(&self, record_name: &str, zone_id_version: &u16) -> xmlrpc::Response {

        let (client, mut request) = self.get_gandi_client("domain.zone.record.list");
        request = request.argument(&self.zone_id.unwrap());
        request = request.argument(zone_id_version);

        #[derive(Debug,RustcEncodable,RustcDecodable)]
        struct Record {
            name: String,
            type_: String,
        }

        let record = Record { name: record_name.to_string(), type_: "A".to_string() };

        request = request.argument(&record);

        request = request.finalize();

        // Horrible hack, because 'type' is a reserved keyword ...
        request.body = request.body.replace("type_", "type");

        client.remote_call(&request).unwrap()
    }

    fn create_new_zone(&self) -> u16 {
        let (client, mut request) = self.get_gandi_client("domain.zone.version.new");
        request = request.argument(&self.zone_id.unwrap());
        request = request.finalize();

        let response = client.remote_call(&request).unwrap();

        let regex = Regex::new(r"<int>([0-9]+)</int>").unwrap();

        let caps = regex.captures(&*response.body).unwrap();

        caps.at(1).unwrap().parse::<u16>().ok().unwrap()
    }

    fn get_record_id(&self, record_name: &str, new_zone_version: &u16) -> u32 {
        let response = &self.get_record_list(record_name, new_zone_version);

        let regex = Regex::new(r"<int>([0-9]+)</int>").unwrap();

        let caps = regex.captures(&*response.body).unwrap();

        caps.at(1).unwrap().parse::<u32>().ok().unwrap()
    }

    fn update_zone_with_record(&self, record_name: &str, ip_addr: &str, new_zone_version: &u16, new_record_id: &u32) {
        let (client, mut request) = self.get_gandi_client("domain.zone.record.update");
        request = request.argument(&self.zone_id.unwrap());
        request = request.argument(new_zone_version);

        #[derive(Debug,RustcEncodable,RustcDecodable)]
        struct NewRecordId { id: u32 };
        request = request.argument(&NewRecordId{ id: *new_record_id });

        #[derive(Debug,RustcEncodable,RustcDecodable)]
        struct Record {
            name: String,
            type_: String,
            value: String,
        }

        let record = Record {
            name: record_name.to_string(),
            type_: "A".to_string(),
            value: ip_addr.to_string(),
        };

        request = request.argument(&record);

        request = request.finalize();

        // Horrible hack, because 'type' is a reserved keyword ...
        request.body = request.body.replace("type_", "type");

        client.remote_call(&request); // ignore response
    }
}

fn main() {
    let matches = App::new("gdu")
        .version(&crate_version!()[..])
        .author("Damien Lecan <dev@dlecan.com>")
        .about("Gandi DNS updater, useful to reflect your dynamic IP address to your Gandi DNS zone file.\nIP address is read from stdin.")
        .args_from_usage(
            "-a --apikey=[apikey] 'Your API key provided by Gandi'
            -d --domain=[domain] 'The domain name whose active zonefile will be updated, e.g. \"domain.com\"'
            -n --dry-run 'Dry run, don't really update Gandi zone file'
            -f --force 'Force new zonefile creation even if IP address isn\'t modified'
            [verbose]... -v 'Verbose mode'")
        .arg(Arg::with_name("record_name")
            .help("Name of the A record to update or create (without domain)")
            .index(1)
            .required(true)
            .multiple(false))
        .get_matches();

    // Init logging to DEBUG only if user requires it
    match matches.occurrences_of("verbose") {
        0 => (),
        1 => {
                println!("Verbose mode");
                env::set_var("RUST_LOG", "DEBUG");
            },
        2 | _=> {
                println!("More verbose mode");
                env::set_var("RUST_LOG", "TRACE");
            }
    }
    env_logger::init().unwrap();

    let apikey = matches.value_of("apikey").unwrap();
    debug!("Using apikey: {}", apikey);

    let domain = matches.value_of("domain").unwrap();
    debug!("Using domain: {}", domain);

    let record_name = matches.value_of("record_name").unwrap();
    debug!("Using record name: {}", record_name);

    let dry_run = matches.is_present("dry-run");
    debug!("Dry run: {}", dry_run);

    let force = matches.is_present("force");
    debug!("Force: {}", force);

    let mut input = String::new();
    match io::stdin().read_line(&mut input) {
        Ok(_) => {
            println!("Input: {}", input);
        }
        Err(error) => println!("error: {}", error),
    }

    let expected_ip_addr = input.trim();

    let mut dns_provider = GandiDNSProvider::new(GANDI_URL_PROD, apikey);

    dns_provider.init(domain);

    let maybe_checked = dns_provider.is_record_already_declared(record_name);

    match maybe_checked {
        Some(ip_addr) => {
            debug!("Record already declared, with IP address: {}", &ip_addr);

            if !force && (&ip_addr == expected_ip_addr) {
                debug!("IP address not modified, no record to update");
            } else {
                debug!("Update record '{}' with IP address '{}'", record_name, &expected_ip_addr);
                let result = dns_provider.update_record(record_name, &expected_ip_addr);
                debug!("End of update process with result: {}", result);
            }
        }
        None => dns_provider.create_record(record_name, expected_ip_addr)
    }

}
<commit_msg>Clear seperation between DNSProvider and GandiRPC apis<commit_after>#[macro_use]
extern crate clap;

#[macro_use]
extern crate log;
extern crate env_logger;

extern crate xmlrpc;
extern crate rustc_serialize;

extern crate regex;

use clap::{Arg, App};
use std::env;
use std::io;
use std::io::prelude::*;

use regex::Regex;

const GANDI_URL_PROD: &'static str = "https://rpc.gandi.net/xmlrpc/";

trait DNSProvider {
    fn init(&mut self, domain: &str);
    fn is_record_already_declared(&self, record_name: &str) -> Option<String>;
    fn update_record(&self, record_name: &str, ip_addr: &str) -> bool;
    fn create_record(&self, record_name: &str, ip_addr: &str);
}

struct GandiDNSProvider<'a> {
    zone_id: Option<u32>,
    zone_id_version: u16,
    gandi_rpc: GandiRPC<'a>,
}

impl<'a> GandiDNSProvider<'a> {

    fn new(gandi_url: &'a str, gandi_apikey: &'a str) -> GandiDNSProvider<'a> {

        let gandi_rpc = GandiRPC {
            xmlrpc_server: gandi_url,
            apikey: gandi_apikey,
        };

        GandiDNSProvider {
            zone_id: None,
            zone_id_version: 0,
            gandi_rpc: gandi_rpc,
        }
    }
}

impl<'a> DNSProvider for GandiDNSProvider<'a> {

    fn init(&mut self, domain: &str) {

        let (client, mut request) = self.gandi_rpc.get_gandi_client("domain.info");
        request = request.argument(&domain.to_string());
        request = request.finalize();

        let response = client.remote_call(&request).unwrap();

        // TODO: fix this ugly code
        // Handle errors
        let zone_id_pos = response.body.find("zone_id").unwrap();
        let (_, body_end) = response.body.split_at(zone_id_pos);
        let first_int_markup = body_end.find("<int>").unwrap();
        let (_, body_end) = body_end.split_at(first_int_markup + "<int>".len());
        let first_end_int_markup = body_end.find("</int>").unwrap();
        let (zone_id, _) = body_end.split_at(first_end_int_markup);
        self.zone_id = zone_id.parse::<u32>().ok();

        debug!("Zone id: {}", self.zone_id.unwrap());
    }

    fn is_record_already_declared(&self, record_name: &str) -> Option<String> {

        let response = &self.gandi_rpc.get_record_list(record_name, &self.zone_id.unwrap(), &self.zone_id_version);

        // Extract already configured IP address
        // We are looking for something like that: <value><string>55.32.210.10</string></value>
        let regex = Regex::new(r"<value><string>([0-9.]*)</string></value>").unwrap();

        let caps = regex.captures(&*response.body);

        caps
            .map_or(None, |caps| caps.at(1))
            .map(|val| val.to_string())
    }

    fn update_record(&self, record_name: &str, ip_addr: &str) -> bool {

        // Create a new zone and get return version

        let new_zone_version = &self.gandi_rpc.create_new_zone(&self.zone_id.unwrap());

        debug!("New zone version: {}", new_zone_version);

        // Extract new record id

        let new_record_id = &self.gandi_rpc.get_record_id(record_name, &self.zone_id.unwrap(), &new_zone_version);

        debug!("New record id: {}", new_record_id);

        // Update zone with the new record

        &self.gandi_rpc.update_zone_with_record(record_name, ip_addr, &self.zone_id.unwrap(), new_zone_version, new_record_id);

        // Activate the new zone
        debug!("Activate version '{}' of the zone '{}'", new_zone_version, &self.zone_id.unwrap());

        let (client, mut request) = self.gandi_rpc.get_gandi_client("domain.zone.version.set");
        request = request.argument(&self.zone_id.unwrap());
        request = request.argument(&new_zone_version);
        request = request.finalize();

        let response = client.remote_call(&request).unwrap();

        let regex = Regex::new(r"<boolean>([0-1]*)</boolean>").unwrap();

        let caps = regex.captures(&*response.body).unwrap();

        let result = caps.at(1).unwrap();

        debug!("Activate version result: {}", result);

        match result {
            "1" => true,
            "0" | _ => false,
        }
    }

    fn create_record(&self, record_name: &str, ip_addr: &str) {
        unimplemented!();
    }
}

struct GandiRPC<'a> {
    xmlrpc_server: &'a str,
    apikey: &'a str,
}

impl<'a> GandiRPC<'a> {

    fn get_gandi_client(&self, rpc_action: &str) -> (xmlrpc::Client, xmlrpc::Request) {
        let client = xmlrpc::Client::new(self.xmlrpc_server);
        let mut request = xmlrpc::Request::new(rpc_action);
        request = request.argument(&self.apikey.to_string());
        (client, request)
    }

    fn get_record_list(&self, record_name: &str, zone_id: &u32, zone_id_version: &u16) -> xmlrpc::Response {

        let (client, mut request) = self.get_gandi_client("domain.zone.record.list");
        request = request.argument(zone_id);
        request = request.argument(zone_id_version);

        #[derive(Debug,RustcEncodable,RustcDecodable)]
        struct Record {
            name: String,
            type_: String,
        }

        let record = Record { name: record_name.to_string(), type_: "A".to_string() };

        request = request.argument(&record);

        request = request.finalize();

        // Horrible hack, because 'type' is a reserved keyword ...
        request.body = request.body.replace("type_", "type");

        client.remote_call(&request).unwrap()
    }

    fn create_new_zone(&self, zone_id: &u32) -> u16 {
        let (client, mut request) = self.get_gandi_client("domain.zone.version.new");
        request = request.argument(zone_id);
        request = request.finalize();

        let response = client.remote_call(&request).unwrap();

        let regex = Regex::new(r"<int>([0-9]+)</int>").unwrap();

        let caps = regex.captures(&*response.body).unwrap();

        caps.at(1).unwrap().parse::<u16>().ok().unwrap()
    }

    fn get_record_id(&self, record_name: &str, zone_id: &u32, new_zone_version: &u16) -> u32 {
        let response = &self.get_record_list(record_name, zone_id, new_zone_version);

        let regex = Regex::new(r"<int>([0-9]+)</int>").unwrap();

        let caps = regex.captures(&*response.body).unwrap();

        caps.at(1).unwrap().parse::<u32>().ok().unwrap()
    }

    fn update_zone_with_record(&self, record_name: &str, ip_addr: &str, zone_id: &u32, new_zone_version: &u16, new_record_id: &u32) {
        let (client, mut request) = self.get_gandi_client("domain.zone.record.update");
        request = request.argument(zone_id);
        request = request.argument(new_zone_version);

        #[derive(Debug,RustcEncodable,RustcDecodable)]
        struct NewRecordId { id: u32 };
        request = request.argument(&NewRecordId{ id: *new_record_id });

        #[derive(Debug,RustcEncodable,RustcDecodable)]
        struct Record {
            name: String,
            type_: String,
            value: String,
        }

        let record = Record {
            name: record_name.to_string(),
            type_: "A".to_string(),
            value: ip_addr.to_string(),
        };

        request = request.argument(&record);

        request = request.finalize();

        // Horrible hack, because 'type' is a reserved keyword ...
        request.body = request.body.replace("type_", "type");

        client.remote_call(&request); // ignore response
    }
}

fn main() {
    let matches = App::new("gdu")
        .version(&crate_version!()[..])
        .author("Damien Lecan <dev@dlecan.com>")
        .about("Gandi DNS updater, useful to reflect your dynamic IP address to your Gandi DNS zone file.\nIP address is read from stdin.")
        .args_from_usage(
            "-a --apikey=[apikey] 'Your API key provided by Gandi'
            -d --domain=[domain] 'The domain name whose active zonefile will be updated, e.g. \"domain.com\"'
            -n --dry-run 'Dry run, don't really update Gandi zone file'
            -f --force 'Force new zonefile creation even if IP address isn\'t modified'
            [verbose]... -v 'Verbose mode'")
        .arg(Arg::with_name("record_name")
            .help("Name of the A record to update or create (without domain)")
            .index(1)
            .required(true)
            .multiple(false))
        .get_matches();

    // Init logging to DEBUG only if user requires it
    match matches.occurrences_of("verbose") {
        0 => (),
        1 => {
                println!("Verbose mode");
                env::set_var("RUST_LOG", "DEBUG");
            },
        2 | _=> {
                println!("More verbose mode");
                env::set_var("RUST_LOG", "TRACE");
            }
    }
    env_logger::init().unwrap();

    let apikey = matches.value_of("apikey").unwrap();
    debug!("Using apikey: {}", apikey);

    let domain = matches.value_of("domain").unwrap();
    debug!("Using domain: {}", domain);

    let record_name = matches.value_of("record_name").unwrap();
    debug!("Using record name: {}", record_name);

    let dry_run = matches.is_present("dry-run");
    debug!("Dry run: {}", dry_run);

    let force = matches.is_present("force");
    debug!("Force: {}", force);

    let mut input = String::new();
    match io::stdin().read_line(&mut input) {
        Ok(_) => {
            println!("Input: {}", input);
        }
        Err(error) => println!("error: {}", error),
    }

    let expected_ip_addr = input.trim();

    let mut dns_provider = GandiDNSProvider::new(GANDI_URL_PROD, apikey);

    dns_provider.init(domain);

    let maybe_checked = dns_provider.is_record_already_declared(record_name);

    match maybe_checked {
        Some(ip_addr) => {
            debug!("Record already declared, with IP address: {}", &ip_addr);

            if !force && (&ip_addr == expected_ip_addr) {
                debug!("IP address not modified, no record to update");
            } else {
                debug!("Update record '{}' with IP address '{}'", record_name, &expected_ip_addr);
                let result = dns_provider.update_record(record_name, &expected_ip_addr);
                debug!("End of update process with result: {}", result);
            }
        }
        None => dns_provider.create_record(record_name, expected_ip_addr)
    }

}
<|endoftext|>


<commit_before>extern crate emailaddress;
extern crate iron;
extern crate lettre;
extern crate openssl;
extern crate rand;
extern crate redis;
extern crate router;
extern crate rustc_serialize;
extern crate serde_json;
extern crate time;
extern crate url;
extern crate urlencoded;

use emailaddress::EmailAddress;
use iron::headers::ContentType;
use iron::middleware::Handler;
use iron::prelude::*;
use iron::status;
use lettre::email::EmailBuilder;
use lettre::transport::smtp::SmtpTransportBuilder;
use lettre::transport::EmailTransport;
use urlencoded::{UrlEncodedBody, UrlEncodedQuery};
use openssl::bn::BigNum;
use openssl::crypto::pkey::PKey;
use openssl::crypto::hash;
use serde_json::builder::ObjectBuilder;
use serde_json::de::from_reader;
use serde_json::value::Value;
use redis::{Client, Commands, RedisResult};
use router::Router;
use rustc_serialize::base64::{self, ToBase64};
use std::collections::HashMap;
use std::error::Error;
use std::fs::File;
use std::io::BufReader;
use time::now_utc;
use url::percent_encoding::{utf8_percent_encode, QUERY_ENCODE_SET};

fn json_response(obj: &Value) -> IronResult<Response> {
    let content = serde_json::to_string(&obj).unwrap();
    let mut rsp = Response::with((status::Ok, content));
    rsp.headers.set(ContentType::json());
    Ok(rsp)
}

struct WelcomeHandler { app: AppConfig }
impl Handler for WelcomeHandler {
    fn handle(&self, _: &mut Request) -> IronResult<Response> {
        json_response(&ObjectBuilder::new()
            .insert("ladaemon", "Welcome")
            .insert("version", &self.app.version)
            .unwrap())
    }
}

struct OIDConfigHandler { app: AppConfig }
impl Handler for OIDConfigHandler {
    fn handle(&self, _: &mut Request) -> IronResult<Response> {
        json_response(&ObjectBuilder::new()
            .insert("issuer", &self.app.base_url)
            .insert("authorization_endpoint",
                    format!("{}/auth", self.app.base_url))
            .insert("jwks_uri", format!("{}/keys.json", self.app.base_url))
            .insert("scopes_supported", vec!["openid", "email"])
            .insert("claims_supported",
                    vec!["aud", "email", "email_verified", "exp", "iat", "iss", "sub"])
            .insert("response_types_supported", vec!["id_token"])
            .insert("response_modes_supported", vec!["form_post"])
            .insert("grant_types_supported", vec!["implicit"])
            .insert("subject_types_supported", vec!["public"])
            .insert("id_token_signing_alg_values_supported", vec!["RS256"])
            .unwrap())
    }
}

#[derive(Clone)]
struct AppConfig {
    version: String,
    base_url: String,
    priv_key: PKey,
    store: Client,
    expire_keys: usize,
    sender: String,
    token_validity: i64,
}

fn json_big_num(n: &BigNum) -> String {
    n.to_vec().to_base64(base64::URL_SAFE)
}

struct KeysHandler { app: AppConfig }
impl Handler for KeysHandler {
    fn handle(&self, _: &mut Request) -> IronResult<Response> {
        let rsa = self.app.priv_key.get_rsa();
        json_response(&ObjectBuilder::new()
            .insert_array("keys", |builder| {
                builder.push_object(|builder| {
                    builder.insert("kty", "RSA")
                        .insert("alg", "RS256")
                        .insert("use", "sig")
                        .insert("kid", "base")
                        .insert("n", json_big_num(&rsa.n().unwrap()))
                        .insert("e", json_big_num(&rsa.e().unwrap()))
                })
            })
            .unwrap())
    }
}

const CODE_CHARS: &'static [char] = &[
    '2', '3', '4', '6', '7', '9', 'a', 'c', 'd', 'e', 'f', 'g', 'h', 'j', 'k',
    'm', 'n', 'p', 'q', 'r', 't', 'v', 'w', 'x', 'y', 'z',
];

struct AuthHandler { app: AppConfig }
impl Handler for AuthHandler {
    fn handle(&self, req: &mut Request) -> IronResult<Response> {

        let chars: String = (0..6).map(|_| CODE_CHARS[rand::random::<usize>() % CODE_CHARS.len()]).collect();
        let params = req.get_ref::<UrlEncodedBody>().unwrap();
        let email_addr = EmailAddress::new(&params.get("login_hint").unwrap()[0]);
        let client_id = &params.get("client_id").unwrap()[0];
        let key = format!("{}:{}", email_addr, client_id);
        let set_res: RedisResult<String> = self.app.store.hset_multiple(key.clone(), &[
            ("code", chars.clone()),
            ("redirect", params.get("redirect_uri").unwrap()[0].clone()),
        ]);
        let exp_res: RedisResult<bool> = self.app.store.expire(key.clone(), self.app.expire_keys);

        let email = EmailBuilder::new()
            .to(email_addr.to_string().as_str())
            .from(self.app.sender.as_str())
            .body(&format!("code: {}", chars))
            .subject("Your login code")
            .build().unwrap();
        let mut mailer = SmtpTransportBuilder::localhost().unwrap().build();
        let result = mailer.send(email);
        mailer.close();

        let href = format!("{}/confirm?email={}&origin={}&code={}",
                           self.app.base_url,
                           utf8_percent_encode(&email_addr.to_string(), QUERY_ENCODE_SET),
                           utf8_percent_encode(client_id, QUERY_ENCODE_SET),
                           utf8_percent_encode(&chars, QUERY_ENCODE_SET));
        let mut obj = ObjectBuilder::new()
            .insert("href", href);
        if !result.is_ok() {
            let error = result.unwrap_err();
            obj = obj.insert("error", error.to_string());
            obj = match error {
                lettre::transport::error::Error::IoError(inner) => {
                    obj.insert("cause", inner.description())
                }
                _ => obj,
            }
        }
        if !set_res.is_ok() {
            obj = obj.insert("hset_multiple", set_res.unwrap_err().to_string());
        }
        if !exp_res.is_ok() {
            obj = obj.insert("expire", exp_res.unwrap_err().to_string());
        }
        json_response(&obj.unwrap())

    }
}

fn create_jwt(key: &PKey, header: &str, payload: &str) -> String {
    let mut input = Vec::<u8>::new();
    input.extend(header.as_bytes().to_base64(base64::URL_SAFE).into_bytes());
    input.push(b'.');
    input.extend(payload.as_bytes().to_base64(base64::URL_SAFE).into_bytes());
    let sha256 = hash::hash(hash::Type::SHA256, &input);
    let sig = key.sign(&sha256);
    input.push(b'.');
    input.extend(sig.to_base64(base64::URL_SAFE).into_bytes());
    String::from_utf8(input).unwrap()
}

const FORWARD_TEMPLATE: &'static str = r#"<!DOCTYPE html>
<html>
  <head>
    <title>Let's Auth</title>
    <script>
      document.addEventListener('DOMContentLoaded', function() {
        document.getElementById('form').submit();
      });
    </script>
  </head>
  <body>
    <form id="form" action="{{ return_url }}" method="post">
      <input type="hidden" name="id_token" value="{{ jwt }}">
    </form>
  </body>
</html>"#;

struct ConfirmHandler { app: AppConfig }
impl Handler for ConfirmHandler {
    fn handle(&self, req: &mut Request) -> IronResult<Response> {

        let params = req.get_ref::<UrlEncodedQuery>().unwrap();
        let email = &params.get("email").unwrap()[0];
        let origin = &params.get("origin").unwrap()[0];
        let key = format!("{}:{}", email, origin);
        let stored: HashMap<String, String> = self.app.store.hgetall(key.clone()).unwrap();

        let req_code = &params.get("code").unwrap()[0];
        if stored.get("code").is_none() || req_code != stored.get("code").unwrap() {
            let mut obj = ObjectBuilder::new().insert("error", "code fail");
            obj = obj.insert("stored", stored);
            return json_response(&obj.unwrap());
        }

        let now = now_utc().to_timespec().sec;
        let payload = ObjectBuilder::new()
            .insert("aud", origin)
            .insert("email", email)
            .insert("email_verified", email)
            .insert("exp", now + self.app.token_validity)
            .insert("iat", now)
            .insert("iss", &self.app.base_url)
            .insert("sub", email)
            .unwrap();
        let headers = ObjectBuilder::new()
            .insert("kid", "base")
            .insert("alg", "RS256")
            .unwrap();
        let jwt = create_jwt(&self.app.priv_key,
                             &serde_json::to_string(&headers).unwrap(),
                             &serde_json::to_string(&payload).unwrap());

        let redirect = stored.get("redirect").unwrap();
        let html = FORWARD_TEMPLATE.replace("{{ return_url }}", redirect)
            .replace("{{ jwt }}", &jwt);

        let mut rsp = Response::with((status::Ok, html));
        rsp.headers.set(ContentType::html());
        Ok(rsp)

    }
}

fn main() {

    let config_file = File::open("test.json").unwrap();
    let config_reader = BufReader::new(config_file);
    let config: Value = from_reader(config_reader).unwrap();

    let key_file_name = config.find("private_key_file").unwrap().as_string().unwrap();
    let priv_key_file = File::open(key_file_name).unwrap();
    let mut reader = BufReader::new(priv_key_file);
    let app = AppConfig {
        version: env!("CARGO_PKG_VERSION").to_string(),
        base_url: config.find("base_url").unwrap().as_string().unwrap().to_string(),
        priv_key: PKey::private_key_from_pem(&mut reader).unwrap(),
        store: Client::open(config.find("redis").unwrap().as_string().unwrap()).unwrap(),
        expire_keys: config.find("expire_keys").unwrap().as_u64().unwrap() as usize,
        sender: "Let's Auth <letsauth@xavamedia.nl>".to_string(),
        token_validity: config.find("token_validity").unwrap().as_i64().unwrap(),
    };

    let mut router = Router::new();
    router.get("/", WelcomeHandler { app: app.clone() });
    router.get("/.well-known/openid-configuration",
               OIDConfigHandler { app: app.clone() });
    router.get("/keys.json", KeysHandler { app: app.clone() });
    router.post("/auth", AuthHandler { app: app.clone() });
    router.get("/confirm", ConfirmHandler { app: app.clone() });
    Iron::new(router).http("0.0.0.0:3333").unwrap();

}
<commit_msg>Improve email content; include full link<commit_after>extern crate emailaddress;
extern crate iron;
extern crate lettre;
extern crate openssl;
extern crate rand;
extern crate redis;
extern crate router;
extern crate rustc_serialize;
extern crate serde_json;
extern crate time;
extern crate url;
extern crate urlencoded;

use emailaddress::EmailAddress;
use iron::headers::ContentType;
use iron::middleware::Handler;
use iron::prelude::*;
use iron::status;
use lettre::email::EmailBuilder;
use lettre::transport::smtp::SmtpTransportBuilder;
use lettre::transport::EmailTransport;
use urlencoded::{UrlEncodedBody, UrlEncodedQuery};
use openssl::bn::BigNum;
use openssl::crypto::pkey::PKey;
use openssl::crypto::hash;
use serde_json::builder::ObjectBuilder;
use serde_json::de::from_reader;
use serde_json::value::Value;
use redis::{Client, Commands, RedisResult};
use router::Router;
use rustc_serialize::base64::{self, ToBase64};
use std::collections::HashMap;
use std::error::Error;
use std::fs::File;
use std::io::BufReader;
use time::now_utc;
use url::percent_encoding::{utf8_percent_encode, QUERY_ENCODE_SET};

fn json_response(obj: &Value) -> IronResult<Response> {
    let content = serde_json::to_string(&obj).unwrap();
    let mut rsp = Response::with((status::Ok, content));
    rsp.headers.set(ContentType::json());
    Ok(rsp)
}

struct WelcomeHandler { app: AppConfig }
impl Handler for WelcomeHandler {
    fn handle(&self, _: &mut Request) -> IronResult<Response> {
        json_response(&ObjectBuilder::new()
            .insert("ladaemon", "Welcome")
            .insert("version", &self.app.version)
            .unwrap())
    }
}

struct OIDConfigHandler { app: AppConfig }
impl Handler for OIDConfigHandler {
    fn handle(&self, _: &mut Request) -> IronResult<Response> {
        json_response(&ObjectBuilder::new()
            .insert("issuer", &self.app.base_url)
            .insert("authorization_endpoint",
                    format!("{}/auth", self.app.base_url))
            .insert("jwks_uri", format!("{}/keys.json", self.app.base_url))
            .insert("scopes_supported", vec!["openid", "email"])
            .insert("claims_supported",
                    vec!["aud", "email", "email_verified", "exp", "iat", "iss", "sub"])
            .insert("response_types_supported", vec!["id_token"])
            .insert("response_modes_supported", vec!["form_post"])
            .insert("grant_types_supported", vec!["implicit"])
            .insert("subject_types_supported", vec!["public"])
            .insert("id_token_signing_alg_values_supported", vec!["RS256"])
            .unwrap())
    }
}

#[derive(Clone)]
struct AppConfig {
    version: String,
    base_url: String,
    priv_key: PKey,
    store: Client,
    expire_keys: usize,
    sender: String,
    token_validity: i64,
}

fn json_big_num(n: &BigNum) -> String {
    n.to_vec().to_base64(base64::URL_SAFE)
}

struct KeysHandler { app: AppConfig }
impl Handler for KeysHandler {
    fn handle(&self, _: &mut Request) -> IronResult<Response> {
        let rsa = self.app.priv_key.get_rsa();
        json_response(&ObjectBuilder::new()
            .insert_array("keys", |builder| {
                builder.push_object(|builder| {
                    builder.insert("kty", "RSA")
                        .insert("alg", "RS256")
                        .insert("use", "sig")
                        .insert("kid", "base")
                        .insert("n", json_big_num(&rsa.n().unwrap()))
                        .insert("e", json_big_num(&rsa.e().unwrap()))
                })
            })
            .unwrap())
    }
}

const CODE_CHARS: &'static [char] = &[
    '2', '3', '4', '6', '7', '9', 'a', 'c', 'd', 'e', 'f', 'g', 'h', 'j', 'k',
    'm', 'n', 'p', 'q', 'r', 't', 'v', 'w', 'x', 'y', 'z',
];

struct AuthHandler { app: AppConfig }
impl Handler for AuthHandler {
    fn handle(&self, req: &mut Request) -> IronResult<Response> {

        let chars: String = (0..6).map(|_| CODE_CHARS[rand::random::<usize>() % CODE_CHARS.len()]).collect();
        let params = req.get_ref::<UrlEncodedBody>().unwrap();
        let email_addr = EmailAddress::new(&params.get("login_hint").unwrap()[0]);
        let client_id = &params.get("client_id").unwrap()[0];
        let key = format!("{}:{}", email_addr, client_id);
        let set_res: RedisResult<String> = self.app.store.hset_multiple(key.clone(), &[
            ("code", chars.clone()),
            ("redirect", params.get("redirect_uri").unwrap()[0].clone()),
        ]);
        let exp_res: RedisResult<bool> = self.app.store.expire(key.clone(), self.app.expire_keys);

        let href = format!("{}/confirm?email={}&origin={}&code={}",
                           self.app.base_url,
                           utf8_percent_encode(&email_addr.to_string(), QUERY_ENCODE_SET),
                           utf8_percent_encode(client_id, QUERY_ENCODE_SET),
                           utf8_percent_encode(&chars, QUERY_ENCODE_SET));
        let email = EmailBuilder::new()
            .to(email_addr.to_string().as_str())
            .from(self.app.sender.as_str())
            .body(&format!("Enter your login code:\n\n{}\n\nOr click this link:\n\n{}",
                           chars, href))
            .subject(&format!("Code: {} - Finish logging in to {}", chars, client_id))
            .build().unwrap();
        let mut mailer = SmtpTransportBuilder::localhost().unwrap().build();
        let result = mailer.send(email);
        mailer.close();

        let mut obj = ObjectBuilder::new()
            .insert("href", href);
        if !result.is_ok() {
            let error = result.unwrap_err();
            obj = obj.insert("error", error.to_string());
            obj = match error {
                lettre::transport::error::Error::IoError(inner) => {
                    obj.insert("cause", inner.description())
                }
                _ => obj,
            }
        }
        if !set_res.is_ok() {
            obj = obj.insert("hset_multiple", set_res.unwrap_err().to_string());
        }
        if !exp_res.is_ok() {
            obj = obj.insert("expire", exp_res.unwrap_err().to_string());
        }
        json_response(&obj.unwrap())

    }
}

fn create_jwt(key: &PKey, header: &str, payload: &str) -> String {
    let mut input = Vec::<u8>::new();
    input.extend(header.as_bytes().to_base64(base64::URL_SAFE).into_bytes());
    input.push(b'.');
    input.extend(payload.as_bytes().to_base64(base64::URL_SAFE).into_bytes());
    let sha256 = hash::hash(hash::Type::SHA256, &input);
    let sig = key.sign(&sha256);
    input.push(b'.');
    input.extend(sig.to_base64(base64::URL_SAFE).into_bytes());
    String::from_utf8(input).unwrap()
}

const FORWARD_TEMPLATE: &'static str = r#"<!DOCTYPE html>
<html>
  <head>
    <title>Let's Auth</title>
    <script>
      document.addEventListener('DOMContentLoaded', function() {
        document.getElementById('form').submit();
      });
    </script>
  </head>
  <body>
    <form id="form" action="{{ return_url }}" method="post">
      <input type="hidden" name="id_token" value="{{ jwt }}">
    </form>
  </body>
</html>"#;

struct ConfirmHandler { app: AppConfig }
impl Handler for ConfirmHandler {
    fn handle(&self, req: &mut Request) -> IronResult<Response> {

        let params = req.get_ref::<UrlEncodedQuery>().unwrap();
        let email = &params.get("email").unwrap()[0];
        let origin = &params.get("origin").unwrap()[0];
        let key = format!("{}:{}", email, origin);
        let stored: HashMap<String, String> = self.app.store.hgetall(key.clone()).unwrap();

        let req_code = &params.get("code").unwrap()[0];
        if stored.get("code").is_none() || req_code != stored.get("code").unwrap() {
            let mut obj = ObjectBuilder::new().insert("error", "code fail");
            obj = obj.insert("stored", stored);
            return json_response(&obj.unwrap());
        }

        let now = now_utc().to_timespec().sec;
        let payload = ObjectBuilder::new()
            .insert("aud", origin)
            .insert("email", email)
            .insert("email_verified", email)
            .insert("exp", now + self.app.token_validity)
            .insert("iat", now)
            .insert("iss", &self.app.base_url)
            .insert("sub", email)
            .unwrap();
        let headers = ObjectBuilder::new()
            .insert("kid", "base")
            .insert("alg", "RS256")
            .unwrap();
        let jwt = create_jwt(&self.app.priv_key,
                             &serde_json::to_string(&headers).unwrap(),
                             &serde_json::to_string(&payload).unwrap());

        let redirect = stored.get("redirect").unwrap();
        let html = FORWARD_TEMPLATE.replace("{{ return_url }}", redirect)
            .replace("{{ jwt }}", &jwt);

        let mut rsp = Response::with((status::Ok, html));
        rsp.headers.set(ContentType::html());
        Ok(rsp)

    }
}

fn main() {

    let config_file = File::open("test.json").unwrap();
    let config_reader = BufReader::new(config_file);
    let config: Value = from_reader(config_reader).unwrap();

    let key_file_name = config.find("private_key_file").unwrap().as_string().unwrap();
    let priv_key_file = File::open(key_file_name).unwrap();
    let mut reader = BufReader::new(priv_key_file);
    let app = AppConfig {
        version: env!("CARGO_PKG_VERSION").to_string(),
        base_url: config.find("base_url").unwrap().as_string().unwrap().to_string(),
        priv_key: PKey::private_key_from_pem(&mut reader).unwrap(),
        store: Client::open(config.find("redis").unwrap().as_string().unwrap()).unwrap(),
        expire_keys: config.find("expire_keys").unwrap().as_u64().unwrap() as usize,
        sender: "Let's Auth <letsauth@xavamedia.nl>".to_string(),
        token_validity: config.find("token_validity").unwrap().as_i64().unwrap(),
    };

    let mut router = Router::new();
    router.get("/", WelcomeHandler { app: app.clone() });
    router.get("/.well-known/openid-configuration",
               OIDConfigHandler { app: app.clone() });
    router.get("/keys.json", KeysHandler { app: app.clone() });
    router.post("/auth", AuthHandler { app: app.clone() });
    router.get("/confirm", ConfirmHandler { app: app.clone() });
    Iron::new(router).http("0.0.0.0:3333").unwrap();

}
<|endoftext|>


<commit_before> #![feature(slicing_syntax)]

extern crate libc;

use std::os;
use std::io::{
    Command,
    process,
};
use std::path::GenericPath;
use std::io::fs::PathExtensions;

mod compiler;

fn has_cache(script_path: &str, cache_path : &str) -> bool {
    let scriptp = Path::new(script_path);
    let cachep = Path::new(cache_path);
    if !scriptp.exists() || !scriptp.is_file() {
        panic!("Can't read script");
    }
    let curr = match scriptp.stat() {
        Ok(stat) => stat,
        Err(_) => { return false },
    };
    if  !cachep.exists() || !cachep.is_file() {
        return false;
    }
    let cached = match cachep.stat() {
        Ok(stat) => stat,
        Err(_) => { return false },
    };
    curr.modified < cached.modified
}
fn get_cache_path(script_name: Path) -> String {
    let stem = script_name.filestem().unwrap();
    format!("/home/dutt/.rusted/{}", String::from_utf8_lossy(stem))
}
fn get_clean_name(script_path: &str) -> String {
    let last_idx = match script_path.rfind(std::path::posix::SEP) {
        Some(idx) => idx,
        None => 0
    };
    String::from_str(script_path[last_idx+1..script_path.len()])
}
fn main() {
    let args = os::args();
    let arg = match args.get(1) {
        Some(a) => a,
        None => panic!("No argument supplied"),
    };
    let cleaned = get_clean_name(arg.as_slice());
    let src = Path::new(String::from_str(cleaned.as_slice()));
    let target = get_cache_path(Path::new(String::from_str(cleaned.as_slice())));
    let sysroot = Path::new("/usr/local");
    let cached = has_cache(arg.as_slice(), target.as_slice());
    if !cached {
        compiler::compile_file(src, Path::new(target.clone()), Some(sysroot));
    }

    let mut cmd = Command::new(target);
    let mut argsvec: Vec<&String> = Vec::new();
    for a in args.iter() {
        argsvec.push(a);
    }
    for i in range(1, argsvec.len()) {
        cmd.arg(argsvec[i]);
    }
    cmd.stdout(process::StdioContainer::InheritFd(libc::STDOUT_FILENO));
    cmd.stderr(process::StdioContainer::InheritFd(libc::STDERR_FILENO));
    cmd.stdin(process::StdioContainer::InheritFd(libc::STDIN_FILENO));
    match cmd.spawn() {
        Ok(_) => (),
        Err(e) => panic!("Failed to execute process: {}", e),
    };
}
<commit_msg>Cleanup, stability<commit_after> #![feature(slicing_syntax)]

extern crate libc;

use std::os;
use std::io::{
    Command,
    process,
};
use std::path::GenericPath;
use std::io::fs::PathExtensions;

mod compiler;

fn has_cache(script_path: &str, cache_path : &str) -> bool {
    let scriptp = Path::new(script_path);
    let cachep = Path::new(cache_path);
    if !scriptp.exists() || !scriptp.is_file() {
        panic!("Can't read script");
    }
    let curr = match scriptp.stat() {
        Ok(stat) => stat,
        Err(_) => { return false },
    };
    if  !cachep.exists() || !cachep.is_file() {
        return false;
    }
    let cached = match cachep.stat() {
        Ok(stat) => stat,
        Err(_) => { return false },
    };
    curr.modified < cached.modified
}
fn get_cache_path(script_name: Path) -> String {
    let file = script_name.filename_str().unwrap();
    let stem = match file.rfind('.') {
        Some(pos) => {
            file[0..pos]
        },
        None => {
            file
        },
    };
    let home = match std::os::homedir() {
        Some(p) => { let s = p.as_str(); String::from_str(s.unwrap()) },
        None => panic!("No home directory identified"),
    };
    let dir = format!("{}/.rusted/", home);
    let dirpath = Path::new(dir);
    match std::io::fs::mkdir_recursive(&dirpath, std::io::USER_RWX) {
        Ok(_) => {},
        Err(e) => panic!("Failed to create cache directory: {}", e),
    }
    format!("{}/{}", dirpath.as_str().unwrap(), stem)
}

fn get_filename(script_path: &str) -> String {
    String::from_str(Path::new(script_path).filename_str().unwrap())
}

fn main() {
    let args = os::args();
    let arg = match args.get(1) {
        Some(a) => a,
        None => panic!("No argument supplied"),
    };
    let cleaned = get_filename(arg.as_slice());
    let src = Path::new(String::from_str(arg.as_slice()));
    let target = get_cache_path(Path::new(cleaned.as_slice()));
    let sysroot = Path::new("/usr/local");
    let cached = has_cache(arg.as_slice(), target.as_slice());
    if !cached {
        compiler::compile_file(src, Path::new(target.as_slice()), Some(sysroot));
    }

    let mut cmd = Command::new(target);
    let mut argsvec: Vec<&String> = Vec::new();
    for a in args.iter() {
        argsvec.push(a);
    }
    for i in range(1, argsvec.len()) {
        cmd.arg(argsvec[i]);
    }
    cmd.stdout(process::StdioContainer::InheritFd(libc::STDOUT_FILENO));
    cmd.stderr(process::StdioContainer::InheritFd(libc::STDERR_FILENO));
    cmd.stdin(process::StdioContainer::InheritFd(libc::STDIN_FILENO));
    match cmd.spawn() {
        Ok(_) => (),
        Err(e) => panic!("Failed to execute process: {}", e),
    };
}
<|endoftext|>