<commit_before>#![feature(duration_float)]
#![feature(euclidean_division)]
#![feature(self_struct_ctor)]
#![feature(try_trait)]

use std::path::PathBuf;

use chrono::Local;

use glium::backend::glutin::headless::Headless;
use glium::glutin::{
    ElementState, Event, HeadlessRendererBuilder, KeyboardInput, VirtualKeyCode, WindowEvent,
};
use glium::Surface;

mod aabb;
mod bvh;
mod camera;
mod color;
mod config;
mod consts;
mod gl_renderer;
mod index_ptr;
mod input;
mod intersect;
mod light;
mod load;
mod material;
mod mesh;
mod obj_load;
mod pt_renderer;
mod scene;
mod stats;
mod texture;
mod triangle;
mod util;
mod vertex;

use self::config::RenderConfig;
use self::gl_renderer::GLRenderer;
use self::input::InputState;
use self::pt_renderer::PTRenderer;

type Float = f64;

fn main() {
    match std::env::args().nth(1).as_ref().map(|s| s.as_str()) {
        Some(_) => benchmark(),
        None => online_render(),
    }
}

fn benchmark() {
    let scenes = [
        "plane",
        "cornell-glossy",
        "cornell-water",
        "indirect",
        "conference",
        "sponza",
    ];
    let config = RenderConfig::benchmark();
    let root_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
    let output_dir = root_dir.join("results");
    std::fs::create_dir_all(output_dir.clone()).unwrap();
    let time_stamp = Local::now().format("%F_%H%M%S").to_string();

    // Initialize an OpenGL context that is needed for post-processing
    let events_loop = glium::glutin::EventsLoop::new();
    // Preferably this wouldn't need use a window at all but alas this is the closest I have gotten.
    // There exists HeadlessContext but that still pops up a window (atleast on Windows).
    let window = glium::glutin::WindowBuilder::new()
        .with_dimensions(glium::glutin::dpi::LogicalSize::new(0.0, 0.0))
        .with_visibility(false)
        .with_decorations(false)
        .with_title("Benchmark");
    let context = glium::glutin::ContextBuilder::new();
    let display = glium::Display::new(window, context, &events_loop).unwrap();

    for scene_name in &scenes {
        stats::new_scene(scene_name);
        let _t = stats::time("Total");
        println!("{}...", scene_name);
        let (scene, camera) = load::load_cpu_scene(scene_name, &config);
        let pt_renderer = PTRenderer::offline_render(&display, &scene, &camera, &config);

        stats::time("Post-process");
        // Save timestamped version in a addition to the default image
        let scene_dir = output_dir.join(scene_name);
        std::fs::create_dir_all(scene_dir.clone()).unwrap();
        let timestamped_image = scene_dir.join(format!("{}_{}.png", scene_name, time_stamp));
        let default_image = output_dir.join(scene_name).with_extension("png");
        pt_renderer.save_image(&timestamped_image);
        pt_renderer.save_image(&default_image);
    }
    let stats_file = output_dir.join(format!("benchmark_{}.txt", time_stamp));
    stats::print_and_save(&stats_file);
}

fn online_render() {
    let mut config = RenderConfig::default();
    let mut events_loop = glium::glutin::EventsLoop::new();
    let window = glium::glutin::WindowBuilder::new()
        .with_dimensions(config.dimensions())
        .with_resizable(false); // TODO: enable resizing
    let context = glium::glutin::ContextBuilder::new().with_depth_buffer(24);
    let display =
        glium::Display::new(window, context, &events_loop).expect("Failed to create display");

    let (mut scene, mut gpu_scene, mut camera) =
        load::load_gpu_scene(VirtualKeyCode::Key1, &display, &config).unwrap();
    let gl_renderer = GLRenderer::new(&display);
    let mut pt_renderer: Option<PTRenderer> = None;

    let mut input = InputState::new();
    let mut quit = false;

    loop {
        let mut target = display.draw();
        target.clear_color_and_depth((0.0, 0.0, 0.0, 1.0), 1.0);
        if let Some(renderer) = &mut pt_renderer {
            renderer.update_and_render(&mut target);
        } else {
            gl_renderer.render(&mut target, &gpu_scene, &camera);
        }
        target.finish().unwrap();

        events_loop.poll_events(|event| {
            input.update(&event);
            match event {
                Event::WindowEvent {
                    event: WindowEvent::KeyboardInput { input, .. },
                    ..
                } => match input {
                    KeyboardInput {
                        state: ElementState::Pressed,
                        virtual_keycode: Some(VirtualKeyCode::Space),
                        ..
                    } => {
                        if pt_renderer.is_some() {
                            pt_renderer = None;
                        } else {
                            pt_renderer =
                                Some(PTRenderer::start_render(&display, &scene, &camera, &config));
                        }
                    }
                    KeyboardInput {
                        state: ElementState::Pressed,
                        virtual_keycode: Some(keycode),
                        ..
                    } => {
                        if pt_renderer.is_none() {
                            if let Some(res) = load::load_gpu_scene(keycode, &display, &config) {
                                scene = res.0;
                                gpu_scene = res.1;
                                camera = res.2;
                            }
                            config.handle_key(keycode);
                        }
                    }
                    _ => (),
                },
                Event::WindowEvent {
                    event: WindowEvent::CloseRequested,
                    ..
                } => quit = true,
                _ => (),
            }
        });
        camera.process_input(&input);
        input.reset_deltas();
        if quit {
            return;
        }
    }
}
<commit_msg>Copy the image to the output dir instead of saving twice<commit_after>#![feature(duration_float)]
#![feature(euclidean_division)]
#![feature(self_struct_ctor)]
#![feature(try_trait)]

use std::path::PathBuf;

use chrono::Local;

use glium::backend::glutin::headless::Headless;
use glium::glutin::{
    ElementState, Event, HeadlessRendererBuilder, KeyboardInput, VirtualKeyCode, WindowEvent,
};
use glium::Surface;

mod aabb;
mod bvh;
mod camera;
mod color;
mod config;
mod consts;
mod gl_renderer;
mod index_ptr;
mod input;
mod intersect;
mod light;
mod load;
mod material;
mod mesh;
mod obj_load;
mod pt_renderer;
mod scene;
mod stats;
mod texture;
mod triangle;
mod util;
mod vertex;

use self::config::RenderConfig;
use self::gl_renderer::GLRenderer;
use self::input::InputState;
use self::pt_renderer::PTRenderer;

type Float = f64;

fn main() {
    match std::env::args().nth(1).as_ref().map(|s| s.as_str()) {
        Some(_) => benchmark(),
        None => online_render(),
    }
}

fn benchmark() {
    let scenes = [
        "plane",
        "cornell-glossy",
        "cornell-water",
        "indirect",
        "conference",
        "sponza",
    ];
    let config = RenderConfig::benchmark();
    let root_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
    let output_dir = root_dir.join("results");
    std::fs::create_dir_all(output_dir.clone()).unwrap();
    let time_stamp = Local::now().format("%F_%H%M%S").to_string();

    // Initialize an OpenGL context that is needed for post-processing
    let events_loop = glium::glutin::EventsLoop::new();
    // Preferably this wouldn't need use a window at all but alas this is the closest I have gotten.
    // There exists HeadlessContext but that still pops up a window (atleast on Windows).
    // TODO: Maybe change this such that the window displays the current render?
    let window = glium::glutin::WindowBuilder::new()
        .with_dimensions(glium::glutin::dpi::LogicalSize::new(0.0, 0.0))
        .with_visibility(false)
        .with_decorations(false)
        .with_title("Benchmark");
    let context = glium::glutin::ContextBuilder::new();
    let display = glium::Display::new(window, context, &events_loop).unwrap();

    for scene_name in &scenes {
        stats::new_scene(scene_name);
        let _t = stats::time("Total");
        println!("{}...", scene_name);
        let (scene, camera) = load::load_cpu_scene(scene_name, &config);
        let pt_renderer = PTRenderer::offline_render(&display, &scene, &camera, &config);

        stats::time("Post-process");
        let scene_dir = output_dir.join(scene_name);
        std::fs::create_dir_all(scene_dir.clone()).unwrap();
        let timestamped_image = scene_dir.join(format!("{}_{}.png", scene_name, time_stamp));
        pt_renderer.save_image(&timestamped_image);
        // Make a copy to the main output directory
        let default_image = output_dir.join(scene_name).with_extension("png");
        std::fs::copy(timestamped_image, default_image).unwrap();
    }
    let stats_file = output_dir.join(format!("benchmark_{}.txt", time_stamp));
    stats::print_and_save(&stats_file);
}

fn online_render() {
    let mut config = RenderConfig::default();
    let mut events_loop = glium::glutin::EventsLoop::new();
    let window = glium::glutin::WindowBuilder::new()
        .with_dimensions(config.dimensions())
        .with_resizable(false); // TODO: enable resizing
    let context = glium::glutin::ContextBuilder::new().with_depth_buffer(24);
    let display =
        glium::Display::new(window, context, &events_loop).expect("Failed to create display");

    let (mut scene, mut gpu_scene, mut camera) =
        load::load_gpu_scene(VirtualKeyCode::Key1, &display, &config).unwrap();
    let gl_renderer = GLRenderer::new(&display);
    let mut pt_renderer: Option<PTRenderer> = None;

    let mut input = InputState::new();
    let mut quit = false;

    loop {
        let mut target = display.draw();
        target.clear_color_and_depth((0.0, 0.0, 0.0, 1.0), 1.0);
        if let Some(renderer) = &mut pt_renderer {
            renderer.update_and_render(&mut target);
        } else {
            gl_renderer.render(&mut target, &gpu_scene, &camera);
        }
        target.finish().unwrap();

        events_loop.poll_events(|event| {
            input.update(&event);
            match event {
                Event::WindowEvent {
                    event: WindowEvent::KeyboardInput { input, .. },
                    ..
                } => match input {
                    KeyboardInput {
                        state: ElementState::Pressed,
                        virtual_keycode: Some(VirtualKeyCode::Space),
                        ..
                    } => {
                        if pt_renderer.is_some() {
                            pt_renderer = None;
                        } else {
                            pt_renderer =
                                Some(PTRenderer::start_render(&display, &scene, &camera, &config));
                        }
                    }
                    KeyboardInput {
                        state: ElementState::Pressed,
                        virtual_keycode: Some(keycode),
                        ..
                    } => {
                        if pt_renderer.is_none() {
                            if let Some(res) = load::load_gpu_scene(keycode, &display, &config) {
                                scene = res.0;
                                gpu_scene = res.1;
                                camera = res.2;
                            }
                            config.handle_key(keycode);
                        }
                    }
                    _ => (),
                },
                Event::WindowEvent {
                    event: WindowEvent::CloseRequested,
                    ..
                } => quit = true,
                _ => (),
            }
        });
        camera.process_input(&input);
        input.reset_deltas();
        if quit {
            return;
        }
    }
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate diesel;
#[macro_use]
extern crate num_derive;

mod config;
mod database;
mod err;
mod feed;
mod path;

use crate::feed::stats::{ListenerAvg, ListenerStatMap, ListenerStats};
use crate::feed::{Feed, FeedNotif};
use anyhow::{anyhow, Context, Result};
use chrono::{DateTime, Datelike, Duration, Local, Timelike, Utc};
use config::Config;
use database::Database;
use diesel::prelude::*;
use once_cell::sync::Lazy;
use parking_lot::{Condvar, Mutex};
use smallvec::SmallVec;
use std::sync::{mpsc, Arc};
use std::thread;

struct CmdOptions {
    reload_config: bool,
}

impl CmdOptions {
    fn from_env() -> Result<Self> {
        let mut args = pico_args::Arguments::from_env();

        if args.contains(["-h", "--help"]) {
            Self::print_help();
        }

        let result = Self {
            reload_config: args.contains(["-r", "--reload"]),
        };

        Ok(result)
    }

    fn print_help() {
        println!(concat!("Usage: ", env!("CARGO_PKG_NAME"), " [OPTIONS]\n"));

        println!("Optional arguments:");
        println!("  -h, --help    show this message");
        println!("  -r, --reload  reload the configuration file on each update");

        std::process::exit(0);
    }
}

fn main() -> Result<()> {
    let args = match CmdOptions::from_env() {
        Ok(args) => args,
        Err(err) => {
            err::error_notif(&err);
            return Err(err);
        }
    };

    let result = run(args);

    if let Err(err) = &result {
        err::error_notif(err);
    }

    result
}

fn run(args: CmdOptions) -> Result<()> {
    let config = {
        let cfg = Config::load_or_new().context("failed to load / create config")?;
        Arc::new(Mutex::new(cfg))
    };

    let db = Database::open().context("failed to open feed database")?;

    let mut listener_stats = ListenerStatMap::with_capacity(200);
    let mut remove_old_feeds_time = Utc::now();

    let event_rx = Event::init_threads(&config).context("failed to init event threads")?;

    loop {
        match event_rx.recv() {
            Ok(Event::RunUpdate) => {
                let cur_time = Utc::now();
                let mut config = config.lock();

                if args.reload_config {
                    match Config::load() {
                        Ok(new) => *config = new,
                        Err(err) => err::error_notif(&err),
                    }
                }

                let result = run_update(&db, &config, &cur_time, &mut listener_stats).and_then(
                    |mut notifs| {
                        FeedNotif::sort_all(&mut notifs, &config);
                        FeedNotif::show_all(&notifs)
                    },
                );

                if let Err(err) = result {
                    err::error_notif(&err);
                }

                if cur_time >= remove_old_feeds_time {
                    ListenerAvg::remove_old_from_db(&db)?;
                    remove_old_feeds_time = cur_time + Duration::hours(12);
                }
            }
            Ok(Event::Exit) => break Ok(()),
            Err(err) => break Err(err.into()),
        }
    }
}

enum Event {
    RunUpdate,
    Exit,
}

impl Event {
    fn init_threads(config: &Arc<Mutex<Config>>) -> Result<mpsc::Receiver<Self>> {
        let (tx, rx) = mpsc::channel();

        Self::spawn_update_thread(tx.clone(), config);
        Self::spawn_signal_handler(tx).context("signal handler spawn failed")?;

        Ok(rx)
    }

    fn spawn_update_thread(
        tx: mpsc::Sender<Self>,
        config: &Arc<Mutex<Config>>,
    ) -> thread::JoinHandle<()> {
        let config = config.clone();

        // This thread should die if something goes horribly wrong, so the uses of unwrap() are intended here
        thread::spawn(move || loop {
            let update_time = {
                let config = config.lock();
                (config.misc.update_time_mins * 60.0) as u64
            };

            if tx.send(Event::RunUpdate).is_err() {
                break;
            }

            thread::sleep(std::time::Duration::from_secs(update_time));
        })
    }

    fn spawn_signal_handler(tx: mpsc::Sender<Self>) -> Result<()> {
        use nix::sys::signal::{signal, SigHandler, Signal};

        static SIG_TRIGGER: Lazy<(Mutex<()>, Condvar)> =
            Lazy::new(|| (Mutex::new(()), Condvar::new()));

        extern "C" fn handle_sig(_: libc::c_int) {
            let (_, cvar) = &*SIG_TRIGGER;
            cvar.notify_one();
        }

        let handler = SigHandler::Handler(handle_sig);
        let sigs = [Signal::SIGHUP, Signal::SIGTERM, Signal::SIGINT];

        unsafe {
            for &sig in &sigs {
                signal(sig, handler)
                    .map_err(|err| anyhow!("failed to register signal handler: {}", err))?;
            }
        }

        thread::spawn(move || {
            let (lock, cvar) = &*SIG_TRIGGER;
            cvar.wait(&mut lock.lock());
            tx.send(Event::Exit).ok();
        });

        Ok(())
    }
}

fn run_update<'a>(
    db: &Database,
    config: &Config,
    cur_time: &DateTime<Utc>,
    listener_stats: &mut ListenerStatMap,
) -> Result<SmallVec<[FeedNotif<'a>; 3]>> {
    use diesel::result::Error;

    let feeds = {
        let mut feeds = Feed::scrape_all(config).context("feed scraping failed")?;
        filter_feeds(config, &mut feeds);
        feeds
    };

    let cur_hour = cur_time.hour() as u8;
    let cur_weekday = Local::today().weekday();

    let mut display = SmallVec::new();

    db.conn()
        .transaction::<_, Error, _>(|| {
            for feed in feeds {
                let stats = listener_stats.entry(feed.id).or_insert_with(|| {
                    ListenerStats::init_from_db(db, cur_hour, feed.id as i32, feed.listeners as f32)
                });

                stats.update(cur_hour, &feed, config, cur_weekday);
                stats.save_to_db(db)?;

                if !stats.should_display_feed(&feed, config) {
                    continue;
                }

                if display.len() > config.misc.show_max as usize {
                    continue;
                }

                display.push(FeedNotif::new(feed, stats));
            }

            Ok(())
        })
        .context("database transaction failed")?;

    Ok(display)
}

fn filter_feeds(config: &Config, feeds: &mut Vec<Feed>) {
    if !config.filters.whitelist.is_empty() {
        feeds.retain(|feed| {
            config
                .filters
                .whitelist
                .iter()
                .any(|entry| entry.matches_feed(feed))
        });
    }

    if !config.filters.blacklist.is_empty() {
        feeds.retain(|feed| {
            config
                .filters
                .blacklist
                .iter()
                .any(|entry| !entry.matches_feed(feed))
        });
    }
}
<commit_msg>Remove unnecessary Result return type from CmdOptions::from_env()<commit_after>#[macro_use]
extern crate diesel;
#[macro_use]
extern crate num_derive;

mod config;
mod database;
mod err;
mod feed;
mod path;

use crate::feed::stats::{ListenerAvg, ListenerStatMap, ListenerStats};
use crate::feed::{Feed, FeedNotif};
use anyhow::{anyhow, Context, Result};
use chrono::{DateTime, Datelike, Duration, Local, Timelike, Utc};
use config::Config;
use database::Database;
use diesel::prelude::*;
use once_cell::sync::Lazy;
use parking_lot::{Condvar, Mutex};
use smallvec::SmallVec;
use std::sync::{mpsc, Arc};
use std::thread;

struct CmdOptions {
    reload_config: bool,
}

impl CmdOptions {
    fn from_env() -> Self {
        let mut args = pico_args::Arguments::from_env();

        if args.contains(["-h", "--help"]) {
            Self::print_help();
        }

        Self {
            reload_config: args.contains(["-r", "--reload"]),
        }
    }

    fn print_help() {
        println!(concat!("Usage: ", env!("CARGO_PKG_NAME"), " [OPTIONS]\n"));

        println!("Optional arguments:");
        println!("  -h, --help    show this message");
        println!("  -r, --reload  reload the configuration file on each update");

        std::process::exit(0);
    }
}

fn main() -> Result<()> {
    let args = CmdOptions::from_env();
    let result = run(args);

    if let Err(err) = &result {
        err::error_notif(err);
    }

    result
}

fn run(args: CmdOptions) -> Result<()> {
    let config = {
        let cfg = Config::load_or_new().context("failed to load / create config")?;
        Arc::new(Mutex::new(cfg))
    };

    let db = Database::open().context("failed to open feed database")?;

    let mut listener_stats = ListenerStatMap::with_capacity(200);
    let mut remove_old_feeds_time = Utc::now();

    let event_rx = Event::init_threads(&config).context("failed to init event threads")?;

    loop {
        match event_rx.recv() {
            Ok(Event::RunUpdate) => {
                let cur_time = Utc::now();
                let mut config = config.lock();

                if args.reload_config {
                    match Config::load() {
                        Ok(new) => *config = new,
                        Err(err) => err::error_notif(&err),
                    }
                }

                let result = run_update(&db, &config, &cur_time, &mut listener_stats).and_then(
                    |mut notifs| {
                        FeedNotif::sort_all(&mut notifs, &config);
                        FeedNotif::show_all(&notifs)
                    },
                );

                if let Err(err) = result {
                    err::error_notif(&err);
                }

                if cur_time >= remove_old_feeds_time {
                    ListenerAvg::remove_old_from_db(&db)?;
                    remove_old_feeds_time = cur_time + Duration::hours(12);
                }
            }
            Ok(Event::Exit) => break Ok(()),
            Err(err) => break Err(err.into()),
        }
    }
}

enum Event {
    RunUpdate,
    Exit,
}

impl Event {
    fn init_threads(config: &Arc<Mutex<Config>>) -> Result<mpsc::Receiver<Self>> {
        let (tx, rx) = mpsc::channel();

        Self::spawn_update_thread(tx.clone(), config);
        Self::spawn_signal_handler(tx).context("signal handler spawn failed")?;

        Ok(rx)
    }

    fn spawn_update_thread(
        tx: mpsc::Sender<Self>,
        config: &Arc<Mutex<Config>>,
    ) -> thread::JoinHandle<()> {
        let config = config.clone();

        // This thread should die if something goes horribly wrong, so the uses of unwrap() are intended here
        thread::spawn(move || loop {
            let update_time = {
                let config = config.lock();
                (config.misc.update_time_mins * 60.0) as u64
            };

            if tx.send(Event::RunUpdate).is_err() {
                break;
            }

            thread::sleep(std::time::Duration::from_secs(update_time));
        })
    }

    fn spawn_signal_handler(tx: mpsc::Sender<Self>) -> Result<()> {
        use nix::sys::signal::{signal, SigHandler, Signal};

        static SIG_TRIGGER: Lazy<(Mutex<()>, Condvar)> =
            Lazy::new(|| (Mutex::new(()), Condvar::new()));

        extern "C" fn handle_sig(_: libc::c_int) {
            let (_, cvar) = &*SIG_TRIGGER;
            cvar.notify_one();
        }

        let handler = SigHandler::Handler(handle_sig);
        let sigs = [Signal::SIGHUP, Signal::SIGTERM, Signal::SIGINT];

        unsafe {
            for &sig in &sigs {
                signal(sig, handler)
                    .map_err(|err| anyhow!("failed to register signal handler: {}", err))?;
            }
        }

        thread::spawn(move || {
            let (lock, cvar) = &*SIG_TRIGGER;
            cvar.wait(&mut lock.lock());
            tx.send(Event::Exit).ok();
        });

        Ok(())
    }
}

fn run_update<'a>(
    db: &Database,
    config: &Config,
    cur_time: &DateTime<Utc>,
    listener_stats: &mut ListenerStatMap,
) -> Result<SmallVec<[FeedNotif<'a>; 3]>> {
    use diesel::result::Error;

    let feeds = {
        let mut feeds = Feed::scrape_all(config).context("feed scraping failed")?;
        filter_feeds(config, &mut feeds);
        feeds
    };

    let cur_hour = cur_time.hour() as u8;
    let cur_weekday = Local::today().weekday();

    let mut display = SmallVec::new();

    db.conn()
        .transaction::<_, Error, _>(|| {
            for feed in feeds {
                let stats = listener_stats.entry(feed.id).or_insert_with(|| {
                    ListenerStats::init_from_db(db, cur_hour, feed.id as i32, feed.listeners as f32)
                });

                stats.update(cur_hour, &feed, config, cur_weekday);
                stats.save_to_db(db)?;

                if !stats.should_display_feed(&feed, config) {
                    continue;
                }

                if display.len() > config.misc.show_max as usize {
                    continue;
                }

                display.push(FeedNotif::new(feed, stats));
            }

            Ok(())
        })
        .context("database transaction failed")?;

    Ok(display)
}

fn filter_feeds(config: &Config, feeds: &mut Vec<Feed>) {
    if !config.filters.whitelist.is_empty() {
        feeds.retain(|feed| {
            config
                .filters
                .whitelist
                .iter()
                .any(|entry| entry.matches_feed(feed))
        });
    }

    if !config.filters.blacklist.is_empty() {
        feeds.retain(|feed| {
            config
                .filters
                .blacklist
                .iter()
                .any(|entry| !entry.matches_feed(feed))
        });
    }
}
<|endoftext|>


<commit_before>#![feature(phase)]

extern crate backbonzo;
extern crate serialize;
extern crate docopt;
extern crate time;
#[phase(plugin)] extern crate docopt_macros;

use docopt::Docopt;
use backbonzo::{init, BackupManager, BonzoError, BonzoResult};

docopt!(Args deriving Show, "
backbonzo

Usage:
  backbonzo OPERATION [options] --key <key>
  backbonzo (-h | --help)

Operations:
  init, backup, restore

Options:
  -s --source=<source>      Source directory [default: ./].
  -d --destination=<dest>   Backup directory [default: /tmp/backbonzo/].
  -k --key=<key>            Encryption key.
  -b --blocksize=<bs>       Size of blocks in megabytes [default: 1].
", arg_OPERATION: Operation, flag_blocksize: uint, flag_key: String)

static DATABASE_FILENAME: &'static str = "index.db3";

#[deriving(Show, Decodable)]
enum Operation {
    Init,
    Backup,
    Restore
}

fn main() {
    let args: Args = Args::docopt().decode().unwrap_or_else(|e| e.exit());

    let source_path = Path::new(args.flag_source);
    let output_path = Path::new(args.flag_destination);
    let block_bytes = 1000 * 1000 * args.flag_blocksize;

    let mut database_path = source_path.clone();
    database_path.push(DATABASE_FILENAME);

    let result = match args.arg_OPERATION {
        Operation::Init    => {
            init(&database_path, args.flag_key)
        },
        Operation::Restore => {
            let manager = BackupManager::new(database_path, source_path, output_path, block_bytes, args.flag_key).ok().unwrap();

            manager.restore(time::get_time().sec as u64)
        },
        Operation::Backup  => {
            let manager = BackupManager::new(database_path, source_path, output_path, block_bytes, args.flag_key).ok().unwrap();

            manager.update()
        }
    };

    handle_result(result);
}

fn handle_result<T>(result: BonzoResult<T>) {
    match result {
        Ok(..)                       => println!("Done!"),
        Err(BonzoError::Database(e)) => println!("Database error: {}", e.message),
        Err(BonzoError::Io(e))       => println!("IO error: {}", e.desc),
        Err(BonzoError::Crypto(..))  => println!("Crypto error!"),
        Err(BonzoError::Other(str))  => println!("Error: {}", str)
    }
}
<commit_msg>add semicolon to item macro<commit_after>#![feature(phase)]

extern crate backbonzo;
extern crate serialize;
extern crate docopt;
extern crate time;
#[phase(plugin)] extern crate docopt_macros;

use docopt::Docopt;
use backbonzo::{init, BackupManager, BonzoError, BonzoResult};

docopt!(Args deriving Show, "
backbonzo

Usage:
  backbonzo OPERATION [options] --key <key>
  backbonzo (-h | --help)

Operations:
  init, backup, restore

Options:
  -s --source=<source>      Source directory [default: ./].
  -d --destination=<dest>   Backup directory [default: /tmp/backbonzo/].
  -k --key=<key>            Encryption key.
  -b --blocksize=<bs>       Size of blocks in megabytes [default: 1].
", arg_OPERATION: Operation, flag_blocksize: uint, flag_key: String);

static DATABASE_FILENAME: &'static str = "index.db3";

#[deriving(Show, Decodable)]
enum Operation {
    Init,
    Backup,
    Restore
}

fn main() {
    let args: Args = Args::docopt().decode().unwrap_or_else(|e| e.exit());

    let source_path = Path::new(args.flag_source);
    let output_path = Path::new(args.flag_destination);
    let block_bytes = 1000 * 1000 * args.flag_blocksize;

    let mut database_path = source_path.clone();
    database_path.push(DATABASE_FILENAME);

    let result = match args.arg_OPERATION {
        Operation::Init    => {
            init(&database_path, args.flag_key)
        },
        Operation::Restore => {
            let manager = BackupManager::new(database_path, source_path, output_path, block_bytes, args.flag_key).ok().unwrap();

            manager.restore(time::get_time().sec as u64)
        },
        Operation::Backup  => {
            let manager = BackupManager::new(database_path, source_path, output_path, block_bytes, args.flag_key).ok().unwrap();

            manager.update()
        }
    };

    handle_result(result);
}

fn handle_result<T>(result: BonzoResult<T>) {
    match result {
        Ok(..)                       => println!("Done!"),
        Err(BonzoError::Database(e)) => println!("Database error: {}", e.message),
        Err(BonzoError::Io(e))       => println!("IO error: {}", e.desc),
        Err(BonzoError::Crypto(..))  => println!("Crypto error!"),
        Err(BonzoError::Other(str))  => println!("Error: {}", str)
    }
}
<|endoftext|>


<commit_before>extern crate clap;
use clap::{Arg, App, AppSettings, SubCommand};
use slog::{LevelFilter, Level, DrainExt, Logger};
use yaml_rust::YamlLoader;
use std::io::prelude::*;
use std::fs::File;
use yaml_rust::{Yaml, YamlEmitter};
use std::error::Error;
use std::path::Path;

#[macro_use]
extern crate slog;
extern crate slog_term;
extern crate slog_stdlog;
extern crate yaml_rust;
extern crate crypto;

#[cfg(test)]
extern crate spectral;

mod model;
mod config;
mod checksum;
mod files;
mod mutate;
mod import;

fn main() {
  let matches: clap::ArgMatches = create_app().get_matches();
  let stream = slog_term::streamer().full().build();
  let verbose: bool = matches.is_present("verbose");
  let force: bool = matches.is_present("force");
  let log = if verbose {
    slog::Logger::root(stream.fuse(), o!())
  } else {
    slog::Logger::root(LevelFilter::new(stream, Level::Info).fuse(), o!())
  };

  let config_file = matches.value_of("config_file").unwrap();
  info!(log, "Starting engine"; "config_file" => config_file);

  let config = match load_config_file(&log, config_file) {
    Ok(content) => content,
    Err(e) => {
      error!(log, "Failed to load config file."; "error" => e.description());
      panic!(1)
    }
  };

  let mut yaml_documents = match YamlLoader::load_from_str(&config) {
    Ok(yaml) => yaml,
    Err(e) => {
      error!(log, "Failed to parse config file."; "error" => e.description());
      panic!(2)
    }
  };
  if yaml_documents.is_empty() {
    let s = "files:";
    yaml_documents = YamlLoader::load_from_str(s).unwrap();
  }

  let maybe_ln_matches = matches.subcommand_matches("ln");
  let maybe_cp_matches = matches.subcommand_matches("cp");
  let maybe_apply_matches = matches.subcommand_matches("apply");
  let maybe_import_matches = matches.subcommand_matches("import");
  if maybe_apply_matches.is_some() {
    let yaml_config = &yaml_documents[0];
    let dot_files: &Yaml = &yaml_config["files"];
    info!(log, "Liftoff! Applying configuration.");
    files::process_dot_files(&log, dot_files, force);
  } else if let Some(ln_matches) = maybe_ln_matches {
    let yaml_config = &yaml_documents[0];
    let link_target = ln_matches.value_of("link_target").unwrap();
    let link_name = ln_matches.value_of("link_name").unwrap();
    let log = log.new(o!("link_target" => link_target.to_string(), "link_name" => link_name.to_string()));
    info!(log, "Liftoff! Adding new link to configuration");
    let new_config = mutate::add_dotfiles_to_config(&log,
                                                   yaml_config,
                                                   &[model::DotFile {
                                                     target: link_name.to_string(),
                                                     source: link_target.to_string(),
                                                     dot_file_type: model::DotFileType::LINK,
                                                   }]);
    write_new_yaml(&log, &new_config, config_file);
  } else if let Some(cp_matches) = maybe_cp_matches {
    let yaml_config = &yaml_documents[0];
    let target = cp_matches.value_of("target").unwrap();
    let source = cp_matches.value_of("source").unwrap();
    let log = log.new(o!("target" => target.to_string(), "source" => source.to_string()));
    info!(log, "Liftoff! Adding new copy to configuration");
    let new_config = mutate::add_dotfiles_to_config(&log,
                                                   yaml_config,
                                                   &[model::DotFile {
                                                     target: target.to_string(),
                                                     source: source.to_string(),
                                                     dot_file_type: model::DotFileType::COPY,
                                                   }]);
    write_new_yaml(&log, &new_config, config_file);
  } else if let Some(import_matches) = maybe_import_matches {
    let  dir =  import_matches.value_of("dir").unwrap();
    let log = log.new(o!("directory" => dir.to_string()));
    let yaml_config = &yaml_documents[0];
    info!(log, "Liftoff! Importing to configuration");
    let dot_files = import::scan_dir(&log, &dir);
    mutate::add_dotfiles_to_config(&log, &yaml_config, &dot_files);
  }
}

fn write_new_yaml(log: &Logger, document: &Yaml, config_file: &str) {
  let mut out_str = String::new();
  {
    let mut emitter = YamlEmitter::new(&mut out_str);
    match emitter.dump(document) {
      Ok(_) => {}
      Err(e) => {
        error!(log, "Failed to write config file."; "error" => format!("{:?}", e));
        return;
      }
    }
  }
  out_str.push_str("\n");
  match write_config_file(config_file, &out_str) {
    Ok(_) => info!(log, "Successfully written configuration"),
    Err(e) => error!(log, "Failed to write config file."; "error" => e.description()),
  };
}

fn write_config_file(file: &str, content: &str) -> Result<(), std::io::Error> {
  let mut file = try!(File::create(file));
  try!(file.write_all(content.as_bytes()));
  Ok(())
}

fn load_config_file(log: &Logger, file: &str) -> Result<String, std::io::Error> {
  let path = Path::new(file);
  if path.exists() {
    let mut file = try!(File::open(file));
    let mut content = String::new();
    try!(file.read_to_string(&mut content));
    Ok(content)
  } else {
    warn!(log, "Configuration doesn't exit"; "file" => file);
    Ok("".to_string())
  }
}

fn create_app<'a>() -> App<'a, 'a> {
  App::new("dotcopter")
    .version("0.1")
    .setting(AppSettings::ColoredHelp)
    .author("Patrick Haun <bomgar85@googlemail.com>")
    .about("manages dotfiles installation")
    .arg(Arg::with_name("force").short("f").long("force").takes_value(false))
    .arg(Arg::with_name("verbose")
           .long("verbose")
           .short("v")
           .help("debug output")
           .required(false)
           .takes_value(false))
    .arg(Arg::with_name("config_file").required(true))
    .subcommand(SubCommand::with_name("apply").about("applies a dotfile configuration"))
    .subcommand(SubCommand::with_name("ln")
                  .about("adds new link to configuration")
                  .arg(Arg::with_name("link_target").required(true))
                  .arg(Arg::with_name("link_name").required(true)))
    .subcommand(SubCommand::with_name("cp")
                  .about("adds new copy to configuration")
                  .arg(Arg::with_name("source").required(true))
                  .arg(Arg::with_name("target").required(true)))
    .subcommand(SubCommand::with_name("import")
                  .about("imports dotfiles from a folder into the configuration")
                  .arg(Arg::with_name("dir").required(true)))
}
<commit_msg>remove some clippy errors<commit_after>extern crate clap;
use clap::{Arg, App, AppSettings, SubCommand};
use slog::{LevelFilter, Level, DrainExt, Logger};
use yaml_rust::YamlLoader;
use std::io::prelude::*;
use std::fs::File;
use yaml_rust::{Yaml, YamlEmitter};
use std::error::Error;
use std::path::Path;

#[macro_use]
extern crate slog;
extern crate slog_term;
extern crate slog_stdlog;
extern crate yaml_rust;
extern crate crypto;

#[cfg(test)]
extern crate spectral;

mod model;
mod config;
mod checksum;
mod files;
mod mutate;
mod import;

fn main() {
  let matches: clap::ArgMatches = create_app().get_matches();
  let stream = slog_term::streamer().full().build();
  let verbose: bool = matches.is_present("verbose");
  let force: bool = matches.is_present("force");
  let log = if verbose {
    slog::Logger::root(stream.fuse(), o!())
  } else {
    slog::Logger::root(LevelFilter::new(stream, Level::Info).fuse(), o!())
  };

  let config_file = matches.value_of("config_file").unwrap();
  info!(log, "Starting engine"; "config_file" => config_file);

  let config = match load_config_file(&log, config_file) {
    Ok(content) => content,
    Err(e) => {
      error!(log, "Failed to load config file."; "error" => e.description());
      panic!(1)
    }
  };

  let mut yaml_documents = match YamlLoader::load_from_str(&config) {
    Ok(yaml) => yaml,
    Err(e) => {
      error!(log, "Failed to parse config file."; "error" => e.description());
      panic!(2)
    }
  };
  if yaml_documents.is_empty() {
    let s = "files:";
    yaml_documents = YamlLoader::load_from_str(s).unwrap();
  }

  let maybe_ln_matches = matches.subcommand_matches("ln");
  let maybe_cp_matches = matches.subcommand_matches("cp");
  let maybe_apply_matches = matches.subcommand_matches("apply");
  let maybe_import_matches = matches.subcommand_matches("import");
  if maybe_apply_matches.is_some() {
    let yaml_config = &yaml_documents[0];
    let dot_files: &Yaml = &yaml_config["files"];
    info!(log, "Liftoff! Applying configuration.");
    files::process_dot_files(&log, dot_files, force);
  } else if let Some(ln_matches) = maybe_ln_matches {
    let yaml_config = &yaml_documents[0];
    let link_target = ln_matches.value_of("link_target").unwrap();
    let link_name = ln_matches.value_of("link_name").unwrap();
    let log = log.new(o!("link_target" => link_target.to_string(), "link_name" => link_name.to_string()));
    info!(log, "Liftoff! Adding new link to configuration");
    let new_config = mutate::add_dotfiles_to_config(&log,
                                                   yaml_config,
                                                   &[model::DotFile {
                                                     target: link_name.to_string(),
                                                     source: link_target.to_string(),
                                                     dot_file_type: model::DotFileType::LINK,
                                                   }]);
    write_new_yaml(&log, &new_config, config_file);
  } else if let Some(cp_matches) = maybe_cp_matches {
    let yaml_config = &yaml_documents[0];
    let target = cp_matches.value_of("target").unwrap();
    let source = cp_matches.value_of("source").unwrap();
    let log = log.new(o!("target" => target.to_string(), "source" => source.to_string()));
    info!(log, "Liftoff! Adding new copy to configuration");
    let new_config = mutate::add_dotfiles_to_config(&log,
                                                   yaml_config,
                                                   &[model::DotFile {
                                                     target: target.to_string(),
                                                     source: source.to_string(),
                                                     dot_file_type: model::DotFileType::COPY,
                                                   }]);
    write_new_yaml(&log, &new_config, config_file);
  } else if let Some(import_matches) = maybe_import_matches {
    let  dir =  import_matches.value_of("dir").unwrap();
    let log = log.new(o!("directory" => dir.to_string()));
    let yaml_config = &yaml_documents[0];
    info!(log, "Liftoff! Importing to configuration");
    let dot_files = import::scan_dir(&log, dir);
    mutate::add_dotfiles_to_config(&log, yaml_config, &dot_files);
  }
}

fn write_new_yaml(log: &Logger, document: &Yaml, config_file: &str) {
  let mut out_str = String::new();
  {
    let mut emitter = YamlEmitter::new(&mut out_str);
    match emitter.dump(document) {
      Ok(_) => {}
      Err(e) => {
        error!(log, "Failed to write config file."; "error" => format!("{:?}", e));
        return;
      }
    }
  }
  out_str.push_str("\n");
  match write_config_file(config_file, &out_str) {
    Ok(_) => info!(log, "Successfully written configuration"),
    Err(e) => error!(log, "Failed to write config file."; "error" => e.description()),
  };
}

fn write_config_file(file: &str, content: &str) -> Result<(), std::io::Error> {
  let mut file = try!(File::create(file));
  try!(file.write_all(content.as_bytes()));
  Ok(())
}

fn load_config_file(log: &Logger, file: &str) -> Result<String, std::io::Error> {
  let path = Path::new(file);
  if path.exists() {
    let mut file = try!(File::open(file));
    let mut content = String::new();
    try!(file.read_to_string(&mut content));
    Ok(content)
  } else {
    warn!(log, "Configuration doesn't exit"; "file" => file);
    Ok("".to_string())
  }
}

fn create_app<'a>() -> App<'a, 'a> {
  App::new("dotcopter")
    .version("0.1")
    .setting(AppSettings::ColoredHelp)
    .author("Patrick Haun <bomgar85@googlemail.com>")
    .about("manages dotfiles installation")
    .arg(Arg::with_name("force").short("f").long("force").takes_value(false))
    .arg(Arg::with_name("verbose")
           .long("verbose")
           .short("v")
           .help("debug output")
           .required(false)
           .takes_value(false))
    .arg(Arg::with_name("config_file").required(true))
    .subcommand(SubCommand::with_name("apply").about("applies a dotfile configuration"))
    .subcommand(SubCommand::with_name("ln")
                  .about("adds new link to configuration")
                  .arg(Arg::with_name("link_target").required(true))
                  .arg(Arg::with_name("link_name").required(true)))
    .subcommand(SubCommand::with_name("cp")
                  .about("adds new copy to configuration")
                  .arg(Arg::with_name("source").required(true))
                  .arg(Arg::with_name("target").required(true)))
    .subcommand(SubCommand::with_name("import")
                  .about("imports dotfiles from a folder into the configuration")
                  .arg(Arg::with_name("dir").required(true)))
}
<|endoftext|>


<commit_before>#![allow(unstable)]
#![feature(plugin)]
extern crate sodiumoxide;
extern crate libc;
extern crate serialize;
extern crate unicode;
extern crate regex;
#[plugin] #[no_link]
extern crate regex_macros;

use regex::NoExpand;

use reader::*;
use controls::*;
use constants::*;
use util::*;
use script::*;
use builtins::*;
use command::*;
use input::*;

use input::InputValue::*;

mod constants;
mod util;
mod termios;
mod signal;
mod controls;
mod input;
mod reader;
mod script;
mod builtins;
mod command;

fn run_line(line:InputValue, term:&mut TermState, env:&mut WashEnv) -> WashArgs {
    // this needs to be re-written, it will be soon
    match line {
        Short(v) | Literal(v) => {
            print!("{}\n", v);
        },
        Long(v) => {
            print!("(");
            for item in v.iter() {
                run_line(item.clone(), term, env);
            }
            print!(")\n");
        },
        Function(n, v) => {
            print!("{}(", n);
            for item in v.iter() {
                run_line(item.clone(), term, env);
            }
            print!(")\n");
        },
        _ => {}
    }
    return WashArgs::Long(vec![]);
}

fn process_line(line:InputValue, term:&mut TermState, env:&mut WashEnv) -> WashArgs {
    // this needs to be re-written, it will be soon
    return WashArgs::Empty;
}

// public so no warnings when we run tests
pub fn main() {
    let mut controls = &mut Controls::new();
    let mut reader = LineReader::new();
    let mut env = WashEnv::new();
    let mut term = TermState::new();
    load_builtins(&mut env);
    term.update_terminal();
    loop {
        controls.flush();
        match reader.read_line() {
            None => {
                if reader.eof {
                    break;
                } else if !reader.line.is_empty() {
                    controls.outc(BEL);
                    reader.restart();
                } else {
                    controls.outc(NL);
                    reader.clear();
                }
            },
            Some(line) => {
                controls.outc(NL);
                match run_line(line, &mut term, &mut env) {
                    WashArgs::Empty => controls.err("Command failed\n"),
                    v => {
                        controls.outs(v.flatten().as_slice());
                        controls.outc(NL);
                    }
                }
                reader.clear();
            }
        }
    }
    controls.outs("\nExiting\n");
    controls.flush();
    term.restore_terminal();
}
<commit_msg>Implement basic function calling<commit_after>#![allow(unstable)]
#![feature(plugin)]
extern crate sodiumoxide;
extern crate libc;
extern crate serialize;
extern crate unicode;
extern crate regex;
#[plugin] #[no_link]
extern crate regex_macros;

use regex::NoExpand;

use reader::*;
use controls::*;
use constants::*;
use util::*;
use script::*;
use builtins::*;
use command::*;
use input::*;

use input::InputValue::*;

mod constants;
mod util;
mod termios;
mod signal;
mod controls;
mod input;
mod reader;
mod script;
mod builtins;
mod command;

fn process_line(line:InputValue, term:&mut TermState, env:&mut WashEnv) -> WashArgs {
    match line {
        Function(n, a) => {
            if env.hasf(&n) {
                let func = WashEnv::getf(env, &n).unwrap();
                let mut args = vec![];
                for item in a.iter() {
                    args.push(process_line(item.clone(), term, env));
                }
                return func(&WashArgs::Long(args), env, term);
            } else {
                return WashArgs::Empty;
            }
        },
        Short(s) => return WashArgs::Flat(s),
        Literal(s) => return WashArgs::Flat(s),
        _ => return WashArgs::Empty
    }
}

// public so no warnings when we run tests
pub fn main() {
    let mut controls = &mut Controls::new();
    let mut reader = LineReader::new();
    let mut env = WashEnv::new();
    let mut term = TermState::new();
    load_builtins(&mut env);
    term.update_terminal();
    loop {
        controls.flush();
        match reader.read_line() {
            None => {
                if reader.eof {
                    break;
                } else if !reader.line.is_empty() {
                    controls.outc(BEL);
                    reader.restart();
                } else {
                    controls.outc(NL);
                    reader.clear();
                }
            },
            Some(line) => {
                controls.outc(NL);
                match process_line(line, &mut term, &mut env) {
                    WashArgs::Empty => controls.err("Command failed\n"),
                    v => {
                        controls.outs(v.flatten().as_slice());
                        controls.outc(NL);
                    }
                }
                reader.clear();
            }
        }
    }
    controls.outs("\nExiting\n");
    controls.flush();
    term.restore_terminal();
}
<|endoftext|>


<commit_before>#[deriving(Show)]
#[deriving(Clone)]
pub enum Doc {
    Nil,
    Append(Box<Doc>, Box<Doc>),
    Nest(uint, Box<Doc>),
    Text(String),
    Break(uint, uint),
    Newline,
    Group(Box<Doc>)
}

impl Doc {
    pub fn append(&self, e:Doc) -> Doc {
        match self {
            &Nil => e,
            x => match e {
                Nil => x.clone(),
                y => Append(box x.clone(), box y)
            }
        }
    }
}

fn concat(ds:&[Doc]) -> Doc {
    ds.iter().fold(Nil, |a, b| Append(box a, box b.clone()))
}

fn int(i:int) -> Doc {
    Text(format!("{}", i))
}

fn char(c:char) -> Doc {
    Text(format!("{}", c))
}

fn bool(b:bool) -> Doc {
    if b {
        Text(String::from_str("true"))
    } else {
        Text(String::from_str("false"))
    }
}

mod mode {
    #[deriving(Clone)]
    pub enum Mode {
        Flat,
        Break
    }
}

fn replicate<A:Clone>(x:A, l:uint) -> Vec<A> {
    let i = 0u;
    let mut res = Vec::new();
    for _ in range(i,l) {
        res.push(x.clone());
    }
    res
}
fn pad_right(c:char, l:uint, str:String) -> String {
    let str_len = str.len();
    if l > str_len {
        let padding = replicate(String::from_chars([c]), l - str_len).concat();
        let mut result = str.clone();
        result.push_str(padding.as_slice());
        result
    } else {
        str
    }
}
fn pad_left(c:char, l:uint, str:String) -> String {
    let str_len = str.len();
    if l > str_len {
      let mut result = replicate(String::from_chars([c]), l - str_len).concat();
      result.push_str(str.as_slice());
      result
    } else {
        str
    }
}

fn nlspace<S>(outs:(|S,String| -> S), s:S, i:uint) -> S {
    outs(s, pad_right(' ', i + String::from_str("\n").len(), String::from_str("\n")))
}

fn spaces<S>(outs:(|S,String| -> S), s:S, i:uint) -> S {
    outs(s, pad_left(' ', i, String::from_str("")))
}

fn fitting(xs:&Vec<(uint,mode::Mode,Doc)>, left:uint) -> bool {
    if left as int >= 0 {
        match xs.as_slice() {
            [] => true,
            [(ref i, ref mode, ref doc), ref rest..] => match doc {
                &Nil => fitting(&rest.to_vec(), left),
                &Append(ref x, ref y) => {
                    let mut ys = [(*i,*mode,*x.clone()), (*i,*mode,*y.clone())].to_vec();
                    ys.push_all(*rest);
                    fitting(&ys,left)
                },
                &Nest(j, ref x) => {
                    let mut ys = [(*i+j,*mode,*x.clone())].to_vec();
                    ys.push_all(*rest);
                    fitting(&ys, left)
                },
                &Text(ref str) => fitting(&rest.to_vec(), left - str.len()),
                &Break(sp, _) => match *mode {
                    mode::Flat => fitting(&rest.to_vec(), left - sp),
                    mode::Break => true
                },
                &Newline => true,
                &Group(ref x) => {
                    let mut ys = [(*i,*mode, *x.clone())].to_vec();
                    ys.push_all(*rest);
                    fitting(&ys, left)
                }
            }
        }
    }
    else {
        false
    }
}

fn best<S:Clone>(w:uint, outs:(|S, String| -> S), s:S, x:Doc) -> S {
    fn go<S:Clone>(w:uint, outs:(|S,String| -> S), s:S, k:uint, xs:Vec<(uint,mode::Mode,Doc)>) -> S {
        match xs.as_slice() {
            [] => s.clone(),
            [(ref i, ref mode, ref doc), ref rest..] => match doc {
                &Nil => go(w, outs, s, k, rest.to_vec()),
                &Append(ref x, ref y) => {
                    let mut rest2 = [(*i,*mode,*x.clone()), (*i,*mode,*y.clone())].to_vec();
                    rest2.push_all(*rest);
                    go(w, outs, s, k, rest2)
                },
                &Nest(j, ref x) => {
                    let mut rest2 = [(*i+j,*mode,*x.clone())].to_vec();
                    rest2.push_all(*rest);
                    go(w, outs, s, k, rest2)
                },
                &Text(ref str) => {
                    let ss = outs(s, str.clone());
                    go(w, outs, ss, (k + str.len()), rest.to_vec())
                },
                &Newline => {
                    let ss = nlspace(|x,y| outs(x,y), s, *i);
                    go(w, outs, ss, *i, rest.to_vec())
                }
                &Break(sp, off) => {
                    match *mode {
                        mode::Flat => {
                            let ss = spaces(|x,y| outs(x,y), s.clone(), sp);
                            go(w, outs, ss, k + sp, rest.to_vec())
                        },
                        mode::Break => {
                            let ss = nlspace(|x,y| outs(x,y), s.clone(), i + off);
                            go(w, outs, ss, i + off, rest.to_vec())
                        }
                    }
                },
                &Group(ref x) => {
                    match *mode {
                        mode::Flat => {
                            let mut rest2 = [(*i,mode::Flat,*x.clone())].to_vec();
                            rest2.push_all(*rest);
                            go(w, outs, s, k, rest2)
                        },
                        mode::Break => {
                            let mut rest2 = [(*i,mode::Flat,*x.clone())].to_vec();
                            rest2.push_all(*rest);
                            if fitting(&rest2, w - k) {
                                go(w, outs, s, k, rest2)
                            } else {
                                let mut rest3 = [(*i,mode::Break,*x.clone())].to_vec();
                                rest3.push_all(*rest);
                                go(w, outs, s, k, rest3)
                            }
                        }
                    }
                }
            }
        }
    }
    go(w, outs, s, 0, [(0,mode::Break,x)].to_vec())
}

impl Doc {
    pub fn to_string(&self, w:uint) -> String {
        let outs = |strs:Vec<String>, s:String| -> Vec<String> {
            let mut welp = [s].to_vec();
            welp.push_all(strs.as_slice()); // horrible
            welp
        };

        let mut strs = best(w, outs, [].to_vec(), self.clone());
        strs.reverse();
        strs.push(String::from_str("\n"));
        strs.concat()
    }
}

fn main() {
}
<commit_msg>Futher minor cleanup<commit_after>#[deriving(Show)]
#[deriving(Clone)]
pub enum Doc {
    Nil,
    Append(Box<Doc>, Box<Doc>),
    Nest(uint, Box<Doc>),
    Text(String),
    Break(uint, uint),
    Newline,
    Group(Box<Doc>)
}

impl Doc {
    pub fn append(&self, e:Doc) -> Doc {
        match self {
            &Nil => e,
            x => match e {
                Nil => x.clone(),
                y => Append(box x.clone(), box y)
            }
        }
    }
}

fn concat(ds:&[Doc]) -> Doc {
    ds.iter().fold(Nil, |a, b| Append(box a, box b.clone()))
}

fn int(i:int) -> Doc {
    Text(format!("{}", i))
}

fn char(c:char) -> Doc {
    Text(format!("{}", c))
}

fn bool(b:bool) -> Doc {
    if b {
        Text(String::from_str("true"))
    } else {
        Text(String::from_str("false"))
    }
}

mod mode {
    #[deriving(Clone)]
    pub enum Mode {
        Flat,
        Break
    }
}

fn replicate<A:Clone>(x:A, l:uint) -> Vec<A> {
    let i = 0u;
    let mut res = Vec::new();
    for _ in range(i,l) {
        res.push(x.clone());
    }
    res
}
fn pad_right(c:char, l:uint, str:String) -> String {
    let str_len = str.len();
    if l > str_len {
        let padding = replicate(String::from_chars([c]), l - str_len).concat();
        let mut res = str.clone();
        res.push_str(padding.as_slice());
        res
    } else {
        str
    }
}
fn pad_left(c:char, l:uint, str:String) -> String {
    let str_len = str.len();
    if l > str_len {
      let mut res = replicate(String::from_chars([c]), l - str_len).concat();
      res.push_str(str.as_slice());
      res
    } else {
        str
    }
}

fn nlspace<S>(outs:|S,String| -> S, s:S, i:uint) -> S {
    outs(s, pad_right(' ', i + String::from_str("\n").len(), String::from_str("\n")))
}

fn spaces<S>(outs:|S,String| -> S, s:S, i:uint) -> S {
    outs(s, pad_left(' ', i, String::from_str("")))
}

fn fitting(xs:&Vec<(uint,mode::Mode,Doc)>, left:uint) -> bool {
    if left as int >= 0 {
        match xs.as_slice() {
            [] => true,
            [(ref i, ref mode, ref doc), ref rest..] => match doc {
                &Nil => fitting(&rest.to_vec(), left),
                &Append(ref x, ref y) => {
                    let mut ys = [(*i,*mode,*x.clone()), (*i,*mode,*y.clone())].to_vec();
                    ys.push_all(*rest);
                    fitting(&ys,left)
                },
                &Nest(j, ref x) => {
                    let mut ys = [(*i+j,*mode,*x.clone())].to_vec();
                    ys.push_all(*rest);
                    fitting(&ys, left)
                },
                &Text(ref str) => fitting(&rest.to_vec(), left - str.len()),
                &Break(sp, _) => match *mode {
                    mode::Flat => fitting(&rest.to_vec(), left - sp),
                    mode::Break => true
                },
                &Newline => true,
                &Group(ref x) => {
                    let mut ys = [(*i,*mode, *x.clone())].to_vec();
                    ys.push_all(*rest);
                    fitting(&ys, left)
                }
            }
        }
    }
    else {
        false
    }
}

fn best<S:Clone>(w:uint, outs:|S, String| -> S, s:S, x:Doc) -> S {
    fn go<S:Clone>(w:uint, outs:|S,String| -> S, s:S, k:uint, xs:Vec<(uint,mode::Mode,Doc)>) -> S {
        match xs.as_slice() {
            [] => s.clone(),
            [(ref i, ref mode, ref doc), ref rest..] => match doc {
                &Nil => go(w, outs, s, k, rest.to_vec()),
                &Append(ref x, ref y) => {
                    let mut rest2 = [(*i,*mode,*x.clone()), (*i,*mode,*y.clone())].to_vec();
                    rest2.push_all(*rest);
                    go(w, outs, s, k, rest2)
                },
                &Nest(j, ref x) => {
                    let mut rest2 = [(*i+j,*mode,*x.clone())].to_vec();
                    rest2.push_all(*rest);
                    go(w, outs, s, k, rest2)
                },
                &Text(ref str) => {
                    let ss = outs(s, str.clone());
                    go(w, outs, ss, (k + str.len()), rest.to_vec())
                },
                &Newline => {
                    let ss = nlspace(|x,y| outs(x,y), s, *i);
                    go(w, outs, ss, *i, rest.to_vec())
                }
                &Break(sp, off) => {
                    match *mode {
                        mode::Flat => {
                            let ss = spaces(|x,y| outs(x,y), s.clone(), sp);
                            go(w, outs, ss, k + sp, rest.to_vec())
                        },
                        mode::Break => {
                            let ss = nlspace(|x,y| outs(x,y), s.clone(), i + off);
                            go(w, outs, ss, i + off, rest.to_vec())
                        }
                    }
                },
                &Group(ref x) => {
                    match *mode {
                        mode::Flat => {
                            let mut rest2 = [(*i,mode::Flat,*x.clone())].to_vec();
                            rest2.push_all(*rest);
                            go(w, outs, s, k, rest2)
                        },
                        mode::Break => {
                            let mut rest2 = [(*i,mode::Flat,*x.clone())].to_vec();
                            rest2.push_all(*rest);
                            if fitting(&rest2, w - k) {
                                go(w, outs, s, k, rest2)
                            } else {
                                let mut rest3 = [(*i,mode::Break,*x.clone())].to_vec();
                                rest3.push_all(*rest);
                                go(w, outs, s, k, rest3)
                            }
                        }
                    }
                }
            }
        }
    }
    go(w, outs, s, 0, [(0,mode::Break,x)].to_vec())
}

impl Doc {
    pub fn to_string(&self, w:uint) -> String {
        let outs = |strs:Vec<String>, s:String| -> Vec<String> {
            let mut welp = [s].to_vec();
            welp.push_all(strs.as_slice()); // horrible
            welp
        };

        let mut strs = best(w, outs, [].to_vec(), self.clone());
        strs.reverse();
        strs.push(String::from_str("\n"));
        strs.concat()
    }
}

fn main() {
}
<|endoftext|>


<commit_before>#![feature(plugin)]
#![plugin(regex_macros)]

extern crate argparse;
extern crate html5ever;
extern crate html5ever_dom_sink;
extern crate hyper;
extern crate iron;
extern crate regex;
extern crate rustc_serialize;
extern crate tempfile;
extern crate tendril;
extern crate url;

mod json;

use argparse::ArgumentParser;
use argparse::Store;
use std::fs::OpenOptions;
use std::io::Read;
use std::io::Write;
use std::process::Command;
use std::process::Stdio;
use std::str::FromStr;

use html5ever::Attribute;
use html5ever::tree_builder::interface::TreeSink;
use html5ever_dom_sink::common::NodeEnum;
use html5ever_dom_sink::rcdom::Handle;
use html5ever_dom_sink::rcdom::RcDom;
use hyper::Client;
use hyper::header::Connection;
use iron::Iron;
use iron::IronResult;
use iron::Request;
use iron::Response;
use iron::headers::ContentType;
use iron::mime::Mime;
use iron::mime::SubLevel;
use iron::mime::TopLevel;
use rustc_serialize::json::Json;
use tempfile::NamedTempFile;
use url::percent_encoding;

use json::JsonPathElement::{Key, Only};

// TODO: consider doing s/en.wikipedia.org/this app's url/ on the HTML before serving it. This
// currently works fine, but might not over HTTPS.
// TODO: Do nested functions work? If so, use those where appropriate.
// TODO: there are some places where I've handrolled try!() equivalents. Fix those.
// TODO: make sure I'm returning Results everywhere, and propagating errors correctly. Remove all
// uses of unwrap() that might panic.
// TODO: The page "Battle_of_Palo_Alto" is truncated. Figure out the best way to debug.

// TODO: return a Result
fn call_wikimedia_api(parameters: Vec<(&str, &str)>) -> String {
    let post_body =
        parameters.into_iter().map(|p| format!("{}={}", p.0, p.1))
        .collect::<Vec<_>>().join("&") + "&format=json";

    let client = Client::new();
    let mut response = client.post("https://en.wikipedia.org/w/api.php")
        .body(&post_body)
        .header(Connection::close())
        .send().unwrap();

    let mut body = String::new();
    response.read_to_string(&mut body).unwrap();
    body
}

// TODO: this name is terrible.
// Returns pairs of (revid, parentid)
fn get_revert_revision_ids(title: &str) -> Result<Vec<(u64, u64)>, String> {
    let json_str = call_wikimedia_api(
        vec![("action", "query"), ("prop", "revisions"), ("titles", title),
             ("rvprop", "comment|ids"), ("rvlimit", "60")]);

    let json = Json::from_str(&json_str).unwrap();
    let revisions = try!(
        json::get_json_array(&json, &[Key("query"), Key("pages"), Only, Key("revisions")]));

    // Filter for revisions that mention vandalism, and that have revid and parentid (which they all
    // should, but merits checking).
    Ok(revisions.into_iter().filter(|revision| {
        match (json::get_json_string(revision, &[Key("comment")]),
               json::get_json_number(revision, &[Key("revid")]),
               json::get_json_number(revision, &[Key("parentid")])) {
            (Ok(ref comment), Ok(_), Ok(_)) => comment.contains("vandal"),
            _ => false,
        }})
       .map(|revision| {
           (json::get_json_number(revision, &[Key("revid")]).unwrap(),
            json::get_json_number(revision, &[Key("parentid")]).unwrap())
       }).collect())
}

fn get_latest_revision_id(title: &str) -> Result<u64, String> {
    let json_str = call_wikimedia_api(
        vec![("action", "query"), ("prop", "revisions"), ("titles", title), ("rvprop", "ids"),
             ("rvlimit", "1")]);
    let json = Json::from_str(&json_str).unwrap();
    json::get_json_number(
        &json, &[Key("query"), Key("pages"), Only, Key("revisions"), Only, Key("revid")])
}

fn get_revision_content(title: &str, id: u64) -> Result<String, String> {
    let json_str = call_wikimedia_api(
        vec![("action", "query"), ("prop", "revisions"), ("titles", title), ("rvprop", "content"),
             ("rvlimit", "1"), ("rvstartid", &id.to_string())]);
    let json = Json::from_str(&json_str).unwrap();
    match json::get_json_string(
        &json, &[Key("query"), Key("pages"), Only, Key("revisions"), Only, Key("*")]) {
        Ok(content) => Ok(content.to_string()),
        Err(msg) => Err(msg),
    }
}

fn get_canonical_title(title: &str) -> Result<String, String> {
    let latest_revision_id = get_latest_revision_id(title).unwrap();
    let page_contents = get_revision_content(title, latest_revision_id).unwrap();

    let regex = regex!(r"#REDIRECT \[\[([^]]+)\]\].*");
    match regex.captures(&page_contents) {
        Some(captures) => get_canonical_title(captures.at(1).unwrap()),
        None => {
            println!("Canonical page title is \"{}\"", title);
            Ok(title.to_string())
        },
    }
}

fn write_to_temp_file(contents: &str) -> NamedTempFile {
    let tempfile = NamedTempFile::new().unwrap();
    let mut file = OpenOptions::new().write(true).open(tempfile.path()).unwrap();
    file.write_all(contents.as_bytes()).unwrap();
    file.flush().unwrap();
    tempfile
}

// TODO: I'm not so sure these parameter names aren't terrible.
fn merge(old: &str, new1: &str, new2: &str) -> Option<String> {
    let old_tempfile = write_to_temp_file(old);
    let new1_tempfile = write_to_temp_file(new1);
    let new2_tempfile = write_to_temp_file(new2);

    let mut process = Command::new("diff3");
    process.arg("-m").args(&[new1_tempfile.path(), old_tempfile.path(), new2_tempfile.path()])
        .stdout(Stdio::piped()).stderr(Stdio::null());
    let output = process.output().unwrap();
    if output.status.success() {
        Some(String::from_utf8(output.stdout).unwrap())
    } else {
        None
    }
}

fn render(title: &str, wikitext: &str) -> Result<String, String> {
    let encoded_wikitext =
        percent_encoding::percent_encode(wikitext.as_bytes(), percent_encoding::QUERY_ENCODE_SET);
    let html = call_wikimedia_api(
        vec![("action", "parse"), ("prop", "text"), ("disablepp", ""), ("contentmodel", "wikitext"),
             ("title", title), ("text", &encoded_wikitext)]);
    // TODO: check return value
    let json = Json::from_str(&html).unwrap();
    match json::get_json_string(&json, &[Key("parse"), Key("text"), Key("*")]) {
        Ok(contents) => Ok(contents.to_string()),
        Err(message) => Err(message),
    }
}

fn has_matching_id(attributes: &Vec<Attribute>, id: &str) -> bool {
    // TODO: could also do this with filter() and is_empty(). Not sure if that would be better.
    for attribute in attributes {
        // TODO: do I seriously have to construct a StrTendril here?
        // There has to be a better way.
        if attribute.name.local.as_slice() == "id" &&
            attribute.value == tendril::StrTendril::from_str(id).unwrap() {
                return true;
            }
    }
    return false;
}

fn find_node_by_id(handle: &Handle, id: &str) -> Result<Handle, String> {
    let node = handle.borrow();
    match node.node {
        NodeEnum::Element(_, ref attributes) if has_matching_id(attributes, id) => Ok(handle.clone()),
        _ => {
            for child in &node.children {
                match find_node_by_id(child, id) {
                    // TODO: this looks weird. Is this an abnormal way to do things?
                    Ok(node) => return Ok(node),
                    _ => continue,
                }
            }
            Err(format!("No node with ID {} found", id))
        },
    }
}

fn replace_node_with_placeholder(original_html: &str, div_id: &str, placeholder: &str) -> Result<String, String> {
    // TODO: check errors
    let html = tendril::StrTendril::from_str(original_html).unwrap();
    let mut dom: RcDom = html5ever::parse(html5ever::one_input(html), Default::default());

    let handle = try!(find_node_by_id(&dom.get_document(), div_id));
    let child_handles =
        (&handle.borrow().children).into_iter().map(|child| child.clone()).collect::<Vec<_>>();
    for child_handle in child_handles {
        dom.remove_from_parent(child_handle);
    }
    dom.append(handle,
               html5ever::tree_builder::interface::NodeOrText::AppendText(
                   tendril::StrTendril::from_str(placeholder).unwrap()));
    let mut serialized: Vec<u8> = vec![];
    html5ever::serialize::serialize(&mut serialized, &dom.document, Default::default());
    // TODO: error handling
    Ok(String::from_utf8(serialized).unwrap())
}

fn get_current_page(title: &str) -> String {
    let client = Client::new();
    let mut res = client.get(
        &format!("https://en.wikipedia.org/wiki/{}", title))
        .header(Connection::close())
        .send().unwrap();

    let mut body = String::new();
    res.read_to_string(&mut body).unwrap();

    body
}

// TODO: this name is terrible.
fn get_page_with_vandalism_restored(title: &str) -> Result<String, String> {
    let canonicalized_title = get_canonical_title(title).unwrap();

    let latest_revid = get_latest_revision_id(&canonicalized_title).unwrap();
    let mut accumulated_contents = get_revision_content(&canonicalized_title, latest_revid).unwrap();
    let mut revisions = vec![];
    for (revert_revid, vandalism_revid) in get_revert_revision_ids(&canonicalized_title).ok().unwrap() {
        let reverted_contents = get_revision_content(&canonicalized_title, revert_revid).unwrap();
        let vandalized_contents = get_revision_content(&canonicalized_title, vandalism_revid).unwrap();
        match merge(&reverted_contents, &vandalized_contents, &accumulated_contents) {
            Some(merged_contents) => {
                // TODO: replace this with a log statement, if there's a good logging framework.
                println!(
                    "For page \"{}\", restored vandalism https://en.wikipedia.org/w/index.php?title={}&diff=prev&oldid={}",
                         &title, &canonicalized_title, revert_revid);
                accumulated_contents = merged_contents;
                revisions.push(revert_revid);
            }
            None => (),
        }
    }

    // TODO: replace this with a log statement, if there's a good logging framework.
    println!("For page \"{}\", restored vandalisms reverted in: {:?}", &title, revisions);

    let rendered_body = render(title, &accumulated_contents).unwrap();
    // Note: "title" rather than "canonicalized_title", so that redirects look right.
    let current_page_contents = get_current_page(&title);
    // TODO: randomize the placeholder string per-request
    let page_contents_with_placeholder =
        replace_node_with_placeholder(
            &current_page_contents, "mw-content-text", "WMW_PLACEHOLDER_TEXT").unwrap();
    Ok(page_contents_with_placeholder.replace("WMW_PLACEHOLDER_TEXT", &rendered_body))
}

fn serve_request(request: &mut Request) -> IronResult<Response> {
    if request.url.path.len() == 2 && request.url.path[0] == "wiki" {
        let mut response = Response::with(
            (iron::status::Ok, get_page_with_vandalism_restored(&request.url.path[1]).unwrap()));
        response.headers.set(ContentType(Mime(TopLevel::Text, SubLevel::Html, vec![])));
        Ok(response)
    } else {
        let client = Client::new();
        // TODO: error handling
        let mut wikipedia_response =
            // TODO: not good enough. Needs to include query string.
            client.get(&format!("https://en.wikipedia.org/{}", request.url.path.join("/")))
            .header(Connection::close())
            .send().unwrap();
        let mut wikipedia_body: Vec<u8> = Vec::new();
        wikipedia_response.read_to_end(&mut wikipedia_body);

        let mut response = Response::with(wikipedia_body);
        response.status = Some(wikipedia_response.status);
        response.headers = wikipedia_response.headers.clone();
        println!("Forwarded request for {} to en.wikipedia.org", request.url.path.join("/"));
        Ok(response)
    }
}

fn main() {
    let mut port = 3000;
    {
        let mut parser = ArgumentParser::new();
        parser.set_description("TODO: Usage description");
        parser.refer(&mut port).add_option(&["-p", "--port"], Store, "The port to serve HTTP on.");
        parser.parse_args_or_exit();
    }
    Iron::new(serve_request).http(("localhost", port)).unwrap();
}

#[cfg(test)]
mod tests {
    use super::{merge, replace_node_with_placeholder};

    #[test]
    fn test_merge_clean() {
        let old = "First line.\n\nSecond line.\n";
        let new1 = "First line.\n\nSecond line changed.\n";
        let new2 = "First line changed.\n\nSecond line.\n";
        assert_eq!("First line changed.\n\nSecond line changed.\n", merge(old, new1, new2).unwrap());
    }

    #[test]
    fn test_merge_conflicting() {
        let old = "First line.\n\nSecond line.\n";
        let new1 = "First line.\n\nSecond line changed one way.\n";
        let new2 = "First line changed.\n\nSecond line changed a different way.\n";
        assert_eq!(None, merge(old, new1, new2));
    }

    #[test]
    fn test_merge_special_characters() {
        let old = "First line.\n\nSecond line.\n";
        let new1 = "First line.\n\nSecond line 𐅃.\n";
        let new2 = "First line さようなら.\n\nSecond line.\n";
        assert_eq!("First line さようなら.\n\nSecond line 𐅃.\n", merge(old, new1, new2).unwrap());
    }

    #[test]
    fn test_replace_html_content() {
        let original_html = "

original text

Other text

";
        let expected_html = "

replaced text

Other text

";
        let processed_html = replace_node_with_placeholder(original_html, "mw-content-text", "replaced text").unwrap();
        assert_eq!(expected_html, processed_html);
    }
}
<commit_msg>Satisfy a TODO: Nest write_to_temp_file and has_matching_id inside the only functinos that use them.<commit_after>#![feature(plugin)]
#![plugin(regex_macros)]

extern crate argparse;
extern crate html5ever;
extern crate html5ever_dom_sink;
extern crate hyper;
extern crate iron;
extern crate regex;
extern crate rustc_serialize;
extern crate tempfile;
extern crate tendril;
extern crate url;

mod json;

use argparse::ArgumentParser;
use argparse::Store;
use std::fs::OpenOptions;
use std::io::Read;
use std::io::Write;
use std::process::Command;
use std::process::Stdio;
use std::str::FromStr;

use html5ever::Attribute;
use html5ever::tree_builder::interface::TreeSink;
use html5ever_dom_sink::common::NodeEnum;
use html5ever_dom_sink::rcdom::Handle;
use html5ever_dom_sink::rcdom::RcDom;
use hyper::Client;
use hyper::header::Connection;
use iron::Iron;
use iron::IronResult;
use iron::Request;
use iron::Response;
use iron::headers::ContentType;
use iron::mime::Mime;
use iron::mime::SubLevel;
use iron::mime::TopLevel;
use rustc_serialize::json::Json;
use tempfile::NamedTempFile;
use url::percent_encoding;

use json::JsonPathElement::{Key, Only};

// TODO: consider doing s/en.wikipedia.org/this app's url/ on the HTML before serving it. This
// currently works fine, but might not over HTTPS.
// TODO: there are some places where I've handrolled try!() equivalents. Fix those.
// TODO: make sure I'm returning Results everywhere, and propagating errors correctly. Remove all
// uses of unwrap() that might panic.
// TODO: The page "Battle_of_Palo_Alto" is truncated. Figure out the best way to debug.

// TODO: return a Result
fn call_wikimedia_api(parameters: Vec<(&str, &str)>) -> String {
    let post_body =
        parameters.into_iter().map(|p| format!("{}={}", p.0, p.1))
        .collect::<Vec<_>>().join("&") + "&format=json";

    let client = Client::new();
    let mut response = client.post("https://en.wikipedia.org/w/api.php")
        .body(&post_body)
        .header(Connection::close())
        .send().unwrap();

    let mut body = String::new();
    response.read_to_string(&mut body).unwrap();
    body
}

// TODO: this name is terrible.
// Returns pairs of (revid, parentid)
fn get_revert_revision_ids(title: &str) -> Result<Vec<(u64, u64)>, String> {
    let json_str = call_wikimedia_api(
        vec![("action", "query"), ("prop", "revisions"), ("titles", title),
             ("rvprop", "comment|ids"), ("rvlimit", "60")]);

    let json = Json::from_str(&json_str).unwrap();
    let revisions = try!(
        json::get_json_array(&json, &[Key("query"), Key("pages"), Only, Key("revisions")]));

    // Filter for revisions that mention vandalism, and that have revid and parentid (which they all
    // should, but merits checking).
    Ok(revisions.into_iter().filter(|revision| {
        match (json::get_json_string(revision, &[Key("comment")]),
               json::get_json_number(revision, &[Key("revid")]),
               json::get_json_number(revision, &[Key("parentid")])) {
            (Ok(ref comment), Ok(_), Ok(_)) => comment.contains("vandal"),
            _ => false,
        }})
       .map(|revision| {
           (json::get_json_number(revision, &[Key("revid")]).unwrap(),
            json::get_json_number(revision, &[Key("parentid")]).unwrap())
       }).collect())
}

fn get_latest_revision_id(title: &str) -> Result<u64, String> {
    let json_str = call_wikimedia_api(
        vec![("action", "query"), ("prop", "revisions"), ("titles", title), ("rvprop", "ids"),
             ("rvlimit", "1")]);
    let json = Json::from_str(&json_str).unwrap();
    json::get_json_number(
        &json, &[Key("query"), Key("pages"), Only, Key("revisions"), Only, Key("revid")])
}

fn get_revision_content(title: &str, id: u64) -> Result<String, String> {
    let json_str = call_wikimedia_api(
        vec![("action", "query"), ("prop", "revisions"), ("titles", title), ("rvprop", "content"),
             ("rvlimit", "1"), ("rvstartid", &id.to_string())]);
    let json = Json::from_str(&json_str).unwrap();
    match json::get_json_string(
        &json, &[Key("query"), Key("pages"), Only, Key("revisions"), Only, Key("*")]) {
        Ok(content) => Ok(content.to_string()),
        Err(msg) => Err(msg),
    }
}

fn get_canonical_title(title: &str) -> Result<String, String> {
    let latest_revision_id = get_latest_revision_id(title).unwrap();
    let page_contents = get_revision_content(title, latest_revision_id).unwrap();

    let regex = regex!(r"#REDIRECT \[\[([^]]+)\]\].*");
    match regex.captures(&page_contents) {
        Some(captures) => get_canonical_title(captures.at(1).unwrap()),
        None => {
            println!("Canonical page title is \"{}\"", title);
            Ok(title.to_string())
        },
    }
}

// TODO: I'm not so sure these parameter names aren't terrible.
fn merge(old: &str, new1: &str, new2: &str) -> Option<String> {
    fn write_to_temp_file(contents: &str) -> NamedTempFile {
        let tempfile = NamedTempFile::new().unwrap();
        let mut file = OpenOptions::new().write(true).open(tempfile.path()).unwrap();
        file.write_all(contents.as_bytes()).unwrap();
        file.flush().unwrap();
        tempfile
    }

    let old_tempfile = write_to_temp_file(old);
    let new1_tempfile = write_to_temp_file(new1);
    let new2_tempfile = write_to_temp_file(new2);

    let mut process = Command::new("diff3");
    process.arg("-m").args(&[new1_tempfile.path(), old_tempfile.path(), new2_tempfile.path()])
        .stdout(Stdio::piped()).stderr(Stdio::null());
    let output = process.output().unwrap();
    if output.status.success() {
        Some(String::from_utf8(output.stdout).unwrap())
    } else {
        None
    }
}

fn render(title: &str, wikitext: &str) -> Result<String, String> {
    let encoded_wikitext =
        percent_encoding::percent_encode(wikitext.as_bytes(), percent_encoding::QUERY_ENCODE_SET);
    let html = call_wikimedia_api(
        vec![("action", "parse"), ("prop", "text"), ("disablepp", ""), ("contentmodel", "wikitext"),
             ("title", title), ("text", &encoded_wikitext)]);
    // TODO: check return value
    let json = Json::from_str(&html).unwrap();
    match json::get_json_string(&json, &[Key("parse"), Key("text"), Key("*")]) {
        Ok(contents) => Ok(contents.to_string()),
        Err(message) => Err(message),
    }
}

fn find_node_by_id(handle: &Handle, id: &str) -> Result<Handle, String> {
    fn has_matching_id(attributes: &Vec<Attribute>, id: &str) -> bool {
        // TODO: could also do this with filter() and is_empty(). Not sure if that would be better.
        for attribute in attributes {
            // TODO: do I seriously have to construct a StrTendril here?
            // There has to be a better way.
            if attribute.name.local.as_slice() == "id" &&
                attribute.value == tendril::StrTendril::from_str(id).unwrap() {
                    return true;
                }
        }
        return false;
    }

    let node = handle.borrow();
    match node.node {
        NodeEnum::Element(_, ref attributes) if has_matching_id(attributes, id) => Ok(handle.clone()),
        _ => {
            for child in &node.children {
                match find_node_by_id(child, id) {
                    // TODO: this looks weird. Is this an abnormal way to do things?
                    Ok(node) => return Ok(node),
                    _ => continue,
                }
            }
            Err(format!("No node with ID {} found", id))
        },
    }
}

fn replace_node_with_placeholder(original_html: &str, div_id: &str, placeholder: &str) -> Result<String, String> {
    // TODO: check errors
    let html = tendril::StrTendril::from_str(original_html).unwrap();
    let mut dom: RcDom = html5ever::parse(html5ever::one_input(html), Default::default());

    let handle = try!(find_node_by_id(&dom.get_document(), div_id));
    let child_handles =
        (&handle.borrow().children).into_iter().map(|child| child.clone()).collect::<Vec<_>>();
    for child_handle in child_handles {
        dom.remove_from_parent(child_handle);
    }
    dom.append(handle,
               html5ever::tree_builder::interface::NodeOrText::AppendText(
                   tendril::StrTendril::from_str(placeholder).unwrap()));
    let mut serialized: Vec<u8> = vec![];
    html5ever::serialize::serialize(&mut serialized, &dom.document, Default::default());
    // TODO: error handling
    Ok(String::from_utf8(serialized).unwrap())
}

fn get_current_page(title: &str) -> String {
    let client = Client::new();
    let mut res = client.get(
        &format!("https://en.wikipedia.org/wiki/{}", title))
        .header(Connection::close())
        .send().unwrap();

    let mut body = String::new();
    res.read_to_string(&mut body).unwrap();

    body
}

// TODO: this name is terrible.
fn get_page_with_vandalism_restored(title: &str) -> Result<String, String> {
    let canonicalized_title = get_canonical_title(title).unwrap();

    let latest_revid = get_latest_revision_id(&canonicalized_title).unwrap();
    let mut accumulated_contents = get_revision_content(&canonicalized_title, latest_revid).unwrap();
    let mut revisions = vec![];
    for (revert_revid, vandalism_revid) in get_revert_revision_ids(&canonicalized_title).ok().unwrap() {
        let reverted_contents = get_revision_content(&canonicalized_title, revert_revid).unwrap();
        let vandalized_contents = get_revision_content(&canonicalized_title, vandalism_revid).unwrap();
        match merge(&reverted_contents, &vandalized_contents, &accumulated_contents) {
            Some(merged_contents) => {
                // TODO: replace this with a log statement, if there's a good logging framework.
                println!(
                    "For page \"{}\", restored vandalism https://en.wikipedia.org/w/index.php?title={}&diff=prev&oldid={}",
                         &title, &canonicalized_title, revert_revid);
                accumulated_contents = merged_contents;
                revisions.push(revert_revid);
            }
            None => (),
        }
    }

    // TODO: replace this with a log statement, if there's a good logging framework.
    println!("For page \"{}\", restored vandalisms reverted in: {:?}", &title, revisions);

    let rendered_body = render(title, &accumulated_contents).unwrap();
    // Note: "title" rather than "canonicalized_title", so that redirects look right.
    let current_page_contents = get_current_page(&title);
    // TODO: randomize the placeholder string per-request
    let page_contents_with_placeholder =
        replace_node_with_placeholder(
            &current_page_contents, "mw-content-text", "WMW_PLACEHOLDER_TEXT").unwrap();
    Ok(page_contents_with_placeholder.replace("WMW_PLACEHOLDER_TEXT", &rendered_body))
}

fn serve_request(request: &mut Request) -> IronResult<Response> {
    if request.url.path.len() == 2 && request.url.path[0] == "wiki" {
        let mut response = Response::with(
            (iron::status::Ok, get_page_with_vandalism_restored(&request.url.path[1]).unwrap()));
        response.headers.set(ContentType(Mime(TopLevel::Text, SubLevel::Html, vec![])));
        Ok(response)
    } else {
        let client = Client::new();
        // TODO: error handling
        let mut wikipedia_response =
            // TODO: not good enough. Needs to include query string.
            client.get(&format!("https://en.wikipedia.org/{}", request.url.path.join("/")))
            .header(Connection::close())
            .send().unwrap();
        let mut wikipedia_body: Vec<u8> = Vec::new();
        wikipedia_response.read_to_end(&mut wikipedia_body);

        let mut response = Response::with(wikipedia_body);
        response.status = Some(wikipedia_response.status);
        response.headers = wikipedia_response.headers.clone();
        println!("Forwarded request for {} to en.wikipedia.org", request.url.path.join("/"));
        Ok(response)
    }
}

fn main() {
    let mut port = 3000;
    {
        let mut parser = ArgumentParser::new();
        parser.set_description("TODO: Usage description");
        parser.refer(&mut port).add_option(&["-p", "--port"], Store, "The port to serve HTTP on.");
        parser.parse_args_or_exit();
    }
    Iron::new(serve_request).http(("localhost", port)).unwrap();
}

#[cfg(test)]
mod tests {
    use super::{merge, replace_node_with_placeholder};

    #[test]
    fn test_merge_clean() {
        let old = "First line.\n\nSecond line.\n";
        let new1 = "First line.\n\nSecond line changed.\n";
        let new2 = "First line changed.\n\nSecond line.\n";
        assert_eq!("First line changed.\n\nSecond line changed.\n", merge(old, new1, new2).unwrap());
    }

    #[test]
    fn test_merge_conflicting() {
        let old = "First line.\n\nSecond line.\n";
        let new1 = "First line.\n\nSecond line changed one way.\n";
        let new2 = "First line changed.\n\nSecond line changed a different way.\n";
        assert_eq!(None, merge(old, new1, new2));
    }

    #[test]
    fn test_merge_special_characters() {
        let old = "First line.\n\nSecond line.\n";
        let new1 = "First line.\n\nSecond line 𐅃.\n";
        let new2 = "First line さようなら.\n\nSecond line.\n";
        assert_eq!("First line さようなら.\n\nSecond line 𐅃.\n", merge(old, new1, new2).unwrap());
    }

    #[test]
    fn test_replace_html_content() {
        let original_html = "

original text

Other text

";
        let expected_html = "

replaced text

Other text

";
        let processed_html = replace_node_with_placeholder(original_html, "mw-content-text", "replaced text").unwrap();
        assert_eq!(expected_html, processed_html);
    }
}
<|endoftext|>


<commit_before>//! # Beamium.
//!
//! Beamium scrap Prometheus endpoint and forward metrics to Warp10.
extern crate bytes;
extern crate cast;
extern crate clap;
extern crate core;
extern crate ctrlc;
extern crate flate2;
extern crate futures;
extern crate hyper;
extern crate hyper_timeout;
extern crate hyper_tls;
extern crate nix;
extern crate regex;
#[macro_use]
extern crate slog;
extern crate slog_async;
#[macro_use]
extern crate slog_scope;
extern crate slog_stream;
extern crate slog_syslog;
extern crate slog_term;
extern crate time;
extern crate tokio_core;
extern crate tokio_timer;
extern crate yaml_rust;

use clap::App;
use std::thread;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::fs;
use std::time::Duration;

mod config;
mod scraper;
mod router;
mod sink;
mod log;

include!("version.rs");

/// Main loop.
fn main() {
    // Setup a bare logger
    log::bootstrap();

    let matches = App::new("beamium")
        .version(&*format!("{} ({})", env!("CARGO_PKG_VERSION"), COMMIT))
        .author("d33d33 <kevin@d33d33.fr>")
        .about("Send Prometheus metrics to Warp10")
        .args_from_usage(
            "-c, --config=[FILE] 'Sets a custom config file'
                              \
                          -v...                'Increase verbosity level (console only)'",
        )
        .get_matches();

    info!("starting");

    // Bootstrap config
    let config_path = matches.value_of("config").unwrap_or("");
    let config = match config::load_config(&config_path) {
        Ok(config) => config,
        Err(err) => {
            crit!("Fail to load config {}: {}", &config_path, err);
            std::process::abort();
        }
    };

    // Setup logging
    match log::log(&config.parameters, matches.occurrences_of("v")) {
        Ok(()) => {}
        Err(err) => {
            crit!("Log setup failure: {}", err);
            std::process::abort();
        }
    }

    // Ensure dirs
    match fs::create_dir_all(&config.parameters.source_dir) {
        Ok(()) => {}
        Err(err) => {
            crit!(
                "Fail to create source directory {}: {}",
                &config.parameters.source_dir,
                err
            );
            std::process::abort();
        }
    };
    match fs::create_dir_all(&config.parameters.sink_dir) {
        Ok(()) => {}
        Err(err) => {
            crit!(
                "Fail to create sink directory {}: {}",
                &config.parameters.source_dir,
                err
            );
            std::process::abort();
        }
    };

    // Synchronisation stuff
    let sigint = Arc::new(AtomicBool::new(false));
    let mut handles = Vec::with_capacity(config.scrapers.len() + 1);

    // Sigint handling
    let r = sigint.clone();
    ctrlc::set_handler(move || {
        r.store(true, Ordering::SeqCst);
    }).expect("Error setting sigint handler");

    // Spawn scrapers
    info!("spawning scrapers");
    for scraper in config.scrapers.clone() {
        let (parameters, sigint) = (config.parameters.clone(), sigint.clone());
        handles.push(thread::spawn(move || {
            slog_scope::scope(
                &slog_scope::logger().new(o!("scraper" => scraper.name.clone())),
                || scraper::scraper(&scraper, &parameters, sigint),
            );
        }));
    }

    // Spawn router
    info!("spawning router");
    {
        let (sinks, labels, parameters, sigint) = (
            config.sinks.clone(),
            config.labels.clone(),
            config.parameters.clone(),
            sigint.clone(),
        );
        handles.push(thread::spawn(move || {
            slog_scope::scope(&slog_scope::logger().new(o!()), || {
                router::router(&sinks, &labels, &parameters, sigint)
            });
        }));
    }

    // Spawn sinks
    info!("spawning sinks");
    let mut sinks: Vec<sink::Sink> = config
        .sinks
        .iter()
        .map(|sink| sink::Sink::new(&sink, &config.parameters))
        .collect();

    sinks.iter_mut().for_each(|s| s.start());

    info!("started");

    // Wait for sigint
    loop {
        thread::sleep(Duration::from_millis(10));

        if sigint.load(Ordering::Relaxed) {
            break;
        }
    }

    info!("shutting down");
    for handle in handles {
        handle.join().unwrap();
    }

    for s in sinks {
        s.stop();
    }
    info!("halted");
}
<commit_msg>feat: validate configuration<commit_after>//! # Beamium.
//!
//! Beamium scrap Prometheus endpoint and forward metrics to Warp10.
extern crate bytes;
extern crate cast;
extern crate clap;
extern crate core;
extern crate ctrlc;
extern crate flate2;
extern crate futures;
extern crate hyper;
extern crate hyper_timeout;
extern crate hyper_tls;
extern crate nix;
extern crate regex;
#[macro_use]
extern crate slog;
extern crate slog_async;
#[macro_use]
extern crate slog_scope;
extern crate slog_stream;
extern crate slog_syslog;
extern crate slog_term;
extern crate time;
extern crate tokio_core;
extern crate tokio_timer;
extern crate yaml_rust;

use clap::App;
use std::thread;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::fs;
use std::time::Duration;

mod config;
mod scraper;
mod router;
mod sink;
mod log;

include!("version.rs");

/// Main loop.
fn main() {
    // Setup a bare logger
    log::bootstrap();

    let matches = App::new("beamium")
        .version(&*format!("{} ({})", env!("CARGO_PKG_VERSION"), COMMIT))
        .author("d33d33 <kevin@d33d33.fr>")
        .about("Send Prometheus metrics to Warp10")
        .args_from_usage(
            "-c, --config=[FILE] 'Sets a custom config file'
                              \
                          -v...                'Increase verbosity level (console only)'
                          -t                   'Test config'",
        )
        .get_matches();

    // Bootstrap config
    let config_path = matches.value_of("config").unwrap_or("");
    let config = match config::load_config(&config_path) {
        Ok(config) => config,
        Err(err) => {
            crit!("Fail to load config {}: {}", &config_path, err);
            std::process::abort();
        }
    };

    if matches.is_present("t") {
        info!("config ok");
        std::process::exit(0);
    }

    info!("starting");

    // Setup logging
    match log::log(&config.parameters, matches.occurrences_of("v")) {
        Ok(()) => {}
        Err(err) => {
            crit!("Log setup failure: {}", err);
            std::process::abort();
        }
    }

    // Ensure dirs
    match fs::create_dir_all(&config.parameters.source_dir) {
        Ok(()) => {}
        Err(err) => {
            crit!(
                "Fail to create source directory {}: {}",
                &config.parameters.source_dir,
                err
            );
            std::process::abort();
        }
    };
    match fs::create_dir_all(&config.parameters.sink_dir) {
        Ok(()) => {}
        Err(err) => {
            crit!(
                "Fail to create sink directory {}: {}",
                &config.parameters.source_dir,
                err
            );
            std::process::abort();
        }
    };

    // Synchronisation stuff
    let sigint = Arc::new(AtomicBool::new(false));
    let mut handles = Vec::with_capacity(config.scrapers.len() + 1);

    // Sigint handling
    let r = sigint.clone();
    ctrlc::set_handler(move || {
        r.store(true, Ordering::SeqCst);
    }).expect("Error setting sigint handler");

    // Spawn scrapers
    info!("spawning scrapers");
    for scraper in config.scrapers.clone() {
        let (parameters, sigint) = (config.parameters.clone(), sigint.clone());
        handles.push(thread::spawn(move || {
            slog_scope::scope(
                &slog_scope::logger().new(o!("scraper" => scraper.name.clone())),
                || scraper::scraper(&scraper, &parameters, sigint),
            );
        }));
    }

    // Spawn router
    info!("spawning router");
    {
        let (sinks, labels, parameters, sigint) = (
            config.sinks.clone(),
            config.labels.clone(),
            config.parameters.clone(),
            sigint.clone(),
        );
        handles.push(thread::spawn(move || {
            slog_scope::scope(&slog_scope::logger().new(o!()), || {
                router::router(&sinks, &labels, &parameters, sigint)
            });
        }));
    }

    // Spawn sinks
    info!("spawning sinks");
    let mut sinks: Vec<sink::Sink> = config
        .sinks
        .iter()
        .map(|sink| sink::Sink::new(&sink, &config.parameters))
        .collect();

    sinks.iter_mut().for_each(|s| s.start());

    info!("started");

    // Wait for sigint
    loop {
        thread::sleep(Duration::from_millis(10));

        if sigint.load(Ordering::Relaxed) {
            break;
        }
    }

    info!("shutting down");
    for handle in handles {
        handle.join().unwrap();
    }

    for s in sinks {
        s.stop();
    }
    info!("halted");
}
<|endoftext|>