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

use std::fs::File;
use std::fs::OpenOptions;
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;
use chrono::Local;

const VERSION: &'static str = env!("CARGO_PKG_VERSION");
const BINARY_MATCHES: [(&str, &str);24] = [ // Global array, so needs an explicit length
  ("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
  ("Bitcoin main chain magic number", r"\xf9\xbe\xb4\xd9")
];

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

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

fn log_to_file(msg: String) {
  let mut file = OpenOptions::new()
              .append(true)
              .create(true)
              .open("portlurker.log").expect("Failed to open local log file for writing");;
  writeln!(file, "{}", msg).unwrap();

}

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(' ');
    result.push_str(&dotline);
  }
  result
}

fn setup() -> App {
  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, file_logging: false };

  let mut config_str = String::new();
  let mut file = match File::open("config.yml") {
      Ok(file) => file,
      Err(e) => { println!("Unable to open configuration file: {}", e.to_string()); exit(-1); },
  };

  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"]["file_logging"].is_badvalue() {
    if config["general"]["file_logging"].as_bool().unwrap() {
      app.file_logging = true;
      println!("Logging to local text file");
    }
  }
  if !config["general"]["sql_logging"].is_badvalue() {
    if config["general"]["sql_logging"].as_bool().unwrap() {
      println!("Logging events to portlurker.sqlite");
      match Connection::open("portlurker.sqlite") {
        Ok(conn) => { app.sql_logging = true;
                      println!("Logging to local SQL database file");
                      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! Logging may not function correctly!"); },
        Err(e) => {
          println!("Failed to open or create database: {}\nContinuing without logging", e.to_string());
        },
      };
    }
  }
  return app; // send our config to the main function
}

fn main() {
  let app: App = setup(); // Print initial UI stuff, then parse the config file, and store the config
  let io_timeout = Duration::new(300, 0); // 5 minutes

  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:");

  // Have to reload config file here - can we improve this?
  let mut config_str = String::new();
  let mut file = match File::open("config.yml") {
      Ok(file) => file,
      Err(e) => { println!("Unable to open configuration file: {}", e.to_string()); exit(-1); },
  };

  file.read_to_string(&mut config_str).unwrap();
  let docs = YamlLoader::load_from_str(&config_str).unwrap();
  let config = &docs[0];
  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)).expect("Port can't be bound (is it in use?)"); // Add more error checking here
        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();

          let current_time = Local::now();
          let formatted_time = format!("{}", current_time.format("%a %d %b %Y - %H:%M.%S"));

          let log_msg = format!("[{}]: CONNECT TCP {} from {}", formatted_time, portno, addr);
          println!("{}", log_msg);
          if app.file_logging {
            log_to_file(log_msg.to_string());
          }
          if app.sql_logging {
            let newdbentry = LoggedConnection {
              remoteip: addr.ip().to_string(),
              remoteport: addr.port(),
              localport: portno
            };
            match Connection::open("portlurker.sqlite") {
              Ok(conn) => { 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! Subsequent logging may also fail.");},
              Err(e) => {
                println!("Failed to open database: {} - Continuing without logging", e.to_string());
                let mut app = app;
                app.sql_logging = false;
              },
            }
          }

          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>Update main.rs<commit_after>extern crate yaml_rust;
extern crate regex;
extern crate rusqlite;
extern crate chrono;

use std::fs::File;
use std::fs::OpenOptions;
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;
use chrono::Local;

const VERSION: &'static str = env!("CARGO_PKG_VERSION");
const BINARY_MATCHES: [(&str, &str);24] = [ // Global array, so needs an explicit length
  ("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
  ("Bitcoin main chain magic number", r"\xf9\xbe\xb4\xd9")
];

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

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

fn log_to_file(msg: String) {
  let mut file = OpenOptions::new()
              .append(true)
              .create(true)
              .open("portlurker.log").expect("Failed to open local log file for writing");;
  writeln!(file, "{}", msg).unwrap();

}

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(' ');
    result.push_str(&dotline);
  }
  result
}

fn setup() -> App {
  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, file_logging: false };

  let mut config_str = String::new();
  let mut file = match File::open("config.yml") {
      Ok(file) => file,
      Err(e) => { println!("Unable to open configuration file: {}", e.to_string()); exit(-1); },
  };

  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"]["file_logging"].is_badvalue() {
    if config["general"]["file_logging"].as_bool().unwrap() {
      app.file_logging = true;
      println!("Logging to local text file");
    }
  }
  if !config["general"]["sql_logging"].is_badvalue() {
    if config["general"]["sql_logging"].as_bool().unwrap() {
      match Connection::open("portlurker.sqlite") {
        Ok(conn) => { app.sql_logging = true;
                      println!("Logging to local SQL database file portlurker.sqlite");
                      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! Logging may not function correctly!"); },
        Err(e) => {
          println!("Enabling SQL logging failed because it was not possible to open or create the database: {}\nContinuing without SQL logging", e.to_string());
        },
      };
    }
  }
  return app; // send our config to the main function
}

fn main() {
  let app: App = setup(); // Print initial UI stuff, then parse the config file, and store the config
  let io_timeout = Duration::new(300, 0); // 5 minutes

  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:");

  // Have to reload config file here - can we improve this?
  let mut config_str = String::new();
  let mut file = match File::open("config.yml") {
      Ok(file) => file,
      Err(e) => { println!("Unable to open configuration file: {}", e.to_string()); exit(-1); },
  };

  file.read_to_string(&mut config_str).unwrap();
  let docs = YamlLoader::load_from_str(&config_str).unwrap();
  let config = &docs[0];
  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)).expect("Port can't be bound (is it in use?)"); // Add more error checking here
        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();

          let current_time = Local::now();
          let formatted_time = format!("{}", current_time.format("%a %d %b %Y - %H:%M.%S"));

          let log_msg = format!("[{}]: CONNECT TCP {} from {}", formatted_time, portno, addr);
          println!("{}", log_msg);
          if app.file_logging {
            log_to_file(log_msg.to_string());
          }
          if app.sql_logging {
            let newdbentry = LoggedConnection {
              remoteip: addr.ip().to_string(),
              remoteport: addr.port(),
              localport: portno
            };
            match Connection::open("portlurker.sqlite") {
              Ok(conn) => { 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! Subsequent logging may also fail.");},
              Err(e) => {
                println!("Failed to open database: {} - Continuing without logging", e.to_string());
                let mut app = app;
                app.sql_logging = false;
              },
            }
          }

          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>//  rpc-perf - RPC Performance Testing
//  Copyright 2015 Twitter, Inc
//
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the License.
//  You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.

#![cfg_attr(feature = "unstable", feature(test))]
#[cfg(feature = "unstable")]
extern crate test;

#[allow(unknown_lints, useless_attribute)]
#[cfg_attr(feature = "cargo-clippy", deny(result_unwrap_used))]
#[macro_use]
extern crate log;
extern crate log_panics;

extern crate bytes;
extern crate byteorder;
extern crate crc;
extern crate getopts;
extern crate mio;
extern crate mpmc;
extern crate pad;
extern crate time;
extern crate rand;
extern crate ratelimit;
extern crate slab;
extern crate tic;
extern crate toml;

mod common;
mod cfgtypes;
mod client;
mod connection;
mod logger;
mod net;
mod options;
mod stats;
mod codec;
mod request;

use client::Client;
use net::InternetProtocol;
use request::{config, workload};
use std::{env, process, thread};

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

pub fn main() {
    let args: Vec<String> = env::args().collect();
    let program = &args[0];
    let opts = options::opts();

    let matches = match opts.parse(&args[1..]) {
        Ok(m) => m,
        Err(f) => {
            println!("ERROR {}", f);
            options::print_usage(program, &opts);
            process::exit(1);
        }
    };

    if matches.opt_present("help") {
        options::print_usage(program, &opts);
        return;
    }

    if matches.opt_present("version") {
        println!("rpc-perf {}", VERSION);
        return;
    }

    // initialize logging
    logger::set_log_level(matches.opt_count("verbose"));
    log_panics::init();

    info!("rpc-perf {} initializing...", VERSION);
    if cfg!(feature = "asm") {
        info!("feature: asm: enabled");
    } else {
        info!("feature: asm: disabled");
    }

    let servers: Vec<String> = matches.opt_strs("server");

    if servers.is_empty() {
        error!("require server parameter");
        options::print_usage(program, &opts);
        return;
    }

    let listen = matches.opt_str("listen");
    let waterfall = matches.opt_str("waterfall");
    let trace = matches.opt_str("trace");

    // Load workload configuration
    let config = match config::load_config(&matches) {
        Ok(cfg) => cfg,
        Err(e) => {
            error!("{}", e);
            return;
        }
    };

    let internet_protocol = match net::choose_layer_3(matches.opt_present("ipv4"),
                                                      matches.opt_present("ipv6")) {
        Ok(i) => i,
        Err(e) => {
            error!("{}", e);
            return;
        }
    };

    print_config(&config, &servers, internet_protocol);

    let stats_receiver = stats::stats_receiver_init(&config, listen, waterfall, trace);

    let mut send_queues = Vec::new();

    let mut client_config = Client::configure();

    client_config
        .set_pool_size(config.poolsize())
        .stats(stats_receiver.get_sender().clone())
        .set_clocksource(stats_receiver.get_clocksource().clone())
        .set_protocol(config.protocol_config.protocol.clone())
        .set_request_timeout(config.request_timeout())
        .set_connect_timeout(config.connect_timeout())
        .set_internet_protocol(internet_protocol);

    for server in servers {
        client_config.add_server(server);
    }

    info!("-----");
    info!("Connecting...");
    for i in 0..config.threads() {
        info!("Client: {}", i);

        let mut client = client_config.clone().build();
        send_queues.push(client.tx());

        let _ = thread::Builder::new()
            .name(format!("client{}", i).to_string())
            .spawn(move || { client.run(); });
    }

    info!("-----");
    info!("Workload:");

    let windows = config.windows();

    workload::launch_workloads(config.protocol_config.workloads,
                               &send_queues,
                               &stats_receiver.get_sender(),
                               &stats_receiver.get_clocksource());

    stats::run(stats_receiver, windows, matches.opt_present("service"));
}

fn print_config(config: &request::BenchmarkConfig,
                servers: &[String],
                internet_protocol: InternetProtocol) {
    info!("-----");
    info!("Config:");
    for server in servers {
        info!("Config: Server: {} Protocol: {}",
              server,
              config.protocol_config.protocol.name());
    }
    info!("Config: IP: {:?} TCP_NODELAY: {}",
          internet_protocol,
          config.tcp_nodelay());
    info!("Config: Threads: {} Poolsize: {}",
          config.threads(),
          config.poolsize());
    info!("Config: Windows: {} Duration: {}",
          config.windows(),
          config.duration());
    match config.request_timeout() {
        Some(v) => {
            info!("Config: Request Timeout: {} ms", v);
        }
        None => {
            info!("Config: Request Timeout: None");
        }
    }
    match config.connect_timeout() {
        Some(v) => {
            info!("Config: Connect Timeout: {} ms", v);
        }
        None => {
            info!("Config: Connect Timeout: None");
        }
    }
}
<commit_msg>cleanup externs<commit_after>//  rpc-perf - RPC Performance Testing
//  Copyright 2015 Twitter, Inc
//
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the License.
//  You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.

#![cfg_attr(feature = "unstable", feature(test))]
#[cfg(feature = "unstable")]
extern crate test;

#[allow(unknown_lints, useless_attribute)]
#[cfg_attr(feature = "cargo-clippy", deny(result_unwrap_used))]

#[macro_use] extern crate log;
extern crate log_panics;
extern crate bytes;
extern crate byteorder;
extern crate crc;
extern crate getopts;
extern crate mio;
extern crate mpmc;
extern crate pad;
extern crate time;
extern crate rand;
extern crate ratelimit;
extern crate slab;
extern crate tic;
extern crate toml;

mod common;
mod cfgtypes;
mod client;
mod connection;
mod logger;
mod net;
mod options;
mod stats;
mod codec;
mod request;

use client::Client;
use net::InternetProtocol;
use request::{config, workload};
use std::{env, process, thread};

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

pub fn main() {
    let args: Vec<String> = env::args().collect();
    let program = &args[0];
    let opts = options::opts();

    let matches = match opts.parse(&args[1..]) {
        Ok(m) => m,
        Err(f) => {
            println!("ERROR {}", f);
            options::print_usage(program, &opts);
            process::exit(1);
        }
    };

    if matches.opt_present("help") {
        options::print_usage(program, &opts);
        return;
    }

    if matches.opt_present("version") {
        println!("rpc-perf {}", VERSION);
        return;
    }

    // initialize logging
    logger::set_log_level(matches.opt_count("verbose"));
    log_panics::init();

    info!("rpc-perf {} initializing...", VERSION);
    if cfg!(feature = "asm") {
        info!("feature: asm: enabled");
    } else {
        info!("feature: asm: disabled");
    }

    let servers: Vec<String> = matches.opt_strs("server");

    if servers.is_empty() {
        error!("require server parameter");
        options::print_usage(program, &opts);
        return;
    }

    let listen = matches.opt_str("listen");
    let waterfall = matches.opt_str("waterfall");
    let trace = matches.opt_str("trace");

    // Load workload configuration
    let config = match config::load_config(&matches) {
        Ok(cfg) => cfg,
        Err(e) => {
            error!("{}", e);
            return;
        }
    };

    let internet_protocol = match net::choose_layer_3(matches.opt_present("ipv4"),
                                                      matches.opt_present("ipv6")) {
        Ok(i) => i,
        Err(e) => {
            error!("{}", e);
            return;
        }
    };

    print_config(&config, &servers, internet_protocol);

    let stats_receiver = stats::stats_receiver_init(&config, listen, waterfall, trace);

    let mut send_queues = Vec::new();

    let mut client_config = Client::configure();

    client_config
        .set_pool_size(config.poolsize())
        .stats(stats_receiver.get_sender().clone())
        .set_clocksource(stats_receiver.get_clocksource().clone())
        .set_protocol(config.protocol_config.protocol.clone())
        .set_request_timeout(config.request_timeout())
        .set_connect_timeout(config.connect_timeout())
        .set_internet_protocol(internet_protocol);

    for server in servers {
        client_config.add_server(server);
    }

    info!("-----");
    info!("Connecting...");
    for i in 0..config.threads() {
        info!("Client: {}", i);

        let mut client = client_config.clone().build();
        send_queues.push(client.tx());

        let _ = thread::Builder::new()
            .name(format!("client{}", i).to_string())
            .spawn(move || { client.run(); });
    }

    info!("-----");
    info!("Workload:");

    let windows = config.windows();

    workload::launch_workloads(config.protocol_config.workloads,
                               &send_queues,
                               &stats_receiver.get_sender(),
                               &stats_receiver.get_clocksource());

    stats::run(stats_receiver, windows, matches.opt_present("service"));
}

fn print_config(config: &request::BenchmarkConfig,
                servers: &[String],
                internet_protocol: InternetProtocol) {
    info!("-----");
    info!("Config:");
    for server in servers {
        info!("Config: Server: {} Protocol: {}",
              server,
              config.protocol_config.protocol.name());
    }
    info!("Config: IP: {:?} TCP_NODELAY: {}",
          internet_protocol,
          config.tcp_nodelay());
    info!("Config: Threads: {} Poolsize: {}",
          config.threads(),
          config.poolsize());
    info!("Config: Windows: {} Duration: {}",
          config.windows(),
          config.duration());
    match config.request_timeout() {
        Some(v) => {
            info!("Config: Request Timeout: {} ms", v);
        }
        None => {
            info!("Config: Request Timeout: None");
        }
    }
    match config.connect_timeout() {
        Some(v) => {
            info!("Config: Connect Timeout: {} ms", v);
        }
        None => {
            info!("Config: Connect Timeout: None");
        }
    }
}
<|endoftext|>


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

use async_std::io;
use async_std::os::unix::io::AsRawFd;
use async_std::task;
use clap::{App, Arg};
use std::convert::TryFrom;
use std::process;

use scout::common::{Result, Text};
use scout::config::Configurator;
use scout::ptty::{self, PTTY};
use scout::supervisor;

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

    trace!("starting main program");

    let args = App::new("scout")
        .version(crate_version!())
        .arg(
            Arg::with_name("full-screen")
                .short("f")
                .long("full-screen")
                .help("Show fuzzy finder in full screen (default)"),
        )
        .arg(
            Arg::with_name("inline")
                .short("i")
                .long("inline")
                .help("Show fuzzy finder under the current line"),
        )
        .arg(
            Arg::with_name("lines")
                .short("l")
                .long("lines")
                .value_name("LINES")
                .takes_value(true)
                .help("Number of lines to display in inline mode, including prompt"),
        )
        .arg(
            Arg::with_name("search")
                .short("s")
                .long("search")
                .value_name("QUERY")
                .takes_value(true)
                .help("Start searching with the given query"),
        )
        .arg(
            Arg::with_name("config")
                .short("c")
                .long("config")
                .value_name("FILE")
                .takes_value(true)
                .help("Sets a custom config file"),
        )
        .get_matches();

    trace!("got args: {:?}", args);

    let res: Result<Option<Text>> = task::block_on(async {
        let mut configurator = Configurator::new();

        match args.value_of("config") {
            Some(config_path) => configurator.from_file(config_path),
            None => configurator.from_default_file(),
        };

        // PTTY = Pseudo Terminal
        let tty = ptty::file().await?;
        let config = configurator.from_ptty(&tty).from_args(&args).build();

        trace!("generated config: {:?}", config);

        // We only need to set up the ptty into noncanonical mode once
        let ptty = PTTY::try_from(tty.as_raw_fd())?;
        ptty.noncanonical_mode()?;

        // Get the list of candidates to filter from the STDIN
        // This list comes most probably from a pipe
        let stdin = io::stdin();

        // The architecture of the app is async, one task (a future) will
        // read from the PTTY while another task will write to it. We can't
        // have one single mutable PTTY reference across these tasks because
        // locking to read would mean blocking printing to the screen.
        // To overcome this we create two different references to the
        // system's PTTY, one for reading and one for writting.
        //
        // Even though the path to the PTTY is the same, reading from and
        // writting to it are fundamentally two different and independent steps.
        //
        // Is this a hack? Most probably, yes. Does it work? Also yes.
        let pttyin = ptty::reader().await?; // to read person's input
        let pttyout = ptty::writer().await?; // to print programs interface

        // The main program's thread will block until the supervisor's task finishes
        // thanks to the `task::block_on` call
        supervisor::run(config, stdin, pttyin, pttyout).await
    });

    trace!("program ended with {:?}", res);

    match res {
        Ok(Some(selection)) => println!("{}", selection),
        Ok(None) => process::exit(130),
        Err(e) => {
            eprintln!("ERROR: {}", e);
            process::exit(1);
        }
    }
}
<commit_msg>Improve CLI help message<commit_after>#[macro_use]
extern crate log;
#[macro_use]
extern crate clap;

use async_std::io;
use async_std::os::unix::io::AsRawFd;
use async_std::task;
use clap::{App, Arg};
use std::convert::TryFrom;
use std::process;

use scout::common::{Result, Text};
use scout::config::Configurator;
use scout::ptty::{self, PTTY};
use scout::supervisor;

const EXTENDED_HELP: &'static str = r#"SUPPORTED KEYS:
    - Enter to select the current highlighted match and print it to STDOUT
    - ^u to clear the prompt
    - ^n or Down arrow key to select the next match
    - ^p or Up arrow key to select the previous match
    - ^e to go to the end of the prompt
    - ^a to go to the beginning of the prompt
    - Left arrow key to move the cursor to the left in the prompt
    - Right arrow key to move the cursor to the right in the prompt
    - ESC to quit without selecting a match

EXAMPLES:
    $ find * -type f | scout

    # Pass an initial query to start filtering right away
    $ find * -type f | scout --search=foo

    # Use a custom config file
    $ find * -type f | scout --config="./config.toml"

    # Select a git branch and check it out with an inline menu
    $ git branch | cut -c 3- | scout -i | xargs git checkout"#;

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

    trace!("starting main program");

    let args = App::new("scout")
        .version(crate_version!())
        .about("Your friendly fuzzy finder")
        .after_help(EXTENDED_HELP)
        .arg(
            Arg::with_name("full-screen")
                .short("f")
                .long("full-screen")
                .help("Show scout in full screen (default)"),
        )
        .arg(
            Arg::with_name("inline")
                .short("i")
                .long("inline")
                .help("Show scout under the current line"),
        )
        .arg(
            Arg::with_name("lines")
                .short("l")
                .long("lines")
                .value_name("LINES")
                .takes_value(true)
                .help("Number of lines to display in inline mode, including prompt"),
        )
        .arg(
            Arg::with_name("search")
                .short("s")
                .long("search")
                .value_name("QUERY")
                .takes_value(true)
                .help("Start searching with the given query"),
        )
        .arg(
            Arg::with_name("config")
                .short("c")
                .long("config")
                .value_name("FILE")
                .takes_value(true)
                .help("Uses a custom config file"),
        )
        .get_matches();

    trace!("got args: {:?}", args);

    let res: Result<Option<Text>> = task::block_on(async {
        let mut configurator = Configurator::new();

        match args.value_of("config") {
            Some(config_path) => configurator.from_file(config_path),
            None => configurator.from_default_file(),
        };

        // PTTY = Pseudo Terminal
        let tty = ptty::file().await?;
        let config = configurator.from_ptty(&tty).from_args(&args).build();

        trace!("generated config: {:?}", config);

        // We only need to set up the ptty into noncanonical mode once
        let ptty = PTTY::try_from(tty.as_raw_fd())?;
        ptty.noncanonical_mode()?;

        // Get the list of candidates to filter from the STDIN
        // This list comes most probably from a pipe
        let stdin = io::stdin();

        // The architecture of the app is async, one task (a future) will
        // read from the PTTY while another task will write to it. We can't
        // have one single mutable PTTY reference across these tasks because
        // locking to read would mean blocking printing to the screen.
        // To overcome this we create two different references to the
        // system's PTTY, one for reading and one for writting.
        //
        // Even though the path to the PTTY is the same, reading from and
        // writting to it are fundamentally two different and independent steps.
        //
        // Is this a hack? Most probably, yes. Does it work? Also yes.
        let pttyin = ptty::reader().await?; // to read person's input
        let pttyout = ptty::writer().await?; // to print programs interface

        // The main program's thread will block until the supervisor's task finishes
        // thanks to the `task::block_on` call
        supervisor::run(config, stdin, pttyin, pttyout).await
    });

    trace!("program ended with {:?}", res);

    match res {
        Ok(Some(selection)) => println!("{}", selection),
        Ok(None) => process::exit(130),
        Err(e) => {
            eprintln!("ERROR: {}", e);
            process::exit(1);
        }
    }
}
<|endoftext|>


<commit_before>extern crate unicode_width;

mod args;
mod matching;
mod screen;
#[cfg(not(windows))] mod ansi;

use std::collections::HashSet;
use std::process;
use std::io::{stdin, BufRead};
use std::cmp::min;
use screen::Screen;
use screen::Key;
use screen::Key::*;
use self::SearchState::*;
use unicode_width::UnicodeWidthStr;

#[cfg(windows)] pub const NEWLINE: &'static str = "\r\n";
#[cfg(not(windows))] pub const NEWLINE: &'static str = "\n";

const VERSION: &'static str = include_str!(concat!(env!("OUT_DIR"), "/version.txt"));
const TIMESTAMP: &'static str = include_str!(concat!(env!("OUT_DIR"), "/timestamp.txt"));
const TARGET: &'static str = include_str!(concat!(env!("OUT_DIR"), "/target.txt"));
const COMMIT: &'static str = include_str!(concat!(env!("OUT_DIR"), "/commit.txt"));

fn main() {
    let args = match args::parse_args() {
        Some(args) => args,
        None => {
            process::exit(1);
        },
    };

    if args.help { return; }

    if args.version {
        if COMMIT == "" {
            println!("heatseeker {} (built {} for {})", VERSION, TIMESTAMP, TARGET);
        } else {
            println!("heatseeker {}-{} (built {} for {})", VERSION, COMMIT, TIMESTAMP, TARGET);
        }
        return;
    }

    let choices = read_choices();
    let initial_search = args.initial_search.clone();
    let choices = choices.iter().map(|x| &x[..]).collect::<Vec<&str>>();
    if args.use_first {
        let matches = matching::compute_matches(&choices, &initial_search, args.filter_only);
        println!("{}", matches.get(0).unwrap_or(&""));
        return;
    } else {
        let desired_rows = if args.full_screen { 999 } else { 20 };
        let selections = event_loop(desired_rows, &choices, &initial_search, args.filter_only);
        print!("{}", selections);
    }
}

fn event_loop(desired_rows: u16, choices: &[&str], initial_search: &str, filter_only: bool) -> String {
    let mut search = Search::new(choices, initial_search.to_string(), filter_only);
    let mut screen = screen::new(desired_rows);

    loop {
        search.recompute_matches();

        match search.state {
            InProgress => draw_screen(screen.as_mut(), &search),
            _ => break,
        }

        let keys = screen.get_buffered_keys();
        for key in &keys {
            handle_key(&mut search, key, screen.visible_choices());
        }
    }

    screen.blank_screen();
    search.get_selections()
}

fn handle_key(search: &mut Search, key: &Key, visible_choices: u16) {
    match *key {
        Char(x) => search.append(x),
        Backspace |
        Control('h') => search.backspace(),
        Control('w') => search.delete_word(),
        Control('u') => search.clear_query(),
        Control('c') |
        Control('g') => search.cancel(),
        Control('t') => { search.toggle_selection(); search.down(visible_choices); },
        Control('p') |
        Up |
        ShiftTab => search.up(visible_choices),
        Control('n') |
        Down |
        Tab => search.down(visible_choices),
        Home => search.home(),
        End => search.end(visible_choices),
        Enter => search.done(),
        PgUp => search.pgup(visible_choices),
        PgDown => search.pgdown(visible_choices),
        _ => {},
    }
}

struct Search<'a> {
    choices: &'a [&'a str],
    query: String,
    matches: Vec<&'a str>,
    stale: bool,
    scroll_offset: usize,
    cursor_index: usize,
    state: SearchState,
    selections: HashSet<String>,
    filter_only: bool,
}

#[derive(PartialEq, Eq)]
enum SearchState {
    InProgress,
    Done,
    Canceled,
}

impl<'a> Search<'a> {
    fn new(choices: &'a [&'a str], initial_search: String, filter_only: bool) -> Search<'a> {
        let matches = choices.to_vec();
        Search {
            choices: choices,
            query: initial_search,
            matches: matches,
            stale: true,
            scroll_offset: 0,
            cursor_index: 0,
            state: InProgress,
            selections: HashSet::new(),
            filter_only: filter_only,
        }
    }

    fn up(&mut self, visible_choices: u16) {
        let match_count = self.matches.len();
        let limit = min(visible_choices as usize - 1, match_count - 1);
        let should_wrap = self.scroll_offset == 0;
        if self.cursor_index == 0 {
            if should_wrap {
                if match_count > visible_choices as usize {
                    self.scroll_offset = match_count - visible_choices as usize;
                    self.cursor_index = visible_choices as usize - 1;
                } else {
                    self.cursor_index = limit;
                }
            } else {
                self.scroll_offset -= 1;
            }
        } else {
            self.cursor_index -= 1;
        }
    }

    fn down(&mut self, visible_choices: u16) {
        let match_count = self.matches.len();
        let limit = min(visible_choices as usize - 1, match_count - 1);
        let should_wrap = self.cursor_index + self.scroll_offset == match_count - 1;
        if self.cursor_index == limit {
            if should_wrap {
                self.cursor_index = 0;
                self.scroll_offset = 0;
            } else {
                self.scroll_offset += 1;
            }
        } else {
            self.cursor_index += 1;
        }
    }

    fn home(&mut self) {
        self.cursor_index = 0;
        self.scroll_offset = 0;
    }

    fn end(&mut self, visible_choices: u16) {
        self.home();
        self.up(visible_choices);
    }

    fn pgup(&mut self, visible_choices: u16) {
        for _ in 0..visible_choices {
            if self.scroll_offset == 0 && self.cursor_index == 0 {
                return;
            }
            self.up(visible_choices);
        }
    }

    fn pgdown(&mut self, visible_choices: u16) {
        for _ in 0..visible_choices {
            if self.scroll_offset + self.cursor_index == self.matches.len() - 1 {
                return;
            }
            self.down(visible_choices);
        }
    }

    fn backspace(&mut self) {
        self.query.pop();
        self.stale = true;
        self.cursor_index = 0;
        self.scroll_offset = 0;
        self.matches = self.choices.to_vec();
    }

    fn delete_word(&mut self) {
        self.stale = true;
        delete_last_word(&mut self.query);
        self.matches = self.choices.to_vec();
    }

    fn append(&mut self, c: char) {
        self.query.push(c);
        self.stale = true;
        self.cursor_index = 0;
        self.scroll_offset = 0;
    }

    fn clear_query(&mut self) {
        self.query.clear();
        self.cursor_index = 0;
        self.scroll_offset = 0;
        self.matches = self.choices.to_vec();
    }

    fn recompute_matches(&mut self) {
        if self.stale {
            self.matches = matching::compute_matches(&self.matches, &self.query, self.filter_only);
            self.stale = false;
        }
    }

    fn toggle_selection(&mut self) {
        self.recompute_matches();
        let selection = self.matches.get(self.scroll_offset + self.cursor_index).unwrap_or(&"").to_string();
        if self.selections.contains(&selection) {
            self.selections.remove(&selection);
        } else {
            self.selections.insert(selection);
        }
    }

    fn get_selections(&mut self) -> String {
        let mut ret = String::new();
        if self.state != Canceled {
            for selection in &self.selections {
                ret.push_str(selection);
                ret.push_str(NEWLINE);
            }
            if ret.is_empty() {
                self.recompute_matches();
                let selection = self.matches.get(self.scroll_offset + self.cursor_index).unwrap_or(&"").to_string();
                ret.push_str(&selection);
            }
        }
        ret
    }

    fn cancel(&mut self) {
        self.state = Canceled
    }

    fn done(&mut self) {
        self.state = Done
    }
}

fn draw_screen(screen: &mut dyn Screen, search: &Search) {
    screen.hide_cursor();
    screen.blank_screen();
    screen.write(&format!("> {} ({}/{} choices){}", search.query, search.matches.len(), search.choices.len(), NEWLINE));

    print_matches(screen, &search.matches, &search.query, search.scroll_offset, search.cursor_index, &search.selections);

    let query_str: &str = &search.query;
    screen.move_cursor_to_prompt_line(2 + UnicodeWidthStr::width(query_str) as u16);
    screen.show_cursor();
}

fn print_matches(screen: &mut dyn Screen, matches: &[&str], query: &str, scroll_offset: usize, cursor_index: usize, selections: &HashSet<String>) {
    let mut i = 1;
    for choice in matches[scroll_offset..].iter() {
        let indices = matching::visual_score(choice, query);
        let max_width = screen.width();
        let mut annotated_choice = choice.to_string();
        if selections.contains(&annotated_choice) {
            annotated_choice.push(' ');
            if cfg!(windows) {
                annotated_choice.push('√');
            } else {
                annotated_choice.push('✓');
            }
        }
        print_match(&annotated_choice, &indices, max_width, &mut |s, highlight| {
            if i == cursor_index + 1 {
                if highlight { screen.write_red_inverted(s); } else { screen.write_inverted(s); }
            } else if highlight {
                screen.write_red(s);
            } else { screen.write(s); }
        });
        if i >= screen.visible_choices() as usize {
            return;
        }
        screen.write(NEWLINE);
        i += 1;
    }
}

fn print_match(choice: &str, indices: &[usize], max_width: u16, writer: &mut dyn FnMut(&str, bool)) {
    #[cfg(windows)] const MARGIN: u16 = 1;
    #[cfg(not(windows))] const MARGIN: u16 = 0;
    let max_width = max_width - MARGIN;
    let chars_in_choice = choice.chars().count();
    let mut chars_to_draw = min(chars_in_choice, max_width as usize);
    while UnicodeWidthStr::width(slice_chars(choice, 0, chars_to_draw)) > max_width as usize {
        chars_to_draw -= 1;
    }
    let mut last_idx = 0;
    for &idx in indices {
        let idx = min(idx, chars_to_draw);
        if last_idx >= chars_to_draw {
            return;
        }
        writer(slice_chars(choice, last_idx, idx), false);
        if idx == chars_to_draw { return }
        writer(slice_chars(choice, idx, idx + 1), true);
        last_idx = idx + 1;
    }
    writer(slice_chars(choice, last_idx, chars_to_draw), false);
}

fn read_choices() -> Vec<String> {
    let stdin = stdin();
    let mut lines = Vec::new();

    let mut stdin = stdin.lock();
    let mut first_error = None;
    let mut suppressed = 0;
    loop {
        let mut s = String::new();
        match stdin.read_line(&mut s) {
            Ok(_) => {
                if s.is_empty() { break; }
                trim(&mut s);
                lines.push(s);
            },
            Err(e) => {
                if first_error.is_some() {
                    suppressed = suppressed + 1;
                } else {
                    first_error = Some(e);
                }
            }
        }
    }
    if first_error.is_some() {
        eprintln!("Warning: Failed to parse one or more lines (\"{}\"); {} additional error(s) suppressed", first_error.unwrap(), suppressed);
    }

    lines
}

pub fn trim(s: &mut String) {
    while let Some(x) = s.pop() {
        if x != '\n' && x != '\r' {
            s.push(x);
            return;
        }
    }
}

fn delete_last_word(s: &mut String) {
    let mut deleted_something = false;
    while let Some(x) = s.pop() {
        if x == ' ' {
            if deleted_something {
                s.push(x);
                return;
            }
        } else {
            deleted_something = true;
        }
    }
}

// This is technically a standard library function, but it's plastered with
// stability warnings and therefore only available on the nightly channel.
fn slice_chars(s: &str, begin: usize, end: usize) -> &str {
    assert!(begin <= end);
    let mut count = 0;
    let mut begin_byte = None;
    let mut end_byte = None;

    // This could be even more efficient by not decoding,
    // only finding the char boundaries
    for (idx, _) in s.char_indices() {
        if count == begin { begin_byte = Some(idx); }
        if count == end { end_byte = Some(idx); break; }
        count += 1;
    }
    if begin_byte.is_none() && count == begin { begin_byte = Some(s.len()) }
    if end_byte.is_none() && count == end { end_byte = Some(s.len()) }

    match (begin_byte, end_byte) {
        (None, _) => panic!("slice_chars: `begin` is beyond end of string"),
        (_, None) => panic!("slice_chars: `end` is beyond end of string"),
        (Some(a), Some(b)) => unsafe { s.get_unchecked(a..b) }
    }
}

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

    #[test]
    fn trim_test() {
        fn should_become(before: &str, after: &str) {
            let mut x = before.to_string();
            trim(&mut x);
            assert_eq!(after.to_string(), x);
        }
        should_become("", "");
        should_become("\n", "");
        should_become("\n\n", "");
        should_become("asdf", "asdf");
        should_become("asdf\n", "asdf");
        should_become("asdf\nasdf\n", "asdf\nasdf");
        should_become("asdf\r\n", "asdf");
    }

    #[test]
    fn delete_word_test() {
        fn should_become(before: &str, after: &str) {
            let mut x = before.to_string();
            delete_last_word(&mut x);
            assert_eq!(after.to_string(), x);
        }
        should_become("", "");
        should_become("a", "");
        should_become("asdf", "");
        should_become("asdf asdf asdf", "asdf asdf ");
        should_become("asdf asdf asdf ", "asdf asdf ");
        should_become("asdf asdf asdf  ", "asdf asdf ");
    }
}
<commit_msg>Enable `RUST_BACKTRACE` on startup<commit_after>extern crate unicode_width;

mod args;
mod matching;
mod screen;
#[cfg(not(windows))] mod ansi;

use std::cmp::min;
use std::collections::HashSet;
use std::env;
use std::io::{stdin, BufRead};
use std::process;
use screen::Screen;
use screen::Key;
use screen::Key::*;
use self::SearchState::*;
use unicode_width::UnicodeWidthStr;

#[cfg(windows)] pub const NEWLINE: &'static str = "\r\n";
#[cfg(not(windows))] pub const NEWLINE: &'static str = "\n";

const VERSION: &'static str = include_str!(concat!(env!("OUT_DIR"), "/version.txt"));
const TIMESTAMP: &'static str = include_str!(concat!(env!("OUT_DIR"), "/timestamp.txt"));
const TARGET: &'static str = include_str!(concat!(env!("OUT_DIR"), "/target.txt"));
const COMMIT: &'static str = include_str!(concat!(env!("OUT_DIR"), "/commit.txt"));

fn main() {
    env::set_var("RUST_BACKTRACE", "1");
    let args = match args::parse_args() {
        Some(args) => args,
        None => {
            process::exit(1);
        },
    };

    if args.help { return; }

    if args.version {
        if COMMIT == "" {
            println!("heatseeker {} (built {} for {})", VERSION, TIMESTAMP, TARGET);
        } else {
            println!("heatseeker {}-{} (built {} for {})", VERSION, COMMIT, TIMESTAMP, TARGET);
        }
        return;
    }

    let choices = read_choices();
    let initial_search = args.initial_search.clone();
    let choices = choices.iter().map(|x| &x[..]).collect::<Vec<&str>>();
    if args.use_first {
        let matches = matching::compute_matches(&choices, &initial_search, args.filter_only);
        println!("{}", matches.get(0).unwrap_or(&""));
        return;
    } else {
        let desired_rows = if args.full_screen { 999 } else { 20 };
        let selections = event_loop(desired_rows, &choices, &initial_search, args.filter_only);
        print!("{}", selections);
    }
}

fn event_loop(desired_rows: u16, choices: &[&str], initial_search: &str, filter_only: bool) -> String {
    let mut search = Search::new(choices, initial_search.to_string(), filter_only);
    let mut screen = screen::new(desired_rows);

    loop {
        search.recompute_matches();

        match search.state {
            InProgress => draw_screen(screen.as_mut(), &search),
            _ => break,
        }

        let keys = screen.get_buffered_keys();
        for key in &keys {
            handle_key(&mut search, key, screen.visible_choices());
        }
    }

    screen.blank_screen();
    search.get_selections()
}

fn handle_key(search: &mut Search, key: &Key, visible_choices: u16) {
    match *key {
        Char(x) => search.append(x),
        Backspace |
        Control('h') => search.backspace(),
        Control('w') => search.delete_word(),
        Control('u') => search.clear_query(),
        Control('c') |
        Control('g') => search.cancel(),
        Control('t') => { search.toggle_selection(); search.down(visible_choices); },
        Control('p') |
        Up |
        ShiftTab => search.up(visible_choices),
        Control('n') |
        Down |
        Tab => search.down(visible_choices),
        Home => search.home(),
        End => search.end(visible_choices),
        Enter => search.done(),
        PgUp => search.pgup(visible_choices),
        PgDown => search.pgdown(visible_choices),
        _ => {},
    }
}

struct Search<'a> {
    choices: &'a [&'a str],
    query: String,
    matches: Vec<&'a str>,
    stale: bool,
    scroll_offset: usize,
    cursor_index: usize,
    state: SearchState,
    selections: HashSet<String>,
    filter_only: bool,
}

#[derive(PartialEq, Eq)]
enum SearchState {
    InProgress,
    Done,
    Canceled,
}

impl<'a> Search<'a> {
    fn new(choices: &'a [&'a str], initial_search: String, filter_only: bool) -> Search<'a> {
        let matches = choices.to_vec();
        Search {
            choices: choices,
            query: initial_search,
            matches: matches,
            stale: true,
            scroll_offset: 0,
            cursor_index: 0,
            state: InProgress,
            selections: HashSet::new(),
            filter_only: filter_only,
        }
    }

    fn up(&mut self, visible_choices: u16) {
        let match_count = self.matches.len();
        let limit = min(visible_choices as usize - 1, match_count - 1);
        let should_wrap = self.scroll_offset == 0;
        if self.cursor_index == 0 {
            if should_wrap {
                if match_count > visible_choices as usize {
                    self.scroll_offset = match_count - visible_choices as usize;
                    self.cursor_index = visible_choices as usize - 1;
                } else {
                    self.cursor_index = limit;
                }
            } else {
                self.scroll_offset -= 1;
            }
        } else {
            self.cursor_index -= 1;
        }
    }

    fn down(&mut self, visible_choices: u16) {
        let match_count = self.matches.len();
        let limit = min(visible_choices as usize - 1, match_count - 1);
        let should_wrap = self.cursor_index + self.scroll_offset == match_count - 1;
        if self.cursor_index == limit {
            if should_wrap {
                self.cursor_index = 0;
                self.scroll_offset = 0;
            } else {
                self.scroll_offset += 1;
            }
        } else {
            self.cursor_index += 1;
        }
    }

    fn home(&mut self) {
        self.cursor_index = 0;
        self.scroll_offset = 0;
    }

    fn end(&mut self, visible_choices: u16) {
        self.home();
        self.up(visible_choices);
    }

    fn pgup(&mut self, visible_choices: u16) {
        for _ in 0..visible_choices {
            if self.scroll_offset == 0 && self.cursor_index == 0 {
                return;
            }
            self.up(visible_choices);
        }
    }

    fn pgdown(&mut self, visible_choices: u16) {
        for _ in 0..visible_choices {
            if self.scroll_offset + self.cursor_index == self.matches.len() - 1 {
                return;
            }
            self.down(visible_choices);
        }
    }

    fn backspace(&mut self) {
        self.query.pop();
        self.stale = true;
        self.cursor_index = 0;
        self.scroll_offset = 0;
        self.matches = self.choices.to_vec();
    }

    fn delete_word(&mut self) {
        self.stale = true;
        delete_last_word(&mut self.query);
        self.matches = self.choices.to_vec();
    }

    fn append(&mut self, c: char) {
        self.query.push(c);
        self.stale = true;
        self.cursor_index = 0;
        self.scroll_offset = 0;
    }

    fn clear_query(&mut self) {
        self.query.clear();
        self.cursor_index = 0;
        self.scroll_offset = 0;
        self.matches = self.choices.to_vec();
    }

    fn recompute_matches(&mut self) {
        if self.stale {
            self.matches = matching::compute_matches(&self.matches, &self.query, self.filter_only);
            self.stale = false;
        }
    }

    fn toggle_selection(&mut self) {
        self.recompute_matches();
        let selection = self.matches.get(self.scroll_offset + self.cursor_index).unwrap_or(&"").to_string();
        if self.selections.contains(&selection) {
            self.selections.remove(&selection);
        } else {
            self.selections.insert(selection);
        }
    }

    fn get_selections(&mut self) -> String {
        let mut ret = String::new();
        if self.state != Canceled {
            for selection in &self.selections {
                ret.push_str(selection);
                ret.push_str(NEWLINE);
            }
            if ret.is_empty() {
                self.recompute_matches();
                let selection = self.matches.get(self.scroll_offset + self.cursor_index).unwrap_or(&"").to_string();
                ret.push_str(&selection);
            }
        }
        ret
    }

    fn cancel(&mut self) {
        self.state = Canceled
    }

    fn done(&mut self) {
        self.state = Done
    }
}

fn draw_screen(screen: &mut dyn Screen, search: &Search) {
    screen.hide_cursor();
    screen.blank_screen();
    screen.write(&format!("> {} ({}/{} choices){}", search.query, search.matches.len(), search.choices.len(), NEWLINE));

    print_matches(screen, &search.matches, &search.query, search.scroll_offset, search.cursor_index, &search.selections);

    let query_str: &str = &search.query;
    screen.move_cursor_to_prompt_line(2 + UnicodeWidthStr::width(query_str) as u16);
    screen.show_cursor();
}

fn print_matches(screen: &mut dyn Screen, matches: &[&str], query: &str, scroll_offset: usize, cursor_index: usize, selections: &HashSet<String>) {
    let mut i = 1;
    for choice in matches[scroll_offset..].iter() {
        let indices = matching::visual_score(choice, query);
        let max_width = screen.width();
        let mut annotated_choice = choice.to_string();
        if selections.contains(&annotated_choice) {
            annotated_choice.push(' ');
            if cfg!(windows) {
                annotated_choice.push('√');
            } else {
                annotated_choice.push('✓');
            }
        }
        print_match(&annotated_choice, &indices, max_width, &mut |s, highlight| {
            if i == cursor_index + 1 {
                if highlight { screen.write_red_inverted(s); } else { screen.write_inverted(s); }
            } else if highlight {
                screen.write_red(s);
            } else { screen.write(s); }
        });
        if i >= screen.visible_choices() as usize {
            return;
        }
        screen.write(NEWLINE);
        i += 1;
    }
}

fn print_match(choice: &str, indices: &[usize], max_width: u16, writer: &mut dyn FnMut(&str, bool)) {
    #[cfg(windows)] const MARGIN: u16 = 1;
    #[cfg(not(windows))] const MARGIN: u16 = 0;
    let max_width = max_width - MARGIN;
    let chars_in_choice = choice.chars().count();
    let mut chars_to_draw = min(chars_in_choice, max_width as usize);
    while UnicodeWidthStr::width(slice_chars(choice, 0, chars_to_draw)) > max_width as usize {
        chars_to_draw -= 1;
    }
    let mut last_idx = 0;
    for &idx in indices {
        let idx = min(idx, chars_to_draw);
        if last_idx >= chars_to_draw {
            return;
        }
        writer(slice_chars(choice, last_idx, idx), false);
        if idx == chars_to_draw { return }
        writer(slice_chars(choice, idx, idx + 1), true);
        last_idx = idx + 1;
    }
    writer(slice_chars(choice, last_idx, chars_to_draw), false);
}

fn read_choices() -> Vec<String> {
    let stdin = stdin();
    let mut lines = Vec::new();

    let mut stdin = stdin.lock();
    let mut first_error = None;
    let mut suppressed = 0;
    loop {
        let mut s = String::new();
        match stdin.read_line(&mut s) {
            Ok(_) => {
                if s.is_empty() { break; }
                trim(&mut s);
                lines.push(s);
            },
            Err(e) => {
                if first_error.is_some() {
                    suppressed = suppressed + 1;
                } else {
                    first_error = Some(e);
                }
            }
        }
    }
    if first_error.is_some() {
        eprintln!("Warning: Failed to parse one or more lines (\"{}\"); {} additional error(s) suppressed", first_error.unwrap(), suppressed);
    }

    lines
}

pub fn trim(s: &mut String) {
    while let Some(x) = s.pop() {
        if x != '\n' && x != '\r' {
            s.push(x);
            return;
        }
    }
}

fn delete_last_word(s: &mut String) {
    let mut deleted_something = false;
    while let Some(x) = s.pop() {
        if x == ' ' {
            if deleted_something {
                s.push(x);
                return;
            }
        } else {
            deleted_something = true;
        }
    }
}

// This is technically a standard library function, but it's plastered with
// stability warnings and therefore only available on the nightly channel.
fn slice_chars(s: &str, begin: usize, end: usize) -> &str {
    assert!(begin <= end);
    let mut count = 0;
    let mut begin_byte = None;
    let mut end_byte = None;

    // This could be even more efficient by not decoding,
    // only finding the char boundaries
    for (idx, _) in s.char_indices() {
        if count == begin { begin_byte = Some(idx); }
        if count == end { end_byte = Some(idx); break; }
        count += 1;
    }
    if begin_byte.is_none() && count == begin { begin_byte = Some(s.len()) }
    if end_byte.is_none() && count == end { end_byte = Some(s.len()) }

    match (begin_byte, end_byte) {
        (None, _) => panic!("slice_chars: `begin` is beyond end of string"),
        (_, None) => panic!("slice_chars: `end` is beyond end of string"),
        (Some(a), Some(b)) => unsafe { s.get_unchecked(a..b) }
    }
}

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

    #[test]
    fn trim_test() {
        fn should_become(before: &str, after: &str) {
            let mut x = before.to_string();
            trim(&mut x);
            assert_eq!(after.to_string(), x);
        }
        should_become("", "");
        should_become("\n", "");
        should_become("\n\n", "");
        should_become("asdf", "asdf");
        should_become("asdf\n", "asdf");
        should_become("asdf\nasdf\n", "asdf\nasdf");
        should_become("asdf\r\n", "asdf");
    }

    #[test]
    fn delete_word_test() {
        fn should_become(before: &str, after: &str) {
            let mut x = before.to_string();
            delete_last_word(&mut x);
            assert_eq!(after.to_string(), x);
        }
        should_become("", "");
        should_become("a", "");
        should_become("asdf", "");
        should_become("asdf asdf asdf", "asdf asdf ");
        should_become("asdf asdf asdf ", "asdf asdf ");
        should_become("asdf asdf asdf  ", "asdf asdf ");
    }
}
<|endoftext|>


<commit_before>extern crate iron;
extern crate hoedown;
extern crate hyper;
extern crate webbrowser;
extern crate router;
extern crate staticfile;

use iron::prelude::*;
use iron::status;

use hoedown::{Markdown, Render};
use hoedown::renderer::html::{Flags, Html};

use hyper::header::{ContentType};
use hyper::mime::{Mime, TopLevel, SubLevel};

use std::io::prelude::*;
use std::path::Path;
use std::fs::File;
use std::fs;

use router::Router;
use staticfile::Static;

fn main() {
    fn make_router() -> Router {
        let mut router = Router::new();

        router.get("/", index, "index");
        router.get("*", file_handler(), "files");

        router
    }

    fn index(_: &mut Request) -> IronResult<Response> {
        let mut html = String::with_capacity(4096);

        render_response(String::from("Cardboard - Index"), &mut html);

        html_response(html)
    }

    fn file_handler() -> Static {
        Static::new(Path::new("public/"))
    }

    fn render_response(title: String, output: &mut String) {
        output.push_str("<title>");
        output.push_str(& title);
        output.push_str("</title>");

        output.push_str("<link rel=\"stylesheet\" type=\"text/css\" href=\"main.css\">");
        output.push_str("");

        render_cards(output);

        output.push_str("");
    }

    fn render_cards(output: &mut String) {
        for card in fs::read_dir(".cardboard/cards").ok().unwrap().flat_map(|dir| dir.map(|e| e.path())) {
            let mut file = File::open(& card).ok().unwrap();
            let mut source = String::new();

            output.push_str("
");

            if file.read_to_string(&mut source).is_ok() {
                let doc = Markdown::new(&source);
                let mut html = Html::new(Flags::empty(), 0);
                let result = html.render(&doc);

                output.push_str(result.to_str().unwrap());
            } else {
                output.push_str("Could not read: ");
                output.push_str(card.to_str().unwrap());
            }

            output.push_str("
");
        }
    }

    fn html_response(content: String) -> IronResult<Response> {
        let mut response = Response::with((status::Ok, content));

        response.headers.set(ContentType(Mime(TopLevel::Text, SubLevel::Html, vec![])));

        Ok(response)
    }

    println!("Listening on port 9000 ...");

    if !webbrowser::open("http://localhost:9000/").is_ok() {
        println!("Open http://localhost:9000 to view your project")
    }

    Iron::new(make_router()).http("localhost:9000").unwrap();
}
<commit_msg>make sure data directory exists<commit_after>extern crate iron;
extern crate hoedown;
extern crate hyper;
extern crate webbrowser;
extern crate router;
extern crate staticfile;

use iron::prelude::*;
use iron::status;

use hoedown::{Markdown, Render};
use hoedown::renderer::html::{Flags, Html};

use hyper::header::{ContentType};
use hyper::mime::{Mime, TopLevel, SubLevel};

use std::io::prelude::*;
use std::path::Path;
use std::fs::File;
use std::fs;

use router::Router;
use staticfile::Static;

fn main() {
    fn make_router() -> Router {
        let mut router = Router::new();

        router.get("/", index, "index");
        router.get("*", file_handler(), "files");

        router
    }

    fn index(_: &mut Request) -> IronResult<Response> {
        let mut html = String::with_capacity(4096);

        render_response(String::from("Cardboard - Index"), &mut html);

        html_response(html)
    }

    fn file_handler() -> Static {
        Static::new(Path::new("public/"))
    }

    fn render_response(title: String, output: &mut String) {
        output.push_str("<title>");
        output.push_str(& title);
        output.push_str("</title>");

        output.push_str("<link rel=\"stylesheet\" type=\"text/css\" href=\"main.css\">");
        output.push_str("");

        render_cards(output);

        output.push_str("");
    }

    fn render_cards(output: &mut String) {
        for card in fs::read_dir(".cardboard/cards").ok().unwrap().flat_map(|dir| dir.map(|e| e.path())) {
            let mut file = File::open(& card).ok().unwrap();
            let mut source = String::new();

            output.push_str("
");

            if file.read_to_string(&mut source).is_ok() {
                let doc = Markdown::new(&source);
                let mut html = Html::new(Flags::empty(), 0);
                let result = html.render(&doc);

                output.push_str(result.to_str().unwrap());
            } else {
                output.push_str("Could not read: ");
                output.push_str(card.to_str().unwrap());
            }

            output.push_str("
");
        }
    }

    fn html_response(content: String) -> IronResult<Response> {
        let mut response = Response::with((status::Ok, content));

        response.headers.set(ContentType(Mime(TopLevel::Text, SubLevel::Html, vec![])));

        Ok(response)
    }

    fn check_data_directory() -> bool {
        let root = ".cardboard";
        let cards = root.to_string() + "/cards";

        (fs::create_dir(root).is_ok() && fs::create_dir(& cards).is_ok()) || Path::new(& cards).exists()
    }

    if check_data_directory() {
        println!("Listening on port 9000 ...");

        if !webbrowser::open("http://localhost:9000/").is_ok() {
            println!("Open http://localhost:9000 to view your project")
        }

        Iron::new(make_router()).http("localhost:9000").unwrap();
    } else {
        println!("Failed to create data directory");
    }
}
<|endoftext|>