<commit_before>use futures::stream::{self, StreamExt};
use futures::Future;
use log::{self, *};
use std::io::Write;
use tokio::task;
use trust_dns_resolver::config::*;
use trust_dns_resolver::lookup::Lookup;
use trust_dns_resolver::proto::rr::RecordType;
use trust_dns_resolver::proto::xfer::DnsRequestOptions;
use trust_dns_resolver::TokioAsyncResolver;

type Result<T> = std::result::Result<T, Box<dyn std::error::Error + Send + Sync>>;

/// Async lookup algorithm -- no work is done
async fn lookup(i: usize) -> Result<Vec<Lookup>> {
    let config = match i {
        1 => ResolverConfig::google(),
        2 => ResolverConfig::cloudflare(),
        _ => ResolverConfig::quad9(),
    };
    let resolver = TokioAsyncResolver::tokio(config, ResolverOpts::default()).await?;
    debug!("Created {} Resolver.", i);

    let futures: Vec<_> = vec![
        RecordType::A,
        RecordType::A,
        RecordType::A,
        RecordType::AAAA,
        RecordType::AAAA,
        RecordType::AAAA,
        RecordType::TXT,
        RecordType::TXT,
        RecordType::TXT,
    ]
    .into_iter()
    .map(|x| make_lookup(i, &resolver, x))
    .collect();

    let mut res = Vec::new();
    let stream = stream::iter(futures);
    let mut buffered_stream = stream.buffer_unordered(2);
    while let Some(f) = buffered_stream.next().await {
        trace!("Received lookup {:?}", f);
        let f = f?;
        res.push(f)
    }

    Ok(res)
}

async fn make_lookup(i: usize, resolver: &TokioAsyncResolver, record_type: RecordType) -> Result<Lookup> {
    let lookup = resolver
        .lookup("www.example.com.", record_type, DnsRequestOptions::default())
        .await?;
    debug!("Received {} Lookup for record type {}.", i, record_type);

    Ok(lookup)
}

/// Create the futures -- no work is done
fn lookup_futures() -> Vec<impl Future<Output = Result<Vec<Lookup>>>> {
    let google = lookup(1);
    let cloudflare = lookup(2);
    let quad6 = lookup(3);

    vec![google, cloudflare, quad6]
}

/// Spawn the futures -- work is started
fn lookups_spawns(lookup_futures: Vec<impl Future<Output = Result<Vec<Lookup>>> + Send + 'static>) -> Vec<task::JoinHandle<Result<Vec<Lookup>>>> {
    lookup_futures.into_iter().map(task::spawn).collect()
}

/// Wait asyncly for the work to finish
async fn lookups(tasks: Vec<task::JoinHandle<Result<Vec<Lookup>>>>) -> Result<Vec<Lookup>> {
    let mut res = Vec::new();
    for t in tasks {
        let r = t.await??;
        res.push(r);
    }

    Ok(res.into_iter().flatten().collect())
}

fn main() -> Result<()> {
    let start = std::time::Instant::now();
    env_logger::Builder::from_default_env()
        .format(move |buf, rec| {
            let t = start.elapsed().as_secs_f32();
            let thread_id_string = format!("{:?}", std::thread::current().id());
            let thread_id = &thread_id_string[9..thread_id_string.len()-1];
            writeln!(buf, "{:.03} [{:5}] ({:}) - {}", t, rec.level(), thread_id, rec.args())
        })
        .init();

    let mut rt = tokio::runtime::Runtime::new()
        .expect("Failed to create runtime.");

    rt.block_on(async {
        let futures = lookup_futures();
        let tasks = lookups_spawns(futures);
        let lookups = lookups(tasks).await;
        match lookups {
            //Ok(lookup) => info!("Lookup: {:#?}", lookup),
            Ok(_) => info!("Done."),
            Err(e) => error!("An error occurred: {}", e),
        }
    });

    Ok(())
}
<commit_msg>Add Thread-Id to logging<commit_after>use futures::stream::{self, StreamExt};
use futures::Future;
use log::{self, *};
use std::io::Write;
use tokio::task;
use trust_dns_resolver::config::*;
use trust_dns_resolver::lookup::Lookup;
use trust_dns_resolver::proto::rr::RecordType;
use trust_dns_resolver::proto::xfer::DnsRequestOptions;
use trust_dns_resolver::TokioAsyncResolver;

type Result<T> = std::result::Result<T, Box<dyn std::error::Error + Send + Sync>>;

/// Async lookup algorithm -- no work is done
async fn lookup(i: usize) -> Result<Vec<Lookup>> {
    let config = match i {
        1 => ResolverConfig::google(),
        2 => ResolverConfig::cloudflare(),
        _ => ResolverConfig::quad9(),
    };
    let resolver = TokioAsyncResolver::tokio(config, ResolverOpts::default()).await?;
    debug!("Created {} Resolver.", i);

    let futures: Vec<_> = vec![
        RecordType::A,
        RecordType::A,
        RecordType::A,
        RecordType::AAAA,
        RecordType::AAAA,
        RecordType::AAAA,
        RecordType::TXT,
        RecordType::TXT,
        RecordType::TXT,
    ]
    .into_iter()
    .map(|x| make_lookup(i, &resolver, x))
    .collect();

    let mut res = Vec::new();
    let stream = stream::iter(futures);
    let mut buffered_stream = stream.buffer_unordered(2);
    while let Some(f) = buffered_stream.next().await {
        trace!("Received lookup {:?}", f);
        let f = f?;
        res.push(f)
    }

    Ok(res)
}

async fn make_lookup(i: usize, resolver: &TokioAsyncResolver, record_type: RecordType) -> Result<Lookup> {
    let lookup = resolver
        .lookup("www.example.com.", record_type, DnsRequestOptions::default())
        .await?;
    debug!("Received {} Lookup for record type {}.", i, record_type);

    Ok(lookup)
}

/// Create the futures -- no work is done
fn lookup_futures() -> Vec<impl Future<Output = Result<Vec<Lookup>>>> {
    let google = lookup(1);
    let cloudflare = lookup(2);
    let quad6 = lookup(3);

    vec![google, cloudflare, quad6]
}

/// Spawn the futures -- work is started
fn lookups_spawns(lookup_futures: Vec<impl Future<Output = Result<Vec<Lookup>>> + Send + 'static>) -> Vec<task::JoinHandle<Result<Vec<Lookup>>>> {
    lookup_futures.into_iter().map(task::spawn).collect()
}

/// Wait asyncly for the work to finish
async fn lookups(tasks: Vec<task::JoinHandle<Result<Vec<Lookup>>>>) -> Result<Vec<Lookup>> {
    let mut res = Vec::new();
    for t in tasks {
        let r = t.await??;
        res.push(r);
    }

    Ok(res.into_iter().flatten().collect())
}

fn main() -> Result<()> {
    let start = std::time::Instant::now();
    env_logger::Builder::from_default_env()
        .format(move |buf, rec| {
            let t = start.elapsed().as_secs_f32();
            let thread_id_string = format!("{:?}", std::thread::current().id());
            let thread_id = &thread_id_string[9..thread_id_string.len()-1];
            writeln!(buf, "{:.03} [{:5}] ({:}) - {}", t, rec.level(), thread_id, rec.args())
        })
        .init();

    let mut rt = tokio::runtime::Runtime::new()
        .expect("Failed to create runtime.");

    rt.block_on(async {
        let futures = lookup_futures();
        debug!("Created all the futures");
        let tasks = lookups_spawns(futures);
        debug!("Spawned all the futures");
        let lookups = lookups(tasks).await;
        debug!("Awaited all the futures");
        match lookups {
            //Ok(lookup) => info!("Lookup: {:#?}", lookup),
            Ok(_) => info!("Done."),
            Err(e) => error!("An error occurred: {}", e),
        }
    });

    Ok(())
}
<|endoftext|>


<commit_before>#[macro_use] extern crate conrod;
extern crate gfx_device_gl;
extern crate graphics;
extern crate piston;
extern crate piston_window;
extern crate sdl2_window;
extern crate vecmath as vecmath_lib;

use std::path::Path;

use conrod::{Labelable, Positionable, Sizeable, Widget};
use piston::window::{AdvancedWindow, WindowSettings};
use piston_window::{Context, G2d, OpenGL, PistonWindow};
use piston::event_loop::*;
use piston::input::*;
use sdl2_window::Sdl2Window;

mod map;
mod vecmath;

/// A transition requested by a game state.
enum StateTransition {
    /// Stay on this state.
    Continue,
    /// End this state and pop it from the stack.
    End,
    /// Pop this state and replace it with this new state.
    Replace(Box<GameState>),
    /// Keep this state, but pause it and push the state on top of the stack.
    Push(Box<GameState>),
}

/// A state the game can be in.
///
/// The application consists of a stack of states, that can each push other
/// states on the stack and/or pop themselves.
/// Examples of states are the main menu, the character selection, or the game
/// itself; starting the game finishes the character selection, but the main
/// menu remains accessible.
trait GameState {
    fn handle_event(&mut self, event: &piston::input::Event<piston::input::Input>, resources: &mut Resources) {}
    fn update(&mut self, dt: f64) -> StateTransition;
    fn draw(&mut self, c: Context, g: &mut G2d);
    fn pause(&mut self) {}
    fn resume(&mut self) {}
}

struct Resources {
    window: PistonWindow<Sdl2Window>,
    width: u32,
    height: u32,
}

struct App {
    states: Vec<Box<GameState>>,
    resources: Resources,
}

impl App {
    fn new() -> App {
        let width = 200;
        let height = 200;

        // Change this to OpenGL::V2_1 if not working.
        let opengl = OpenGL::V3_2;

        // Create an SDL2 window.
        let window: PistonWindow<Sdl2Window> = WindowSettings::new(
                "fluffy-fiesta",
                [width, height]
            )
            .opengl(opengl)
            .build()
            .unwrap();

        let mut app = App {
            states: Vec::new(),
            resources: Resources {
                window: window,
            width: width,
            height: height,
            },
        };
        let game = Game::new(&mut app.resources);
        app.states.push(Box::new(game));
        app
    }

    fn run(&mut self) {
        loop {
            let transition: StateTransition = if let Some(state) = self.states.last_mut() {
                Self::handle_state(&mut self.resources, state)
            } else {
                break
            };
            match transition {
                StateTransition::Continue => panic!("App::run() got Transition::Continue"),
                StateTransition::End => {
                    self.states.pop().expect("Transition::End with no states");
                }
                StateTransition::Replace(state) => {
                    self.states.pop().expect("Transition::Replace with no states");
                    self.states.push(state);
                }
                StateTransition::Push(state) => {
                    self.states.push(state);
                }
            }
        }
    }

    fn handle_state(resources: &mut Resources, state: &mut Box<GameState>) -> StateTransition {
        state.resume();

        let mut events = resources.window.events();
        while let Some(event) = events.next(&mut resources.window) {
            // Handle generic event
            state.handle_event(&event, resources);

            // Call update method
            if let Some(u) = event.update_args() {
                let transition = state.update(u.dt);
                match transition {
                    StateTransition::Continue => {},
                    t => {
                        state.pause();
                        return t;
                    }
                }
            }

            // Call draw method
            resources.window.draw_2d(&event, |c, g| state.draw(c, g));
        }
        state.pause();
        StateTransition::End
    }
}

fn main() {
    let mut app = App::new();
    app.run();
}

widget_ids!(struct GameWidgetIds { canvas, resume, quit });

struct Game {
    rotation: f64,  // Rotation for the square.
    ui: conrod::Ui,
    widget_ids: GameWidgetIds,
    capture: bool,
    ui_on: bool,
    quit: bool,
    // FIXME: Spelling out gfx_device_gl (and direct dep) shouldn't be necessary
    image_map: conrod::image::Map<piston_window::Texture<gfx_device_gl::Resources>>,
    text_texture_cache: conrod::backend::piston_window::GlyphCache,
}

impl Game {
    fn new(resources: &mut Resources) -> Game {
        // Construct our `Ui`.
        let mut ui = conrod::UiBuilder::new().build();

        // Generate the widget identifiers.
        let ids = GameWidgetIds::new(ui.widget_id_generator());

        // Add a `Font` to the `Ui`'s `font::Map` from file.
        let font_path = Path::new("assets/NotoSans-Regular.ttf");
        assert!(font_path.exists());
        ui.fonts.insert_from_file(font_path).unwrap();

        // Create a texture to use for efficiently caching text on the GPU.
        let text_texture_cache =
            conrod::backend::piston_window::GlyphCache::new(&mut resources.window,
                                                            resources.width,
                                                            resources.height);

        // The image map describing each of our widget->image mappings (in our case, none).
        let image_map = conrod::image::Map::new();

        Game {
            rotation: 0.0,
            ui: ui,
            widget_ids: ids,
            capture: false,
            ui_on: false,
            quit: false,
            image_map: image_map,
            text_texture_cache: text_texture_cache,
        }
    }
}

impl GameState for Game {
    fn handle_event(&mut self, event: &piston::input::Event<piston::input::Input>, resources: &mut Resources) {
        // Convert the piston event to a conrod event.
        if let Some(ce) = conrod::backend::piston_window::convert_event(event.clone(), &mut resources.window) {
            self.ui.handle_event(ce);
        }

        if self.ui_on {
            let ui = &mut self.ui.set_widgets();

            // Create a background canvas upon which we'll place the button.
            conrod::widget::Canvas::new().floating(true).w_h(100.0, 70.0).pad(10.0).middle()
                .set(self.widget_ids.canvas, ui);

            // Draw the buttons.
            if conrod::widget::Button::new()
                .mid_top_of(self.widget_ids.canvas)
                .w_h(80.0, 20.0)
                .label("Resume")
                .set(self.widget_ids.resume, ui)
                .was_clicked()
            {
                self.ui_on = false;
            }

            if conrod::widget::Button::new()
                .down(10.0)
                .w_h(80.0, 20.0)
                .label("Quit")
                .set(self.widget_ids.quit, ui)
                .was_clicked()
            {
                self.quit = true;
            }
        }

        if let Some(Button::Mouse(button)) = event.press_args() {
            println!("Pressed mouse button '{:?}'", button);
        }

        if let Some(Button::Keyboard(key)) = event.press_args() {
            println!("Pressed key '{:?}'", key);
            if key == Key::Escape {
                println!("ui on");
                self.ui_on = true;
                if self.capture {
                    self.capture = false;
                    resources.window.set_capture_cursor(false);
                }
            } else if !self.ui_on && key == Key::C {
                self.capture = !self.capture;
                resources.window.set_capture_cursor(self.capture);
                println!("capture {}", if self.capture { "on" } else { "off" });
            }
        }
    }

    fn update(&mut self, dt: f64) -> StateTransition {
        // Rotate 2 radians per second.
        self.rotation += 2.0 * dt;

        if self.quit {
            StateTransition::End
        } else {
            StateTransition::Continue
        }
    }

    fn draw(&mut self, c: Context, g: &mut G2d) {
        use graphics::*;

        const GREEN: [f32; 4] = [0.0, 1.0, 0.0, 1.0];
        const RED:   [f32; 4] = [1.0, 0.0, 0.0, 1.0];

        let square = rectangle::square(0.0, 0.0, 150.0);
        let rotation = self.rotation;
        let (x, y) = (100.0, 100.0);

        // Clear the screen.
        clear(GREEN, g);

        let transform = c.transform.trans(x, y)
                                   .rot_rad(rotation)
                                   .trans(-75.0, -75.0);

        // Draw a box rotating around the middle of the screen.
        rectangle(RED, square, transform, g);

        if self.ui_on {
            let primitives = self.ui.draw();
            fn texture_from_image<T>(img: &T) -> &T { img };
            conrod::backend::piston_window::draw(c, g, primitives,
                                                 &mut self.text_texture_cache,
                                                 &self.image_map,
                                                 texture_from_image);
        }
    }
}
<commit_msg>Improve states API<commit_after>#[macro_use] extern crate conrod;
extern crate gfx_device_gl;
extern crate graphics;
extern crate piston;
extern crate piston_window;
extern crate sdl2_window;
extern crate vecmath as vecmath_lib;

use std::path::Path;

use conrod::{Labelable, Positionable, Sizeable, Widget};
use piston::window::{AdvancedWindow, WindowSettings};
use piston_window::{Context, G2d, OpenGL, PistonWindow};
use piston::event_loop::*;
use piston::input::*;
use sdl2_window::Sdl2Window;

mod map;
mod vecmath;

/// A transition requested by a game state.
enum StateTransition {
    /// Stay on this state.
    Continue,
    /// End this state and pop it from the stack.
    End,
    /// Pop this state and replace it with this new state.
    Replace(Box<GameState>),
    /// Keep this state, but pause it and push the state on top of the stack.
    Push(Box<GameState>),
}

/// A state the game can be in.
///
/// The application consists of a stack of states, that can each push other
/// states on the stack and/or pop themselves.
/// Examples of states are the main menu, the character selection, or the game
/// itself; starting the game finishes the character selection, but the main
/// menu remains accessible.
trait GameState {
    fn handle_event(&mut self, event: &piston::input::Event<piston::input::Input>,
                    resources: &mut Resources) -> StateTransition
    {
        StateTransition::Continue
    }
    fn update(&mut self, dt: f64) -> StateTransition;
    fn draw(&mut self, c: Context, g: &mut G2d);
    fn pause(&mut self, resources: &mut Resources) {}
    fn resume(&mut self, resources: &mut Resources) {}
}

struct Resources {
    window: PistonWindow<Sdl2Window>,
    width: u32,
    height: u32,
}

struct App {
    states: Vec<Box<GameState>>,
    resources: Resources,
}

impl App {
    fn new() -> App {
        let width = 200;
        let height = 200;

        // Change this to OpenGL::V2_1 if not working.
        let opengl = OpenGL::V3_2;

        // Create an SDL2 window.
        let window: PistonWindow<Sdl2Window> = WindowSettings::new(
                "fluffy-fiesta",
                [width, height]
            )
            .opengl(opengl)
            .build()
            .unwrap();

        let mut app = App {
            states: Vec::new(),
            resources: Resources {
                window: window,
            width: width,
            height: height,
            },
        };
        let game = Game::new(&mut app.resources);
        app.states.push(Box::new(game));
        app
    }

    fn run(&mut self) {
        loop {
            let transition: StateTransition = if let Some(state) = self.states.last_mut() {
                Self::handle_state(&mut self.resources, state)
            } else {
                break
            };
            match transition {
                StateTransition::Continue => panic!("App::run() got Transition::Continue"),
                StateTransition::End => {
                    self.states.pop().expect("Transition::End with no states");
                }
                StateTransition::Replace(state) => {
                    self.states.pop().expect("Transition::Replace with no states");
                    self.states.push(state);
                }
                StateTransition::Push(state) => {
                    self.states.push(state);
                }
            }
        }
    }

    fn handle_state(resources: &mut Resources, state: &mut Box<GameState>) -> StateTransition {
        state.resume(resources);

        let mut events = resources.window.events();
        while let Some(event) = events.next(&mut resources.window) {
            // Handle generic event
            let transition = state.handle_event(&event, resources);
            match transition {
                StateTransition::Continue => {},
                t => {
                    state.pause(resources);
                    return t;
                }
            }

            // Call update method
            if let Some(u) = event.update_args() {
                let transition = state.update(u.dt);
                match transition {
                    StateTransition::Continue => {},
                    t => {
                        state.pause(resources);
                        return t;
                    }
                }
            }

            // Call draw method
            resources.window.draw_2d(&event, |c, g| state.draw(c, g));
        }
        state.pause(resources);
        StateTransition::End
    }
}

fn main() {
    let mut app = App::new();
    app.run();
}

widget_ids!(struct GameWidgetIds { canvas, resume, quit });

struct Game {
    rotation: f64,  // Rotation for the square.
    ui: conrod::Ui,
    widget_ids: GameWidgetIds,
    capture: bool,
    ui_on: bool,
    // FIXME: Spelling out gfx_device_gl (and direct dep) shouldn't be necessary
    image_map: conrod::image::Map<piston_window::Texture<gfx_device_gl::Resources>>,
    text_texture_cache: conrod::backend::piston_window::GlyphCache,
}

impl Game {
    fn new(resources: &mut Resources) -> Game {
        // Construct our `Ui`.
        let mut ui = conrod::UiBuilder::new().build();

        // Generate the widget identifiers.
        let ids = GameWidgetIds::new(ui.widget_id_generator());

        // Add a `Font` to the `Ui`'s `font::Map` from file.
        let font_path = Path::new("assets/NotoSans-Regular.ttf");
        assert!(font_path.exists());
        ui.fonts.insert_from_file(font_path).unwrap();

        // Create a texture to use for efficiently caching text on the GPU.
        let text_texture_cache =
            conrod::backend::piston_window::GlyphCache::new(&mut resources.window,
                                                            resources.width,
                                                            resources.height);

        // The image map describing each of our widget->image mappings (in our case, none).
        let image_map = conrod::image::Map::new();

        Game {
            rotation: 0.0,
            ui: ui,
            widget_ids: ids,
            capture: false,
            ui_on: false,
            image_map: image_map,
            text_texture_cache: text_texture_cache,
        }
    }
}

impl GameState for Game {
    fn handle_event(&mut self, event: &piston::input::Event<piston::input::Input>,
                    resources: &mut Resources) -> StateTransition
    {
        // Convert the piston event to a conrod event.
        if let Some(ce) = conrod::backend::piston_window::convert_event(event.clone(), &mut resources.window) {
            self.ui.handle_event(ce);
        }

        if self.ui_on {
            let ui = &mut self.ui.set_widgets();

            // Create a background canvas upon which we'll place the button.
            conrod::widget::Canvas::new().floating(true).w_h(100.0, 70.0).pad(10.0).middle()
                .set(self.widget_ids.canvas, ui);

            // Draw the buttons.
            if conrod::widget::Button::new()
                .mid_top_of(self.widget_ids.canvas)
                .w_h(80.0, 20.0)
                .label("Resume")
                .set(self.widget_ids.resume, ui)
                .was_clicked()
            {
                self.ui_on = false;
            }

            if conrod::widget::Button::new()
                .down(10.0)
                .w_h(80.0, 20.0)
                .label("Quit")
                .set(self.widget_ids.quit, ui)
                .was_clicked()
            {
                return StateTransition::End;
            }
        }

        if let Some(Button::Mouse(button)) = event.press_args() {
            println!("Pressed mouse button '{:?}'", button);
        }

        if let Some(Button::Keyboard(key)) = event.press_args() {
            println!("Pressed key '{:?}'", key);
            if key == Key::Escape {
                println!("ui on");
                self.ui_on = true;
                if self.capture {
                    self.capture = false;
                    resources.window.set_capture_cursor(false);
                }
            } else if !self.ui_on && key == Key::C {
                self.capture = !self.capture;
                resources.window.set_capture_cursor(self.capture);
                println!("capture {}", if self.capture { "on" } else { "off" });
            }
        }

        StateTransition::Continue
    }

    fn update(&mut self, dt: f64) -> StateTransition {
        // Rotate 2 radians per second.
        self.rotation += 2.0 * dt;

        StateTransition::Continue
    }

    fn draw(&mut self, c: Context, g: &mut G2d) {
        use graphics::*;

        const GREEN: [f32; 4] = [0.0, 1.0, 0.0, 1.0];
        const RED:   [f32; 4] = [1.0, 0.0, 0.0, 1.0];

        let square = rectangle::square(0.0, 0.0, 150.0);
        let rotation = self.rotation;
        let (x, y) = (100.0, 100.0);

        // Clear the screen.
        clear(GREEN, g);

        let transform = c.transform.trans(x, y)
                                   .rot_rad(rotation)
                                   .trans(-75.0, -75.0);

        // Draw a box rotating around the middle of the screen.
        rectangle(RED, square, transform, g);

        if self.ui_on {
            let primitives = self.ui.draw();
            fn texture_from_image<T>(img: &T) -> &T { img };
            conrod::backend::piston_window::draw(c, g, primitives,
                                                 &mut self.text_texture_cache,
                                                 &self.image_map,
                                                 texture_from_image);
        }
    }
}
<|endoftext|>


<commit_before>use csv;
use indicatif;
use lazy_static;
use prettytable;
use regex;
use rusqlite;
use rustyline;
use std::env;
use std::error::Error;
use std::fs::File;

fn _normalize_col(col: &str) -> String {
    lazy_static::lazy_static! {
        static ref RE: regex::Regex = regex::Regex::new(r"\(.*?\)").unwrap();
    }
    RE.replace_all(col, "")
        .to_lowercase()
        .trim()
        .replace(" ", "_")
        .replace(".", "_")
        .replace("-", "_")
        .replace("/", "_")
        .replace("?", "")
}

fn _create_table(db: &mut rusqlite::Connection, table_name: &str, cols: &[String]) {
    let create_columns = cols
        .iter()
        .map(|c| format!("{} varchar", c))
        .collect::<Vec<String>>()
        .join(", ");
    db.execute(
        &format!("CREATE TABLE {} ({})", table_name, create_columns),
        &[] as &[&rusqlite::types::ToSql],
    )
    .unwrap();
}

fn _load_table_from_path(
    db: &mut rusqlite::Connection,
    table_name: &str,
    path: String,
) -> Vec<String> {
    let mut num_rows = 0;
    let f = File::open(path).unwrap();
    let file_size = f.metadata().unwrap().len();
    let mut reader = csv::Reader::from_reader(f);

    let normalized_cols =
        reader
            .headers()
            .unwrap()
            .iter()
            .map(_normalize_col)
            .fold(vec![], |mut v, orig_col| {
                let mut col = orig_col.clone();
                let mut i = 1;
                while v.contains(&col) {
                    col = format!("{}_{}", orig_col, i);
                    i += 1
                }
                v.push(col);
                v
            });
    _create_table(db, table_name, &normalized_cols);

    let insert_query = format!(
        "INSERT INTO {} VALUES ({})",
        table_name,
        normalized_cols
            .iter()
            .map(|_| "?")
            .collect::<Vec<_>>()
            .join(", ")
    );
    let pb = indicatif::ProgressBar::new(file_size);
    pb.set_style(
        indicatif::ProgressStyle::default_bar()
            .template("[{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({eta})")
            .progress_chars("#>-"),
    );
    let mut records = reader.records();
    let tx = db.transaction().unwrap();
    {
        let mut stmt = tx.prepare(&insert_query).unwrap();
        while let Some(row) = records.next() {
            let res = stmt.execute(&row.unwrap());
            assert!(res.is_ok());

            num_rows += 1;
            if num_rows % 10000 == 0 {
                pb.set_position(records.reader().position().byte())
            }
        }
    }
    tx.commit().unwrap();
    pb.finish();

    println!(
        "Loaded {} rows into {}({})",
        num_rows,
        table_name,
        normalized_cols.join(", "),
    );
    normalized_cols
}

struct FromAnySqlType {
    value: String,
}

impl rusqlite::types::FromSql for FromAnySqlType {
    fn column_result(
        value: rusqlite::types::ValueRef,
    ) -> Result<FromAnySqlType, rusqlite::types::FromSqlError> {
        let result = match value {
            rusqlite::types::ValueRef::Null => "null".to_string(),
            rusqlite::types::ValueRef::Integer(v) => v.to_string(),
            rusqlite::types::ValueRef::Real(v) => v.to_string(),
            rusqlite::types::ValueRef::Text(v) => v.to_string(),
            rusqlite::types::ValueRef::Blob(v) => String::from_utf8(v.to_vec()).unwrap(),
        };
        Ok(FromAnySqlType { value: result })
    }
}

fn _print_table(conn: &mut rusqlite::Connection, line: &str) {
    let mut table = prettytable::Table::new();
    table.set_format(*prettytable::format::consts::FORMAT_NO_LINESEP_WITH_TITLE);
    let mut stmt = match conn.prepare(&line) {
        Ok(stmt) => stmt,
        Err(e) => {
            println!("{}", e.description());
            return;
        }
    };

    let mut title_row = prettytable::Row::new(vec![]);
    for col in stmt.column_names() {
        title_row.add_cell(prettytable::Cell::new(col));
    }
    table.set_titles(title_row);

    let mut results = stmt.query(&[] as &[&rusqlite::types::ToSql]).unwrap();
    while let Ok(Some(r)) = results.next() {
        let mut row = prettytable::Row::new(vec![]);
        for i in 0..r.column_count() {
            let cell: FromAnySqlType = r.get(i).unwrap();
            row.add_cell(prettytable::Cell::new(&cell.value));
        }
        table.add_row(row);
    }
    table.printstd();
}

struct SimpleWordCompleter {
    words: Vec<String>,
}

static BREAK_CHARS: [u8; 4] = [b' ', b'(', b')', b','];
impl SimpleWordCompleter {
    fn new(words: Vec<String>) -> SimpleWordCompleter {
        SimpleWordCompleter { words }
    }
}

impl rustyline::Helper for SimpleWordCompleter {}

impl rustyline::hint::Hinter for SimpleWordCompleter {
    fn hint(&self, _line: &str, _pos: usize, _ctx: &rustyline::Context) -> Option<String> {
        None
    }
}

impl rustyline::highlight::Highlighter for SimpleWordCompleter {}

impl rustyline::completion::Completer for SimpleWordCompleter {
    type Candidate = String;

    fn complete(
        &self,
        line: &str,
        pos: usize,
        _ctx: &rustyline::Context,
    ) -> rustyline::Result<(usize, Vec<String>)> {
        let (start, word) = rustyline::completion::extract_word(line, pos, None, &BREAK_CHARS);

        let matches = self
            .words
            .iter()
            .filter(|w| w.starts_with(word))
            .cloned()
            .collect();
        Ok((start, matches))
    }
}

fn main() {
    let mut paths = env::args().skip(1);

    let mut conn = rusqlite::Connection::open_in_memory().unwrap();

    let mut base_words = vec![
        "distinct", "select", "from", "group", "by", "order", "where", "count", "limit", "offset",
    ]
    .iter()
    .map(|s| s.to_string())
    .collect::<Vec<String>>();

    if paths.len() == 1 {
        let mut col_names = _load_table_from_path(&mut conn, "t", paths.next().unwrap());
        base_words.append(&mut col_names);
    } else {
        for (idx, path) in paths.enumerate() {
            let mut col_names = _load_table_from_path(&mut conn, &format!("t{}", idx + 1), path);
            base_words.append(&mut col_names);
        }
    }

    let completer = SimpleWordCompleter::new(base_words);
    let mut rl = rustyline::Editor::new();
    rl.set_helper(Some(completer));
    loop {
        match rl.readline("> ") {
            Ok(line) => {
                if line.trim().is_empty() {
                    continue;
                }
                rl.add_history_entry(line.clone());
                _print_table(&mut conn, &line);
            }
            Err(rustyline::error::ReadlineError::Interrupted) => {
                println!("Interrupted");
                continue;
            }
            Err(rustyline::error::ReadlineError::Eof) => {
                break;
            }
            Err(err) => {
                println!("Error: {}", err);
                break;
            }
        }
    }
}
<commit_msg>Improve error message quality on insert failure<commit_after>use csv;
use indicatif;
use lazy_static;
use prettytable;
use regex;
use rusqlite;
use rustyline;
use std::env;
use std::error::Error;
use std::fs::File;

fn _normalize_col(col: &str) -> String {
    lazy_static::lazy_static! {
        static ref RE: regex::Regex = regex::Regex::new(r"\(.*?\)").unwrap();
    }
    RE.replace_all(col, "")
        .to_lowercase()
        .trim()
        .replace(" ", "_")
        .replace(".", "_")
        .replace("-", "_")
        .replace("/", "_")
        .replace("?", "")
}

fn _create_table(db: &mut rusqlite::Connection, table_name: &str, cols: &[String]) {
    let create_columns = cols
        .iter()
        .map(|c| format!("{} varchar", c))
        .collect::<Vec<String>>()
        .join(", ");
    db.execute(
        &format!("CREATE TABLE {} ({})", table_name, create_columns),
        &[] as &[&rusqlite::types::ToSql],
    )
    .unwrap();
}

fn _load_table_from_path(
    db: &mut rusqlite::Connection,
    table_name: &str,
    path: String,
) -> Vec<String> {
    let mut num_rows = 0;
    let f = File::open(path).unwrap();
    let file_size = f.metadata().unwrap().len();
    let mut reader = csv::Reader::from_reader(f);

    let normalized_cols =
        reader
            .headers()
            .unwrap()
            .iter()
            .map(_normalize_col)
            .fold(vec![], |mut v, orig_col| {
                let mut col = orig_col.clone();
                let mut i = 1;
                while v.contains(&col) {
                    col = format!("{}_{}", orig_col, i);
                    i += 1
                }
                v.push(col);
                v
            });
    _create_table(db, table_name, &normalized_cols);

    let insert_query = format!(
        "INSERT INTO {} VALUES ({})",
        table_name,
        normalized_cols
            .iter()
            .map(|_| "?")
            .collect::<Vec<_>>()
            .join(", ")
    );
    let pb = indicatif::ProgressBar::new(file_size);
    pb.set_style(
        indicatif::ProgressStyle::default_bar()
            .template("[{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({eta})")
            .progress_chars("#>-"),
    );
    let mut records = reader.records();
    let tx = db.transaction().unwrap();
    {
        let mut stmt = tx.prepare(&insert_query).unwrap();
        while let Some(row) = records.next() {
            stmt.execute(&row.unwrap()).unwrap();

            num_rows += 1;
            if num_rows % 10000 == 0 {
                pb.set_position(records.reader().position().byte())
            }
        }
    }
    tx.commit().unwrap();
    pb.finish();

    println!(
        "Loaded {} rows into {}({})",
        num_rows,
        table_name,
        normalized_cols.join(", "),
    );
    normalized_cols
}

struct FromAnySqlType {
    value: String,
}

impl rusqlite::types::FromSql for FromAnySqlType {
    fn column_result(
        value: rusqlite::types::ValueRef,
    ) -> Result<FromAnySqlType, rusqlite::types::FromSqlError> {
        let result = match value {
            rusqlite::types::ValueRef::Null => "null".to_string(),
            rusqlite::types::ValueRef::Integer(v) => v.to_string(),
            rusqlite::types::ValueRef::Real(v) => v.to_string(),
            rusqlite::types::ValueRef::Text(v) => v.to_string(),
            rusqlite::types::ValueRef::Blob(v) => String::from_utf8(v.to_vec()).unwrap(),
        };
        Ok(FromAnySqlType { value: result })
    }
}

fn _print_table(conn: &mut rusqlite::Connection, line: &str) {
    let mut table = prettytable::Table::new();
    table.set_format(*prettytable::format::consts::FORMAT_NO_LINESEP_WITH_TITLE);
    let mut stmt = match conn.prepare(&line) {
        Ok(stmt) => stmt,
        Err(e) => {
            println!("{}", e.description());
            return;
        }
    };

    let mut title_row = prettytable::Row::new(vec![]);
    for col in stmt.column_names() {
        title_row.add_cell(prettytable::Cell::new(col));
    }
    table.set_titles(title_row);

    let mut results = stmt.query(&[] as &[&rusqlite::types::ToSql]).unwrap();
    while let Ok(Some(r)) = results.next() {
        let mut row = prettytable::Row::new(vec![]);
        for i in 0..r.column_count() {
            let cell: FromAnySqlType = r.get(i).unwrap();
            row.add_cell(prettytable::Cell::new(&cell.value));
        }
        table.add_row(row);
    }
    table.printstd();
}

struct SimpleWordCompleter {
    words: Vec<String>,
}

static BREAK_CHARS: [u8; 4] = [b' ', b'(', b')', b','];
impl SimpleWordCompleter {
    fn new(words: Vec<String>) -> SimpleWordCompleter {
        SimpleWordCompleter { words }
    }
}

impl rustyline::Helper for SimpleWordCompleter {}

impl rustyline::hint::Hinter for SimpleWordCompleter {
    fn hint(&self, _line: &str, _pos: usize, _ctx: &rustyline::Context) -> Option<String> {
        None
    }
}

impl rustyline::highlight::Highlighter for SimpleWordCompleter {}

impl rustyline::completion::Completer for SimpleWordCompleter {
    type Candidate = String;

    fn complete(
        &self,
        line: &str,
        pos: usize,
        _ctx: &rustyline::Context,
    ) -> rustyline::Result<(usize, Vec<String>)> {
        let (start, word) = rustyline::completion::extract_word(line, pos, None, &BREAK_CHARS);

        let matches = self
            .words
            .iter()
            .filter(|w| w.starts_with(word))
            .cloned()
            .collect();
        Ok((start, matches))
    }
}

fn main() {
    let mut paths = env::args().skip(1);

    let mut conn = rusqlite::Connection::open_in_memory().unwrap();

    let mut base_words = vec![
        "distinct", "select", "from", "group", "by", "order", "where", "count", "limit", "offset",
    ]
    .iter()
    .map(|s| s.to_string())
    .collect::<Vec<String>>();

    if paths.len() == 1 {
        let mut col_names = _load_table_from_path(&mut conn, "t", paths.next().unwrap());
        base_words.append(&mut col_names);
    } else {
        for (idx, path) in paths.enumerate() {
            let mut col_names = _load_table_from_path(&mut conn, &format!("t{}", idx + 1), path);
            base_words.append(&mut col_names);
        }
    }

    let completer = SimpleWordCompleter::new(base_words);
    let mut rl = rustyline::Editor::new();
    rl.set_helper(Some(completer));
    loop {
        match rl.readline("> ") {
            Ok(line) => {
                if line.trim().is_empty() {
                    continue;
                }
                rl.add_history_entry(line.clone());
                _print_table(&mut conn, &line);
            }
            Err(rustyline::error::ReadlineError::Interrupted) => {
                println!("Interrupted");
                continue;
            }
            Err(rustyline::error::ReadlineError::Eof) => {
                break;
            }
            Err(err) => {
                println!("Error: {}", err);
                break;
            }
        }
    }
}
<|endoftext|>


<commit_before>extern crate alpm;
extern crate clap;
extern crate curl;
extern crate env_logger;
extern crate itertools;
extern crate json;
#[macro_use]
extern crate log;
extern crate select;

use clap::{Arg, App};
use curl::easy::Easy;
use itertools::Itertools;
use select::document::Document;
use select::predicate::Name;
use std::cmp::Ordering;
use std::collections::HashMap;
use std::collections::hash_map::Entry::{Occupied, Vacant};
use std::default::Default;
use std::process::exit;
use std::str;

#[derive(Debug, Clone)]
struct ASA {
    cve: Vec<String>,
    version: Option<String>,
}

impl Default for ASA {
    fn default() -> ASA {
        ASA {
            cve: vec![],
            version: None,
        }
    }
}

#[derive(Debug)]
struct Options {
    format: Option<String>,
    quiet: u64,
    upgradable_only: bool,
}

impl Default for Options {
    fn default() -> Options {
        Options {
            format: None,
            quiet: 0,
            upgradable_only: false
        }
    }
}

fn main() {
    env_logger::init().unwrap();

    let args = App::new("arch-audit")
                        .version("0.1.4")
                        .arg(Arg::with_name("dbpath")
                             .short("b")
                             .long("dbpath")
                             .takes_value(true)
                             .help("Set an alternate database location"))
                        .arg(Arg::with_name("format")
                             .short("f")
                             .long("format")
                             .takes_value(true)
                             .help("Specify a format to control the output. Placeholders are %n (pkgname) and %c (CVEs)"))
                        .arg(Arg::with_name("quiet")
                             .short("q")
                             .long("quiet")
                             .multiple(true)
                             .help("Show only vulnerable package names and their versions"))
                        .arg(Arg::with_name("upgradable")
                             .short("u")
                             .long("upgradable")
                             .help("Show only packages that have already been fixed"))
                        .get_matches();

    let options = Options {
        format: { match args.value_of("format") {
            Some(f) => { Some(f.to_string()) },
            None => { None },
            } },
        quiet: args.occurrences_of("quiet"),
        upgradable_only: args.is_present("upgradable"),
    };

    let mut wikipage = String::new();
    {
        info!("Downloading CVE wiki page...");
        let wikipage_url = "https://wiki.archlinux.org/api.php?format=json&action=parse&page=CVE&section=5";

        let mut easy = Easy::new();
        easy.url(wikipage_url).unwrap();
        let mut transfer = easy.transfer();
        transfer.write_function(|data| {
            wikipage.push_str(str::from_utf8(data).unwrap());
            Ok(data.len())
        }).unwrap();
        match transfer.perform() {
            Ok(_) => {},
            Err(_) => { println!("Cannot fetch data, please check your network connection!"); exit(1) },
        };
    }

    let mut infos: HashMap<String, Vec<_>> = HashMap::new();
    {
        let json = json::parse(wikipage.as_str()).unwrap();
        let document = Document::from(json["parse"]["text"]["*"].as_str().unwrap());

        for tr in document.find(Name("tbody")).find(Name("tr")).iter() {
            let tds = tr.find(Name("td"));

            match tds.first() {
                Some(td) => {
                    let mut next = tds.next().next();
                    let pkgname = next.first().unwrap().text().trim().to_string();
                    next = next.next().next().next().next().next().next();
                    let info = ASA {
                        cve: td.text().split_whitespace().filter(|s| s.starts_with("CVE")).map(|s| s.to_string()).collect(),
                        version: {
                            let v = next.first().unwrap().text().trim().to_string();
                            if !v.is_empty() && v != "?".to_string() && v != "-".to_string() { Some(v) } else { None }
                        },
                    };
                    next = next.next().next().next().next();
                    let status = next.first().unwrap().text().trim().to_string();

                    if !status.starts_with("Not Affected") {
                        match infos.entry(pkgname) {
                            Occupied(c) => { c.into_mut() },
                            Vacant(c) => { c.insert(vec![]) },
                        }.push(info);
                    }
                },
                None => {},
            };
        }
    }

    let pacman = match args.value_of("dbpath") {
        Some(path) => { alpm::Alpm::with_dbpath(path.to_string()).unwrap() },
        None => { alpm::Alpm::new().unwrap() },
    };

    for (pkg, cves) in infos {
        let pkg_cves = get_asas_for_package_by_version(&pacman, &pkg, &cves);

        let (relevant_cves, version) = cves_with_relevant_version(pkg_cves, &pacman);

        if !relevant_cves.is_empty() {
            print_asa(&options, &pkg, relevant_cves, version);
        }
    }
}

/// gets a map from version to ASAs for a specific package, returns list of CVEs for that
/// package and a version, which is the most updated
fn cves_with_relevant_version(pkg_cves: HashMap<Option<String>, Vec<ASA>>, pacman: &alpm::Alpm) -> (Vec<String>, Option<String>) {
    let mut cves: Vec<String> = vec![];
    // the newest version, which fixes all/most ? cves
    let mut version: Option<String> = None;
    for (v, asas) in pkg_cves {
        for asa in asas {
            for cve in asa.cve {
                cves.push(cve);
            }
        }
        match v {
            Some(ref v) => {
                version = match version {
                    None => Some(v.clone()),
                    Some(ref vv) => {
                        match pacman.vercmp(v.clone(), vv.clone()).unwrap() {
                            Ordering::Greater => Some(v.clone()),
                            _ => version.clone()
                        }
                    }
                }
            }
            _ => {}
        }
    }
    (cves, version)
}

#[test]
fn test_cves_with_relevant_version() {
    let mut map = HashMap::new();
    map.insert(Some("1.0".to_string()), vec![ASA{cve: vec!["a".to_string(), "b".to_string()], version: Some("1.0".to_string())}]);
    map.insert(Some("2.0".to_string()), vec![ASA{cve: vec!["c".to_string()], version: Some("2.0".to_string())}]);
    map.insert(Some("3.0".to_string()), vec![ASA{cve: vec![], version: Some("3.0".to_string())}]);

    let pacman = alpm::Alpm::new().unwrap();
    let (mut cves, version) = cves_with_relevant_version(map, &pacman);

    cves.sort();
    assert_eq!(cves, vec!["a".to_string(), "b".to_string(), "c".to_string()]);
    assert_eq!(version, Some("3.0".to_string()));
}

/// creates a map between version and ASAs from a list of ASAs
fn get_asas_for_package_by_version(pacman: &alpm::Alpm, pkg: &String, cves: &Vec<ASA>) -> HashMap<Option<String>, Vec<ASA>> {
    let mut pkg_cves: HashMap<Option<String>, Vec<ASA>> = HashMap::new();
    match pacman.query_package_version(pkg.clone()) {
        Ok(v) => {
            info!("Found installed version {} for package {}", v, pkg);
            for cve in cves {
                match cve.version {
                    Some(ref version) => {
                        info!("Comparing with fixed version {}", version);
                        match pacman.vercmp(v.clone(), version.clone()).unwrap() {
                            Ordering::Less => {
                                match pkg_cves.entry(Some(version.clone())) {
                                    Occupied(c) => { c.into_mut() },
                                    Vacant(c) => { c.insert(vec![]) },
                                }.push(cve.clone());
                            },
                            _ => {},
                        };
                    },
                    None => {
                        match pkg_cves.entry(None) {
                            Occupied(c) => { c.into_mut() },
                            Vacant(c) => { c.insert(vec![]) },
                        }.push(cve.clone());
                    },
                };
            };
        },
        Err(_) => { debug!("Package {} not installed", pkg) },
    }
    pkg_cves
}

fn print_asa(options: &Options, pkgname: &String, cve: Vec<String>, version: Option<String>) {
    let msg = format!("Package {} is affected by {:?}", pkgname, cve);

    match version {
        Some(v) => {
            if options.quiet == 1 {
                println!("{}>={}", pkgname, v);
            } else if options.quiet >= 2 {
                println!("{}", pkgname);
            } else {
                match options.format {
                    Some(ref f) => { println!("{}", f.replace("%n", pkgname).replace("%c", cve.iter().join(",").as_str())) },
                    None => { println!("{}. Update to {}!", msg, v) },
                }
            }
        }
        None => {
            if !options.upgradable_only {
                if options.quiet > 0 {
                    println!("{}", pkgname);
                } else {
                    match options.format {
                        Some(ref f) => { println!("{}", f.replace("%n", pkgname).replace("%c", cve.iter().join(",").as_str())) },
                        None => { println!("{}. VULNERABLE!", msg) },
                    }
                }
            }
        }
    }
}
<commit_msg>rustfmt<commit_after>extern crate alpm;
extern crate clap;
extern crate curl;
extern crate env_logger;
extern crate itertools;
extern crate json;
#[macro_use]
extern crate log;
extern crate select;

use clap::{Arg, App};
use curl::easy::Easy;
use itertools::Itertools;
use select::document::Document;
use select::predicate::Name;
use std::cmp::Ordering;
use std::collections::HashMap;
use std::collections::hash_map::Entry::{Occupied, Vacant};
use std::default::Default;
use std::process::exit;
use std::str;

#[derive(Debug, Clone)]
struct ASA {
    cve: Vec<String>,
    version: Option<String>,
}

impl Default for ASA {
    fn default() -> ASA {
        ASA {
            cve: vec![],
            version: None,
        }
    }
}

#[derive(Debug)]
struct Options {
    format: Option<String>,
    quiet: u64,
    upgradable_only: bool,
}

impl Default for Options {
    fn default() -> Options {
        Options {
            format: None,
            quiet: 0,
            upgradable_only: false,
        }
    }
}

fn main() {
    env_logger::init().unwrap();

    let args = App::new("arch-audit")
        .version("0.1.4")
        .arg(Arg::with_name("dbpath")
            .short("b")
            .long("dbpath")
            .takes_value(true)
            .help("Set an alternate database location"))
        .arg(Arg::with_name("format")
            .short("f")
            .long("format")
            .takes_value(true)
            .help("Specify a format to control the output. Placeholders are %n (pkgname) and %c \
                   (CVEs)"))
        .arg(Arg::with_name("quiet")
            .short("q")
            .long("quiet")
            .multiple(true)
            .help("Show only vulnerable package names and their versions"))
        .arg(Arg::with_name("upgradable")
            .short("u")
            .long("upgradable")
            .help("Show only packages that have already been fixed"))
        .get_matches();

    let options = Options {
        format: {
            match args.value_of("format") {
                Some(f) => Some(f.to_string()),
                None => None,
            }
        },
        quiet: args.occurrences_of("quiet"),
        upgradable_only: args.is_present("upgradable"),
    };

    let mut wikipage = String::new();
    {
        info!("Downloading CVE wiki page...");
        let wikipage_url = "https://wiki.archlinux.org/api.\
                            php?format=json&action=parse&page=CVE&section=5";

        let mut easy = Easy::new();
        easy.url(wikipage_url).unwrap();
        let mut transfer = easy.transfer();
        transfer.write_function(|data| {
                wikipage.push_str(str::from_utf8(data).unwrap());
                Ok(data.len())
            })
            .unwrap();
        match transfer.perform() {
            Ok(_) => {}
            Err(_) => {
                println!("Cannot fetch data, please check your network connection!");
                exit(1)
            }
        };
    }

    let mut infos: HashMap<String, Vec<_>> = HashMap::new();
    {
        let json = json::parse(wikipage.as_str()).unwrap();
        let document = Document::from(json["parse"]["text"]["*"].as_str().unwrap());

        for tr in document.find(Name("tbody")).find(Name("tr")).iter() {
            let tds = tr.find(Name("td"));

            match tds.first() {
                Some(td) => {
                    let mut next = tds.next().next();
                    let pkgname = next.first().unwrap().text().trim().to_string();
                    next = next.next().next().next().next().next().next();
                    let info = ASA {
                        cve: td.text()
                            .split_whitespace()
                            .filter(|s| s.starts_with("CVE"))
                            .map(|s| s.to_string())
                            .collect(),
                        version: {
                            let v = next.first().unwrap().text().trim().to_string();
                            if !v.is_empty() && v != "?".to_string() && v != "-".to_string() {
                                Some(v)
                            } else {
                                None
                            }
                        },
                    };
                    next = next.next().next().next().next();
                    let status = next.first().unwrap().text().trim().to_string();

                    if !status.starts_with("Not Affected") {
                        match infos.entry(pkgname) {
                                Occupied(c) => c.into_mut(),
                                Vacant(c) => c.insert(vec![]),
                            }
                            .push(info);
                    }
                }
                None => {}
            };
        }
    }

    let pacman = match args.value_of("dbpath") {
        Some(path) => alpm::Alpm::with_dbpath(path.to_string()).unwrap(),
        None => alpm::Alpm::new().unwrap(),
    };

    for (pkg, cves) in infos {
        let pkg_cves = get_asas_for_package_by_version(&pacman, &pkg, &cves);

        let (relevant_cves, version) = cves_with_relevant_version(pkg_cves, &pacman);

        if !relevant_cves.is_empty() {
            print_asa(&options, &pkg, relevant_cves, version);
        }
    }
}

/// gets a map from version to ASAs for a specific package, returns list of CVEs for that
/// package and a version, which is the most updated
fn cves_with_relevant_version(pkg_cves: HashMap<Option<String>, Vec<ASA>>,
                              pacman: &alpm::Alpm)
                              -> (Vec<String>, Option<String>) {
    let mut cves: Vec<String> = vec![];
    // the newest version, which fixes all/most ? cves
    let mut version: Option<String> = None;
    for (v, asas) in pkg_cves {
        for asa in asas {
            for cve in asa.cve {
                cves.push(cve);
            }
        }
        match v {
            Some(ref v) => {
                version = match version {
                    None => Some(v.clone()),
                    Some(ref vv) => {
                        match pacman.vercmp(v.clone(), vv.clone()).unwrap() {
                            Ordering::Greater => Some(v.clone()),
                            _ => version.clone(),
                        }
                    }
                }
            }
            _ => {}
        }
    }
    (cves, version)
}

#[test]
fn test_cves_with_relevant_version() {
    let mut map = HashMap::new();
    map.insert(Some("1.0".to_string()),
               vec![ASA {
                        cve: vec!["a".to_string(), "b".to_string()],
                        version: Some("1.0".to_string()),
                    }]);
    map.insert(Some("2.0".to_string()),
               vec![ASA {
                        cve: vec!["c".to_string()],
                        version: Some("2.0".to_string()),
                    }]);
    map.insert(Some("3.0".to_string()),
               vec![ASA {
                        cve: vec![],
                        version: Some("3.0".to_string()),
                    }]);

    let pacman = alpm::Alpm::new().unwrap();
    let (mut cves, version) = cves_with_relevant_version(map, &pacman);

    cves.sort();
    assert_eq!(cves,
               vec!["a".to_string(), "b".to_string(), "c".to_string()]);
    assert_eq!(version, Some("3.0".to_string()));
}

/// creates a map between version and ASAs from a list of ASAs
fn get_asas_for_package_by_version(pacman: &alpm::Alpm,
                                   pkg: &String,
                                   cves: &Vec<ASA>)
                                   -> HashMap<Option<String>, Vec<ASA>> {
    let mut pkg_cves: HashMap<Option<String>, Vec<ASA>> = HashMap::new();
    match pacman.query_package_version(pkg.clone()) {
        Ok(v) => {
            info!("Found installed version {} for package {}", v, pkg);
            for cve in cves {
                match cve.version {
                    Some(ref version) => {
                        info!("Comparing with fixed version {}", version);
                        match pacman.vercmp(v.clone(), version.clone()).unwrap() {
                            Ordering::Less => {
                                match pkg_cves.entry(Some(version.clone())) {
                                        Occupied(c) => c.into_mut(),
                                        Vacant(c) => c.insert(vec![]),
                                    }
                                    .push(cve.clone());
                            }
                            _ => {}
                        };
                    }
                    None => {
                        match pkg_cves.entry(None) {
                                Occupied(c) => c.into_mut(),
                                Vacant(c) => c.insert(vec![]),
                            }
                            .push(cve.clone());
                    }
                };
            }
        }
        Err(_) => debug!("Package {} not installed", pkg),
    }
    pkg_cves
}

fn print_asa(options: &Options, pkgname: &String, cve: Vec<String>, version: Option<String>) {
    let msg = format!("Package {} is affected by {:?}", pkgname, cve);

    match version {
        Some(v) => {
            if options.quiet == 1 {
                println!("{}>={}", pkgname, v);
            } else if options.quiet >= 2 {
                println!("{}", pkgname);
            } else {
                match options.format {
                    Some(ref f) => {
                        println!("{}",
                                 f.replace("%n", pkgname)
                                     .replace("%c", cve.iter().join(",").as_str()))
                    }
                    None => println!("{}. Update to {}!", msg, v),
                }
            }
        }
        None => {
            if !options.upgradable_only {
                if options.quiet > 0 {
                    println!("{}", pkgname);
                } else {
                    match options.format {
                        Some(ref f) => {
                            println!("{}",
                                     f.replace("%n", pkgname)
                                         .replace("%c", cve.iter().join(",").as_str()))
                        }
                        None => println!("{}. VULNERABLE!", msg),
                    }
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>//! Joins two blocks of text, such that row 1 of the first block will be
//! appended to row 1 of the second block and so on.
//!
//! Input is read from stdin until EOF, output to stdout. The the first half of
//! the input lines will become the first block.

use std::io::prelude::*;
use std::io::stdin;
use std::error::Error;
use std::env;

fn main() {
    run().unwrap();
}

fn run() -> Result<(), Box<Error>> {
    let separator: String = env::args().nth(1).unwrap_or("".to_string());
    let mut buffer = String::new();
    try!(stdin().read_to_string(&mut buffer));
    let lines: Vec<&str> = buffer.lines().collect();
    let lines_per_block = lines.len() / 2;
    let first = &lines[0..lines_per_block];
    let second = &lines[lines_per_block..];
    for (&left, &right) in first.iter().zip(second.iter()) {
        println!("{}{}{}", left, separator, right);
    }
    Ok(())
}
<commit_msg>Refactor: simplify a bit<commit_after>//! Joins two blocks of text, such that row 1 of the first block will be
//! appended to row 1 of the second block and so on.
//!
//! Input is read from stdin until EOF, output to stdout. The the first half of
//! the input lines will become the first block.

use std::io::prelude::*;
use std::io::stdin;
use std::error::Error;
use std::env;

fn main() {
    run().unwrap();
}

fn run() -> Result<(), Box<Error>> {
    let separator: String = env::args().nth(1).unwrap_or("".to_string());
    let mut buffer = String::new();
    try!(stdin().read_to_string(&mut buffer));
    let lines: Vec<&str> = buffer.lines().collect();
    let (first, second) = lines.split_at(lines.len() / 2);
    for (&left, &right) in first.iter().zip(second.iter()) {
        println!("{}{}{}", left, separator, right);
    }
    Ok(())
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use proc_macro2::Span;
use std::iter::Peekable;
use std::{fmt::Display, sync::Arc};
use syn::{parse_quote, Ident, PathSegment, TypePath};

use crate::known_types::known_types;

pub(crate) fn make_ident<S: AsRef<str>>(id: S) -> Ident {
    Ident::new(id.as_ref(), Span::call_site())
}

/// Newtype wrapper for a C++ namespace.
#[derive(Debug, PartialEq, PartialOrd, Eq, Hash, Clone)]
#[allow(clippy::rc_buffer)]
pub struct Namespace(Arc<Vec<String>>);

impl Namespace {
    pub(crate) fn new() -> Self {
        Self(Arc::new(Vec::new()))
    }

    #[must_use]
    pub(crate) fn push(&self, segment: String) -> Self {
        let mut bigger = (*self.0).clone();
        bigger.push(segment);
        Namespace(Arc::new(bigger))
    }

    pub(crate) fn is_empty(&self) -> bool {
        self.0.is_empty()
    }

    pub(crate) fn iter(&self) -> impl Iterator<Item = &String> {
        self.0.iter()
    }

    #[cfg(test)]
    pub(crate) fn from_user_input(input: &str) -> Self {
        Self(Arc::new(input.split("::").map(|x| x.to_string()).collect()))
    }

    pub(crate) fn depth(&self) -> usize {
        self.0.len()
    }

    pub(crate) fn to_display_suffix(&self) -> String {
        if self.is_empty() {
            String::new()
        } else {
            format!(" (in namespace {})", self)
        }
    }
}

impl Display for Namespace {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(&self.0.join("::"))
    }
}

impl<'a> IntoIterator for &'a Namespace {
    type Item = &'a String;

    type IntoIter = std::slice::Iter<'a, String>;

    fn into_iter(self) -> Self::IntoIter {
        self.0.iter()
    }
}

/// Any time we store a qualified name, we should use this. Stores the type
/// and its namespace. Namespaces should be stored without any
/// 'bindgen::root' prefix; that means a type not in any C++
/// namespace should have an empty namespace segment list.
/// Some types have names that change as they flow through the
/// autocxx pipeline. e.g. you start with std::string
/// and end up with CxxString. This TypeName type can store
/// either. It doesn't directly have functionality to convert
/// from one to the other; `replace_type_path_without_arguments`
/// does that.
#[derive(Debug, PartialEq, PartialOrd, Eq, Hash, Clone)]
pub struct QualifiedName(Namespace, String);

impl QualifiedName {
    /// From a TypePath which starts with 'root'
    pub(crate) fn from_type_path(typ: &TypePath) -> Self {
        let mut seg_iter = typ.path.segments.iter().peekable();
        let first_seg = seg_iter.next().unwrap().ident.clone();
        if first_seg == "root" {
            // This is a C++ type prefixed with a namespace,
            // e.g. std::string or something the user has defined.
            Self::from_segments(seg_iter) // all but 'root'
        } else {
            // This is actually a Rust type e.g.
            // std::os::raw::c_ulong. Start iterating from the beginning again.
            Self::from_segments(typ.path.segments.iter().peekable())
        }
    }

    fn from_segments<'a, T: Iterator<Item = &'a PathSegment>>(mut seg_iter: Peekable<T>) -> Self {
        let mut ns = Namespace::new();
        while let Some(seg) = seg_iter.next() {
            if seg_iter.peek().is_some() {
                ns = ns.push(seg.ident.to_string());
            } else {
                return Self(ns, seg.ident.to_string());
            }
        }
        unreachable!()
    }

    /// Create from a type encountered in the code.
    pub(crate) fn new(ns: &Namespace, id: Ident) -> Self {
        Self(ns.clone(), id.to_string())
    }

    /// Create from user input, e.g. a name in an AllowPOD directive.
    pub(crate) fn new_from_user_input(id: &str) -> Self {
        let mut seg_iter = id.split("::").peekable();
        let mut ns = Namespace::new();
        while let Some(seg) = seg_iter.next() {
            if seg_iter.peek().is_some() {
                ns = ns.push(seg.to_string());
            } else {
                return Self(ns, seg.to_string());
            }
        }
        unreachable!()
    }

    /// Return the actual type name, without any namespace
    /// qualification. Avoid unless you have a good reason.
    pub(crate) fn get_final_ident(&self) -> &str {
        &self.1
    }

    pub(crate) fn get_namespace(&self) -> &Namespace {
        &self.0
    }

    pub(crate) fn get_bindgen_path_idents(&self) -> Vec<Ident> {
        ["bindgen", "root"]
            .iter()
            .map(make_ident)
            .chain(self.ns_segment_iter().map(make_ident))
            .chain(std::iter::once(make_ident(self.get_final_ident().clone())))
            .collect()
    }

    /// Output the fully-qualified C++ name of this type.
    pub(crate) fn to_cpp_name(&self) -> String {
        let special_cpp_name = known_types().special_cpp_name(&self);
        match special_cpp_name {
            Some(name) => name,
            None => {
                let mut s = String::new();
                for seg in &self.0 {
                    s.push_str(&seg);
                    s.push_str("::");
                }
                s.push_str(&self.1);
                s
            }
        }
    }

    pub(crate) fn to_type_path(&self) -> TypePath {
        if let Some(known_type_path) = known_types().known_type_type_path(self) {
            known_type_path
        } else {
            let root = "root".to_string();
            let segs = std::iter::once(&root)
                .chain(self.ns_segment_iter())
                .chain(std::iter::once(&self.1))
                .map(make_ident);
            parse_quote! {
                #(#segs)::*
            }
        }
    }

    /// Iterator over segments in the namespace of this type.
    pub(crate) fn ns_segment_iter(&self) -> impl Iterator<Item = &String> {
        self.0.iter()
    }

    pub(crate) fn is_cvoid(&self) -> bool {
        self.to_cpp_name() == "void"
    }
}

impl Display for QualifiedName {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        for seg in &self.0 {
            f.write_str(&seg)?;
            f.write_str("::")?;
        }
        f.write_str(&self.1)
    }
}

#[cfg(test)]
mod tests {
    use super::QualifiedName;

    #[test]
    fn test_ints() {
        assert_eq!(
            QualifiedName::new_from_user_input("i8").to_cpp_name(),
            "int8_t"
        );
        assert_eq!(
            QualifiedName::new_from_user_input("u64").to_cpp_name(),
            "uint64_t"
        );
    }
}
<commit_msg>Clippy.<commit_after>// Copyright 2020 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use proc_macro2::Span;
use std::iter::Peekable;
use std::{fmt::Display, sync::Arc};
use syn::{parse_quote, Ident, PathSegment, TypePath};

use crate::known_types::known_types;

pub(crate) fn make_ident<S: AsRef<str>>(id: S) -> Ident {
    Ident::new(id.as_ref(), Span::call_site())
}

/// Newtype wrapper for a C++ namespace.
#[derive(Debug, PartialEq, PartialOrd, Eq, Hash, Clone)]
#[allow(clippy::rc_buffer)]
pub struct Namespace(Arc<Vec<String>>);

impl Namespace {
    pub(crate) fn new() -> Self {
        Self(Arc::new(Vec::new()))
    }

    #[must_use]
    pub(crate) fn push(&self, segment: String) -> Self {
        let mut bigger = (*self.0).clone();
        bigger.push(segment);
        Namespace(Arc::new(bigger))
    }

    pub(crate) fn is_empty(&self) -> bool {
        self.0.is_empty()
    }

    pub(crate) fn iter(&self) -> impl Iterator<Item = &String> {
        self.0.iter()
    }

    #[cfg(test)]
    pub(crate) fn from_user_input(input: &str) -> Self {
        Self(Arc::new(input.split("::").map(|x| x.to_string()).collect()))
    }

    pub(crate) fn depth(&self) -> usize {
        self.0.len()
    }

    pub(crate) fn to_display_suffix(&self) -> String {
        if self.is_empty() {
            String::new()
        } else {
            format!(" (in namespace {})", self)
        }
    }
}

impl Display for Namespace {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(&self.0.join("::"))
    }
}

impl<'a> IntoIterator for &'a Namespace {
    type Item = &'a String;

    type IntoIter = std::slice::Iter<'a, String>;

    fn into_iter(self) -> Self::IntoIter {
        self.0.iter()
    }
}

/// Any time we store a qualified name, we should use this. Stores the type
/// and its namespace. Namespaces should be stored without any
/// 'bindgen::root' prefix; that means a type not in any C++
/// namespace should have an empty namespace segment list.
/// Some types have names that change as they flow through the
/// autocxx pipeline. e.g. you start with std::string
/// and end up with CxxString. This TypeName type can store
/// either. It doesn't directly have functionality to convert
/// from one to the other; `replace_type_path_without_arguments`
/// does that.
#[derive(Debug, PartialEq, PartialOrd, Eq, Hash, Clone)]
pub struct QualifiedName(Namespace, String);

impl QualifiedName {
    /// From a TypePath which starts with 'root'
    pub(crate) fn from_type_path(typ: &TypePath) -> Self {
        let mut seg_iter = typ.path.segments.iter().peekable();
        let first_seg = seg_iter.next().unwrap().ident.clone();
        if first_seg == "root" {
            // This is a C++ type prefixed with a namespace,
            // e.g. std::string or something the user has defined.
            Self::from_segments(seg_iter) // all but 'root'
        } else {
            // This is actually a Rust type e.g.
            // std::os::raw::c_ulong. Start iterating from the beginning again.
            Self::from_segments(typ.path.segments.iter().peekable())
        }
    }

    fn from_segments<'a, T: Iterator<Item = &'a PathSegment>>(mut seg_iter: Peekable<T>) -> Self {
        let mut ns = Namespace::new();
        while let Some(seg) = seg_iter.next() {
            if seg_iter.peek().is_some() {
                ns = ns.push(seg.ident.to_string());
            } else {
                return Self(ns, seg.ident.to_string());
            }
        }
        unreachable!()
    }

    /// Create from a type encountered in the code.
    pub(crate) fn new(ns: &Namespace, id: Ident) -> Self {
        Self(ns.clone(), id.to_string())
    }

    /// Create from user input, e.g. a name in an AllowPOD directive.
    pub(crate) fn new_from_user_input(id: &str) -> Self {
        let mut seg_iter = id.split("::").peekable();
        let mut ns = Namespace::new();
        while let Some(seg) = seg_iter.next() {
            if seg_iter.peek().is_some() {
                ns = ns.push(seg.to_string());
            } else {
                return Self(ns, seg.to_string());
            }
        }
        unreachable!()
    }

    /// Return the actual type name, without any namespace
    /// qualification. Avoid unless you have a good reason.
    pub(crate) fn get_final_ident(&self) -> &str {
        &self.1
    }

    pub(crate) fn get_namespace(&self) -> &Namespace {
        &self.0
    }

    pub(crate) fn get_bindgen_path_idents(&self) -> Vec<Ident> {
        ["bindgen", "root"]
            .iter()
            .map(make_ident)
            .chain(self.ns_segment_iter().map(make_ident))
            .chain(std::iter::once(make_ident(self.get_final_ident())))
            .collect()
    }

    /// Output the fully-qualified C++ name of this type.
    pub(crate) fn to_cpp_name(&self) -> String {
        let special_cpp_name = known_types().special_cpp_name(&self);
        match special_cpp_name {
            Some(name) => name,
            None => {
                let mut s = String::new();
                for seg in &self.0 {
                    s.push_str(&seg);
                    s.push_str("::");
                }
                s.push_str(&self.1);
                s
            }
        }
    }

    pub(crate) fn to_type_path(&self) -> TypePath {
        if let Some(known_type_path) = known_types().known_type_type_path(self) {
            known_type_path
        } else {
            let root = "root".to_string();
            let segs = std::iter::once(&root)
                .chain(self.ns_segment_iter())
                .chain(std::iter::once(&self.1))
                .map(make_ident);
            parse_quote! {
                #(#segs)::*
            }
        }
    }

    /// Iterator over segments in the namespace of this type.
    pub(crate) fn ns_segment_iter(&self) -> impl Iterator<Item = &String> {
        self.0.iter()
    }

    pub(crate) fn is_cvoid(&self) -> bool {
        self.to_cpp_name() == "void"
    }
}

impl Display for QualifiedName {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        for seg in &self.0 {
            f.write_str(&seg)?;
            f.write_str("::")?;
        }
        f.write_str(&self.1)
    }
}

#[cfg(test)]
mod tests {
    use super::QualifiedName;

    #[test]
    fn test_ints() {
        assert_eq!(
            QualifiedName::new_from_user_input("i8").to_cpp_name(),
            "int8_t"
        );
        assert_eq!(
            QualifiedName::new_from_user_input("u64").to_cpp_name(),
            "uint64_t"
        );
    }
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

//! Structured specification of command-line arguments.
//!
//! This module specifies all command-line arguments that RRG offers and exposes
//! functions for parsing them into a high-level structure.
//!
//! Ideally, only one instance of this high-level structure should ever be
//! created (using the [`from_args`] function). Then this instance should be
//! shared through the entire lifetime of a program and explicitly passed to
//! functions that care about it.
//!
//! [`from_args`]: fn.from_args.html

use std::time::Duration;
use std::path::PathBuf;

use structopt::StructOpt;

#[derive(StructOpt)]
#[structopt(name = "RRG", about = "A GRR agent rewritten in Rust.")]
pub struct Opts {
    /// A level of log verbosity.
    #[structopt(long="log-verbosity", name="LEVEL", default_value="info",
                help="Specifies the level of log verbosity")]
    pub log_verbosity: Verbosity,

    /// A standard stream to log into.
    #[structopt(long="log-stream", name="STREAM",
                help="Enables logging to the specified standard stream")]
    pub log_stream: Option<Stream>,

    /// A path to the file to log into.
    #[structopt(long="log-file", name="FILE",
                help="Enables logging to the specified file")]
    pub log_file: Option<PathBuf>,

    /// A frequency of heartbeat messages to send to the Fleetspeak client.
    #[structopt(long="heartbeat-rate", name="DURATION", default_value="5s",
                parse(try_from_str = humantime::parse_duration),
                help="Specifies the frequency of heartbeat messages")]
    pub heartbeat_rate: Duration,
}

/// Parses command-line arguments.
///
/// This is a just a convenience function intended to be used as a shortcut for
/// creating instances of [`Opts`]. Ideally, it should be called only once in
/// the entire lifetime of the agent.
///
/// [`Opts`]: struct.Opts.html
pub fn from_args() -> Opts {
    Opts::from_args()
}

/// A type representing level of log verbosity.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Verbosity {
    level: log::LevelFilter,
}

impl Verbosity {

    /// Yields a corresponding log filter.
    pub fn level(&self) -> log::LevelFilter {
        self.level
    }
}

impl std::str::FromStr for Verbosity {

    type Err = String; // TODO.

    fn from_str(string: &str) -> std::result::Result<Verbosity, String> {
        use log::LevelFilter::*;

        let level = match string {
            "quiet" => Off,
            "error" => Error,
            "warn" => Warn,
            "info" => Info,
            "debug" => Debug,
            "trace" => Trace,
            _ => return Err(format!("invalid verbosity choice '{}'", string)),
        };

        Ok(Verbosity {
            level: level,
        })
    }
}

/// A type listing different options for logging to standard streams.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Stream {
    mode: simplelog::TerminalMode,
}

impl Stream {

    /// Yields a corresponding terminal mode.
    pub fn mode(&self) -> simplelog::TerminalMode {
        self.mode
    }
}

impl std::str::FromStr for Stream {

    type Err = String; // TODO.

    fn from_str(string: &str) -> std::result::Result<Stream, String> {
        use simplelog::TerminalMode::*;

        let mode = match string {
            "stdout" => Stdout,
            "stderr" => Stderr,
            "mixed" => Mixed,
            _ => return Err(format!("invalid stream choice '{}'", string)),
        };

        Ok(Stream {
            mode: mode,
        })
    }
}
<commit_msg>Improve error handling in option parsing.<commit_after>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

//! Structured specification of command-line arguments.
//!
//! This module specifies all command-line arguments that RRG offers and exposes
//! functions for parsing them into a high-level structure.
//!
//! Ideally, only one instance of this high-level structure should ever be
//! created (using the [`from_args`] function). Then this instance should be
//! shared through the entire lifetime of a program and explicitly passed to
//! functions that care about it.
//!
//! [`from_args`]: fn.from_args.html

use std::time::Duration;
use std::path::PathBuf;

use structopt::StructOpt;

#[derive(StructOpt)]
#[structopt(name = "RRG", about = "A GRR agent rewritten in Rust.")]
pub struct Opts {
    /// A level of log verbosity.
    #[structopt(long="log-verbosity", name="LEVEL", default_value="info",
                help="Specifies the level of log verbosity")]
    pub log_verbosity: Verbosity,

    /// A standard stream to log into.
    #[structopt(long="log-stream", name="STREAM",
                help="Enables logging to the specified standard stream")]
    pub log_stream: Option<Stream>,

    /// A path to the file to log into.
    #[structopt(long="log-file", name="FILE",
                help="Enables logging to the specified file")]
    pub log_file: Option<PathBuf>,

    /// A frequency of heartbeat messages to send to the Fleetspeak client.
    #[structopt(long="heartbeat-rate", name="DURATION", default_value="5s",
                parse(try_from_str = humantime::parse_duration),
                help="Specifies the frequency of heartbeat messages")]
    pub heartbeat_rate: Duration,
}

/// Parses command-line arguments.
///
/// This is a just a convenience function intended to be used as a shortcut for
/// creating instances of [`Opts`]. Ideally, it should be called only once in
/// the entire lifetime of the agent.
///
/// [`Opts`]: struct.Opts.html
pub fn from_args() -> Opts {
    Opts::from_args()
}

/// A type representing level of log verbosity.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Verbosity {
    level: log::LevelFilter,
}

impl Verbosity {

    /// Yields a corresponding log filter.
    pub fn level(&self) -> log::LevelFilter {
        self.level
    }
}

impl std::str::FromStr for Verbosity {

    type Err = ParseVerbosityError;

    fn from_str(string: &str) -> std::result::Result<Verbosity, ParseVerbosityError> {
        use log::LevelFilter::*;

        let level = match string {
            "quiet" => Off,
            "error" => Error,
            "warn" => Warn,
            "info" => Info,
            "debug" => Debug,
            "trace" => Trace,
            _ => return Err(ParseVerbosityError::new(string)),
        };

        Ok(Verbosity {
            level: level,
        })
    }
}

/// A type listing different options for logging to standard streams.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Stream {
    mode: simplelog::TerminalMode,
}

impl Stream {

    /// Yields a corresponding terminal mode.
    pub fn mode(&self) -> simplelog::TerminalMode {
        self.mode
    }
}

impl std::str::FromStr for Stream {

    type Err = ParseStreamError;

    fn from_str(string: &str) -> std::result::Result<Stream, ParseStreamError> {
        use simplelog::TerminalMode::*;

        let mode = match string {
            "stdout" => Stdout,
            "stderr" => Stderr,
            "mixed" => Mixed,
            _ => return Err(ParseStreamError::new(string)),
        };

        Ok(Stream {
            mode: mode,
        })
    }
}

/// An error type for failures related to parsing of verbosity levels.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ParseVerbosityError {
    string: String,
}

impl ParseVerbosityError {

    /// Constructs a new error indicating failure of parsing given string.
    fn new<S: Into<String>>(string: S) -> ParseVerbosityError {
        ParseVerbosityError {
            string: string.into(),
        }
    }
}

impl std::fmt::Display for ParseVerbosityError {

    fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(fmt, "invalid verbosity level: '{}'", self.string)
    }
}

impl std::error::Error for ParseVerbosityError {

    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        None
    }
}

/// An error type for failures related to parsing names of standard streams.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ParseStreamError {
    string: String,
}

impl ParseStreamError {

    /// Constructs a new error indicating failure of parsing given string.
    fn new<S: Into<String>>(string: S) -> ParseStreamError {
        ParseStreamError {
            string: string.into(),
        }
    }
}

impl std::fmt::Display for ParseStreamError {

    fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(fmt, "invalid stream name: '{}'", self.string)
    }
}

impl std::error::Error for ParseStreamError {

    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        None
    }
}
<|endoftext|>


<commit_before>extern crate libc;

use std::mem;
use std::os::unix::io::{FromRawFd, IntoRawFd, RawFd};
use std::fs::File;
use std::process::Stdio;

#[derive(Debug)]
pub struct Handle {
    // The struct *owns* this file descriptor, and will close it in drop().
    fd: RawFd,
}

impl Handle {
    pub fn to_file(self) -> File {
        unsafe { File::from_raw_fd(self.into_raw_fd()) }
    }

    pub fn to_stdio(self) -> Stdio {
        unsafe { Stdio::from_raw_fd(self.into_raw_fd()) }
    }
}

impl Clone for Handle {
    fn clone(&self) -> Self {
        let new_fd = unsafe { libc::dup(self.fd) };
        assert!(new_fd >= 0);
        unsafe { FromRawFd::from_raw_fd(new_fd) }
    }
}

impl FromRawFd for Handle {
    unsafe fn from_raw_fd(fd: RawFd) -> Self {
        Handle{fd: fd}
    }
}

impl IntoRawFd for Handle {
    fn into_raw_fd(self) -> RawFd {
        let fd = self.fd;
        mem::forget(self);  // prevent drop() from closing the fd
        fd
    }
}

impl Drop for Handle {
    fn drop(&mut self) {
        let error = unsafe { libc::close(self.fd) };
        assert_eq!(error, 0);
    }
}

// (read, write)
// TODO: error handling
pub fn open_pipe() -> (Handle, Handle) {
    unsafe {
        let mut pipes = [0, 0];
        let error = libc::pipe(pipes.as_mut_ptr());
        assert_eq!(error, 0);
        // prevent child processes from inheriting these by default
        for fd in &pipes {
            let ret = libc::ioctl(*fd, libc::FIOCLEX);
            assert_eq!(ret, 0);
        }
        (Handle::from_raw_fd(pipes[0]), Handle::from_raw_fd(pipes[1]))
    }
}

<commit_msg>test for pipe<commit_after>extern crate libc;

use std::mem;
use std::os::unix::io::{FromRawFd, IntoRawFd, RawFd};
use std::fs::File;
use std::process::Stdio;

#[derive(Debug)]
pub struct Handle {
    // The struct *owns* this file descriptor, and will close it in drop().
    fd: RawFd,
}

impl Handle {
    pub fn to_file(self) -> File {
        unsafe { File::from_raw_fd(self.into_raw_fd()) }
    }

    pub fn to_stdio(self) -> Stdio {
        unsafe { Stdio::from_raw_fd(self.into_raw_fd()) }
    }
}

impl Clone for Handle {
    fn clone(&self) -> Self {
        let new_fd = unsafe { libc::dup(self.fd) };
        assert!(new_fd >= 0);
        unsafe { FromRawFd::from_raw_fd(new_fd) }
    }
}

impl FromRawFd for Handle {
    unsafe fn from_raw_fd(fd: RawFd) -> Self {
        Handle{fd: fd}
    }
}

impl IntoRawFd for Handle {
    fn into_raw_fd(self) -> RawFd {
        let fd = self.fd;
        mem::forget(self);  // prevent drop() from closing the fd
        fd
    }
}

impl Drop for Handle {
    fn drop(&mut self) {
        let error = unsafe { libc::close(self.fd) };
        assert_eq!(error, 0);
    }
}

// (read, write)
// TODO: error handling
pub fn open_pipe() -> (Handle, Handle) {
    unsafe {
        let mut pipes = [0, 0];
        let error = libc::pipe(pipes.as_mut_ptr());
        assert_eq!(error, 0);
        // prevent child processes from inheriting these by default
        for fd in &pipes {
            let ret = libc::ioctl(*fd, libc::FIOCLEX);
            assert_eq!(ret, 0);
        }
        (Handle::from_raw_fd(pipes[0]), Handle::from_raw_fd(pipes[1]))
    }
}

#[cfg(test)]
mod test {
    use super::open_pipe;
    use std::io::prelude::*;

    #[test]
    fn it_works() {
        let (r, w) = open_pipe();
        let mut r_file = r.to_file();
        let mut w_file = w.to_file();
        w_file.write_all(b"some stuff").unwrap();
        drop(w_file);
        let mut output = Vec::new();
        r_file.read_to_end(&mut output).unwrap();
    }
}
<|endoftext|>