<commit_before>extern crate clap;
#[macro_use]
extern crate glium;
extern crate rand;

use std::collections::{HashSet, VecDeque};

use glium::Surface;
use glium::glutin;

/// Maximum frames per second.
const FPS: u32 = 30;

#[derive(Clone, Copy, Debug, PartialEq)]
enum Color {
    Red,
    Yellow,
    Green,
    LightBrown,
    Purple,
    Cyan,
    Blue,
    Fuchsia,
}

const COLORS: &[Color] = &[
    Color::Red,
    Color::Yellow,
    Color::Green,
    Color::LightBrown,
    Color::Purple,
    Color::Cyan,
    Color::Blue,
    Color::Fuchsia,
];

impl rand::Rand for Color {
    fn rand<R: rand::Rng>(rng: &mut R) -> Self {
        *rng.choose(COLORS).unwrap()
    }
}

struct Grid {
    width: u8,
    height: u8,
    num_colors: u8,

    cells: Vec<Color>,
    population: Vec<u16>,

    max_clicks: u16,
    num_clicks: u16,
}

impl Grid {
    pub fn new(num_colors: u8, size: u8) -> Self {
        use rand::distributions::{IndependentSample, Range};
        let between = Range::new(0, num_colors as usize);
        let mut rng = rand::thread_rng();

        let cells: Vec<Color> = (0..size as u64 * size as u64)
            .map(|_| COLORS[between.ind_sample(&mut rng)])
            .collect();
        let max_clicks = 25 * 2 * size as u16 * num_colors as u16 / (2 * 14 * 6);
        let mut population = vec![0; COLORS.len()];
        for &c in &cells {
            population[c as usize] += 1;
        }

        Self {
            width: size,
            height: size,
            num_colors,

            cells,
            population,

            max_clicks,
            num_clicks: 0,
        }
    }

    fn index(&self, row: u8, column: u8) -> usize {
        self.width as usize * row as usize + column as usize
    }

    fn current_color(&self) -> Color {
        self.cells[0]
    }

    fn solved(&self) -> bool {
        let mut colors_present = 0;
        for &x in &self.population {
            if x > 0 {
                colors_present += 1;
            }
        }
        colors_present == 1
    }

    pub fn click(&mut self, row: u8, column: u8) -> bool {
        if row >= self.height || column >= self.width {
            return false;
        }

        let i = self.index(row, column);
        let new_color = self.cells[i];
        if self.current_color() == new_color {
            return false;
        }

        self.num_clicks += 1;
        self.flood(new_color);
        self.solved()
    }

    /// Flood the grid from the top left cell and return the number of cell which are connected to
    /// the top left cell by cells of indentical color.
    fn flood(&mut self, new_color: Color) {
        let current_color = self.current_color();

        let rows = self.height as usize;
        let columns = self.width as usize;

        let mut visited = HashSet::new();
        let mut queue = VecDeque::new();
        queue.push_back(0);

        while let Some(i) = queue.pop_front() {
            if self.cells[i] == new_color {
                continue;
            }
            let c = self.cells[i];
            self.population[c as usize] -= 1;
            self.cells[i] = new_color;
            self.population[new_color as usize] += 1;
            visited.insert(i);

            if i % columns > 0 && self.cells[i - 1] == current_color &&
                !visited.contains(&(i - 1))
            {
                queue.push_back(i - 1);
            }
            if i % columns < columns - 1 && self.cells[i + 1] == current_color &&
                !visited.contains(&(i + 1))
            {
                queue.push_back(i + 1);
            }
            if i >= columns && self.cells[i - columns] == current_color &&
                !visited.contains(&(i - columns))
            {
                queue.push_back(i - columns);
            }
            if i < (rows - 1) * columns && self.cells[i + columns] == current_color &&
                !visited.contains(&(i + columns))
            {
                queue.push_back(i + columns);
            }
        }
    }

    pub fn aspect_ratio(&self) -> f32 {
        self.height as f32 / self.width as f32
    }

    pub fn reset(&mut self) {
        use rand::distributions::{IndependentSample, Range};
        let between = Range::new(0, self.num_colors as usize);
        let mut rng = rand::thread_rng();

        self.cells.iter_mut().for_each(|x| {
            *x = COLORS[between.ind_sample(&mut rng)]
        });

        self.population = vec![0; COLORS.len()];
        for &c in &self.cells {
            self.population[c as usize] += 1;
        }
        self.num_clicks = 0;
    }
}

#[derive(Copy, Clone)]
struct Vertex {
    color: f32,
    position: [f32; 2],
}

implement_vertex!(Vertex, color, position);

const TITLE: &'static str = "Flood-It";

const VERTEX_SHADER: &str = r#"
    #version 140
    in vec2 position;
    in float color;
    uniform mat4 matrix;
    out float c;
    void main() {
        c = color;
        gl_Position = matrix * vec4(position, 0.0, 1.0);
    }
"#;

/// Map Color::* to RGB values.
const FRAGMENT_SHADER: &str = r#"
    #version 140
    in float c;
    out vec4 color;
    void main() {
        if (c <= 0.0)
            color = vec4(1.0, 0.0, 0.0, 1.0);
        else if (c <= 1.0)
            color = vec4(1.0, 1.0, 0.0, 1.0);
        else if (c <= 2.0)
            color = vec4(0.0, 0.69, 0.0, 1.0);
        else if (c <= 3.0)
            color = vec4(1.0, 0.8, 0.4, 1.0);
        else if (c <= 4.0)
            color = vec4(0.5, 0.0, 0.5, 1.0);
        else if (c <= 5.0)
            color = vec4(0.0, 1.0, 1.0, 1.0);
        else if (c <= 6.0)
            color = vec4(0.0, 0.0, 1.0, 1.0);
        else if (c <= 7.0)
            color = vec4(1.0, 0.0, 1.0, 1.0);
        else
            // Should never happen
            color = vec4(1.0, 1.0, 1.0, 1.0);
    }
"#;

fn parse_args() -> (u8, u8) {
    use clap::{App, Arg};

    let matches = App::new(TITLE)
        .author(env!("CARGO_PKG_AUTHORS"))
        .version(env!("CARGO_PKG_VERSION"))
        .arg(
            Arg::with_name("colors")
                .takes_value(true)
                .help("The number of different colors")
                .default_value("6"),
        )
        .arg(
            Arg::with_name("size")
                .takes_value(true)
                .help("The height and width of the grid")
                .default_value("14"),
        )
        .get_matches();

    let colors = {
        let tmp: u8 = matches
            .value_of("colors")
            .and_then(|x| x.parse().ok())
            .unwrap_or(6);
        if tmp < 3 || tmp > 8 {
            panic!("Flood-It only supports 3 through 8 (inclusive) colors.");
        } else {
            tmp
        }
    };
    let size: u8 = {
        let tmp: u8 = matches
            .value_of("size")
            .and_then(|x| x.parse().ok())
            .unwrap_or(14);
        if tmp < 2 {
            panic!("Flood-It needs a grid of at least 2x2 cells.");
        } else {
            tmp
        }
    };

    (colors, size)
}

fn main() {
    use glium::glutin::{ContextBuilder, ElementState, EventsLoop, KeyboardInput, WindowBuilder};
    use glutin::{Event, WindowEvent};

    let (colors, size) = parse_args();

    let mut events_loop = EventsLoop::new();
    let window = WindowBuilder::new().with_title(TITLE);
    let context = ContextBuilder::new();
    let display = glium::Display::new(window, context, &events_loop).unwrap();

    let program = glium::Program::from_source(&display, VERTEX_SHADER, FRAGMENT_SHADER, None)
        .unwrap();

    let indices = glium::index::NoIndices(glium::index::PrimitiveType::TrianglesList);

    let params = glium::DrawParameters {
        backface_culling: glium::draw_parameters::BackfaceCullingMode::CullCounterClockwise,
        ..Default::default()
    };

    let mut grid = Grid::new(colors, size);
    let grid_aspect_ratio = grid.aspect_ratio();

    let mut cursor_position = (0.0, 0.0);

    main_loop(|| {
        let mut target = display.draw();
        target.clear_color(0.0, 0.0, 0.0, 1.0);
        let (width, height) = target.get_dimensions();
        let screen_aspect_ratio = height as f32 / width as f32;

        let (ratio_x, ratio_y) = if screen_aspect_ratio < grid_aspect_ratio {
            (screen_aspect_ratio, 1.0)
        } else {
            (1.0, screen_aspect_ratio.recip())
        };

        let uniforms =
            uniform! {
            matrix: [
                [ratio_x, 0.0, 0.0, 0.0],
                [0.0, -ratio_y, 0.0, 0.0],
                [0.0, 0.0, 1.0, 0.0],
                [0.0, 0.0, 0.0, 1.0_f32],
            ]
        };

        let mut triangles = vec![];

        for i in 0..grid.cells.len() {
            let x = (i % grid.width as usize) as f32;
            let y = (i / grid.width as usize) as f32;

            let cell_size = 2.0 / grid.height.max(grid.width) as f32;

            let mut quad = generate_rectangle_vertices(
                x * cell_size - 1.0,
                y * cell_size - 1.0,
                (x + 1.0) * cell_size - 1.0,
                (y + 1.0) * cell_size - 1.0,
                grid.cells[i],
            );
            triangles.append(&mut quad);
        }

        let vertex_buffer = glium::VertexBuffer::new(&display, &triangles).unwrap();
        target
            .draw(&vertex_buffer, &indices, &program, &uniforms, &params)
            .unwrap();

        target.finish().unwrap();

        let mut closed = false;
        events_loop.poll_events(|event| match event {
            Event::WindowEvent { event, .. } => {
                match event {
                    WindowEvent::Closed => closed = true,
                    WindowEvent::MouseMoved { position, .. } =>
                        cursor_position = position,
                    WindowEvent::KeyboardInput { input: KeyboardInput{virtual_keycode: Some(key), state: ElementState::Released, ..}, ..} => {
                        use glium::glutin::VirtualKeyCode;
                        match key {
                            VirtualKeyCode::Space | VirtualKeyCode::N | VirtualKeyCode::R =>
                                if grid.solved() {
                                    grid.reset();
                                }
                            _ => (),
                        }
                    }
                    WindowEvent::MouseInput {
                        state: glium::glutin::ElementState::Released,
                        button: glium::glutin::MouseButton::Left,
                        ..
                    } => {
                        let offsets = {
                            let offset = (width as f64 - height as f64).abs() / 2.0;
                            if screen_aspect_ratio < grid_aspect_ratio {
                                (offset, 0.0)
                            } else {
                                (0.0, offset)
                            }
                        };
                        if cursor_position.0 - offsets.0 >= 0.0 &&
                            cursor_position.1 - offsets.1 >= 0.0
                        {
                            let column = ((cursor_position.0 - offsets.0) * grid.width as f64 /
                                              (width as f64 - 2.0 * offsets.0))
                                .floor() as u8;
                            let row = ((cursor_position.1 - offsets.1) * grid.height as f64 /
                                           (height as f64 - 2.0 * offsets.1))
                                .floor() as u8;
                            if grid.click(row, column) {
                                    if grid.num_clicks <= grid.max_clicks {
                                        println!("You win! You used {} out of {} available moves.", grid.num_clicks, grid.max_clicks);
                                    } else {
                                        println!("You lose. You took {} moves but should have finished in {}.", grid.num_clicks, grid.max_clicks);
                                }
                            }
                        }
                    }
                    _ => (),
                }
            }
            _ => (),
        });
        !closed
    });
}

fn generate_rectangle_vertices(
    left: f32,
    bottom: f32,
    right: f32,
    top: f32,
    color: Color,
) -> Vec<Vertex> {
    let color = color as u32 as f32;
    vec![
        Vertex {
            color,
            position: [left, bottom],
        },
        Vertex {
            color,
            position: [right, bottom],
        },
        Vertex {
            color,
            position: [left, top],
        },
        Vertex {
            color,
            position: [right, bottom],
        },
        Vertex {
            color,
            position: [right, top],
        },
        Vertex {
            color,
            position: [left, top],
        },
    ]
}

fn main_loop<F: FnMut() -> bool>(mut callback: F) {
    use std::time::{Duration, Instant};
    use std::thread;

    let one_frame = Duration::new(0, 10 ^ 9 / FPS + 1);
    let mut last_frame = Instant::now();

    while callback() {
        let now = Instant::now();
        let remaining = one_frame.checked_sub(now - last_frame).unwrap_or_default();
        last_frame = now;
        thread::sleep(remaining);
    }
}
<commit_msg>Render using just one quad and a texture<commit_after>extern crate clap;
#[macro_use]
extern crate glium;
extern crate rand;

use std::collections::{HashSet, VecDeque};

use glium::Surface;
use glium::glutin;

/// Maximum frames per second.
const FPS: u32 = 30;

#[derive(Clone, Copy, Debug, PartialEq)]
enum Color {
    Red,
    Yellow,
    Green,
    LightBrown,
    Purple,
    Cyan,
    Blue,
    Fuchsia,
}

const COLORS: &[Color] = &[
    Color::Red,
    Color::Yellow,
    Color::Green,
    Color::LightBrown,
    Color::Purple,
    Color::Cyan,
    Color::Blue,
    Color::Fuchsia,
];

impl Color {
    fn to_rgb(self) -> (u8, u8, u8) {
        match self {
            Color::Red => (255, 0, 0),
            Color::Yellow => (255, 255, 0),
            Color::Green => (0, 176, 0),
            Color::LightBrown => (255, 204, 102),
            Color::Purple => (128, 0, 128),
            Color::Cyan => (0, 255, 255),
            Color::Blue => (0, 0, 255),
            Color::Fuchsia => (255, 0, 255),
        }
    }
}

impl rand::Rand for Color {
    fn rand<R: rand::Rng>(rng: &mut R) -> Self {
        *rng.choose(COLORS).unwrap()
    }
}

struct Grid {
    width: u8,
    height: u8,
    num_colors: u8,

    cells: Vec<Color>,
    population: Vec<u16>,

    max_clicks: u16,
    num_clicks: u16,
}

impl Grid {
    pub fn new(num_colors: u8, size: u8) -> Self {
        use rand::distributions::{IndependentSample, Range};
        let between = Range::new(0, num_colors as usize);
        let mut rng = rand::thread_rng();

        let cells: Vec<Color> = (0..size as u64 * size as u64)
            .map(|_| COLORS[between.ind_sample(&mut rng)])
            .collect();
        let max_clicks = 25 * 2 * size as u16 * num_colors as u16 / (2 * 14 * 6);
        let mut population = vec![0; COLORS.len()];
        for &c in &cells {
            population[c as usize] += 1;
        }

        Self {
            width: size,
            height: size,
            num_colors,

            cells,
            population,

            max_clicks,
            num_clicks: 0,
        }
    }

    fn index(&self, row: u8, column: u8) -> usize {
        self.width as usize * row as usize + column as usize
    }

    fn current_color(&self) -> Color {
        self.cells[0]
    }

    fn solved(&self) -> bool {
        let mut colors_present = 0;
        for &x in &self.population {
            if x > 0 {
                colors_present += 1;
            }
        }
        colors_present == 1
    }

    pub fn click(&mut self, row: u8, column: u8) -> bool {
        if row >= self.height || column >= self.width {
            return false;
        }

        let i = self.index(row, column);
        let new_color = self.cells[i];
        if self.current_color() == new_color {
            return false;
        }

        self.num_clicks += 1;
        self.flood(new_color);
        self.solved()
    }

    /// Flood the grid from the top left cell and return the number of cell which are connected to
    /// the top left cell by cells of indentical color.
    fn flood(&mut self, new_color: Color) {
        let current_color = self.current_color();

        let rows = self.height as usize;
        let columns = self.width as usize;

        let mut visited = HashSet::new();
        let mut queue = VecDeque::new();
        queue.push_back(0);

        while let Some(i) = queue.pop_front() {
            if self.cells[i] == new_color {
                continue;
            }
            let c = self.cells[i];
            self.population[c as usize] -= 1;
            self.cells[i] = new_color;
            self.population[new_color as usize] += 1;
            visited.insert(i);

            if i % columns > 0 && self.cells[i - 1] == current_color &&
                !visited.contains(&(i - 1))
            {
                queue.push_back(i - 1);
            }
            if i % columns < columns - 1 && self.cells[i + 1] == current_color &&
                !visited.contains(&(i + 1))
            {
                queue.push_back(i + 1);
            }
            if i >= columns && self.cells[i - columns] == current_color &&
                !visited.contains(&(i - columns))
            {
                queue.push_back(i - columns);
            }
            if i < (rows - 1) * columns && self.cells[i + columns] == current_color &&
                !visited.contains(&(i + columns))
            {
                queue.push_back(i + columns);
            }
        }
    }

    pub fn aspect_ratio(&self) -> f32 {
        self.height as f32 / self.width as f32
    }

    pub fn reset(&mut self) {
        use rand::distributions::{IndependentSample, Range};
        let between = Range::new(0, self.num_colors as usize);
        let mut rng = rand::thread_rng();

        self.cells.iter_mut().for_each(|x| {
            *x = COLORS[between.ind_sample(&mut rng)]
        });

        self.population = vec![0; COLORS.len()];
        for &c in &self.cells {
            self.population[c as usize] += 1;
        }
        self.num_clicks = 0;
    }
}

#[derive(Copy, Clone)]
struct Vertex {
    position: [f32; 2],
}

implement_vertex!(Vertex, position);

const TITLE: &str = "Flood-It";

const VERTEX_SHADER: &str = r#"
    #version 140

    in  vec2 position;
    out vec2 v_tex_coords;

    uniform vec2 size;
    uniform mat4 matrix;

    void main() {
        v_tex_coords = 0.5 * (position + vec2(1.0, 1.0));
        gl_Position = matrix * vec4(position, 0.0, 1.0);
    }
"#;

/// Map Color::* to RGB values.
const FRAGMENT_SHADER: &str = r#"
    #version 140

    in  vec2 v_tex_coords;
    out vec4 color;

    uniform sampler2D colors;

    void main() {
        color = texture(colors, v_tex_coords);
    }
"#;

fn parse_args() -> (u8, u8) {
    use clap::{App, Arg};

    let matches = App::new(TITLE)
        .author(env!("CARGO_PKG_AUTHORS"))
        .version(env!("CARGO_PKG_VERSION"))
        .arg(
            Arg::with_name("colors")
                .takes_value(true)
                .help("The number of different colors")
                .default_value("6"),
        )
        .arg(
            Arg::with_name("size")
                .takes_value(true)
                .help("The height and width of the grid")
                .default_value("14"),
        )
        .get_matches();

    let colors = {
        let tmp: u8 = matches
            .value_of("colors")
            .and_then(|x| x.parse().ok())
            .unwrap_or(6);
        if tmp < 3 || tmp > 8 {
            panic!("Flood-It only supports 3 through 8 (inclusive) colors.");
        } else {
            tmp
        }
    };
    let size: u8 = {
        let tmp: u8 = matches
            .value_of("size")
            .and_then(|x| x.parse().ok())
            .unwrap_or(14);
        if tmp < 2 {
            panic!("Flood-It needs a grid of at least 2x2 cells.");
        } else {
            tmp
        }
    };

    (colors, size)
}

fn main() {
    use std::borrow::Cow;
    use glium::glutin::{ContextBuilder, ElementState, EventsLoop, KeyboardInput, WindowBuilder};
    use glium::glutin::{Event, WindowEvent};
    use glium::texture::{RawImage2d, Texture2d, Texture2dDataSink};

    let (colors, size) = parse_args();

    let mut events_loop = EventsLoop::new();
    let window = WindowBuilder::new().with_title(TITLE);
    let context = ContextBuilder::new();
    let display = glium::Display::new(window, context, &events_loop).unwrap();

    let program = glium::Program::from_source(&display, VERTEX_SHADER, FRAGMENT_SHADER, None)
        .unwrap();

    let indices = glium::index::NoIndices(glium::index::PrimitiveType::TrianglesList);

    let params = glium::DrawParameters {
        backface_culling: glium::draw_parameters::BackfaceCullingMode::CullCounterClockwise,
        ..Default::default()
    };

    let mut grid = Grid::new(colors, size);
    let grid_aspect_ratio = grid.aspect_ratio();

    let mut cursor_position = (0.0, 0.0);

    let rect = generate_rectangle_vertices(-1.0, -1.0, 1.0, 1.0);
    let vertex_buffer = glium::VertexBuffer::new(&display, &rect).unwrap();

    main_loop(|| {
        let mut target = display.draw();
        target.clear_color(0.0, 0.0, 0.0, 1.0);
        let (width, height) = target.get_dimensions();
        let screen_aspect_ratio = height as f32 / width as f32;

        let (ratio_x, ratio_y) = if screen_aspect_ratio < grid_aspect_ratio {
            (screen_aspect_ratio, 1.0)
        } else {
            (1.0, screen_aspect_ratio.recip())
        };

        // Create a texture representation of the colored grid
        let cell_colors: Vec<_> = grid.cells.iter().map(|x| x.to_rgb()).collect();
        let cell_image = RawImage2d::from_raw(
            Cow::from(cell_colors),
            grid.width as u32,
            grid.height as u32,
        );
        let cell_texture = Texture2d::new(&display, cell_image).unwrap();

        let uniforms =
            uniform! {
                size: (grid.width as f32, grid.height as f32),
                colors: cell_texture.sampled()
                    //.wrap_function(glium::uniforms::SamplerWrapFunction::Clamp)
                    .magnify_filter(glium::uniforms::MagnifySamplerFilter::Nearest),
                matrix: [
                    [ratio_x, 0.0, 0.0, 0.0],
                    [0.0, -ratio_y, 0.0, 0.0],
                    [0.0, 0.0, 1.0, 0.0],
                    [0.0, 0.0, 0.0, 1.0_f32],
                ]
        };
        target
            .draw(&vertex_buffer, &indices, &program, &uniforms, &params)
            .unwrap();

        target.finish().unwrap();

        let mut closed = false;
        events_loop.poll_events(|event| match event {
            Event::WindowEvent { event, .. } => {
                match event {
                    WindowEvent::Closed => closed = true,
                    WindowEvent::MouseMoved { position, .. } =>
                        cursor_position = position,
                    WindowEvent::KeyboardInput { input: KeyboardInput{virtual_keycode: Some(key), state: ElementState::Released, ..}, ..} => {
                        use glium::glutin::VirtualKeyCode;
                        match key {
                            VirtualKeyCode::Space | VirtualKeyCode::N | VirtualKeyCode::R =>
                                if grid.solved() {
                                    grid.reset();
                                }
                            _ => (),
                        }
                    }
                    WindowEvent::MouseInput {
                        state: glium::glutin::ElementState::Released,
                        button: glium::glutin::MouseButton::Left,
                        ..
                    } => {
                        let offsets = {
                            let offset = (width as f64 - height as f64).abs() / 2.0;
                            if screen_aspect_ratio < grid_aspect_ratio {
                                (offset, 0.0)
                            } else {
                                (0.0, offset)
                            }
                        };
                        if cursor_position.0 - offsets.0 >= 0.0 &&
                            cursor_position.1 - offsets.1 >= 0.0
                        {
                            let column = ((cursor_position.0 - offsets.0) * grid.width as f64 /
                                              (width as f64 - 2.0 * offsets.0))
                                .floor() as u8;
                            let row = ((cursor_position.1 - offsets.1) * grid.height as f64 /
                                           (height as f64 - 2.0 * offsets.1))
                                .floor() as u8;
                            if grid.click(row, column) {
                                    if grid.num_clicks <= grid.max_clicks {
                                        println!("You win! You used {} out of {} available moves.", grid.num_clicks, grid.max_clicks);
                                    } else {
                                        println!("You lose. You took {} moves but should have finished in {}.", grid.num_clicks, grid.max_clicks);
                                }
                            }
                        }
                    }
                    _ => (),
                }
            }
            _ => (),
        });
        !closed
    });
}

fn generate_rectangle_vertices(left: f32, bottom: f32, right: f32, top: f32) -> Vec<Vertex> {
    vec![
        Vertex { position: [left, bottom] },
        Vertex { position: [right, bottom] },
        Vertex { position: [left, top] },
        Vertex { position: [right, bottom] },
        Vertex { position: [right, top] },
        Vertex { position: [left, top] },
    ]
}

fn main_loop<F: FnMut() -> bool>(mut callback: F) {
    use std::time::{Duration, Instant};
    use std::thread;

    let one_frame = Duration::new(0, 10 ^ 9 / FPS + 1);
    let mut last_frame = Instant::now();

    while callback() {
        let now = Instant::now();
        let remaining = one_frame.checked_sub(now - last_frame).unwrap_or_default();
        last_frame = now;
        thread::sleep(remaining);
    }
}
<|endoftext|>


<commit_before>//! Muxed. A tmux project manager with no runtime dependencies.
extern crate common;
extern crate docopt;
extern crate edit;
extern crate load;
extern crate new;
extern crate snapshot;

use common::args::Args;
use docopt::Docopt;
use std::env;
use std::process::exit;

#[macro_export]
macro_rules! try_or_err (
    ($expr: expr) => ({
        match $expr {
            Ok(val) => val,
            Err(e) => {
                println!("Muxed ran in to a problem: {}", e);
                exit(1);
            }
        }
    })
);

static DISALLOWED_PROJECT_NAMES: [&str; 3] = ["new", "edit", "load"];

static USAGE: &str = "
Usage:
    muxed [options] <project>
    muxed edit [options] <project>
    muxed new [options] <project>
    muxed snapshot [options] <project>
    muxed load [options] <project>
    muxed (-h | --help)
    muxed (-v | --version)

Flags:
    -d                  If you want to create a muxed session without connecting to it
    -f                  Overwrite existing file if one exists
    --debug             Prints debug information while executing (project opening only)
    -h, --help          Prints help information
    -v, --version       Prints version information

Options:
    -p <project_dir>    The directory your project config files live in. Defaults to ~/.muxed/
    -t <session>        The name of the running TMUX session to codify

Args:
    <project>           The name of your project to open

Subcommands:
    edit <project>                   Edit an existing project file
    new <project>                    To create a new project file
    snapshot -t <session> <project>  Capture a running session and create a config file for it
";

/// The main execution method.
/// Currently accepts a single option. The option represents a configuration
/// file in the same naming format. Given a project file name `projectName.yml`
/// in the `~/.muxed/` directory.
///
/// # Examples
///
/// ~/.muxed/projectName.yml
///
/// ```
/// root: ~/projects/muxed/
/// windows:
///     - cargo: "cargo build"
///     - vim: "vim ."
///     - git: ""
/// ```
///
/// You can run the command:
///
/// ```
/// $ ./muxed projectName
/// ```
pub fn main() {
    let args: Args = Docopt::new(USAGE)
        .and_then(|d| d.deserialize())
        .unwrap_or_else(|e| e.exit());

    if args.flag_v {
        println!("Muxed {}", env!("CARGO_PKG_VERSION"));
        exit(0);
    };

    if args.cmd_new {
        try_or_err!(new::exec(args))
    } else if args.cmd_edit {
        try_or_err!(edit::exec(args))
    } else if args.cmd_snapshot {
        try_or_err!(snapshot::exec(args))
    } else if DISALLOWED_PROJECT_NAMES.contains(&args.arg_project.as_ref()) {
        println!(
            "Tried to call sub-command {} without a project. Please specify a project name.",
            args.arg_project
        );
        exit(1);
    } else {
        try_or_err!(load::exec(args))
    }
}
<commit_msg>Load also isnt an original sub-command<commit_after>//! Muxed. A tmux project manager with no runtime dependencies.
extern crate common;
extern crate docopt;
extern crate edit;
extern crate load;
extern crate new;
extern crate snapshot;

use common::args::Args;
use docopt::Docopt;
use std::env;
use std::process::exit;

#[macro_export]
macro_rules! try_or_err (
    ($expr: expr) => ({
        match $expr {
            Ok(val) => val,
            Err(e) => {
                println!("Muxed ran in to a problem: {}", e);
                exit(1);
            }
        }
    })
);

static DISALLOWED_PROJECT_NAMES: [&str; 2] = ["new", "edit"];

static USAGE: &str = "
Usage:
    muxed [options] <project>
    muxed edit [options] <project>
    muxed new [options] <project>
    muxed snapshot [options] <project>
    muxed (-h | --help)
    muxed (-v | --version)

Flags:
    -d                  If you want to create a muxed session without connecting to it
    -f                  Overwrite existing file if one exists
    --debug             Prints debug information while executing (project opening only)
    -h, --help          Prints help information
    -v, --version       Prints version information

Options:
    -p <project_dir>    The directory your project config files live in. Defaults to ~/.muxed/
    -t <session>        The name of the running TMUX session to codify

Args:
    <project>           The name of your project to open

Subcommands:
    edit <project>                   Edit an existing project file
    new <project>                    To create a new project file
    snapshot -t <session> <project>  Capture a running session and create a config file for it
";

/// The main execution method.
/// Currently accepts a single option. The option represents a configuration
/// file in the same naming format. Given a project file name `projectName.yml`
/// in the `~/.muxed/` directory.
///
/// # Examples
///
/// ~/.muxed/projectName.yml
///
/// ```
/// root: ~/projects/muxed/
/// windows:
///     - cargo: "cargo build"
///     - vim: "vim ."
///     - git: ""
/// ```
///
/// You can run the command:
///
/// ```
/// $ ./muxed projectName
/// ```
pub fn main() {
    let args: Args = Docopt::new(USAGE)
        .and_then(|d| d.deserialize())
        .unwrap_or_else(|e| e.exit());

    if args.flag_v {
        println!("Muxed {}", env!("CARGO_PKG_VERSION"));
        exit(0);
    };

    if args.cmd_new {
        try_or_err!(new::exec(args))
    } else if args.cmd_edit {
        try_or_err!(edit::exec(args))
    } else if args.cmd_snapshot {
        try_or_err!(snapshot::exec(args))
    } else if DISALLOWED_PROJECT_NAMES.contains(&args.arg_project.as_ref()) {
        println!(
            "Tried to call sub-command {} without a project. Please specify a project name.",
            args.arg_project
        );
        exit(1);
    } else {
        try_or_err!(load::exec(args))
    }
}
<|endoftext|>


<commit_before>#![recursion_limit = "256"]
#![feature(proc_macro_hygiene, decl_macro)]
#![allow(proc_macro_derive_resolution_fallback)]

extern crate ape;
extern crate app_dirs;
extern crate base64;
extern crate core;
extern crate crypto;
#[macro_use]
extern crate diesel;
#[macro_use]
extern crate diesel_migrations;
#[macro_use]
extern crate error_chain;
extern crate getopts;
extern crate id3;
extern crate image;
extern crate lewton;
extern crate metaflac;
extern crate mp3_duration;
extern crate rand;
extern crate regex;
extern crate reqwest;
extern crate ring;
#[macro_use]
extern crate rocket;
extern crate rocket_contrib;
extern crate rustfm_scrobble;
extern crate serde;
#[macro_use]
extern crate serde_derive;
extern crate serde_json;
extern crate toml;
#[macro_use]
extern crate log;
extern crate simplelog;

#[cfg(windows)]
extern crate uuid;
#[cfg(windows)]
extern crate winapi;

#[cfg(unix)]
extern crate unix_daemonize;

#[cfg(unix)]
use std::fs::File;
#[cfg(unix)]
use std::io::prelude::*;
#[cfg(unix)]
use unix_daemonize::{daemonize_redirect, ChdirMode};

use core::ops::Deref;
use crate::errors::*;
use getopts::Options;
use simplelog::{Level, LevelFilter, SimpleLogger, TermLogger};
use std::path::Path;
use std::sync::Arc;

mod api;
#[cfg(test)]
mod api_tests;
mod config;
mod db;
mod ddns;
mod errors;
mod index;
mod lastfm;
mod metadata;
mod playlist;
mod serve;
mod server;
mod swagger;
#[cfg(test)]
mod test;
mod thumbnails;
mod ui;
mod user;
mod utils;
mod vfs;
mod web;

static LOG_CONFIG: simplelog::Config = simplelog::Config {
	time: Some(Level::Error),
	level: Some(Level::Error),
	target: Some(Level::Error),
	location: Some(Level::Error),
	time_format: None,
};

fn main() {
	if let Err(ref e) = run() {
		println!("Error: {}", e);

		for e in e.iter().skip(1) {
			println!("caused by: {}", e);
		}
		if let Some(backtrace) = e.backtrace() {
			println!("backtrace: {:?}", backtrace);
		}
		::std::process::exit(1);
	}
}

#[cfg(unix)]
fn daemonize(options: &getopts::Matches) -> Result<()> {
	if options.opt_present("f") {
		return Ok(());
	}
	let mut log_file = utils::get_data_root()?;
	log_file.push("polaris.log");
	let pid = match daemonize_redirect(Some(&log_file), Some(&log_file), ChdirMode::NoChdir) {
		Ok(p) => p,
		Err(_) => bail!(ErrorKind::DaemonError),
	};
	let mut pid_path = utils::get_data_root()?;
	pid_path.push("polaris.pid");
	let mut file = File::create(pid_path)?;
	file.write_all(pid.to_string().as_bytes())?;
	Ok(())
}

#[cfg(unix)]
fn init_log(log_level: LevelFilter, options: &getopts::Matches) -> Result<()> {
	if options.opt_present("f") {
		if let Err(e) = TermLogger::init(log_level, LOG_CONFIG) {
			bail!("Error starting terminal logger: {}", e);
		};
	} else {
		if let Err(e) = SimpleLogger::init(log_level, LOG_CONFIG) {
			bail!("Error starting simple logger: {}", e);
		}
	}
	Ok(())
}

#[cfg(windows)]
fn init_log(log_level: LevelFilter, _: &getopts::Matches) -> Result<()> {
	if TermLogger::init(log_level, LOG_CONFIG).is_err() {
		if let Err(e) = SimpleLogger::init(log_level, LOG_CONFIG) {
			bail!("Error starting simple logger: {}", e);
		}
	};
	Ok(())
}

fn run() -> Result<()> {
	// Parse CLI options
	let args: Vec<String> = std::env::args().collect();
	let mut options = Options::new();
	options.optopt("c", "config", "set the configuration file", "FILE");
	options.optopt("p", "port", "set polaris to run on a custom port", "PORT");
	options.optopt("d", "database", "set the path to index database", "FILE");
	options.optopt("w", "web", "set the path to web client files", "DIRECTORY");
	options.optopt("s", "swagger", "set the path to swagger files", "DIRECTORY");
	options.optopt(
		"l",
		"log",
		"set the log level to a value between 0 (off) and 3 (debug)",
		"LEVEL",
	);

	#[cfg(unix)]
	options.optflag(
		"f",
		"foreground",
		"run polaris in the foreground instead of daemonizing",
	);

	options.optflag("h", "help", "print this help menu");

	let matches = options.parse(&args[1..])?;

	if matches.opt_present("h") {
		let program = args[0].clone();
		let brief = format!("Usage: {} [options]", program);
		print!("{}", options.usage(&brief));
		return Ok(());
	}

	let log_level = match matches.opt_str("l").as_ref().map(String::as_ref) {
		Some("0") => LevelFilter::Off,
		Some("1") => LevelFilter::Error,
		Some("2") => LevelFilter::Info,
		Some("3") => LevelFilter::Debug,
		_ => LevelFilter::Info,
	};

	init_log(log_level, &matches)?;

	#[cfg(unix)]
	daemonize(&matches)?;

	// Init DB
	info!("Starting up database");
	let db_path = matches.opt_str("d");
	let mut default_db_path = utils::get_data_root()?;
	default_db_path.push("db.sqlite");
	let db_path = db_path
		.map(|n| Path::new(n.as_str()).to_path_buf())
		.unwrap_or(default_db_path);
	let db = Arc::new(db::DB::new(&db_path)?);

	// Parse config
	info!("Parsing configuration");
	let config_file_name = matches.opt_str("c");
	let config_file_path = config_file_name.map(|p| Path::new(p.as_str()).to_path_buf());
	if let Some(path) = config_file_path {
		let config = config::parse_toml_file(&path)?;
		config::overwrite(db.deref(), &config)?;
	}
	let config = config::read(db.deref())?;

	// Init index
	info!("Initializing index");
	let command_sender = index::init(db.clone());

	// Trigger auto-indexing
	let db_auto_index = db.clone();
	let command_sender_auto_index = command_sender.clone();
	std::thread::spawn(move || {
		index::self_trigger(db_auto_index.deref(), &command_sender_auto_index);
	});

	// API mount target
	let prefix_url = config.prefix_url.unwrap_or_else(|| "".to_string());
	let api_url = format!("/{}api", &prefix_url);
	info!("Mounting API on {}", api_url);

	// Web client mount target
	let web_dir_name = matches.opt_str("w");
	let mut default_web_dir = utils::get_data_root()?;
	default_web_dir.push("web");
	let web_dir_path = web_dir_name
		.map(|n| Path::new(n.as_str()).to_path_buf())
		.unwrap_or(default_web_dir);
	info!("Static files location is {}", web_dir_path.display());
	let web_url = format!("/{}", &prefix_url);
	info!("Mounting web client files on {}", web_url);

	// Swagger files mount target
	let swagger_dir_name = matches.opt_str("s");
	let mut default_swagger_dir = utils::get_data_root()?;
	default_swagger_dir.push("swagger");
	let swagger_dir_path = swagger_dir_name
		.map(|n| Path::new(n.as_str()).to_path_buf())
		.unwrap_or(default_swagger_dir);
	info!("Swagger files location is {}", swagger_dir_path.display());
	let swagger_url = format!("/{}swagger", &prefix_url);
	info!("Mounting swagger files on {}", swagger_url);

	// Start server
	info!("Starting up server");
	let port: u16 = matches
		.opt_str("p")
		.unwrap_or_else(|| "5050".to_owned())
		.parse()
		.or(Err("invalid port number"))?;

	let server = server::get_server(
		port,
		&api_url,
		&web_url,
		&web_dir_path,
		&swagger_url,
		&swagger_dir_path,
		db.clone(),
		command_sender,
	)?;
	std::thread::spawn(move || {
		server.launch();
	});

	// Start DDNS updates
	let db_ddns = db.clone();
	std::thread::spawn(move || {
		ddns::run(db_ddns.deref());
	});

	// Run UI
	ui::run();

	info!("Shutting down server");
	Ok(())
}
<commit_msg>Fix logging in docker.<commit_after>#![recursion_limit = "256"]
#![feature(proc_macro_hygiene, decl_macro)]
#![allow(proc_macro_derive_resolution_fallback)]

extern crate ape;
extern crate app_dirs;
extern crate base64;
extern crate core;
extern crate crypto;
#[macro_use]
extern crate diesel;
#[macro_use]
extern crate diesel_migrations;
#[macro_use]
extern crate error_chain;
extern crate getopts;
extern crate id3;
extern crate image;
extern crate lewton;
extern crate metaflac;
extern crate mp3_duration;
extern crate rand;
extern crate regex;
extern crate reqwest;
extern crate ring;
#[macro_use]
extern crate rocket;
extern crate rocket_contrib;
extern crate rustfm_scrobble;
extern crate serde;
#[macro_use]
extern crate serde_derive;
extern crate serde_json;
extern crate toml;
#[macro_use]
extern crate log;
extern crate simplelog;

#[cfg(windows)]
extern crate uuid;
#[cfg(windows)]
extern crate winapi;

#[cfg(unix)]
extern crate unix_daemonize;

#[cfg(unix)]
use std::fs::File;
#[cfg(unix)]
use std::io::prelude::*;
#[cfg(unix)]
use unix_daemonize::{daemonize_redirect, ChdirMode};

use core::ops::Deref;
use crate::errors::*;
use getopts::Options;
use simplelog::{Level, LevelFilter, SimpleLogger, TermLogger};
use std::path::Path;
use std::sync::Arc;

mod api;
#[cfg(test)]
mod api_tests;
mod config;
mod db;
mod ddns;
mod errors;
mod index;
mod lastfm;
mod metadata;
mod playlist;
mod serve;
mod server;
mod swagger;
#[cfg(test)]
mod test;
mod thumbnails;
mod ui;
mod user;
mod utils;
mod vfs;
mod web;

static LOG_CONFIG: simplelog::Config = simplelog::Config {
	time: Some(Level::Error),
	level: Some(Level::Error),
	target: Some(Level::Error),
	location: Some(Level::Error),
	time_format: None,
};

fn main() {
	if let Err(ref e) = run() {
		println!("Error: {}", e);

		for e in e.iter().skip(1) {
			println!("caused by: {}", e);
		}
		if let Some(backtrace) = e.backtrace() {
			println!("backtrace: {:?}", backtrace);
		}
		::std::process::exit(1);
	}
}

#[cfg(unix)]
fn daemonize(options: &getopts::Matches) -> Result<()> {
	if options.opt_present("f") {
		return Ok(());
	}
	let mut log_file = utils::get_data_root()?;
	log_file.push("polaris.log");
	let pid = match daemonize_redirect(Some(&log_file), Some(&log_file), ChdirMode::NoChdir) {
		Ok(p) => p,
		Err(_) => bail!(ErrorKind::DaemonError),
	};
	let mut pid_path = utils::get_data_root()?;
	pid_path.push("polaris.pid");
	let mut file = File::create(pid_path)?;
	file.write_all(pid.to_string().as_bytes())?;
	Ok(())
}

#[cfg(unix)]
fn init_log(log_level: LevelFilter, options: &getopts::Matches) -> Result<()> {
	if options.opt_present("f") {
		if let Err(_) = TermLogger::init(log_level, LOG_CONFIG) {
            if let Err(e) = SimpleLogger::init(log_level, LOG_CONFIG) {
                bail!("Error starting terminal logger: {}", e);
            };
		};
	} else {
		if let Err(e) = SimpleLogger::init(log_level, LOG_CONFIG) {
			bail!("Error starting simple logger: {}", e);
		}
	}
	Ok(())
}

#[cfg(windows)]
fn init_log(log_level: LevelFilter, _: &getopts::Matches) -> Result<()> {
	if TermLogger::init(log_level, LOG_CONFIG).is_err() {
		if let Err(e) = SimpleLogger::init(log_level, LOG_CONFIG) {
			bail!("Error starting simple logger: {}", e);
		}
	};
	Ok(())
}

fn run() -> Result<()> {
	// Parse CLI options
	let args: Vec<String> = std::env::args().collect();
	let mut options = Options::new();
	options.optopt("c", "config", "set the configuration file", "FILE");
	options.optopt("p", "port", "set polaris to run on a custom port", "PORT");
	options.optopt("d", "database", "set the path to index database", "FILE");
	options.optopt("w", "web", "set the path to web client files", "DIRECTORY");
	options.optopt("s", "swagger", "set the path to swagger files", "DIRECTORY");
	options.optopt(
		"l",
		"log",
		"set the log level to a value between 0 (off) and 3 (debug)",
		"LEVEL",
	);

	#[cfg(unix)]
	options.optflag(
		"f",
		"foreground",
		"run polaris in the foreground instead of daemonizing",
	);

	options.optflag("h", "help", "print this help menu");

	let matches = options.parse(&args[1..])?;

	if matches.opt_present("h") {
		let program = args[0].clone();
		let brief = format!("Usage: {} [options]", program);
		print!("{}", options.usage(&brief));
		return Ok(());
	}

	let log_level = match matches.opt_str("l").as_ref().map(String::as_ref) {
		Some("0") => LevelFilter::Off,
		Some("1") => LevelFilter::Error,
		Some("2") => LevelFilter::Info,
		Some("3") => LevelFilter::Debug,
		_ => LevelFilter::Info,
	};

	init_log(log_level, &matches)?;

	#[cfg(unix)]
	daemonize(&matches)?;

	// Init DB
	info!("Starting up database");
	let db_path = matches.opt_str("d");
	let mut default_db_path = utils::get_data_root()?;
	default_db_path.push("db.sqlite");
	let db_path = db_path
		.map(|n| Path::new(n.as_str()).to_path_buf())
		.unwrap_or(default_db_path);
	let db = Arc::new(db::DB::new(&db_path)?);

	// Parse config
	info!("Parsing configuration");
	let config_file_name = matches.opt_str("c");
	let config_file_path = config_file_name.map(|p| Path::new(p.as_str()).to_path_buf());
	if let Some(path) = config_file_path {
		let config = config::parse_toml_file(&path)?;
		config::overwrite(db.deref(), &config)?;
	}
	let config = config::read(db.deref())?;

	// Init index
	info!("Initializing index");
	let command_sender = index::init(db.clone());

	// Trigger auto-indexing
	let db_auto_index = db.clone();
	let command_sender_auto_index = command_sender.clone();
	std::thread::spawn(move || {
		index::self_trigger(db_auto_index.deref(), &command_sender_auto_index);
	});

	// API mount target
	let prefix_url = config.prefix_url.unwrap_or_else(|| "".to_string());
	let api_url = format!("/{}api", &prefix_url);
	info!("Mounting API on {}", api_url);

	// Web client mount target
	let web_dir_name = matches.opt_str("w");
	let mut default_web_dir = utils::get_data_root()?;
	default_web_dir.push("web");
	let web_dir_path = web_dir_name
		.map(|n| Path::new(n.as_str()).to_path_buf())
		.unwrap_or(default_web_dir);
	info!("Static files location is {}", web_dir_path.display());
	let web_url = format!("/{}", &prefix_url);
	info!("Mounting web client files on {}", web_url);

	// Swagger files mount target
	let swagger_dir_name = matches.opt_str("s");
	let mut default_swagger_dir = utils::get_data_root()?;
	default_swagger_dir.push("swagger");
	let swagger_dir_path = swagger_dir_name
		.map(|n| Path::new(n.as_str()).to_path_buf())
		.unwrap_or(default_swagger_dir);
	info!("Swagger files location is {}", swagger_dir_path.display());
	let swagger_url = format!("/{}swagger", &prefix_url);
	info!("Mounting swagger files on {}", swagger_url);

	// Start server
	info!("Starting up server");
	let port: u16 = matches
		.opt_str("p")
		.unwrap_or_else(|| "5050".to_owned())
		.parse()
		.or(Err("invalid port number"))?;

	let server = server::get_server(
		port,
		&api_url,
		&web_url,
		&web_dir_path,
		&swagger_url,
		&swagger_dir_path,
		db.clone(),
		command_sender,
	)?;
	std::thread::spawn(move || {
		server.launch();
	});

	// Start DDNS updates
	let db_ddns = db.clone();
	std::thread::spawn(move || {
		ddns::run(db_ddns.deref());
	});

	// Run UI
	ui::run();

	info!("Shutting down server");
	Ok(())
}
<|endoftext|>


<commit_before>#![feature(phase)]
extern crate cargo;
extern crate docopt;
#[phase(plugin)] extern crate docopt_macros;
extern crate graphviz;
extern crate serialize;

use cargo::core::{Resolve, SourceId, PackageId};
use docopt::{Error, FlagParser};
use graphviz as dot;
use std::io::File;
use std::str;

docopt!(Flags, "
Generate a graph of package dependencies in graphviz format

Usage: cargo dot [options]
       cargo dot --help

Options:
    -h, --help         Show this message
    -V, --version      Print version info and exit
    --lock-file=FILE   Specify location of input file, default \"Cargo.lock\"
    --dot-file=FILE    Output to file, default prints to stdout
    --source-labels    Use sources for the label instead of package names
")

fn main() {
    let config = docopt::Config { version: Some("0.1.0".to_string()) , ..docopt::DEFAULT_CONFIG };
    let flags: Flags = FlagParser::parse_conf(config).unwrap_or_else(|e| e.exit());
    let lock_file  = unless_empty(flags.flag_lock_file, "Cargo.lock");
    let dot_f_flag = if flags.flag_dot_file.is_empty() { None } else { Some(flags.flag_dot_file) };
    let source_labels = flags.flag_source_labels;

    let lock_file = Path::new(lock_file);
    let project_dir = Path::new(lock_file.dirname());
    let project_dir = std::os::make_absolute(&project_dir);
    let src_id = SourceId::for_path(&project_dir).unwrap();
    let resolved = cargo::ops::load_lockfile(&lock_file, &src_id)
                     .unwrap_or_else(|e| exit_with(e.description().as_slice()))
                     .unwrap_or_else(||  exit_with("Lock file not found."));

    let mut graph = Graph::with_root(resolved.root(), source_labels);
    graph.add_dependencies(&resolved);

    match dot_f_flag {
        None           => graph.render_to(&mut std::io::stdio::stdout()),
        Some(dot_file) => graph.render_to(&mut File::create(&Path::new(dot_file)))
    };

}

fn exit_with(s: &str) -> ! {
    fail!("Error parsing lockfile: {}", s)
}

fn unless_empty(s: String, default: &str) -> String {
    if s.is_empty() {
        default.to_string()
    } else {
        s
    }
}

pub type Nd = uint;
pub type Ed = (uint, uint);
pub struct Graph<'a> {
    nodes: Vec<&'a PackageId>,
    edges: Vec<Ed>,
    source_labels: bool
}

impl<'a> Graph<'a> {
    pub fn with_root(root: &PackageId, source_labels: bool) -> Graph {
        Graph { nodes: vec![root], edges: vec![], source_labels: source_labels }
    }

    pub fn add_dependencies(&mut self, resolved: &'a Resolve) {
        for crat in resolved.iter() {
            match resolved.deps(crat) {
                Some(mut crate_deps) => {
                    let idl = self.find_or_add(crat);
                    for dep in crate_deps {
                        let idr = self.find_or_add(dep);
                        self.edges.push((idl, idr));
                    };
                },
                None => { }
            }
        }
    }

    fn find_or_add(&mut self, new: &'a PackageId) -> uint {
        for (i, id) in self.nodes.iter().enumerate() {
            if *id == new {
                return i
            }
        }
        self.nodes.push(new);
        self.nodes.len() - 1
    }

    pub fn render_to<W:Writer>(&'a self, output: &mut W) {
        dot::render(self, output).unwrap()
    }
}

impl<'a> dot::Labeller<'a, Nd, Ed> for Graph<'a> {
    fn graph_id(&self) -> dot::Id<'a> {
        dot::Id::new("example3")
    }
    fn node_id(&self, n: &Nd) -> dot::Id {
        dot::Id::new(format!("N{:u}", *n))
    }
    fn node_label<'a>(&'a self, i: &Nd) -> dot::LabelText<'a> {
        if !self.source_labels {
            dot::LabelStr(str::Slice(self.nodes[*i].get_name()))
        } else {
            dot::LabelStr(str::Owned(self.nodes[*i].get_source_id().get_url().to_string()))
        }
    }
}

impl<'a> dot::GraphWalk<'a, Nd, Ed> for Graph<'a> {
    fn nodes(&self) -> dot::Nodes<'a,Nd> {
        range(0, self.nodes.len()).collect()
    }
    fn edges(&'a self) -> dot::Edges<'a,Ed> {
        dot::maybe_owned_vec::Borrowed(self.edges.as_slice())
    }
    fn source(&self, &(s, _): &Ed) -> Nd { s }
    fn target(&self, &(_, t): &Ed) -> Nd { t }
}
<commit_msg>Name output graph "dependencies"<commit_after>#![feature(phase)]
extern crate cargo;
extern crate docopt;
#[phase(plugin)] extern crate docopt_macros;
extern crate graphviz;
extern crate serialize;

use cargo::core::{Resolve, SourceId, PackageId};
use docopt::{Error, FlagParser};
use graphviz as dot;
use std::io::File;
use std::str;

docopt!(Flags, "
Generate a graph of package dependencies in graphviz format

Usage: cargo dot [options]
       cargo dot --help

Options:
    -h, --help         Show this message
    -V, --version      Print version info and exit
    --lock-file=FILE   Specify location of input file, default \"Cargo.lock\"
    --dot-file=FILE    Output to file, default prints to stdout
    --source-labels    Use sources for the label instead of package names
")

fn main() {
    let config = docopt::Config { version: Some("0.1.0".to_string()) , ..docopt::DEFAULT_CONFIG };
    let flags: Flags = FlagParser::parse_conf(config).unwrap_or_else(|e| e.exit());
    let lock_file  = unless_empty(flags.flag_lock_file, "Cargo.lock");
    let dot_f_flag = if flags.flag_dot_file.is_empty() { None } else { Some(flags.flag_dot_file) };
    let source_labels = flags.flag_source_labels;

    let lock_file = Path::new(lock_file);
    let project_dir = Path::new(lock_file.dirname());
    let project_dir = std::os::make_absolute(&project_dir);
    let src_id = SourceId::for_path(&project_dir).unwrap();
    let resolved = cargo::ops::load_lockfile(&lock_file, &src_id)
                     .unwrap_or_else(|e| exit_with(e.description().as_slice()))
                     .unwrap_or_else(||  exit_with("Lock file not found."));

    let mut graph = Graph::with_root(resolved.root(), source_labels);
    graph.add_dependencies(&resolved);

    match dot_f_flag {
        None           => graph.render_to(&mut std::io::stdio::stdout()),
        Some(dot_file) => graph.render_to(&mut File::create(&Path::new(dot_file)))
    };

}

fn exit_with(s: &str) -> ! {
    fail!("Error parsing lockfile: {}", s)
}

fn unless_empty(s: String, default: &str) -> String {
    if s.is_empty() {
        default.to_string()
    } else {
        s
    }
}

pub type Nd = uint;
pub type Ed = (uint, uint);
pub struct Graph<'a> {
    nodes: Vec<&'a PackageId>,
    edges: Vec<Ed>,
    source_labels: bool
}

impl<'a> Graph<'a> {
    pub fn with_root(root: &PackageId, source_labels: bool) -> Graph {
        Graph { nodes: vec![root], edges: vec![], source_labels: source_labels }
    }

    pub fn add_dependencies(&mut self, resolved: &'a Resolve) {
        for crat in resolved.iter() {
            match resolved.deps(crat) {
                Some(mut crate_deps) => {
                    let idl = self.find_or_add(crat);
                    for dep in crate_deps {
                        let idr = self.find_or_add(dep);
                        self.edges.push((idl, idr));
                    };
                },
                None => { }
            }
        }
    }

    fn find_or_add(&mut self, new: &'a PackageId) -> uint {
        for (i, id) in self.nodes.iter().enumerate() {
            if *id == new {
                return i
            }
        }
        self.nodes.push(new);
        self.nodes.len() - 1
    }

    pub fn render_to<W:Writer>(&'a self, output: &mut W) {
        dot::render(self, output).unwrap()
    }
}

impl<'a> dot::Labeller<'a, Nd, Ed> for Graph<'a> {
    fn graph_id(&self) -> dot::Id<'a> {
        dot::Id::new("dependencies")
    }
    fn node_id(&self, n: &Nd) -> dot::Id {
        dot::Id::new(format!("N{:u}", *n))
    }
    fn node_label<'a>(&'a self, i: &Nd) -> dot::LabelText<'a> {
        if !self.source_labels {
            dot::LabelStr(str::Slice(self.nodes[*i].get_name()))
        } else {
            dot::LabelStr(str::Owned(self.nodes[*i].get_source_id().get_url().to_string()))
        }
    }
}

impl<'a> dot::GraphWalk<'a, Nd, Ed> for Graph<'a> {
    fn nodes(&self) -> dot::Nodes<'a,Nd> {
        range(0, self.nodes.len()).collect()
    }
    fn edges(&'a self) -> dot::Edges<'a,Ed> {
        dot::maybe_owned_vec::Borrowed(self.edges.as_slice())
    }
    fn source(&self, &(s, _): &Ed) -> Nd { s }
    fn target(&self, &(_, t): &Ed) -> Nd { t }
}
<|endoftext|>


<commit_before>extern crate bindgen;
#[macro_use]
extern crate log;
extern crate docopt;
#[macro_use]
extern crate rustc_serialize;
extern crate env_logger;

use bindgen::{Builder, LinkType};
use std::io::{self, Write};
use std::fs::File;
use std::process::exit;

const USAGE: &'static str = "
Generate C bindings for Rust.

Usage:
  bindgen [options] <file> [-- <clang-args>...]
  bindgen (-h | --help)

Options:
  <clang-args>                 Options passed directly to clang.
  -h, --help                   Display this help message.
  --link=<library>             Link to a dynamic library, can be provided multiple times.
                               <library> is in the format `[kind=]lib`, where `kind` is
                               one of `static`, `dynamic` or `framework`.
  --output=<output>            Write bindings to <output> (- is stdout).
                               [default: -]
  --match=<name>               Only output bindings for definitions from files
                               whose name contains <name>
                               If multiple -match options are provided, files
                               matching any rule are bound to.
  --builtins                   Output bindings for builtin definitions
                               (for example __builtin_va_list)
  --emit-clang-ast             Output the ast (for debugging purposes)
  --override-enum-type=<type>  Override enum type, type name could be
                                 uchar
                                 schar
                                 ushort
                                 sshort
                                 uint
                                 sint
                                 ulong
                                 slong
                                 ulonglong
                                 slonglong
  --use-core                  Use `core` as a base crate for `Option` and such.
                              See also `--ctypes-prefix`.
  --ctypes-prefix=<prefix>    Use this prefix for all the types in the generated
                              code.
                              [default: std::os::raw]
  --remove-prefix=<prefix>    Prefix to remove from all the symbols, like
                              `libfoo_`. The removal is case-insensitive.
  --no-derive-debug           Disable `derive(Debug)` for all generated types.
  --no-rust-enums             Convert C enums to Rust constants instead of enums.
  --dont-convert-floats       Disables the convertion of C `float` and `double`
                              to Rust `f32` and `f64`.
";

#[derive(Debug, RustcDecodable)]
struct Args {
    arg_file: String,
    arg_clang_args: Vec<String>,
    flag_link: Option<String>,
    flag_output: String,
    flag_match: Option<String>,
    flag_builtins: bool,
    flag_emit_clang_ast: bool,
    flag_override_enum_type: String,
    flag_ctypes_prefix: String,
    flag_use_core: bool,
    flag_remove_prefix: Option<String>,
    // TODO: allow finer control.
    flag_no_derive_debug: bool,
    flag_no_rust_enums: bool,
    flag_dont_convert_floats: bool,
}

fn args_to_opts(args: Args) -> Builder<'static> {
    let mut builder = Builder::new(args.arg_file);
    builder.emit_ast(args.flag_emit_clang_ast)
           .ctypes_prefix(args.flag_ctypes_prefix
                              .split("::")
                              .map(String::from)
                              .collect::<Vec<_>>())
           .use_core(args.flag_use_core)
           .derive_debug(!args.flag_no_derive_debug)
           .rust_enums(!args.flag_no_rust_enums)
           .override_enum_ty(args.flag_override_enum_type);
    for arg in args.arg_clang_args {
        builder.clang_arg(arg);
    }
    if let Some(s) = args.flag_match {
        builder.match_pat(s);
    }
    if let Some(s) = args.flag_remove_prefix {
        builder.remove_prefix(s);
    }
    if args.flag_builtins {
        builder.builtins();
    }
    if args.flag_dont_convert_floats {
        builder.dont_convert_floats();
    }
    if let Some(link) = args.flag_link {
        let mut parts = link.split('=');
        let (lib, kind) = match (parts.next(), parts.next()) {
            (Some(lib), None) => (lib, LinkType::Dynamic),
            (Some(kind), Some(lib)) => {
                (lib,
                 match kind {
                    "static" => LinkType::Static,
                    "dynamic" => LinkType::Dynamic,
                    "framework" => LinkType::Framework,
                    _ => {
                        println!("Link type unknown: {}", kind);
                        exit(1);
                    }
                })
            }
            _ => {
                println!("Wrong link format: {}", link);
                exit(1);
            }
        };
        builder.link(lib, kind);
    }
    builder
}

fn get_output(o: &str) -> Box<Write> {
    if o == "-" {
        Box::new(io::stdout())
    } else {
        Box::new(File::create(o).expect(&format!("\"{}\" unwritable", o)))
    }
}

pub fn main() {
    env_logger::LogBuilder::new().parse("bindgen=warn").init().unwrap();

    let args: Args = docopt::Docopt::new(USAGE)
                         .and_then(|d| d.decode())
                         .unwrap_or_else(|e| e.exit());
    debug!("{:?}", args);

    let output = get_output(&args.flag_output);

    let builder = args_to_opts(args);
    debug!("{:?}", builder);

    match builder.generate() {
        Ok(bindings) => {
            match bindings.write(output) {
                Ok(()) => (),
                Err(e) => {
                    error!("Unable to write bindings to file. {}", e);
                    exit(-1);
                }
            }
        }
        Err(()) => exit(-1),
    }
}
<commit_msg>Bring back multiple matches support. Closes #362.<commit_after>extern crate bindgen;
#[macro_use]
extern crate log;
extern crate docopt;
#[macro_use]
extern crate rustc_serialize;
extern crate env_logger;

use bindgen::{Builder, LinkType};
use std::io::{self, Write};
use std::fs::File;
use std::process::exit;

const USAGE: &'static str = "
Generate C bindings for Rust.

Usage:
  bindgen [options] <file> [-- <clang-args>...]
  bindgen [options] (--match=<name> ...) <file> [-- <clang-args>...]
  bindgen (-h | --help)

Options:
  <clang-args>                 Options passed directly to clang.
  -h, --help                   Display this help message.
  --link=<library>             Link to a dynamic library, can be provided multiple times.
                               <library> is in the format `[kind=]lib`, where `kind` is
                               one of `static`, `dynamic` or `framework`.
  --output=<output>            Write bindings to <output> (- is stdout).
                               [default: -]
  --match=<name>               Only output bindings for definitions from files
                               whose name contains <name>
                               If multiple -match options are provided, files
                               matching any rule are bound to.
  --builtins                   Output bindings for builtin definitions
                               (for example __builtin_va_list)
  --emit-clang-ast             Output the ast (for debugging purposes)
  --override-enum-type=<type>  Override enum type, type name could be
                                 uchar
                                 schar
                                 ushort
                                 sshort
                                 uint
                                 sint
                                 ulong
                                 slong
                                 ulonglong
                                 slonglong
  --use-core                  Use `core` as a base crate for `Option` and such.
                              See also `--ctypes-prefix`.
  --ctypes-prefix=<prefix>    Use this prefix for all the types in the generated
                              code.
                              [default: std::os::raw]
  --remove-prefix=<prefix>    Prefix to remove from all the symbols, like
                              `libfoo_`. The removal is case-insensitive.
  --no-derive-debug           Disable `derive(Debug)` for all generated types.
  --no-rust-enums             Convert C enums to Rust constants instead of enums.
  --dont-convert-floats       Disables the convertion of C `float` and `double`
                              to Rust `f32` and `f64`.
";

#[derive(Debug, RustcDecodable)]
struct Args {
    arg_file: String,
    arg_clang_args: Vec<String>,
    flag_link: Option<String>,
    flag_output: String,
    flag_match: Vec<String>,
    flag_builtins: bool,
    flag_emit_clang_ast: bool,
    flag_override_enum_type: String,
    flag_ctypes_prefix: String,
    flag_use_core: bool,
    flag_remove_prefix: Option<String>,
    // TODO: allow finer control.
    flag_no_derive_debug: bool,
    flag_no_rust_enums: bool,
    flag_dont_convert_floats: bool,
}

fn args_to_opts(args: Args) -> Builder<'static> {
    let mut builder = Builder::new(args.arg_file);
    builder.emit_ast(args.flag_emit_clang_ast)
           .ctypes_prefix(args.flag_ctypes_prefix
                              .split("::")
                              .map(String::from)
                              .collect::<Vec<_>>())
           .use_core(args.flag_use_core)
           .derive_debug(!args.flag_no_derive_debug)
           .rust_enums(!args.flag_no_rust_enums)
           .override_enum_ty(args.flag_override_enum_type);
    for arg in args.arg_clang_args {
        builder.clang_arg(arg);
    }
    for flag_match in args.flag_match {
        builder.match_pat(flag_match);
    }
    if let Some(s) = args.flag_remove_prefix {
        builder.remove_prefix(s);
    }
    if args.flag_builtins {
        builder.builtins();
    }
    if args.flag_dont_convert_floats {
        builder.dont_convert_floats();
    }
    if let Some(link) = args.flag_link {
        let mut parts = link.split('=');
        let (lib, kind) = match (parts.next(), parts.next()) {
            (Some(lib), None) => (lib, LinkType::Dynamic),
            (Some(kind), Some(lib)) => {
                (lib,
                 match kind {
                    "static" => LinkType::Static,
                    "dynamic" => LinkType::Dynamic,
                    "framework" => LinkType::Framework,
                    _ => {
                        println!("Link type unknown: {}", kind);
                        exit(1);
                    }
                })
            }
            _ => {
                println!("Wrong link format: {}", link);
                exit(1);
            }
        };
        builder.link(lib, kind);
    }
    builder
}

fn get_output(o: &str) -> Box<Write> {
    if o == "-" {
        Box::new(io::stdout())
    } else {
        Box::new(File::create(o).expect(&format!("\"{}\" unwritable", o)))
    }
}

pub fn main() {
    env_logger::LogBuilder::new().parse("bindgen=warn").init().unwrap();

    let args: Args = docopt::Docopt::new(USAGE)
                         .and_then(|d| d.decode())
                         .unwrap_or_else(|e| e.exit());
    debug!("{:?}", args);

    let output = get_output(&args.flag_output);

    let builder = args_to_opts(args);
    debug!("{:?}", builder);

    match builder.generate() {
        Ok(bindings) => {
            match bindings.write(output) {
                Ok(()) => (),
                Err(e) => {
                    error!("Unable to write bindings to file. {}", e);
                    exit(-1);
                }
            }
        }
        Err(()) => exit(-1),
    }
}
<|endoftext|>


<commit_before>use std::io::ErrorKind;
use std::process::Command;
use structopt::StructOpt;

mod commands;

/// Personal CLI.
#[derive(StructOpt)]
enum Commands {
    /// Reproduce my computer's configuration.
    Setup,

    /// Upgrade what's installed on my computer.
    Upgrade,

    /// Clean up my computer's state.
    End,

    /// Start a study session.
    Study,
}

fn main() {
    match Command::new("which").arg("killall").output() {
        Ok(_) => (),
        Err(error) => match error.kind() {
            ErrorKind::NotFound => panic!("A Unix environment is required."),
            other_error => panic!("There was a problem: {:?}", other_error),
        },
    };

    match Command::new("brew").output() {
        Ok(_) => (),
        Err(error) => match error.kind() {
            ErrorKind::NotFound => panic!("Homebrew is required."),
            other_error => panic!("There was a problem: {:?}", other_error),
        },
    }

    match Commands::from_args() {
        Commands::Setup => commands::setup(),
        Commands::Upgrade => commands::upgrade(),
        Commands::End => commands::end(),
        Commands::Study => commands::study(),
    }

    println!("Finished.");
}
<commit_msg>Remove prerequisite checks covered by being able to run the code<commit_after>use structopt::StructOpt;

mod commands;

/// Personal CLI.
#[derive(StructOpt)]
enum Commands {
    /// Reproduce my computer's configuration.
    Setup,

    /// Upgrade what's installed on my computer.
    Upgrade,

    /// Clean up my computer's state.
    End,

    /// Start a study session.
    Study,
}

fn main() {
    match Commands::from_args() {
        Commands::Setup => commands::setup(),
        Commands::Upgrade => commands::upgrade(),
        Commands::End => commands::end(),
        Commands::Study => commands::study(),
    }

    println!("Finished.");
}
<|endoftext|>


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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

  loop { thread::sleep(Duration::new(60, 0)); }
}
<commit_msg>Added improved configuration file error handling<commit_after>extern crate yaml_rust;
extern crate regex;
extern crate rusqlite;

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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