<commit_before>use gl;
use gl::types::GLint;
use std;
use std::mem;
use sdl2;
use sdl2::event;
use sdl2::event::Event;
use sdl2::video;
use stopwatch::TimerSet;
use yaglw::gl_context::GLContext;
use yaglw::shader::Shader;
use yaglw::texture::{BufferTexture, TextureUnit};
use yaglw::vertex_buffer::{ArrayHandle};

pub const WINDOW_WIDTH: u32 = 800;
pub const WINDOW_HEIGHT: u32 = 800;

pub fn main() {
  let timers = TimerSet::new();

  let mut sdl = sdl2::init().everything().unwrap();
  let window = make_window(&sdl);

  let _sdl_gl_context = window.gl_create_context().unwrap();

  // Load the OpenGL function pointers.
  gl::load_with(|s| unsafe {
    mem::transmute(video::gl_get_proc_address(s))
  });

  let mut gl = unsafe {
    GLContext::new()
  };
  let gl = &mut gl;

  match gl.get_error() {
    gl::NO_ERROR => {},
    err => {
      println!("OpenGL error 0x{:x} in setup", err);
      return;
    },
  }

  let mut shader = make_shader(gl);
  shader.use_shader(gl);

  let mut heightmap =
    BufferTexture::new(gl, gl::R32F, WINDOW_WIDTH as usize);
  for i in 0..WINDOW_WIDTH {
    let h = i as f32 / WINDOW_HEIGHT as f32 * 2.0 - 1.0;
    heightmap.buffer.push(gl, &[h]);
  }

  {
    let mut bind = |name, id| {
      let unit: TextureUnit = Default::default();
      unsafe {
        gl::ActiveTexture(unit.gl_id());
        gl::BindTexture(gl::TEXTURE_BUFFER, id);
      }
      let loc = shader.get_uniform_location(name);
      unsafe {
        gl::Uniform1i(loc, unit.glsl_id as GLint);
      }
    };

    bind("heightmap", heightmap.buffer.byte_buffer.handle.gl_id);
  }

  let empty_vao = ArrayHandle::new(gl);

  unsafe {
    gl::BindVertexArray(empty_vao.gl_id);
  }

  let mut event_pump = sdl.event_pump();

  while process_events(&mut event_pump) {
    timers.time("draw", || {
      gl.clear_buffer();
      unsafe {
        gl::DrawArrays(gl::TRIANGLE_STRIP, 0, 4);
      }
      // swap buffers
      window.gl_swap_window();
    });

    std::thread::sleep_ms(10);
  }

  timers.print();
}

fn make_shader<'a, 'b:'a>(
  gl: &'a GLContext,
) -> Shader<'b> {
  let vertex_shader: String = format!("
    #version 330 core

    out vec4 world_pos;

    void main() {{
      if (gl_VertexID == 0) {{
        world_pos = vec4(-1, 1, 0, 1);
      }} else if (gl_VertexID == 1) {{
        world_pos = vec4(-1, -1, 0, 1);
      }} else if (gl_VertexID == 2) {{
        world_pos = vec4(1, 1, 0, 1);
      }} else {{
        world_pos = vec4(1, -1, 0, 1);
      }}
      gl_Position = world_pos;
    }}
  ");

  let fragment_shader: String = format!("
    #version 330 core

    uniform samplerBuffer heightmap;

    in vec4 world_pos;

    layout(location=0) out vec4 frag_color;

    void main() {{
      int x = int(round(gl_FragCoord.x));
      float h = texelFetch(heightmap, x).r;
      if (world_pos.y <= h) {{
        frag_color = vec4(1, 0, 0, 1);
      }} else {{
        frag_color = vec4(0, 0, 0, 1);
      }}
    }}
  ");

  let components = vec!(
    (gl::VERTEX_SHADER, vertex_shader),
    (gl::FRAGMENT_SHADER, fragment_shader),
  );

  Shader::new(gl, components.into_iter())
}

fn make_window(sdl: &sdl2::Sdl) -> video::Window {
  video::gl_attr::set_context_profile(video::GLProfile::Core);
  video::gl_attr::set_context_version(3, 3);

  // Open the window as fullscreen at the current resolution.
  let mut window =
    video::WindowBuilder::new(
      &sdl,
      "Hello, Mandelbrot",
      WINDOW_WIDTH, WINDOW_HEIGHT,
    );

  let window = window.position_centered();
  window.opengl();

  window.build().unwrap()
}

fn process_events<'a>(
  event_pump: &mut event::EventPump,
) -> bool {
  for event in event_pump.poll_iter() {
    match event {
      Event::Quit {..} => {
        return false;
      },
      Event::AppTerminating {..} => {
        return false;
      },
      _ => {},
    }
  }

  true
}
<commit_msg>cleanups<commit_after>use gl;
use gl::types::GLint;
use std;
use std::mem;
use sdl2;
use sdl2::event;
use sdl2::event::Event;
use sdl2::video;
use stopwatch::TimerSet;
use yaglw::gl_context::GLContext;
use yaglw::shader::Shader;
use yaglw::texture::{BufferTexture, TextureUnit};
use yaglw::vertex_buffer::{ArrayHandle};

pub const WINDOW_WIDTH: u32 = 800;
pub const WINDOW_HEIGHT: u32 = 800;

pub fn main() {
  let timers = TimerSet::new();

  let mut sdl = sdl2::init().everything().unwrap();
  let window = make_window(&sdl);

  let _sdl_gl_context = window.gl_create_context().unwrap();

  // Load the OpenGL function pointers.
  gl::load_with(|s| unsafe {
    mem::transmute(video::gl_get_proc_address(s))
  });

  let mut gl = unsafe {
    GLContext::new()
  };
  let gl = &mut gl;

  match gl.get_error() {
    gl::NO_ERROR => {},
    err => {
      println!("OpenGL error 0x{:x} in setup", err);
      return;
    },
  }

  let mut shader = make_shader(gl);
  shader.use_shader(gl);

  let heightmap = make_heightmap(gl);

  {
    let mut bind = |name, id| {
      let unit: TextureUnit = Default::default();
      unsafe {
        gl::ActiveTexture(unit.gl_id());
        gl::BindTexture(gl::TEXTURE_BUFFER, id);
      }
      let loc = shader.get_uniform_location(name);
      unsafe {
        gl::Uniform1i(loc, unit.glsl_id as GLint);
      }
    };

    bind("heightmap", heightmap.buffer.byte_buffer.handle.gl_id);
  }

  let empty_vao = ArrayHandle::new(gl);

  unsafe {
    gl::BindVertexArray(empty_vao.gl_id);
  }

  let mut event_pump = sdl.event_pump();

  while process_events(&mut event_pump) {
    timers.time("draw", || {
      gl.clear_buffer();
      unsafe {
        gl::DrawArrays(gl::TRIANGLE_STRIP, 0, 4);
      }
      // swap buffers
      window.gl_swap_window();
    });

    std::thread::sleep_ms(10);
  }

  timers.print();
}

fn make_shader<'a, 'b:'a>(
  gl: &'a GLContext,
) -> Shader<'b> {
  let vertex_shader: String = format!("
    #version 330 core

    out vec4 world_pos;

    void main() {{
      if (gl_VertexID == 0) {{
        world_pos = vec4(-1, 1, 0, 1);
      }} else if (gl_VertexID == 1) {{
        world_pos = vec4(-1, -1, 0, 1);
      }} else if (gl_VertexID == 2) {{
        world_pos = vec4(1, 1, 0, 1);
      }} else {{
        world_pos = vec4(1, -1, 0, 1);
      }}
      gl_Position = world_pos;
    }}
  ");

  let fragment_shader: String = format!("
    #version 330 core

    uniform samplerBuffer heightmap;

    in vec4 world_pos;

    layout(location=0) out vec4 frag_color;

    void main() {{
      int x = int(round(gl_FragCoord.x));
      float h = texelFetch(heightmap, x).r;
      if (world_pos.y <= h) {{
        frag_color = vec4(1, 0, 0, 1);
      }} else {{
        frag_color = vec4(0, 0, 0, 1);
      }}
    }}
  ");

  let components = vec!(
    (gl::VERTEX_SHADER, vertex_shader),
    (gl::FRAGMENT_SHADER, fragment_shader),
  );

  Shader::new(gl, components.into_iter())
}

fn make_window(sdl: &sdl2::Sdl) -> video::Window {
  video::gl_attr::set_context_profile(video::GLProfile::Core);
  video::gl_attr::set_context_version(3, 3);

  // Open the window as fullscreen at the current resolution.
  let mut window =
    video::WindowBuilder::new(
      &sdl,
      "Hello, Mandelbrot",
      WINDOW_WIDTH, WINDOW_HEIGHT,
    );

  let window = window.position_centered();
  window.opengl();

  window.build().unwrap()
}

fn make_heightmap<'a, 'b:'a>(
  gl: &'a mut GLContext,
) -> BufferTexture<'b, f32> {
  let mut ram_heightmap = [0.0; WINDOW_WIDTH as usize];

  for i in 0..WINDOW_WIDTH as usize {
    let h = i as f32 / WINDOW_HEIGHT as f32 * 2.0 - 1.0;
    ram_heightmap[i] = h;
  }

  let mut vram_heightmap = BufferTexture::new(gl, gl::R32F, WINDOW_WIDTH as usize);
  vram_heightmap.buffer.push(gl, &ram_heightmap);
  vram_heightmap
}

fn process_events<'a>(
  event_pump: &mut event::EventPump,
) -> bool {
  for event in event_pump.poll_iter() {
    match event {
      Event::Quit {..} => {
        return false;
      },
      Event::AppTerminating {..} => {
        return false;
      },
      _ => {},
    }
  }

  true
}
<|endoftext|>


<commit_before>// Copyright (c) 2016, 2018 Nikita Pekin and the ferru contributors
// See the README.md file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Ferru is a lightweight, blazing fast static site generator.
//!
//! This binary serves as a frontend to the Ferru static site generation
//! functionality provided by the ferru library.

#![deny(missing_docs)]
#![deny(non_camel_case_types)]
#![cfg_attr(feature = "nightly", feature(plugin))]
#![cfg_attr(feature = "clippy", plugin(clippy))]

#[macro_use]
extern crate clap;
extern crate env_logger;
extern crate ferru;
#[macro_use]
extern crate log;
extern crate mustache;

use clap::{App, ArgMatches};
use ferru::Config;

fn main() {
    env_logger::init()
        .expect("Failed to initialize env_logger");

    // Create the CLI.
    let yml = load_yaml!("../cli.yml");
    let matches = App::from_yaml(yml).get_matches();

    // Run the program.
    run(matches);
}

fn run(m: ArgMatches) {
    // Match the raw subcommand, and get its sub-matches "m".
    if let (name, Some(sub_m)) = m.subcommand() {
        let config = config_from_matches(sub_m);
        match name {
            "build" => ferru::build(&config).unwrap(),
            _ => unreachable!(),
        }
    }
}

fn config_from_matches<'a>(matches: &'a ArgMatches) -> Config<'a> {
    Config {
        source_directory: matches.value_of("source"),
        dest_directory: matches.value_of("destination"),
    }
}
<commit_msg>impr: remove unused crate<commit_after>// Copyright (c) 2016, 2018 Nikita Pekin and the ferru contributors
// See the README.md file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Ferru is a lightweight, blazing fast static site generator.
//!
//! This binary serves as a frontend to the Ferru static site generation
//! functionality provided by the ferru library.

#![deny(missing_docs)]
#![deny(non_camel_case_types)]
#![cfg_attr(feature = "nightly", feature(plugin))]
#![cfg_attr(feature = "clippy", plugin(clippy))]

#[macro_use]
extern crate clap;
extern crate env_logger;
extern crate ferru;
extern crate mustache;

use clap::{App, ArgMatches};
use ferru::Config;

fn main() {
    env_logger::init()
        .expect("Failed to initialize env_logger");

    // Create the CLI.
    let yml = load_yaml!("../cli.yml");
    let matches = App::from_yaml(yml).get_matches();

    // Run the program.
    run(matches);
}

fn run(m: ArgMatches) {
    // Match the raw subcommand, and get its sub-matches "m".
    if let (name, Some(sub_m)) = m.subcommand() {
        let config = config_from_matches(sub_m);
        match name {
            "build" => ferru::build(&config).unwrap(),
            _ => unreachable!(),
        }
    }
}

fn config_from_matches<'a>(matches: &'a ArgMatches) -> Config<'a> {
    Config {
        source_directory: matches.value_of("source"),
        dest_directory: matches.value_of("destination"),
    }
}
<|endoftext|>


<commit_before>extern crate cpuid;

#[cfg(not(test))]
fn main () {
    println!("cpuid is present: {}", cpuid::is_present());
    println!("cpuid version: {}", cpuid::version());
    let info = cpuid::identify();
    match info {
        Err(err) => println!("cpuid error: {}", err),
        Ok(info) => {
            println!("Found: {} CPU", info.vendor);
            println!("Processor model is: {}", info.codename);
            println!("The full brand string is: {}", info.brand);
            println!("The processor has {} cores and {} logical processors", info.num_cores, info.num_logical_cpus);
            match cpuid::clock_frequency() {
                Some(frequency) => println!("CPU speed: {} MHz", frequency),
                None => println!("Couldn't get CPU speed."),
            };
        }
    }
}
<commit_msg>Updated main.rs from readme.<commit_after>extern crate cpuid;

#[cfg(not(test))]
fn main () {
    println!("cpuid is present: {}", cpuid::is_present());
    println!("cpuid version: {}", cpuid::version());
    match cpuid::identify() {
        Ok(info) => {
            println!("Found: {} CPU, model: {}", info.vendor, info.codename);
            println!("The full brand string is: {}", info.brand);
            println!("The processor has {} cores and {} logical processors", info.num_cores, info.num_logical_cpus);
            if info.has_feature(cpuid::FloatingPointUnit) {
                println!("Floating point unit available");
            }
        },
        Err(err) => println!("cpuid error: {}", err),
    };
    match cpuid::clock_frequency() {
        Some(frequency) => println!("CPU speed: {} MHz", frequency),
        None => println!("Couldn't get CPU speed."),
    };
}
<|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("/etc/portlurker.sqlite").expect("Failed to open or create database! (Disk space? Permissions?");
      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>Fixed a bug due to debug test left in the code<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"),
    ("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)); }
}
<|endoftext|>


<commit_before>#[macro_use] extern crate clap;
extern crate tempdir;
extern crate rand;
extern crate base32;
#[macro_use] extern crate serde_json;

use std::fs;

use clap::{Arg, SubCommand};

use tempdir::TempDir;

struct CmdErr {
    message: String,
    secondary: Option<String>,
}

trait ToCmdErr<T> {
    fn check(self, message: &str) -> Result<T, CmdErr>;
}

impl<T, E: ::std::fmt::Display> ToCmdErr<T> for Result<T, E> {
    fn check(self, message: &str) -> Result<T, CmdErr> {
        self.map_err(|e| CmdErr {
            message: message.to_owned(),
            secondary: Some(format!("{}", e)),
        })
    }
}

impl<T> ToCmdErr<T> for Option<T> {
    fn check(self, message: &str) -> Result<T, CmdErr> {
        match self {
            Some(t) => Ok(t),
            None => Err(CmdErr {
                message: message.to_owned(),
                secondary: None,
            }),
        }
    }
}

impl ToCmdErr<()> for bool {
    fn check(self, message: &str) -> Result<(), CmdErr> {
        if self {
            Ok(())
        } else {
            Err(CmdErr {
                message: message.to_owned(),
                secondary: None,
            })
        }
    }
}

fn cmd_run(test: bool, image: &str, rootless: bool, volumes: &[(&str, &str)]) -> Result<Option<i32>, CmdErr> {
    let tmp_dir = TempDir::new("shipc").check("Cannot create temporary directory")?;
    eprintln!("info: Temporary directory: {:?}", tmp_dir.path());
    eprint!("info: Running image {}", image);
    if rootless {
        eprint!(" in rootless mode");
    }
    eprintln!();
    for &(ref orig, ref dest) in volumes {
        eprintln!("info:   with volume {}:{}", orig, dest);
    }
    let image_path = fs::canonicalize(::std::path::Path::new(image)).check("Image path is invalid")?;
    let image_metadata = fs::metadata(image).check("Invalid image path")?;
    let mut image_dir_path = image_path.to_path_buf();

    // ** Prepare image **

    // Support for .tar.gz images: un-tar them in the temporary directory
    if image_metadata.file_type().is_file() {
        image_path.file_name().expect("Filename is expected for a file")
            .to_string_lossy().ends_with(".tar.gz").check("File is not a tarball")?;
        image_dir_path = tmp_dir.path().join("image");
        fs::create_dir(&image_dir_path).check("Internal error")?;
        eprintln!("info: Uncompressing image into {}", image_dir_path.to_string_lossy());
        let tar_output = ::std::process::Command::new("tar")
            .arg("--strip-components")
            .arg("1")
            .arg("-C")
            .arg(&image_dir_path)
            .arg("-xf")
            .arg(image_path)
            .output()
            .check("Cannot execute tar")?;
        if !tar_output.status.success() {
            return Err(CmdErr {
                message: "Failed to un-tar image".to_owned(),
                secondary: Some(String::from_utf8_lossy(&tar_output.stderr).into_owned()),
            });
        }
    }

    // ** Unpack image

    let bundle = "bundle";

    let bundle_path = tmp_dir.path().join(bundle);
    eprintln!("info: Unpacking image {} to {}", image_dir_path.to_string_lossy(), bundle_path.to_string_lossy());

    let mut umoci_command = ::std::process::Command::new("umoci");
    umoci_command
        .current_dir(tmp_dir.path())
        .arg("unpack")
        .arg("--image")
        .arg(&image_dir_path);
    if rootless {
        umoci_command.arg("--rootless");
    }
    umoci_command.arg(bundle);

    let umoci_output = umoci_command.output()
        .check("Cannot run umoci")?;

    if !umoci_output.status.success() {
        return Err(CmdErr {
            message: "Failed to unpack image with umoci".to_owned(),
            secondary: Some(String::from_utf8_lossy(&umoci_output.stderr).into_owned()),
        });
    }

    // ** Edit spec (volumes and command) **

    let container_name = {
        use rand::Rng;
        let bytes: [u8; 8] = unsafe { ::std::mem::transmute(rand::thread_rng().gen::<u64>().to_be()) };
        ::base32::encode(::base32::Alphabet::RFC4648 { padding: false }, &bytes[..]).to_lowercase()
    };

    let config_json_path = bundle_path.join("config.json");
    let json_str = {
        use std::io::Read;
        let mut out = String::with_capacity(1 << 12);
        let mut file = fs::File::open(&config_json_path).check("Internal error: invalid bundle")?;
        file.read_to_string(&mut out).check("Internal error: invalid bundle")?;
        out
    };

    let spec = {
        let mut root: ::serde_json::Value = ::serde_json::from_str(&json_str).check("Internal error: invalid bundle spec")?;
        {
            let hostname = root.get_mut("hostname").check("Internal error: invalid bundle spec")?;
            *hostname = json!(container_name);
        }
        {
            let mounts_value = root.get_mut("mounts").check("Internal error: invalid bundle spec")?;
            let mounts = mounts_value.as_array_mut().check("Internal error: invalid bundle spec")?;
            for &(ref orig, ref dest) in volumes {
                let orig_path = fs::canonicalize(orig).check("Invalid volume source")?;
                mounts.push(json!({
                    "source": orig_path,
                    "destination": dest,
                    "type": "bind",
                    "options": ["rbind", "rw"],
                }));
            }
        }
        root
    };

    {
        use std::io::Write;
        let mut file = fs::OpenOptions::new().write(true).truncate(true).open(&config_json_path).check("Internal error")?;
        let spec_json_str = ::serde_json::ser::to_string_pretty(&spec).check("Internal error: invalid json")?;
        file.write_all(spec_json_str.as_bytes()).check("Internal error: cannot write spec")?;
    }

    // ** Spawn container **

    let mut runc_command = ::std::process::Command::new("runc");
    runc_command
        .current_dir(bundle_path);
    if rootless {
        let root_dir_path = tmp_dir.path().join("root");
        fs::create_dir(&root_dir_path).check("Internal error")?;
        runc_command
            .arg("--root")
            .arg(&root_dir_path);
    }
    runc_command
        .arg("run")
        .arg(&container_name);

    eprintln!("info: Starting container with name {}", &container_name);
    eprintln!("info: If runc is successful, the next line will be inside the container");

    if test {
        eprintln!("test: Press enter to continue (temporary dir cleaned up afterwards)");
        use std::io::BufRead;
        let stdin = std::io::stdin();
        stdin.lock().lines().next();
        eprintln!("test: Test run successful (container not started)");
        return Ok(Some(0));
    }

    let runc_status = runc_command.status().check("Failed to start the container")?;

    eprintln!("info: Container exited with {}", runc_status);

    Ok(runc_status.code())
}

fn fail(message: &str) {
    eprintln!("error: {}", message);
    eprintln!("error: Run with --help for usage");
    ::std::process::exit(-1);
}

fn check(condition: bool, message: &str) {
    if !condition {
        fail(message);
    }
}

fn main() {
    let app = app_from_crate!()
        .about("Unpack and run oci images")
        .arg(Arg::with_name("test")
             .long("test")
             .help("Run in test mode (do not spawn the container)")
             .hidden(true))
        .subcommand(SubCommand::with_name("run")
                    .about("Run an oci image")
                    .arg(Arg::with_name("image")
                         .required(true))
                    .arg(Arg::with_name("rootless")
                         .help("Run in rootless mode")
                         .long("rootless"))
                    .arg(Arg::with_name("volume")
                         .help("Rebind volumes")
                         .long("volume")
                         .short("v")
                         .takes_value(true)
                         .value_name("VOLUME")
                         .multiple(true)));
    let matches = app.clone().get_matches();

    let test = matches.is_present("test");

    let result = match matches.subcommand() {
        ("run", Some(sub_matches)) => {
            let rootless: bool = sub_matches.is_present("rootless");
            let image = sub_matches.value_of("image").unwrap();
            let volumes = sub_matches.values_of("volume").map(|vol_strs| {
                vol_strs.map(|s| {
                    let vol_arr: Vec<&str> = s.split(":").collect();
                    check(vol_arr.len() == 2, "Volume should be origin:destination");
                    let orig = vol_arr[0];
                    let dest = vol_arr[1];
                    (orig, dest)
                }).collect::<Vec<_>>()
            }).unwrap_or(Vec::new());
            cmd_run(test, image, rootless, &volumes[..])
        },
        _ => {
            let mut out = ::std::io::stdout();
            app.write_help(&mut out).ok().expect("error: Failed to print usage to stdout");
            println!();
            ::std::process::exit(-1);
        },
    };
    match result {
        Ok(exit_code) => ::std::process::exit(exit_code.unwrap_or(-1)),
        Err(err) => {
            eprintln!("error: {}", err.message);
            if let Some(secondary) = err.secondary {
                eprintln!("error: {}", secondary);
            }
            eprintln!("error: Run with --help for usage");
        },
    }
}
<commit_msg>shipc run: change command<commit_after>#[macro_use] extern crate clap;
extern crate tempdir;
extern crate rand;
extern crate base32;
#[macro_use] extern crate serde_json;

use std::fs;

use clap::{Arg, SubCommand};

use tempdir::TempDir;

struct CmdErr {
    message: String,
    secondary: Option<String>,
}

trait ToCmdErr<T> {
    fn check(self, message: &str) -> Result<T, CmdErr>;
}

impl<T, E: ::std::fmt::Display> ToCmdErr<T> for Result<T, E> {
    fn check(self, message: &str) -> Result<T, CmdErr> {
        self.map_err(|e| CmdErr {
            message: message.to_owned(),
            secondary: Some(format!("{}", e)),
        })
    }
}

impl<T> ToCmdErr<T> for Option<T> {
    fn check(self, message: &str) -> Result<T, CmdErr> {
        match self {
            Some(t) => Ok(t),
            None => Err(CmdErr {
                message: message.to_owned(),
                secondary: None,
            }),
        }
    }
}

impl ToCmdErr<()> for bool {
    fn check(self, message: &str) -> Result<(), CmdErr> {
        if self {
            Ok(())
        } else {
            Err(CmdErr {
                message: message.to_owned(),
                secondary: None,
            })
        }
    }
}

fn cmd_run(test: bool, image: &str, rootless: bool, volumes: &[(&str, &str)], command: Option<Vec<&str>>) -> Result<Option<i32>, CmdErr> {
    let tmp_dir = TempDir::new("shipc").check("Cannot create temporary directory")?;
    eprintln!("info: Temporary directory: {:?}", tmp_dir.path());
    eprint!("info: Running image {}", image);
    if rootless {
        eprint!(" in rootless mode");
    }
    eprintln!();
    for &(ref orig, ref dest) in volumes {
        eprintln!("info:   with volume {}:{}", orig, dest);
    }
    let image_path = fs::canonicalize(::std::path::Path::new(image)).check("Image path is invalid")?;
    let image_metadata = fs::metadata(image).check("Invalid image path")?;
    let mut image_dir_path = image_path.to_path_buf();

    // ** Prepare image **

    // Support for .tar.gz images: un-tar them in the temporary directory
    if image_metadata.file_type().is_file() {
        image_path.file_name().expect("Filename is expected for a file")
            .to_string_lossy().ends_with(".tar.gz").check("File is not a tarball")?;
        image_dir_path = tmp_dir.path().join("image");
        fs::create_dir(&image_dir_path).check("Internal error")?;
        eprintln!("info: Uncompressing image into {}", image_dir_path.to_string_lossy());
        let tar_output = ::std::process::Command::new("tar")
            .arg("--strip-components")
            .arg("1")
            .arg("-C")
            .arg(&image_dir_path)
            .arg("-xf")
            .arg(image_path)
            .output()
            .check("Cannot execute tar")?;
        if !tar_output.status.success() {
            return Err(CmdErr {
                message: "Failed to un-tar image".to_owned(),
                secondary: Some(String::from_utf8_lossy(&tar_output.stderr).into_owned()),
            });
        }
    }

    // ** Unpack image

    let bundle = "bundle";

    let bundle_path = tmp_dir.path().join(bundle);
    eprintln!("info: Unpacking image {} to {}", image_dir_path.to_string_lossy(), bundle_path.to_string_lossy());

    let mut umoci_command = ::std::process::Command::new("umoci");
    umoci_command
        .current_dir(tmp_dir.path())
        .arg("unpack")
        .arg("--image")
        .arg(&image_dir_path);
    if rootless {
        umoci_command.arg("--rootless");
    }
    umoci_command.arg(bundle);

    let umoci_output = umoci_command.output()
        .check("Cannot run umoci")?;

    if !umoci_output.status.success() {
        return Err(CmdErr {
            message: "Failed to unpack image with umoci".to_owned(),
            secondary: Some(String::from_utf8_lossy(&umoci_output.stderr).into_owned()),
        });
    }

    // ** Edit spec (volumes and command) **

    let container_name = {
        use rand::Rng;
        let bytes: [u8; 8] = unsafe { ::std::mem::transmute(rand::thread_rng().gen::<u64>().to_be()) };
        ::base32::encode(::base32::Alphabet::RFC4648 { padding: false }, &bytes[..]).to_lowercase()
    };

    let config_json_path = bundle_path.join("config.json");
    let json_str = {
        use std::io::Read;
        let mut out = String::with_capacity(1 << 12);
        let mut file = fs::File::open(&config_json_path).check("Internal error: invalid bundle")?;
        file.read_to_string(&mut out).check("Internal error: invalid bundle")?;
        out
    };

    let spec = {
        let mut root: ::serde_json::Value = ::serde_json::from_str(&json_str).check("Internal error: invalid bundle spec")?;
        {
            let hostname = root.get_mut("hostname").check("Internal error: invalid bundle spec")?;
            *hostname = json!(container_name);
        }
        {
            let mounts_value = root.get_mut("mounts").check("Internal error: invalid bundle spec")?;
            let mounts = mounts_value.as_array_mut().check("Internal error: invalid bundle spec")?;
            for &(ref orig, ref dest) in volumes {
                let orig_path = fs::canonicalize(orig).check("Invalid volume source")?;
                mounts.push(json!({
                    "source": orig_path,
                    "destination": dest,
                    "type": "bind",
                    "options": ["rbind", "rw"],
                }));
            }
        }
        if let Some(command_args) = command {
            let process = root.get_mut("process").check("Internal error: invalid bundle spec")?;
            let args = process.get_mut("args").check("Internal error: invalid bundle spec")?;
            *args = json!(command_args);
        }
        root
    };

    {
        use std::io::Write;
        let mut file = fs::OpenOptions::new().write(true).truncate(true).open(&config_json_path).check("Internal error")?;
        let spec_json_str = ::serde_json::ser::to_string_pretty(&spec).check("Internal error: invalid json")?;
        file.write_all(spec_json_str.as_bytes()).check("Internal error: cannot write spec")?;
    }

    // ** Spawn container **

    let mut runc_command = ::std::process::Command::new("runc");
    runc_command
        .current_dir(bundle_path);
    if rootless {
        let root_dir_path = tmp_dir.path().join("root");
        fs::create_dir(&root_dir_path).check("Internal error")?;
        runc_command
            .arg("--root")
            .arg(&root_dir_path);
    }
    runc_command
        .arg("run")
        .arg(&container_name);

    eprintln!("info: Starting container with name {}", &container_name);
    eprintln!("info: If runc is successful, the next line will be inside the container");

    if test {
        eprintln!("test: Press enter to continue (temporary dir cleaned up afterwards)");
        use std::io::BufRead;
        let stdin = std::io::stdin();
        stdin.lock().lines().next();
        eprintln!("test: Test run successful (container not started)");
        return Ok(Some(0));
    }

    let runc_status = runc_command.status().check("Failed to start the container")?;

    eprintln!("info: Container exited with {}", runc_status);

    Ok(runc_status.code())
}

fn fail(message: &str) {
    eprintln!("error: {}", message);
    eprintln!("error: Run with --help for usage");
    ::std::process::exit(-1);
}

fn check(condition: bool, message: &str) {
    if !condition {
        fail(message);
    }
}

fn main() {
    let app = app_from_crate!()
        .about("Unpack and run oci images")
        .arg(Arg::with_name("test")
             .long("test")
             .help("Run in test mode (do not spawn the container)")
             .hidden(true))
        .subcommand(SubCommand::with_name("run")
                    .arg(Arg::with_name("image")
                         .required(true))
                    .arg(Arg::with_name("command")
                         .multiple(true))
                    .arg(Arg::with_name("rootless")
                         .help("Run in rootless mode")
                         .long("rootless"))
                    .arg(Arg::with_name("volume")
                         .help("Rebind volumes")
                         .long("volume")
                         .short("v")
                         .takes_value(true)
                         .value_name("volume")
                         .number_of_values(1)
                         .multiple(true))
                    );
    let matches = app.clone().get_matches();

    let test = matches.is_present("test");

    let result = match matches.subcommand() {
        ("run", Some(sub_matches)) => {
            let rootless: bool = sub_matches.is_present("rootless");
            let image = sub_matches.value_of("image").unwrap();
            let volumes = sub_matches.values_of("volume").map(|vol_strs| {
                vol_strs.map(|s| {
                    let vol_arr: Vec<&str> = s.split(":").collect();
                    check(vol_arr.len() == 2, "Volume should be origin:destination");
                    let orig = vol_arr[0];
                    let dest = vol_arr[1];
                    (orig, dest)
                }).collect::<Vec<_>>()
            }).unwrap_or(Vec::new());
            cmd_run(test, image, rootless, &volumes[..], sub_matches.values_of("command").map(|x| x.collect()))
        },
        _ => {
            let mut out = ::std::io::stdout();
            app.write_help(&mut out).ok().expect("error: Failed to print usage to stdout");
            println!();
            ::std::process::exit(-1);
        },
    };
    match result {
        Ok(exit_code) => ::std::process::exit(exit_code.unwrap_or(-1)),
        Err(err) => {
            eprintln!("error: {}", err.message);
            if let Some(secondary) = err.secondary {
                eprintln!("error: {}", secondary);
            }
            eprintln!("error: Run with --help for usage");
        },
    }
}
<|endoftext|>


<commit_before>const VERSION: &'static str = env!("CARGO_PKG_VERSION");
const AUTHORS: &'static str = env!("CARGO_PKG_AUTHORS");

static FORMAT: &'static str = "{request-time}: {method} {uri} {status} {response-time}";

#[macro_use]
extern crate clap;
use clap::{Arg, ArgMatches, App, SubCommand};

extern crate toml;
extern crate rustc_serialize;
extern crate regex;

extern crate iron;
use iron::prelude::*;
use iron::Handler;
use iron::method::Method;
use std::str::FromStr;

extern crate router;
use router::Router;

extern crate staticfile;
use staticfile::Static;

extern crate mount;
use mount::Mount;

#[macro_use]
extern crate log;
extern crate simplelog;
use simplelog::{SimpleLogger, LogLevelFilter};

extern crate logger;
use logger::Logger;
use logger::format::Format;

mod config;
use config::{Config, ServerConfig, RequestConfig, ResponseConfig};
use std::collections::HashMap;
use std::ascii::AsciiExt;
mod error;
use error::HaxoniteError;
mod handlers;
use handlers::*;
mod utils;

use std::path::Path;

fn main() {
    let _ = SimpleLogger::init(LogLevelFilter::Info, simplelog::Config::default());

    match run() {
        Ok(_) => {}
        Err(err) => {
            error!("{}!", err);
            std::process::exit(-1);
        }
    }
}

fn run() -> Result<(), HaxoniteError> {
    let matches = define_command_line_options().get_matches();
    if let Some(matches) = matches.subcommand_matches("new") {
        let generate_full_project = matches.is_present("full");
        if let Some(project_name) = matches.value_of("project_name") {
            try!(utils::create_new_project(project_name, generate_full_project));
        }
        return Ok(());
    }

    try!(serve(matches));

    Ok(())
}

fn serve(matches: ArgMatches) -> Result<(), HaxoniteError> {
    let config_file = matches.value_of("config_file").unwrap_or("config.toml");
    let config_content = try!(config::read_config(config_file));
    let config: Config = toml::decode_str(config_content.as_ref()).unwrap(); // TODO: remove unwrap to handle malformed config
    debug!("Using config: {:?}!", config);

    let server_config = config.server.unwrap_or_else(ServerConfig::default);
    let host = match value_t!(matches, "host", String) { // TODO: rewrite these matches
        Ok(host) => host,
        Err(_) => server_config.host.unwrap_or_else(config::default_host),
    };
    let port_number = match value_t!(matches, "port_number", u16) {
        Ok(port_number) => port_number,
        Err(_) => server_config.port.unwrap_or_else(config::default_port),
    };

    // Initialize Iron's router and pass it to config processing function to define routes
    let mut router = Router::new();
    let mut mount = Mount::new();
    match config.requests {
        Some(requests) => {
            try!(process_config_requests(&requests, &mut mount, &mut router));
        }
        None => return Err(HaxoniteError::NoRequestDefined),
    }
    mount.mount("/", router);

    // Insert logger middlewares before and after router
    let mut chain = Chain::new(mount);
    let format = Format::new(FORMAT);
    let (logger_before, logger_after) = Logger::new(Some(format.unwrap()));
    chain.link_before(logger_before);
    chain.link_after(logger_after);

    // Initialize Iron to process requests
    let _iron = try!(Iron::new(chain).http((host.as_ref(), port_number)));
    info!("Haxonite running on port: {}!", port_number);
    Ok(())
}

fn process_config_requests(requests: &HashMap<String, RequestConfig>,
                           mut mount: &mut Mount,
                           mut router: &mut Router)
                           -> Result<(), HaxoniteError> {
    info!("Processing config!");
    for (request_name, request_config) in requests {
        info!("Processing config for {}: {:?}!",
              request_name,
              request_config);
        try!(define_route(request_name.clone(),
                          request_config.clone(),
                          &mut mount,
                          &mut router));
    }
    Ok(())
}

fn define_route(request_name: String, request_config: RequestConfig, mount: &mut Mount, router: &mut Router) -> Result<(), HaxoniteError> {
    let method = request_config.method.unwrap_or_else(config::default_method);
    let path = match request_config.path {
        Some(path) => path,
        None => return Err(HaxoniteError::NoPathDefined(request_name.clone())),
    };

    let type_ = request_config.type_.unwrap_or_else(config::default_type);
    let response_configs = match request_config.responses {
        Some(response_configs) => response_configs,
        None => return Err(HaxoniteError::NoResponseDefined(request_name.clone())),
    };
    let handler: Box<Handler> = match type_.as_ref() {
        "single" => {
            let response_config = try!(response_configs.first()
                .ok_or(HaxoniteError::NoResponseDefined(request_name.clone())));
            Box::new(try!(get_single_response_handler(request_name.clone(), response_config.clone())))
        }
        "random" => {
            let mut handler = RandomResponse::new();
            for response_config in &response_configs {
                let weight = response_config.weight.unwrap_or_else(config::default_weight);
                handler.add_handler(weight,
                                    try!(get_single_response_handler(request_name.clone(), response_config.clone())));
            }
            Box::new(handler)
        }
        "roundrobin" => {
            let mut handler = IteratingResponse::new("roundrobin");
            for response_config in &response_configs {
                handler.add_handler(try!(get_single_response_handler(request_name.clone(), response_config.clone())));
            }
            Box::new(handler)
        }
        "chain" => {
            let mut handler = IteratingResponse::new("chain");
            for response_config in &response_configs {
                handler.add_handler(try!(get_single_response_handler(request_name.clone(), response_config.clone())));
            }
            Box::new(handler)
        }
        "static" => {
            let response_config = try!(response_configs.first()
                .ok_or(HaxoniteError::NoResponseDefined(request_name.clone())));
            let response = match response_config.response {
                Some(ref response) => response,
                None => return Err(HaxoniteError::NoResponseDefined(request_name.clone())),
            };
            let handler = Static::new(Path::new(&response));
            Box::new(handler)
        }
        _ => return Err(HaxoniteError::InvalidType(request_name.clone())),
    };

    match Method::from_str(method.to_ascii_uppercase().as_ref()) {
        Ok(Method::Extension(_)) |
        Err(_) => return Err(HaxoniteError::InvalidHTTPMethod(request_name.clone())),
        Ok(a) => {
            match type_.as_ref() {
                "static" => {
                    info!("Mounting static for: {} using {} type of handler.\n",
                          path,
                          type_);
                    mount.mount(path.as_ref(), handler);
                }
                _ => {
                    info!("Defining route for: {} using {} type of handler.\n",
                          path,
                          type_);
                    router.route(a, path.clone(), handler, request_name);
                }
            }
        }
    };
    Ok(())
}

fn get_single_response_handler(request_name: String, response_config: ResponseConfig) -> Result<SingleResponse, HaxoniteError> {
    let headers = process_headers(response_config.headers);
    let status = match response_config.status {
        Some(status) if status >= 100 && status <= 599 => status, // switch to Range's contains() when it will be stable
        Some(status) => return Err(HaxoniteError::InvalidHTTPStatus(request_name.clone(), status)),
        None => config::default_status(),
    };
    let response = match response_config.response {
        Some(response) => response,
        None => return Err(HaxoniteError::NoResponseDefined(request_name.clone())),
    };
    if !Path::new(&response).exists() {
        return Err(HaxoniteError::ResponseDoesNotExist(request_name.clone()));
    }
    Ok(SingleResponse::new(status, headers, response, response_config.delay))
}

fn process_headers(headers: Option<Vec<String>>) -> Vec<Vec<String>> {
    // TODO: This method should be optimized to not compile regex on every run,
    // not match on header when the default is used (it's valid)
    use regex::Regex;
    let re = Regex::new(r"([\w-]+):\s+([\w/]+)").unwrap();

    let headers = headers.unwrap_or_else(config::default_headers);
    headers.iter()
        .filter(|header| {
            if re.is_match(header) {
                true
            } else {
                warn!("Header {} is invalid (does not match \": \"). Skipping.",
                      header);
                false
            }
        })
        .map(|header| {
            header.split(": ")
                .map(// TODO: This is quite stupid and could be done with &str,
                     // but I decided to use String to not fight borrow checker for now...
                     |hdr| hdr.to_string())
                .collect()
        })
        .collect()
}

fn define_command_line_options<'a, 'b>() -> App<'a, 'b> {
    App::new("Haxonite")
        .version(VERSION)
        .author(AUTHORS)
        .about("Easy API mocking")
        .arg(Arg::with_name("host")
            .short("h")
            .long("host")
            .value_name("HOST")
            .help("Run Haxonite on the host. Default: localhost.")
            .takes_value(true))
        .arg(Arg::with_name("port_number")
            .short("p")
            .long("port")
            .value_name("PORT")
            .help("Run Haxonite on the specified port. Default: 4000.")
            .takes_value(true))
        .arg(Arg::with_name("config_file")
            .short("c")
            .long("config")
            .value_name("FILE")
            .help("Sets a custom config file. Default: config.toml.")
            .takes_value(true))
        .subcommand(SubCommand::with_name("new")
            .about("Create new Haxonite project")
            .version(VERSION)
            .author(AUTHORS)
            .arg(Arg::with_name("full")
                .help("Use to create a project with full config.toml file")
                .short("f")
                .long("full"))
            .arg(Arg::with_name("project_name")
                .help("Name of the project (will be used as a directory name)")
                .index(1)
                .required(true)))
}
<commit_msg>#6: Add serve subcommand function<commit_after>const VERSION: &'static str = env!("CARGO_PKG_VERSION");
const AUTHORS: &'static str = env!("CARGO_PKG_AUTHORS");

static FORMAT: &'static str = "{request-time}: {method} {uri} {status} {response-time}";

#[macro_use]
extern crate clap;
use clap::{Arg, ArgMatches, App, SubCommand};

extern crate toml;
extern crate rustc_serialize;
extern crate regex;

extern crate iron;
use iron::prelude::*;
use iron::Handler;
use iron::method::Method;
use std::str::FromStr;

extern crate router;
use router::Router;

extern crate staticfile;
use staticfile::Static;

extern crate mount;
use mount::Mount;

#[macro_use]
extern crate log;
extern crate simplelog;
use simplelog::{SimpleLogger, LogLevelFilter};

extern crate logger;
use logger::Logger;
use logger::format::Format;

mod config;
use config::{Config, ServerConfig, RequestConfig, ResponseConfig};
use std::collections::HashMap;
use std::ascii::AsciiExt;
mod error;
use error::HaxoniteError;
mod handlers;
use handlers::*;
mod utils;

use std::path::Path;

fn main() {
    let _ = SimpleLogger::init(LogLevelFilter::Info, simplelog::Config::default());

    match run() {
        Ok(_) => {}
        Err(err) => {
            error!("{}!", err);
            std::process::exit(-1);
        }
    }
}

fn run() -> Result<(), HaxoniteError> {
    let matches = define_command_line_options().get_matches();
    if let Some(matches) = matches.subcommand_matches("new") {
        let generate_full_project = matches.is_present("full");
        if let Some(project_name) = matches.value_of("project_name") {
            try!(utils::create_new_project(project_name, generate_full_project));
        }
        return Ok(());
    }

    if let Some(matches) = matches.subcommand_matches("new") {
        try!(serve(matches));
    }
    Ok(())
}

fn serve(matches: &ArgMatches) -> Result<(), HaxoniteError> {
    let config_file = matches.value_of("config_file").unwrap_or("config.toml");
    let config_content = try!(config::read_config(config_file));
    let config: Config = toml::decode_str(config_content.as_ref()).unwrap(); // TODO: remove unwrap to handle malformed config
    debug!("Using config: {:?}!", config);

    let server_config = config.server.unwrap_or_else(ServerConfig::default);
    let host = match value_t!(matches, "host", String) { // TODO: rewrite these matches
        Ok(host) => host,
        Err(_) => server_config.host.unwrap_or_else(config::default_host),
    };
    let port_number = match value_t!(matches, "port_number", u16) {
        Ok(port_number) => port_number,
        Err(_) => server_config.port.unwrap_or_else(config::default_port),
    };

    // Initialize Iron's router and pass it to config processing function to define routes
    let mut router = Router::new();
    let mut mount = Mount::new();
    match config.requests {
        Some(requests) => {
            try!(process_config_requests(&requests, &mut mount, &mut router));
        }
        None => return Err(HaxoniteError::NoRequestDefined),
    }
    mount.mount("/", router);

    // Insert logger middlewares before and after router
    let mut chain = Chain::new(mount);
    let format = Format::new(FORMAT);
    let (logger_before, logger_after) = Logger::new(Some(format.unwrap()));
    chain.link_before(logger_before);
    chain.link_after(logger_after);

    // Initialize Iron to process requests
    let _iron = try!(Iron::new(chain).http((host.as_ref(), port_number)));
    info!("Haxonite running on port: {}!", port_number);
    Ok(())
}

fn process_config_requests(requests: &HashMap<String, RequestConfig>,
                           mut mount: &mut Mount,
                           mut router: &mut Router)
                           -> Result<(), HaxoniteError> {
    info!("Processing config!");
    for (request_name, request_config) in requests {
        info!("Processing config for {}: {:?}!",
              request_name,
              request_config);
        try!(define_route(request_name.clone(),
                          request_config.clone(),
                          &mut mount,
                          &mut router));
    }
    Ok(())
}

fn define_route(request_name: String, request_config: RequestConfig, mount: &mut Mount, router: &mut Router) -> Result<(), HaxoniteError> {
    let method = request_config.method.unwrap_or_else(config::default_method);
    let path = match request_config.path {
        Some(path) => path,
        None => return Err(HaxoniteError::NoPathDefined(request_name.clone())),
    };

    let type_ = request_config.type_.unwrap_or_else(config::default_type);
    let response_configs = match request_config.responses {
        Some(response_configs) => response_configs,
        None => return Err(HaxoniteError::NoResponseDefined(request_name.clone())),
    };
    let handler: Box<Handler> = match type_.as_ref() {
        "single" => {
            let response_config = try!(response_configs.first()
                .ok_or(HaxoniteError::NoResponseDefined(request_name.clone())));
            Box::new(try!(get_single_response_handler(request_name.clone(), response_config.clone())))
        }
        "random" => {
            let mut handler = RandomResponse::new();
            for response_config in &response_configs {
                let weight = response_config.weight.unwrap_or_else(config::default_weight);
                handler.add_handler(weight,
                                    try!(get_single_response_handler(request_name.clone(), response_config.clone())));
            }
            Box::new(handler)
        }
        "roundrobin" => {
            let mut handler = IteratingResponse::new("roundrobin");
            for response_config in &response_configs {
                handler.add_handler(try!(get_single_response_handler(request_name.clone(), response_config.clone())));
            }
            Box::new(handler)
        }
        "chain" => {
            let mut handler = IteratingResponse::new("chain");
            for response_config in &response_configs {
                handler.add_handler(try!(get_single_response_handler(request_name.clone(), response_config.clone())));
            }
            Box::new(handler)
        }
        "static" => {
            let response_config = try!(response_configs.first()
                .ok_or(HaxoniteError::NoResponseDefined(request_name.clone())));
            let response = match response_config.response {
                Some(ref response) => response,
                None => return Err(HaxoniteError::NoResponseDefined(request_name.clone())),
            };
            let handler = Static::new(Path::new(&response));
            Box::new(handler)
        }
        _ => return Err(HaxoniteError::InvalidType(request_name.clone())),
    };

    match Method::from_str(method.to_ascii_uppercase().as_ref()) {
        Ok(Method::Extension(_)) |
        Err(_) => return Err(HaxoniteError::InvalidHTTPMethod(request_name.clone())),
        Ok(a) => {
            match type_.as_ref() {
                "static" => {
                    info!("Mounting static for: {} using {} type of handler.\n",
                          path,
                          type_);
                    mount.mount(path.as_ref(), handler);
                }
                _ => {
                    info!("Defining route for: {} using {} type of handler.\n",
                          path,
                          type_);
                    router.route(a, path.clone(), handler, request_name);
                }
            }
        }
    };
    Ok(())
}

fn get_single_response_handler(request_name: String, response_config: ResponseConfig) -> Result<SingleResponse, HaxoniteError> {
    let headers = process_headers(response_config.headers);
    let status = match response_config.status {
        Some(status) if status >= 100 && status <= 599 => status, // switch to Range's contains() when it will be stable
        Some(status) => return Err(HaxoniteError::InvalidHTTPStatus(request_name.clone(), status)),
        None => config::default_status(),
    };
    let response = match response_config.response {
        Some(response) => response,
        None => return Err(HaxoniteError::NoResponseDefined(request_name.clone())),
    };
    if !Path::new(&response).exists() {
        return Err(HaxoniteError::ResponseDoesNotExist(request_name.clone()));
    }
    Ok(SingleResponse::new(status, headers, response, response_config.delay))
}

fn process_headers(headers: Option<Vec<String>>) -> Vec<Vec<String>> {
    // TODO: This method should be optimized to not compile regex on every run,
    // not match on header when the default is used (it's valid)
    use regex::Regex;
    let re = Regex::new(r"([\w-]+):\s+([\w/]+)").unwrap();

    let headers = headers.unwrap_or_else(config::default_headers);
    headers.iter()
        .filter(|header| {
            if re.is_match(header) {
                true
            } else {
                warn!("Header {} is invalid (does not match \": \"). Skipping.",
                      header);
                false
            }
        })
        .map(|header| {
            header.split(": ")
                .map(// TODO: This is quite stupid and could be done with &str,
                     // but I decided to use String to not fight borrow checker for now...
                     |hdr| hdr.to_string())
                .collect()
        })
        .collect()
}

fn define_command_line_options<'a, 'b>() -> App<'a, 'b> {
    App::new("Haxonite")
        .version(VERSION)
        .author(AUTHORS)
        .about("Easy API mocking")
        .subcommand(SubCommand::with_name("serve")
            .about("Starts Haxonite server")
            .version(VERSION)
            .author(AUTHORS)
            .arg(Arg::with_name("host")
                 .short("h")
                 .long("host")
                 .value_name("HOST")
                 .help("Run Haxonite on the host. Default: localhost.")
                 .takes_value(true))
            .arg(Arg::with_name("port_number")
                 .short("p")
                 .long("port")
                 .value_name("PORT")
                 .help("Run Haxonite on the specified port. Default: 4000.")
                 .takes_value(true))
            .arg(Arg::with_name("config_file")
                 .short("c")
                 .long("config")
                 .value_name("FILE")
                 .help("Sets a custom config file. Default: config.toml.")
                 .takes_value(true)))
        .subcommand(SubCommand::with_name("watch")
            .about("Starts Haxonite server with dynamic reload on change")
            .version(VERSION)
            .author(AUTHORS)
            .arg(Arg::with_name("host")
                 .short("h")
                 .long("host")
                 .value_name("HOST")
                 .help("Run Haxonite on the host. Default: localhost.")
                 .takes_value(true))
            .arg(Arg::with_name("port_number")
                 .short("p")
                 .long("port")
                 .value_name("PORT")
                 .help("Run Haxonite on the specified port. Default: 4000.")
                 .takes_value(true))
            .arg(Arg::with_name("config_file")
                 .short("c")
                 .long("config")
                 .value_name("FILE")
                 .help("Sets a custom config file. Default: config.toml.")
                 .takes_value(true)))
        .subcommand(SubCommand::with_name("new")
            .about("Create new Haxonite project")
            .version(VERSION)
            .author(AUTHORS)
            .arg(Arg::with_name("full")
                .help("Use to create a project with full config.toml file")
                .short("f")
                .long("full"))
            .arg(Arg::with_name("project_name")
                .help("Name of the project (will be used as a directory name)")
                .index(1)
                .required(true)))
}
<|endoftext|>


<commit_before>use clap::{crate_version, App, Arg};
use std::io;
use std::process;
use table::{reader, writer};

const DESCRIPTION: &str = "A command to print ASCII table from stdin";
const USAGE_TEMPLATE: &str = r#"
Usage:
    {usage}

Flags:
{flags}

Options:
{options}
"#;

fn main() {
    let matches = App::new("table")
        .version(crate_version!())
        .author("Naoto Kaneko <naoty.k@gmail.com>")
        .about(DESCRIPTION)
        .template(USAGE_TEMPLATE.trim())
        .version_short("v")
        .arg(
            Arg::with_name("format")
                .short("f")
                .long("format")
                .takes_value(true)
                .default_value("")
                .value_name("FORMAT")
                .help("Config input/output data format"),
        )
        .arg(
            Arg::with_name("header")
                .short("H")
                .long("header")
                .help("Prints table with headers"),
        )
        .get_matches();

    let reader: &mut dyn table::Read =
        &mut reader::CsvReader::new(io::stdin(), b',', matches.is_present("header"));
    let writer: &mut dyn table::Write = &mut writer::AsciiWriter::new(io::stdout());
    let result = reader
        .read()
        .and_then(|table| writer.write(table))
        .and_then(|_| writer.flush());

    match result {
        Ok(_) => {
            process::exit(0);
        }
        Err(error) => {
            eprintln!("{:?}", error);
            process::exit(1);
        }
    }
}
<commit_msg>Switch reader and writer with -f option<commit_after>use clap::{crate_version, App, Arg};
use std::io;
use std::process;
use table::{reader, writer};

const DESCRIPTION: &str = "A command to print ASCII table from stdin";
const USAGE_TEMPLATE: &str = r#"
Usage:
    {usage}

Flags:
{flags}

Options:
{options}
"#;

fn main() {
    let matches = App::new("table")
        .version(crate_version!())
        .author("Naoto Kaneko <naoty.k@gmail.com>")
        .about(DESCRIPTION)
        .template(USAGE_TEMPLATE.trim())
        .version_short("v")
        .arg(
            Arg::with_name("format")
                .short("f")
                .long("format")
                .takes_value(true)
                .default_value("")
                .value_name("FORMAT")
                .help("Config input/output data format"),
        )
        .arg(
            Arg::with_name("header")
                .short("H")
                .long("header")
                .help("Prints table with headers"),
        )
        .get_matches();

    let format = matches.value_of("format").unwrap_or_default();
    let mut tokens = format.split(":").take(2);

    let mut csv_reader = reader::CsvReader::new(io::stdin(), b',', matches.is_present("header"));
    let mut tsv_reader = reader::CsvReader::new(io::stdin(), b'\t', matches.is_present("header"));
    let reader: &mut dyn table::Read = match tokens.next() {
        Some("csv") => &mut csv_reader,
        _ => &mut tsv_reader,
    };

    let mut ascii_writer = writer::AsciiWriter::new(io::stdout());
    let writer: &mut dyn table::Write = match tokens.next() {
        _ => &mut ascii_writer,
    };

    let result = reader
        .read()
        .and_then(|table| writer.write(table))
        .and_then(|_| writer.flush());

    match result {
        Ok(_) => {
            process::exit(0);
        }
        Err(error) => {
            eprintln!("{:?}", error);
            process::exit(1);
        }
    }
}
<|endoftext|>