<commit_before>#![feature(io, slice_patterns)]
use std::fs;
use std::io;
use std::io::{Read, Write};
use std::path::Path;

const MEMSIZE: usize = 4096;

enum ParserResult {
    Something(Op),
    Nothing,
    EOF
}

use ParserResult::*;

#[derive(Debug)]
enum Op {
    Add(u8),
    Mov(usize),
    In,
    Out,
    Loop(OpStream)
}

use Op::*;

#[derive(Debug)]
struct OpStream {
    ops: Vec<Op>
}

impl OpStream {
    fn add(&mut self, op: Op) {
        self.ops.push(op);
    }

    fn optimize(&mut self) {
        let mut i = 0;
        while i < self.ops.len() {
            match &self.ops[i..] {
                [Add(a), Add(b), ..] => {
                    self.ops[i] = Add(a.wrapping_add(b));
                    self.ops.remove(i + 1);
                }
                [Mov(a), Mov(b), ..] => {
                    self.ops[i] = Mov(a.wrapping_add(b));
                    self.ops.remove(i + 1);
                }
                [Add(0), ..] | [Mov(0), ..] => {
                    self.ops.remove(i);
                }
                _ => i += 1
            }
        }
        for op in &mut self.ops[..] {
            if let &mut Loop(ref mut stream) = op {
                stream.optimize();
            }
        }
    }

    fn get(&self) -> &[Op] {
        &self.ops[..]
    }

    fn new() -> OpStream {
        OpStream { ops: Vec::new() }
    }
}

struct State {
    index: usize,
    memory: [u8; MEMSIZE]
}

impl State {
    fn new() -> State {
        State { index: 0, memory: [0; MEMSIZE] }
    }

    fn peek(&self) -> u8 {
        self.memory[self.index % MEMSIZE]
    }

    fn poke(&mut self, value: u8) {
        self.memory[self.index % MEMSIZE] = value;
    }

    fn step(&mut self, op: &Op) -> bool {
        match op {
            &Add(i) => {
                // XXX: Ugly expression due to lexical scoping of borrow,
                //      cf. https://github.com/rust-lang/rust/issues/6393
                let x = self.peek();
                self.poke(x.wrapping_add(i));
            }
            &Mov(n) => {
                self.index = self.index.wrapping_add(n);
            }
            &In => {
                let mut c = vec![0u8];
                if io::stdin().read(&mut c).unwrap() == 0 {
                    return false;
                }
                self.poke(c[0]);
            }
            &Out => {
                io::stdout().write(&vec![self.peek()]).unwrap();
            }
            &Loop(ref ops) => {
                while self.peek() != 0 {
                    if !self.run(ops.get()) {
                        return false;
                    }
                }
            }
        }
        return true;
    }

    fn run(&mut self, ops: &[Op]) -> bool {
        for op in ops {
            if !self.step(op) {
                return false;
            }
        }
        return true;
    }
}

fn main() {
    let filenames: Vec<String> = std::iter::FromIterator::from_iter(std::env::args());

    for filename in &filenames[1..] {
        match reader(&Path::new(filename)) {
            Ok(reader) => {
                let mut chars = reader.chars().peekable();

                let mut opstream = OpStream::new();
                loop {
                    match parse(&mut chars) {
                        Something(op) => opstream.add(op),
                        EOF => break,
                        _ => {}
                    }
                }

                opstream.optimize();

                State::new().run(opstream.get());
            }
            Err(e) => {
                writeln!(&mut io::stderr(), "Error while processing {}: {}", filename, e).unwrap();
            }
        }
    }
}

fn parse<T: io::Read>(mut chars: &mut std::iter::Peekable<io::Chars<T>>) -> ParserResult {
    match chars.next() {
        Some(Ok(c)) => match c {
            '+' => Something(Add(1)),
            '-' => Something(Add(-1)),
            '<' => Something(Mov(-1)),
            '>' => Something(Mov(1)),
            ',' => Something(In),
            '.' => Something(Out),
            '[' => {
                let mut childstream = OpStream::new();
                while match chars.peek() {
                    Some(&Ok(c)) if c != ']' => true,
                    _ => false
                } {
                    match parse(&mut chars) {
                        Something(op) => childstream.add(op),
                        Nothing => {},
                        EOF => panic!()
                    }
                }
                chars.next();
                Something(Loop(childstream))
            },
            _ => Nothing  // other characters
        },
        _ => EOF  // XXX: really?
    }
}

fn reader(path: &Path) -> Result<io::BufReader<fs::File>, io::Error> {
    Ok(io::BufReader::new(try!(fs::File::open(path))))
}

#[cfg(test)]
mod tests {
    use super::{State, OpStream};
    use super::Op::*;

    macro_rules! assert_let(
        ($test:pat, $value:expr, $then:block) => (
            if let $test = $value {
                $then
            } else {
                panic!("{:?} doesn't match {}", $value, stringify!($test))
            }
        );
        ($test:pat, $value:expr) => (
            assert_let!($test, $value, {})
        );
    );

    #[test]
    fn state_peek() {
        let mut state = State::new();
        state.memory[state.index] = 23;
        assert_eq!(23, state.peek());
        state.index = 5;
        state.memory[state.index] = 42;
        assert_eq!(42, state.peek());
    }

    #[test]
    fn state_poke() {
        let mut state = State::new();
        state.poke(23);
        assert_eq!(23, state.memory[state.index]);
        state.index = 5;
        state.poke(42);
        assert_eq!(42, state.memory[state.index]);
    }

    #[test]
    fn state_step_add() {
        let mut state = State::new();
        state.step(&Add(23));
        assert_eq!(23, state.peek());
        state.step(&Add(42));
        assert_eq!(65, state.peek());
        state.step(&Add(-66));
        assert_eq!(-1, state.peek());
    }

    #[test]
    fn state_step_mov() {
        let mut state = State::new();
        state.step(&Mov(1));
        assert_eq!(1, state.index);
        state.step(&Mov(42));
        assert_eq!(43, state.index);
        state.step(&Mov(-1));
        assert_eq!(42, state.index);
    }

    #[test]
    fn state_step_loop() {
        let mut state = State::new();
        state.poke(23);
        state.step(&Loop(OpStream { ops: vec![Add(1)] }));
        assert_eq!(0, state.peek());
    }

    #[test]
    fn opstream_new_empty() {
        let opstream = OpStream::new();
        assert_eq!(0, opstream.ops.len());
    }

    #[test]
    fn opstream_optimize() {
        let mut opstream = OpStream::new();
        opstream.add(Mov(1));
        opstream.add(Mov(1));
        opstream.add(Add(1));
        opstream.add(Add(-1));
        opstream.add(Add(-1));
        opstream.add(Mov(1));
        opstream.add(Mov(-1));
        let mut opstream2 = OpStream::new();
        opstream2.add(Mov(2));
        opstream2.add(Mov(3));
        opstream.add(Loop(opstream2));
        opstream.optimize();

        assert_let!([Mov(2), Add(-1), Loop(ref s2)], &opstream.ops[..], {
            assert_let!([Mov(5)], &s2.ops[..]);
        });
    }
}
<commit_msg>refactor: move recursive optimization<commit_after>#![feature(io, slice_patterns)]
use std::fs;
use std::io;
use std::io::{Read, Write};
use std::path::Path;

const MEMSIZE: usize = 4096;

enum ParserResult {
    Something(Op),
    Nothing,
    EOF
}

use ParserResult::*;

#[derive(Debug)]
enum Op {
    Add(u8),
    Mov(usize),
    In,
    Out,
    Loop(OpStream)
}

use Op::*;

#[derive(Debug)]
struct OpStream {
    ops: Vec<Op>
}

impl OpStream {
    fn add(&mut self, op: Op) {
        self.ops.push(op);
    }

    fn optimize(&mut self) {
        let mut i = 0;
        while i < self.ops.len() {
            match &self.ops[i..] {
                [Add(a), Add(b), ..] => {
                    self.ops[i] = Add(a.wrapping_add(b));
                    self.ops.remove(i + 1);
                }
                [Mov(a), Mov(b), ..] => {
                    self.ops[i] = Mov(a.wrapping_add(b));
                    self.ops.remove(i + 1);
                }
                [Add(0), ..] | [Mov(0), ..] => {
                    self.ops.remove(i);
                }
                [Loop(_), ..] => {
                    if let &mut Loop(ref mut stream) = &mut self.ops[i] {
                        stream.optimize();
                    }
                    i += 1
                }
                _ => i += 1
            }
        }
    }

    fn get(&self) -> &[Op] {
        &self.ops[..]
    }

    fn new() -> OpStream {
        OpStream { ops: Vec::new() }
    }
}

struct State {
    index: usize,
    memory: [u8; MEMSIZE]
}

impl State {
    fn new() -> State {
        State { index: 0, memory: [0; MEMSIZE] }
    }

    fn peek(&self) -> u8 {
        self.memory[self.index % MEMSIZE]
    }

    fn poke(&mut self, value: u8) {
        self.memory[self.index % MEMSIZE] = value;
    }

    fn step(&mut self, op: &Op) -> bool {
        match op {
            &Add(i) => {
                // XXX: Ugly expression due to lexical scoping of borrow,
                //      cf. https://github.com/rust-lang/rust/issues/6393
                let x = self.peek();
                self.poke(x.wrapping_add(i));
            }
            &Mov(n) => {
                self.index = self.index.wrapping_add(n);
            }
            &In => {
                let mut c = vec![0u8];
                if io::stdin().read(&mut c).unwrap() == 0 {
                    return false;
                }
                self.poke(c[0]);
            }
            &Out => {
                io::stdout().write(&vec![self.peek()]).unwrap();
            }
            &Loop(ref ops) => {
                while self.peek() != 0 {
                    if !self.run(ops.get()) {
                        return false;
                    }
                }
            }
        }
        return true;
    }

    fn run(&mut self, ops: &[Op]) -> bool {
        for op in ops {
            if !self.step(op) {
                return false;
            }
        }
        return true;
    }
}

fn main() {
    let filenames: Vec<String> = std::iter::FromIterator::from_iter(std::env::args());

    for filename in &filenames[1..] {
        match reader(&Path::new(filename)) {
            Ok(reader) => {
                let mut chars = reader.chars().peekable();

                let mut opstream = OpStream::new();
                loop {
                    match parse(&mut chars) {
                        Something(op) => opstream.add(op),
                        EOF => break,
                        _ => {}
                    }
                }

                opstream.optimize();

                State::new().run(opstream.get());
            }
            Err(e) => {
                writeln!(&mut io::stderr(), "Error while processing {}: {}", filename, e).unwrap();
            }
        }
    }
}

fn parse<T: io::Read>(mut chars: &mut std::iter::Peekable<io::Chars<T>>) -> ParserResult {
    match chars.next() {
        Some(Ok(c)) => match c {
            '+' => Something(Add(1)),
            '-' => Something(Add(-1)),
            '<' => Something(Mov(-1)),
            '>' => Something(Mov(1)),
            ',' => Something(In),
            '.' => Something(Out),
            '[' => {
                let mut childstream = OpStream::new();
                while match chars.peek() {
                    Some(&Ok(c)) if c != ']' => true,
                    _ => false
                } {
                    match parse(&mut chars) {
                        Something(op) => childstream.add(op),
                        Nothing => {},
                        EOF => panic!()
                    }
                }
                chars.next();
                Something(Loop(childstream))
            },
            _ => Nothing  // other characters
        },
        _ => EOF  // XXX: really?
    }
}

fn reader(path: &Path) -> Result<io::BufReader<fs::File>, io::Error> {
    Ok(io::BufReader::new(try!(fs::File::open(path))))
}

#[cfg(test)]
mod tests {
    use super::{State, OpStream};
    use super::Op::*;

    macro_rules! assert_let(
        ($test:pat, $value:expr, $then:block) => (
            if let $test = $value {
                $then
            } else {
                panic!("{:?} doesn't match {}", $value, stringify!($test))
            }
        );
        ($test:pat, $value:expr) => (
            assert_let!($test, $value, {})
        );
    );

    #[test]
    fn state_peek() {
        let mut state = State::new();
        state.memory[state.index] = 23;
        assert_eq!(23, state.peek());
        state.index = 5;
        state.memory[state.index] = 42;
        assert_eq!(42, state.peek());
    }

    #[test]
    fn state_poke() {
        let mut state = State::new();
        state.poke(23);
        assert_eq!(23, state.memory[state.index]);
        state.index = 5;
        state.poke(42);
        assert_eq!(42, state.memory[state.index]);
    }

    #[test]
    fn state_step_add() {
        let mut state = State::new();
        state.step(&Add(23));
        assert_eq!(23, state.peek());
        state.step(&Add(42));
        assert_eq!(65, state.peek());
        state.step(&Add(-66));
        assert_eq!(-1, state.peek());
    }

    #[test]
    fn state_step_mov() {
        let mut state = State::new();
        state.step(&Mov(1));
        assert_eq!(1, state.index);
        state.step(&Mov(42));
        assert_eq!(43, state.index);
        state.step(&Mov(-1));
        assert_eq!(42, state.index);
    }

    #[test]
    fn state_step_loop() {
        let mut state = State::new();
        state.poke(23);
        state.step(&Loop(OpStream { ops: vec![Add(1)] }));
        assert_eq!(0, state.peek());
    }

    #[test]
    fn opstream_new_empty() {
        let opstream = OpStream::new();
        assert_eq!(0, opstream.ops.len());
    }

    #[test]
    fn opstream_optimize() {
        let mut opstream = OpStream::new();
        opstream.add(Mov(1));
        opstream.add(Mov(1));
        opstream.add(Add(1));
        opstream.add(Add(-1));
        opstream.add(Add(-1));
        opstream.add(Mov(1));
        opstream.add(Mov(-1));
        let mut opstream2 = OpStream::new();
        opstream2.add(Mov(2));
        opstream2.add(Mov(3));
        opstream.add(Loop(opstream2));
        opstream.optimize();

        assert_let!([Mov(2), Add(-1), Loop(ref s2)], &opstream.ops[..], {
            assert_let!([Mov(5)], &s2.ops[..]);
        });
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2013-2014 Sandstorm Development Group, Inc. and contributors
// Licensed under the MIT License:
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

//! # Cap'n Proto Schema Compiler Plugin Executable
//!
//! [See this.](http://kentonv.github.io/capnproto/otherlang.html#how-to-write-compiler-plugins)
//!
//!

#![feature(globs)]

#![crate_name="capnpc-rust"]
#![crate_type = "bin"]

extern crate capnp;
extern crate capnpc;

pub fn main() {
    //! Generate Rust code according to a `schema_capnp::code_generator_request` read from stdin.

    match ::capnpc::codegen::main(&mut ::std::io::stdin()) {
        Ok(()) => {}
        Err(e) => {
            std::os::set_exit_status(1);
            println!("error: {}", e)
        }
    }
}
<commit_msg>oops, missed this one<commit_after>// Copyright (c) 2013-2014 Sandstorm Development Group, Inc. and contributors
// Licensed under the MIT License:
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

//! # Cap'n Proto Schema Compiler Plugin Executable
//!
//! [See this.](http://kentonv.github.io/capnproto/otherlang.html#how-to-write-compiler-plugins)
//!
//!

#![crate_name="capnpc-rust"]
#![crate_type = "bin"]

extern crate capnp;
extern crate capnpc;

pub fn main() {
    //! Generate Rust code according to a `schema_capnp::code_generator_request` read from stdin.

    match ::capnpc::codegen::main(&mut ::std::io::stdin()) {
        Ok(()) => {}
        Err(e) => {
            std::os::set_exit_status(1);
            println!("error: {}", e)
        }
    }
}
<|endoftext|>


<commit_before>//! A simple evolutionary algorithm written in Rust.

extern crate rand;

use rand::Rng;
use std::collections::HashMap;

fn main() {
    let target = "METHINKS IT IS LIKE A WEASEL".to_string();
    println!("{}", target);

    let mut parent = String::new();
    generate_first_sentence(&mut parent);

    println!("{}", parent);

    let nb_copy = 400;
    let mutation_rate = 0.05;
    let mut counter = 0;

    while fitness(&target, &parent) != 0 {
        counter += 1;

        let mut sentences = HashMap::new();
        let mut f_min = 30;

        for _ in 0..nb_copy {
            let sentence = mutate(&mut parent, mutation_rate);
            let f = fitness(&target, &sentence);

            sentences.insert(f, sentence);

            if f < f_min {
                f_min = f
            }
        }

        if fitness(&target, &parent) > f_min {
            match sentences.get(&f_min) {
                Some(s) => {
                    parent = s.clone();
                    println!("{} : {}", parent, counter);
                }
                None => println!("Error, fitness minimum but no sentence."),
            }
        }
    }
}

fn fitness(target: &String, sentence: &String) -> u32 {
    let mut fitness = 0;

    for (c1, c2) in target.chars().zip(sentence.chars()) {
        if c1 != c2 {
            fitness += 1
        };
    }

    fitness
}

fn mutate(sentence: &mut String, mutation_rate: f64) -> String {
    let mut rng = rand::thread_rng();
    let mut mutation = String::new();

    for c in sentence.chars() {
        match mutation_rate > rng.gen_range(0.0, 1.0) {
            false => mutation.push(c),               // no mutation
            true => mutation.push(random_char()),   // mutation
        }
    }
    mutation
}

fn generate_first_sentence(parent: &mut String) {
    for _ in 0..28 {
        parent.push(random_char());
    }
}

fn random_char() -> char {
    match rand::thread_rng().gen_range('A' as u8, '\\' as u8) as char {
        '[' => ' ',
        c => c,
    }
}
<commit_msg>Taking &str rather than &String is idiomatic<commit_after>//! A simple evolutionary algorithm written in Rust.

extern crate rand;

use rand::Rng;
use std::collections::HashMap;

fn main() {
    let target = "METHINKS IT IS LIKE A WEASEL".to_string();
    println!("{}", target);

    let mut parent = String::new();
    generate_first_sentence(&mut parent);

    println!("{}", parent);

    let nb_copy = 400;
    let mutation_rate = 0.05;
    let mut counter = 0;

    while fitness(&target, &parent) != 0 {
        counter += 1;

        let mut sentences = HashMap::new();
        let mut f_min = 30;

        for _ in 0..nb_copy {
            let sentence = mutate(&mut parent, mutation_rate);
            let f = fitness(&target, &sentence);

            sentences.insert(f, sentence);

            if f < f_min {
                f_min = f
            }
        }

        if fitness(&target, &parent) > f_min {
            match sentences.get(&f_min) {
                Some(s) => {
                    parent = s.clone();
                    println!("{} : {}", parent, counter);
                }
                None => println!("Error, fitness minimum but no sentence."),
            }
        }
    }
}

fn fitness(target: &str, sentence: &str) -> u32 {
    let mut fitness = 0;

    for (c1, c2) in target.chars().zip(sentence.chars()) {
        if c1 != c2 {
            fitness += 1
        };
    }

    fitness
}

fn mutate(sentence: &str, mutation_rate: f64) -> String {
    let mut rng = rand::thread_rng();
    let mut mutation = String::new();

    for c in sentence.chars() {
        match mutation_rate > rng.gen_range(0.0, 1.0) {
            false => mutation.push(c),               // no mutation
            true => mutation.push(random_char()),   // mutation
        }
    }

    mutation
}

fn generate_first_sentence(parent: &mut String) {
    for _ in 0..28 {
        parent.push(random_char());
    }
}

fn random_char() -> char {
    match rand::thread_rng().gen_range('A' as u8, '\\' as u8) as char {
        '[' => ' ',
        c => c,
    }
}
<|endoftext|>


<commit_before>// Copyright 2017 CoreOS, 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.

//! update-ssh-keys
//!
//! this command allows users of container linux to administer ssh keys

#[macro_use]
extern crate clap;
#[macro_use]
extern crate error_chain;
extern crate users;

extern crate update_ssh_keys;

use clap::{Arg, App};
use users::get_current_username;

use std::fs::File;

use update_ssh_keys::*;
use update_ssh_keys::errors::*;

#[derive(Clone, Debug)]
struct Config {
    user: String,
    command: Command,
}
#[derive(Clone, Debug)]
enum Command {
    Add {
        name: String,
        force: bool,
        replace: bool,
        stdin: bool,
        keyfiles: Vec<String>,
    },
    Delete { name: String },
    Disable { name: String },
    List,
    Sync,
}

quick_main!(run);

fn run() -> Result<()> {
    let config = config()
        .chain_err(|| "command line configuration")?;

    let user = users::get_user_by_name(&config.user)
        .ok_or_else(|| format!("failed to find user with name '{}'", config.user))?;

    let mut aks =  AuthorizedKeys::open(user, true)
        .chain_err(|| format!("failed to open authorized keys directory for user '{}'", config.user))?;

    match config.command {
        Command::Add{name, force, replace, stdin, keyfiles} => {
            let keys = if stdin {
                // read the keys from stdin
                AuthorizedKeys::read_keys(std::io::stdin())?
            } else {
                // keys are in provided files
                let mut keys = vec![];
                for keyfile in keyfiles {
                    let file = File::open(&keyfile)
                        .chain_err(|| format!("failed to open keyfile '{:?}'", keyfile))?;
                    keys.append(&mut AuthorizedKeys::read_keys(file)?);
                }
                keys
            };
            let res = aks.add_keys(&name, keys, replace, force);
            match res {
                Err(Error(ErrorKind::KeysDisabled(name), _)) => println!("Skipping add {} for {}, disabled.", name, config.user),
                Err(Error(ErrorKind::KeysExist(_), _)) => println!("Skipping add {} for {}, already exists.", name, config.user),
                _ => res.chain_err(|| "failed to add keys")?,
            }
        },
        Command::Delete{name} => aks.remove_keys(&name),
        Command::Disable{name} => aks.disable_keys(&name),
        Command::List => {
            let keys = aks.get_all_keys();
            println!("All keys for {}:", config.user);
            for (name, keyset) in keys {
                println!("{}:", name);
                for key in &keyset.keys {
                    println!("{}", key.to_fingerprint_string())
                }
            }
        },
        Command::Sync => {},
    }

    aks.write()
        .chain_err(|| "failed to update authorized keys directory")?;
    aks.sync()
        .chain_err(|| "failed to update authorized keys")?;

    println!("Updated {:?}", aks.file);

    Ok(())
}

fn config() -> Result<Config> {
    // get the default user by figuring out the current user if the current user
    // is root (or doesn't exist) then use core
    let default_user = get_current_username().map_or("core".into(), |u| {
        if u == "root" {
            "core".into()
        } else {
            u
        }
    });

    // setup cli
    // TODO(sdemos): help should print out the default user
    // tried format! but it doesn't seem to like multi-line strings?
    let matches = App::new("update-ssh-keys")
        .version(crate_version!())
        .help("Usage: update-ssh-keys [-l] [-u user] [-a name file1... | -d name]
Options:
    -u USER     Update the given user's authorized_keys file [{}]
    -a NAME     Add the given keys, using the given name to identify them.
    -A NAME     Add the given keys, even if it was disabled with '-D'
    -n          When adding, don't replace an existing key with the given name.
    -d NAME     Delete keys identified by the given name.
    -D NAME     Disable the given set from being added with '-a'
    -l          List the names and number of keys currently installed.
    -h          This ;-)

This tool provides a consistent way for different systems to add ssh public
keys to a given user account, usually the default '${UPDATE_USER}' user.
If -a, -A, -d, nor -D are provided then the authorized_keys file is simply
regenerated using the existing keys.

With the -a option keys may be provided as files on the command line. If no
files are provided with the -a option the keys will be read from stdin.")
        .arg(Arg::with_name("user")
             .short("u")
             .help("Update the given user's authorized_keys file.")
             .takes_value(true))
        .arg(Arg::with_name("no-replace")
             .short("n")
             .help("When adding, don't replace an existing key with the given name."))
        .arg(Arg::with_name("list")
             .short("l")
             .help("List the names and number of keys currently installed."))
        .arg(Arg::with_name("add")
             .short("a")
             .help("Add the given keys, using the given name to identify them.")
             .takes_value(true))
        .arg(Arg::with_name("add-force")
             .short("A")
             .help("Add the given keys, even if it was disabled with '-D'.")
             .takes_value(true))
        .arg(Arg::with_name("delete")
             .short("d")
             .help("Delete keys identified by the given name.")
             .takes_value(true))
        .arg(Arg::with_name("disable")
             .short("D")
             .help("Disable the given set from being added with '-a'.")
             .takes_value(true))
        .arg(Arg::with_name("keys")
             .multiple(true))
        .get_matches();

    let command = matches.value_of("add").map(|name| Command::Add{
            name: name.into(),
            force: false,
            replace: !matches.is_present("no-replace"),
            stdin: !matches.is_present("keys"),
            keyfiles: matches.values_of("keys").map(|vals| vals.map(|s| s.into()).collect::<Vec<_>>()).unwrap_or_default(),
        })
        .or_else(|| matches.value_of("add-force").map(|name| Command::Add{
            name: name.into(),
            force: true,
            replace: !matches.is_present("no-replace"),
            stdin: !matches.is_present("keys"),
            keyfiles: matches.values_of("keys").map(|vals| vals.map(|s| s.into()).collect::<Vec<_>>()).unwrap_or_default(),
        }))
        .or_else(|| matches.value_of("delete").map(|name| Command::Delete{name: name.into()}))
        .or_else(|| matches.value_of("disable").map(|name| Command::Disable{name: name.into()}))
        .unwrap_or(if matches.is_present("list") { Command::List } else { Command::Sync });

    let user = matches.value_of("user")
        .map_or(default_user, String::from);

    Ok(Config {
        user: user,
        command: command,
    })
}
<commit_msg>main: remove one spurious todo<commit_after>// Copyright 2017 CoreOS, 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.

//! update-ssh-keys
//!
//! this command allows users of container linux to administer ssh keys

#[macro_use]
extern crate clap;
#[macro_use]
extern crate error_chain;
extern crate users;

extern crate update_ssh_keys;

use clap::{Arg, App};
use users::get_current_username;

use std::fs::File;

use update_ssh_keys::*;
use update_ssh_keys::errors::*;

#[derive(Clone, Debug)]
struct Config {
    user: String,
    command: Command,
}
#[derive(Clone, Debug)]
enum Command {
    Add {
        name: String,
        force: bool,
        replace: bool,
        stdin: bool,
        keyfiles: Vec<String>,
    },
    Delete { name: String },
    Disable { name: String },
    List,
    Sync,
}

quick_main!(run);

fn run() -> Result<()> {
    let config = config()
        .chain_err(|| "command line configuration")?;

    let user = users::get_user_by_name(&config.user)
        .ok_or_else(|| format!("failed to find user with name '{}'", config.user))?;

    let mut aks =  AuthorizedKeys::open(user, true)
        .chain_err(|| format!("failed to open authorized keys directory for user '{}'", config.user))?;

    match config.command {
        Command::Add{name, force, replace, stdin, keyfiles} => {
            let keys = if stdin {
                // read the keys from stdin
                AuthorizedKeys::read_keys(std::io::stdin())?
            } else {
                // keys are in provided files
                let mut keys = vec![];
                for keyfile in keyfiles {
                    let file = File::open(&keyfile)
                        .chain_err(|| format!("failed to open keyfile '{:?}'", keyfile))?;
                    keys.append(&mut AuthorizedKeys::read_keys(file)?);
                }
                keys
            };
            let res = aks.add_keys(&name, keys, replace, force);
            match res {
                Err(Error(ErrorKind::KeysDisabled(name), _)) => println!("Skipping add {} for {}, disabled.", name, config.user),
                Err(Error(ErrorKind::KeysExist(_), _)) => println!("Skipping add {} for {}, already exists.", name, config.user),
                _ => res.chain_err(|| "failed to add keys")?,
            }
        },
        Command::Delete{name} => aks.remove_keys(&name),
        Command::Disable{name} => aks.disable_keys(&name),
        Command::List => {
            let keys = aks.get_all_keys();
            println!("All keys for {}:", config.user);
            for (name, keyset) in keys {
                println!("{}:", name);
                for key in &keyset.keys {
                    println!("{}", key.to_fingerprint_string())
                }
            }
        },
        Command::Sync => {},
    }

    aks.write()
        .chain_err(|| "failed to update authorized keys directory")?;
    aks.sync()
        .chain_err(|| "failed to update authorized keys")?;

    println!("Updated {:?}", aks.file);

    Ok(())
}

fn config() -> Result<Config> {
    // get the default user by figuring out the current user if the current user
    // is root (or doesn't exist) then use core
    let default_user = get_current_username().map_or("core".into(), |u| {
        if u == "root" {
            "core".into()
        } else {
            u
        }
    });

    // setup cli
    let matches = App::new("update-ssh-keys")
        .version(crate_version!())
        .help(format!(
r#"Usage: update-ssh-keys [-l] [-u user] [-a name file1... | -d name]
Options:
    -u USER     Update the given user's authorized_keys file [{0}]
    -a NAME     Add the given keys, using the given name to identify them.
    -A NAME     Add the given keys, even if it was disabled with '-D'
    -n          When adding, don't replace an existing key with the given name.
    -d NAME     Delete keys identified by the given name.
    -D NAME     Disable the given set from being added with '-a'
    -l          List the names and number of keys currently installed.
    -h          This ;-)

This tool provides a consistent way for different systems to add ssh public
keys to a given user account, usually the default current user.
If -a, -A, -d, nor -D are provided then the authorized_keys file is simply
regenerated using the existing keys.

With the -a option keys may be provided as files on the command line. If no
files are provided with the -a option the keys will be read from stdin."#, default_user).as_ref())
        .arg(Arg::with_name("user")
             .short("u")
             .help("Update the given user's authorized_keys file.")
             .takes_value(true))
        .arg(Arg::with_name("no-replace")
             .short("n")
             .help("When adding, don't replace an existing key with the given name."))
        .arg(Arg::with_name("list")
             .short("l")
             .help("List the names and number of keys currently installed."))
        .arg(Arg::with_name("add")
             .short("a")
             .help("Add the given keys, using the given name to identify them.")
             .takes_value(true))
        .arg(Arg::with_name("add-force")
             .short("A")
             .help("Add the given keys, even if it was disabled with '-D'.")
             .takes_value(true))
        .arg(Arg::with_name("delete")
             .short("d")
             .help("Delete keys identified by the given name.")
             .takes_value(true))
        .arg(Arg::with_name("disable")
             .short("D")
             .help("Disable the given set from being added with '-a'.")
             .takes_value(true))
        .arg(Arg::with_name("keys")
             .multiple(true))
        .get_matches();

    let command = matches.value_of("add").map(|name| Command::Add{
            name: name.into(),
            force: false,
            replace: !matches.is_present("no-replace"),
            stdin: !matches.is_present("keys"),
            keyfiles: matches.values_of("keys").map(|vals| vals.map(|s| s.into()).collect::<Vec<_>>()).unwrap_or_default(),
        })
        .or_else(|| matches.value_of("add-force").map(|name| Command::Add{
            name: name.into(),
            force: true,
            replace: !matches.is_present("no-replace"),
            stdin: !matches.is_present("keys"),
            keyfiles: matches.values_of("keys").map(|vals| vals.map(|s| s.into()).collect::<Vec<_>>()).unwrap_or_default(),
        }))
        .or_else(|| matches.value_of("delete").map(|name| Command::Delete{name: name.into()}))
        .or_else(|| matches.value_of("disable").map(|name| Command::Disable{name: name.into()}))
        .unwrap_or(if matches.is_present("list") { Command::List } else { Command::Sync });

    let user = matches.value_of("user")
        .map_or(default_user, String::from);

    Ok(Config {
        user: user,
        command: command,
    })
}
<|endoftext|>


<commit_before>//! [한국 러스트 사용자 그룹 홈페이지][rust-kr] 소스코드
//! --------
//!
//! ### 사용법
//! ```bash
//! git clone https://github.com/rust-kr/rust-kr.org.wiki.git docs
//! cargo run
//! # 웹 브라우저에서 http://localhost:8000 를 열어보세요
//! ```
//!
//! [rust-kr]: http://rust-kr.org

// 한국 러스트 사용자 그룹 홈페이지
// Copyright (C) 2016  Hyeon Kim
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

#[macro_use(itry)] extern crate iron;
#[macro_use(router)] extern crate router;
extern crate staticfile;
extern crate handlebars_iron as hbs;
extern crate pulldown_cmark as cmark;

use std::collections::BTreeMap;
use std::fmt;
use iron::prelude::*;
use iron::status;
use iron::middleware::{Handler, AfterMiddleware};
use staticfile::Static;
use hbs::{Template, HandlebarsEngine, DirectorySource};

//
// Configs
//
/// 서버가 소켓을 열 주소
const ADDR: &'static str = "localhost:8000";
/// 마크다운 문서가 담긴 주소
const DOCS_DIR: &'static str = "docs";

/// Entry point
fn main() {
    let mode = get_mode();

    //
    //      User sends request
    //               ⇩
    //                                '/public' 디렉토리에서 해당 파일이 있는지 먼저 찾아본다.
    //    try static file serving     있을경우 이를 전송하고, 없을경우 아래의 main handler로
    //   (only in development mode)   넘어간다. production mode로 실행될 경우, 이 동작이 없어지므로
    //                                nginx나 apache의 정적파일서빙 기능을 써야한다.
    //               ⇩
    //        "main handlers"         index, page, all_docs 등의 함수로 요청을 처리한다
    //               ⇩
    //       404 page handler         위의 두개가 모두 실패하였을경우 404페이지를 렌더링한다
    //               ⇩
    //     handlebar templating
    //

    let mut c = Chain::new(router! {
        get "/" => index,
        get "/pages/:name" => page,
        get "/pages/_pages" => all_docs,
    });

    // Static file serving (only in development mode)
    if mode == Mode::Development {
        c.link_around(move |main: Box<Handler>| -> Box<Handler> {
            Box::new(move |req: &mut Request| -> IronResult<Response> {
                Static::new("public/").handle(req)
                    .or_else(|_| main.handle(req))
            })
        });
    }
    // 404 page handler
    c.link_after(catch(|mut err: IronError| {
        if err.response.body.is_some() { return Err(err); }
        if err.response.status != Some(status::NotFound) { return Err(err); }

        // TODO: 좀더 내용 채우기
        let mut data = BTreeMap::new();
        data.insert("content", r#"

404

This is not the web page you are looking for.

        "#);
        err.response.set_mut(Template::new("default", data));
        Err(err)
    }));
    // Handlebar templating
    c.link_after({
        let mut hbr = HandlebarsEngine::new();
        hbr.add(Box::new(DirectorySource::new("templates", ".hbs")));
        if let Err(r) = hbr.reload() {
            use std::error::Error;
            panic!("{}", r.description());
        }
        hbr
    });

    println!("\nServer running at \x1b[32mhttp://{}\x1b[0m in \x1b[33m{}\x1b[0m mode\n", ADDR, mode);
    Iron::new(c).http(ADDR).unwrap();
}

/// Root page handler
fn index(_: &mut Request) -> IronResult<Response> {
    read_docs("Home")
}

/// Wiki page handler
fn page(req: &mut Request) -> IronResult<Response> {
    use router::Router;

    // Parse URL
    let name = req.extensions.get::<Router>().unwrap().find("name").unwrap_or("Home");
    read_docs(name)
}

/// "All documents" page handler
fn all_docs(_: &mut Request) -> IronResult<Response> {
    use std::fs::read_dir;

    let mut paths: Vec<_> = itry!(read_dir(DOCS_DIR), status::NotFound)
        .filter_map(|f| f.ok())
        .map(|f| f.path())
        .filter(|p| p.is_file())
        .filter(|p| p.as_os_str().to_str()
            .map(|s| s.ends_with(".md"))
            .unwrap_or(false))
        .collect();
    paths.sort();

    let pages = paths.iter()
        .filter_map(|p| p.file_stem())
        .filter_map(|stem| stem.to_str());

    // TODO: 러스트로 이짓하지 말고 Handlebar로 대체
    let mut html = "
".to_string();
    for page in pages {
        use std::fmt::Write;
        write!(&mut html, r#"
{0}
"#, page).unwrap();
    }
    html.push_str("
");

    // Fill in to the template
    let mut data = BTreeMap::new();
    data.insert("content", html);
    Ok(Response::with((status::Ok, Template::new("default", data))))
}

/// `name`을 인자로 받아, `docs/<name>.md` 파일을 렌더링하여 `IronResult<Response>`로 반환합니다.
fn read_docs(name: &str) -> IronResult<Response> {
    use std::fs::File;
    use std::io::Read;
    use cmark::Parser;
    use cmark::html::push_html;

    // Read file
    let path = format!("{}/{}.md", DOCS_DIR, name);
    let mut file = itry!(File::open(&path), status::NotFound);
    let mut md = String::new();
    itry!(file.read_to_string(&mut md), status::ServiceUnavailable);

    // Parse markdown
    let parser = Parser::new(&md);
    let mut html = String::new();
    push_html(&mut html, parser);

    // Fill in to the template
    let mut data = BTreeMap::new();
    data.insert("content", html);
    Ok(Response::with((status::Ok, Template::new("default", data))))
}

/// 현재 프로그램이 개발모드로 실행되는지, 실서비스 모드로 실행되는지 여부를 저장하는 타입.
///
/// `get_mode()` 함수로 값을 얻어올 수 있다.
#[derive(PartialEq)]
enum Mode { Production, Development }

impl fmt::Display for Mode {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.write_str(match *self {
            Mode::Production => "production",
            Mode::Development => "development",
        })
    }
}

/// `RUST_KR` 환경변수가 `"production"`일경우 Mode::Production, 그 외의경우 Mode::Development를 반환한다.
fn get_mode() -> Mode {
    use std::env;
    use Mode::*;

    env::var("RUST_KR")
        .map(|s| if s == "production" { Production } else { Development })
        .unwrap_or(Development)
}

/// Helper struct for handle 404 page of rust-kr
struct RustkrHandler<F> { handler: F }
/// Helper function for handle 404 page of rust-kr
fn catch<F>(handler: F) -> RustkrHandler<F> { RustkrHandler { handler: handler } }
impl<F> AfterMiddleware for RustkrHandler<F>
    where F: Send + Sync + 'static + Fn(IronError) -> IronResult<Response> {
    fn catch(&self, _: &mut Request, err: IronError) -> IronResult<Response> {
        (self.handler)(err)
    }
}
<commit_msg>Change the bind address to '[::]'<commit_after>//! [한국 러스트 사용자 그룹 홈페이지][rust-kr] 소스코드
//! --------
//!
//! ### 사용법
//! ```bash
//! git clone https://github.com/rust-kr/rust-kr.org.wiki.git docs
//! cargo run
//! # 웹 브라우저에서 http://localhost:8000 를 열어보세요
//! ```
//!
//! [rust-kr]: http://rust-kr.org

// 한국 러스트 사용자 그룹 홈페이지
// Copyright (C) 2016  Hyeon Kim
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

#[macro_use(itry)] extern crate iron;
#[macro_use(router)] extern crate router;
extern crate staticfile;
extern crate handlebars_iron as hbs;
extern crate pulldown_cmark as cmark;

use std::collections::BTreeMap;
use std::fmt;
use iron::prelude::*;
use iron::status;
use iron::middleware::{Handler, AfterMiddleware};
use staticfile::Static;
use hbs::{Template, HandlebarsEngine, DirectorySource};

//
// Configs
//
/// 서버가 소켓을 열 주소
const ADDR: &'static str = "[::]:8000";
/// 마크다운 문서가 담긴 주소
const DOCS_DIR: &'static str = "docs";

/// Entry point
fn main() {
    let mode = get_mode();

    //
    //      User sends request
    //               ⇩
    //                                '/public' 디렉토리에서 해당 파일이 있는지 먼저 찾아본다.
    //    try static file serving     있을경우 이를 전송하고, 없을경우 아래의 main handler로
    //   (only in development mode)   넘어간다. production mode로 실행될 경우, 이 동작이 없어지므로
    //                                nginx나 apache의 정적파일서빙 기능을 써야한다.
    //               ⇩
    //        "main handlers"         index, page, all_docs 등의 함수로 요청을 처리한다
    //               ⇩
    //       404 page handler         위의 두개가 모두 실패하였을경우 404페이지를 렌더링한다
    //               ⇩
    //     handlebar templating
    //

    let mut c = Chain::new(router! {
        get "/" => index,
        get "/pages/:name" => page,
        get "/pages/_pages" => all_docs,
    });

    // Static file serving (only in development mode)
    if mode == Mode::Development {
        c.link_around(move |main: Box<Handler>| -> Box<Handler> {
            Box::new(move |req: &mut Request| -> IronResult<Response> {
                Static::new("public/").handle(req)
                    .or_else(|_| main.handle(req))
            })
        });
    }
    // 404 page handler
    c.link_after(catch(|mut err: IronError| {
        if err.response.body.is_some() { return Err(err); }
        if err.response.status != Some(status::NotFound) { return Err(err); }

        // TODO: 좀더 내용 채우기
        let mut data = BTreeMap::new();
        data.insert("content", r#"

404

This is not the web page you are looking for.

        "#);
        err.response.set_mut(Template::new("default", data));
        Err(err)
    }));
    // Handlebar templating
    c.link_after({
        let mut hbr = HandlebarsEngine::new();
        hbr.add(Box::new(DirectorySource::new("templates", ".hbs")));
        if let Err(r) = hbr.reload() {
            use std::error::Error;
            panic!("{}", r.description());
        }
        hbr
    });

    println!("\nServer running at \x1b[32mhttp://{}\x1b[0m in \x1b[33m{}\x1b[0m mode\n", ADDR, mode);
    Iron::new(c).http(ADDR).unwrap();
}

/// Root page handler
fn index(_: &mut Request) -> IronResult<Response> {
    read_docs("Home")
}

/// Wiki page handler
fn page(req: &mut Request) -> IronResult<Response> {
    use router::Router;

    // Parse URL
    let name = req.extensions.get::<Router>().unwrap().find("name").unwrap_or("Home");
    read_docs(name)
}

/// "All documents" page handler
fn all_docs(_: &mut Request) -> IronResult<Response> {
    use std::fs::read_dir;

    let mut paths: Vec<_> = itry!(read_dir(DOCS_DIR), status::NotFound)
        .filter_map(|f| f.ok())
        .map(|f| f.path())
        .filter(|p| p.is_file())
        .filter(|p| p.as_os_str().to_str()
            .map(|s| s.ends_with(".md"))
            .unwrap_or(false))
        .collect();
    paths.sort();

    let pages = paths.iter()
        .filter_map(|p| p.file_stem())
        .filter_map(|stem| stem.to_str());

    // TODO: 러스트로 이짓하지 말고 Handlebar로 대체
    let mut html = "
".to_string();
    for page in pages {
        use std::fmt::Write;
        write!(&mut html, r#"
{0}
"#, page).unwrap();
    }
    html.push_str("
");

    // Fill in to the template
    let mut data = BTreeMap::new();
    data.insert("content", html);
    Ok(Response::with((status::Ok, Template::new("default", data))))
}

/// `name`을 인자로 받아, `docs/<name>.md` 파일을 렌더링하여 `IronResult<Response>`로 반환합니다.
fn read_docs(name: &str) -> IronResult<Response> {
    use std::fs::File;
    use std::io::Read;
    use cmark::Parser;
    use cmark::html::push_html;

    // Read file
    let path = format!("{}/{}.md", DOCS_DIR, name);
    let mut file = itry!(File::open(&path), status::NotFound);
    let mut md = String::new();
    itry!(file.read_to_string(&mut md), status::ServiceUnavailable);

    // Parse markdown
    let parser = Parser::new(&md);
    let mut html = String::new();
    push_html(&mut html, parser);

    // Fill in to the template
    let mut data = BTreeMap::new();
    data.insert("content", html);
    Ok(Response::with((status::Ok, Template::new("default", data))))
}

/// 현재 프로그램이 개발모드로 실행되는지, 실서비스 모드로 실행되는지 여부를 저장하는 타입.
///
/// `get_mode()` 함수로 값을 얻어올 수 있다.
#[derive(PartialEq)]
enum Mode { Production, Development }

impl fmt::Display for Mode {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.write_str(match *self {
            Mode::Production => "production",
            Mode::Development => "development",
        })
    }
}

/// `RUST_KR` 환경변수가 `"production"`일경우 Mode::Production, 그 외의경우 Mode::Development를 반환한다.
fn get_mode() -> Mode {
    use std::env;
    use Mode::*;

    env::var("RUST_KR")
        .map(|s| if s == "production" { Production } else { Development })
        .unwrap_or(Development)
}

/// Helper struct for handle 404 page of rust-kr
struct RustkrHandler<F> { handler: F }
/// Helper function for handle 404 page of rust-kr
fn catch<F>(handler: F) -> RustkrHandler<F> { RustkrHandler { handler: handler } }
impl<F> AfterMiddleware for RustkrHandler<F>
    where F: Send + Sync + 'static + Fn(IronError) -> IronResult<Response> {
    fn catch(&self, _: &mut Request, err: IronError) -> IronResult<Response> {
        (self.handler)(err)
    }
}
<|endoftext|>


<commit_before>use std::fs::OpenOptions;
use std::io::BufReader;
use std::io::BufWriter;

struct Repo {
	status:u32,
	files: Option<Vec<String>>,
	dirs: Option<Vec<String>>,
}

impl Repo {
	fn new() -> Self {
		Repo {status:0,files:None,dirs:None}
	}
}

fn init(repo_instance:&mut Repo) {
	repo_instance.status +=1;
	println!("Initialized a graphin repository");
// create directory here;
}

fn status(repo_instance:&Repo) {
	if (repo_instance.status == 0) {
		println!("Fatal Error: Not a graphin repository");
	} else {
	println!("Working directory clean");
	}
// show status of directory here
}

fn add() {
	println!("Added to graphin repository");
// code to add file to tracker index
}

fn branch() {
	println!("Branches:");
// code to add for branches
}

fn main() {

	let mut master:Repo = Repo::new();
	match &std::env::args().nth(1).unwrap()[..] {
	"init" => {init(&mut master);},
	"status" => status(&master),
	"add *.*" => add(),
	"branch" => branch(),
	_ => println!("Command Not Found"),
}
}
<commit_msg>updated main.rs<commit_after>use std::fs::OpenOptions;
use std::io::BufReader;
use std::io::BufWriter;
use std::fs;

struct Repo {
	status:u32,
	files: Option<Vec<String>>,
	dirs: Option<Vec<String>>,
}

impl Repo {
	fn new() -> Self {
		Repo {status:0,files:None,dirs:None}
	}
}

fn init(repo_instance:&mut Repo) {
	repo_instance.status +=1;
    let metadata = OpenOptions::new()
            .read(true)
            .write(true)
            .create(true)
            .open(".graphin");
    println!("Initialized an empty graphin repository");
}

fn status(repo_instance:&Repo) {

	match fs::metadata(".graphin") {
        Ok(metadata) => {
            if !(metadata.is_file()) {
                println!("Not a valid graphin repository");
            } else {
                // code to add file to tracker index
                // if changes in blob object, notify else
                println!("Working directory clean", );
    }
        },
        Err(_) => {println!("Not a valid graphin repository");},
    }
// show status of directory here
}

fn add(repo_instance:&mut Repo) {

    match fs::metadata(".graphin") {
        Ok(metadata) => {
            if !(metadata.is_file()) {
                println!("Not a valid graphin repository");

            } else {
                // code to add file to tracker index
                // if changes in blob object, notify else
                let file = std::env::args().nth(2).unwrap();
                println!("{:?} added to master branch", file);
                println!("Working directory clean", );

    }
        },
        Err(_) => {println!("Not a valid graphin repository");},
    }

}

fn branch() {
	println!("Branches:");
// code to add for branches
}

fn main() {
	let mut master:Repo = Repo::new();
	match &std::env::args().nth(1).unwrap()[..] {
	"init" => {init(&mut master);},
	"status" => status(&master),
	"add" => add(&mut master),
	"branch" => branch(),
	_ => println!("Command Not Found"),
}
}
<|endoftext|>


<commit_before>#![feature(globs)]

extern crate libc;
extern crate neovim;
extern crate rgtk;
extern crate serialize;

use neovim::nvim_main;
use rgtk::*;
use std::collections::HashSet;

mod projects;
mod ui;
mod utils;

mod ffi {
    use libc::c_int;

    extern "C" {
        pub fn fork () -> c_int;
        pub fn kill (pid: c_int, sig: c_int);
    }
}

fn start_gui(pty: &mut gtk::VtePty, pid: i32) {
    gtk::init();

    // constants

    let width = 1242;
    let height = 768;
    let editor_height = ((height as f32) * 0.75) as i32;

    // create the window

    let mut window = gtk::Window::new(gtk::WindowType::TopLevel).unwrap();
    window.set_title("SolidOak");
    window.set_window_position(gtk::WindowPosition::Center);
    window.set_default_size(width, height);

    window.connect(gtk::signals::DeleteEvent::new(|_| {
        unsafe { ffi::kill(pid, 15); }
        gtk::main_quit();
        true
    }));

    // create the panes

    let new_button = gtk::Button::new_with_label("New Project").unwrap();
    let import_button = gtk::Button::new_with_label("Import").unwrap();
    let rename_button = gtk::Button::new_with_label("Rename").unwrap();
    let remove_button = gtk::Button::new_with_label("Remove").unwrap();

    let mut proj_btns = gtk::Box::new(gtk::Orientation::Horizontal, 0).unwrap();
    proj_btns.set_size_request(-1, -1);
    proj_btns.add(&new_button);
    proj_btns.add(&import_button);
    proj_btns.add(&rename_button);
    proj_btns.add(&remove_button);

    let mut proj_tree = gtk::TreeView::new().unwrap();
    let selection = proj_tree.get_selection().unwrap();
    let column_types = [glib::ffi::g_type_string, glib::ffi::g_type_string];
    let store = gtk::TreeStore::new(&column_types).unwrap();
    let model = store.get_model().unwrap();
    proj_tree.set_model(&model);
    proj_tree.set_headers_visible(false);

    let mut scroll_pane = gtk::ScrolledWindow::new(None, None).unwrap();
    scroll_pane.add(&proj_tree);

    let column = gtk::TreeViewColumn::new().unwrap();
    let cell = gtk::CellRendererText::new().unwrap();
    column.pack_start(&cell, true);
    column.add_attribute(&cell, "text", 0);
    proj_tree.append_column(&column);

    let mut proj_pane = gtk::Box::new(gtk::Orientation::Vertical, 0).unwrap();
    proj_pane.set_size_request(-1, -1);
    proj_pane.pack_start(&proj_btns, false, true, 0);
    proj_pane.pack_start(&scroll_pane, true, true, 0);

    let mut editor_pane = gtk::VteTerminal::new().unwrap();
    editor_pane.set_size_request(-1, editor_height);
    editor_pane.set_pty(pty);
    editor_pane.watch_child(pid);

    let run_button = gtk::Button::new_with_label("Run").unwrap();
    let build_button = gtk::Button::new_with_label("Build").unwrap();
    let test_button = gtk::Button::new_with_label("Test").unwrap();
    let clean_button = gtk::Button::new_with_label("Clean").unwrap();
    let stop_button = gtk::Button::new_with_label("Stop").unwrap();

    let mut build_btns = gtk::Box::new(gtk::Orientation::Horizontal, 0).unwrap();
    build_btns.set_size_request(-1, -1);
    build_btns.add(&run_button);
    build_btns.add(&build_button);
    build_btns.add(&test_button);
    build_btns.add(&clean_button);
    build_btns.add(&stop_button);

    let build_term = gtk::VteTerminal::new().unwrap();

    let mut build_pane = gtk::Box::new(gtk::Orientation::Vertical, 0).unwrap();
    build_pane.pack_start(&build_btns, false, true, 0);
    build_pane.pack_start(&build_term, true, true, 0);

    let mut content = gtk::Box::new(gtk::Orientation::Vertical, 0).unwrap();
    content.pack_start(&editor_pane, false, true, 0);
    content.pack_start(&build_pane, true, true, 0);

    let mut hbox = gtk::Box::new(gtk::Orientation::Horizontal, 0).unwrap();
    hbox.pack_start(&proj_pane, false, true, 0);
    hbox.pack_start(&content, true, true, 0);
    window.add(&hbox);

    // populate the project tree

    let mut state = ::utils::State{
        projects: HashSet::new(),
        expansions: HashSet::new(),
        selection: None,
        tree_model: &model,
        tree_store: &store,
        tree_selection: &selection,
        rename_button: &rename_button,
        remove_button: &remove_button,
    };

    ::utils::create_data_dir();
    ::utils::read_prefs(&mut state);
    ::ui::update_project_tree(&mut state, &mut proj_tree);

    // connect to the signals

    new_button.connect(gtk::signals::Clicked::new(|| {
        ::projects::new_project(&mut state, &mut proj_tree);
    }));
    import_button.connect(gtk::signals::Clicked::new(|| {
        ::projects::import_project(&mut state, &mut proj_tree);
    }));
    rename_button.connect(gtk::signals::Clicked::new(|| {
        ::projects::rename_file(&mut state);
    }));
    remove_button.connect(gtk::signals::Clicked::new(|| {
        ::projects::remove_item(&mut state);
    }));
    selection.connect(gtk::signals::Changed::new(|| {
        ::projects::update_selection(&mut state);
    }));
    proj_tree.connect(gtk::signals::RowCollapsed::new(|iter_raw, _| {
        let iter = gtk::TreeIter::wrap_pointer(iter_raw);
        ::projects::remove_expansion(&mut state, &iter);
    }));
    proj_tree.connect(gtk::signals::RowExpanded::new(|iter_raw, _| {
        let iter = gtk::TreeIter::wrap_pointer(iter_raw);
        ::projects::add_expansion(&mut state, &iter);
    }));

    // show the window

    window.show_all();
    gtk::main();
}

fn main() {
    let mut pty = gtk::VtePty::new().unwrap();
    let pid = unsafe { ffi::fork() };

    if pid > 0 {
        start_gui(&mut pty, pid);
    } else {
        pty.child_setup();
        nvim_main();
    }
}
<commit_msg>Start work on communication between processes<commit_after>#![feature(globs)]

extern crate libc;
extern crate neovim;
extern crate rgtk;
extern crate serialize;

use neovim::nvim_main;
use rgtk::*;
use std::collections::HashSet;
use std::{io, str, time};

mod projects;
mod ui;
mod utils;

mod ffi {
    pub use libc::{c_char, c_uchar, c_int, c_void, uint64_t};
    pub use libc::funcs::posix88::unistd::{close, pipe, read, write};
    pub use libc::types::os::arch::c95::size_t;

    extern "C" {
        pub fn fork () -> c_int;
        pub fn kill (pid: c_int, sig: c_int);
        pub fn channel_from_fds (read_fd: c_int, write_fd: c_int) -> uint64_t;
        pub fn channel_subscribe(id: uint64_t, event: *const c_char);
    }
}

fn gui_main(
    pty: &mut gtk::VtePty,
    pid: ffi::c_int,
    read_fd: ffi::c_int,
    write_fd: ffi::c_int)
{
    gtk::init();

    // constants

    let width = 1242;
    let height = 768;
    let editor_height = ((height as f32) * 0.75) as i32;

    // create the window

    let mut window = gtk::Window::new(gtk::WindowType::TopLevel).unwrap();
    window.set_title("SolidOak");
    window.set_window_position(gtk::WindowPosition::Center);
    window.set_default_size(width, height);

    window.connect(gtk::signals::DeleteEvent::new(|_| {
        unsafe {
            ffi::close(read_fd);
            ffi::close(write_fd);
            ffi::kill(pid, 15);
        }
        gtk::main_quit();
        true
    }));

    // create the panes

    let new_button = gtk::Button::new_with_label("New Project").unwrap();
    let import_button = gtk::Button::new_with_label("Import").unwrap();
    let rename_button = gtk::Button::new_with_label("Rename").unwrap();
    let remove_button = gtk::Button::new_with_label("Remove").unwrap();

    let mut proj_btns = gtk::Box::new(gtk::Orientation::Horizontal, 0).unwrap();
    proj_btns.set_size_request(-1, -1);
    proj_btns.add(&new_button);
    proj_btns.add(&import_button);
    proj_btns.add(&rename_button);
    proj_btns.add(&remove_button);

    let mut proj_tree = gtk::TreeView::new().unwrap();
    let selection = proj_tree.get_selection().unwrap();
    let column_types = [glib::ffi::g_type_string, glib::ffi::g_type_string];
    let store = gtk::TreeStore::new(&column_types).unwrap();
    let model = store.get_model().unwrap();
    proj_tree.set_model(&model);
    proj_tree.set_headers_visible(false);

    let mut scroll_pane = gtk::ScrolledWindow::new(None, None).unwrap();
    scroll_pane.add(&proj_tree);

    let column = gtk::TreeViewColumn::new().unwrap();
    let cell = gtk::CellRendererText::new().unwrap();
    column.pack_start(&cell, true);
    column.add_attribute(&cell, "text", 0);
    proj_tree.append_column(&column);

    let mut proj_pane = gtk::Box::new(gtk::Orientation::Vertical, 0).unwrap();
    proj_pane.set_size_request(-1, -1);
    proj_pane.pack_start(&proj_btns, false, true, 0);
    proj_pane.pack_start(&scroll_pane, true, true, 0);

    let mut editor_pane = gtk::VteTerminal::new().unwrap();
    editor_pane.set_size_request(-1, editor_height);
    editor_pane.set_pty(pty);
    editor_pane.watch_child(pid);

    let run_button = gtk::Button::new_with_label("Run").unwrap();
    let build_button = gtk::Button::new_with_label("Build").unwrap();
    let test_button = gtk::Button::new_with_label("Test").unwrap();
    let clean_button = gtk::Button::new_with_label("Clean").unwrap();
    let stop_button = gtk::Button::new_with_label("Stop").unwrap();

    let mut build_btns = gtk::Box::new(gtk::Orientation::Horizontal, 0).unwrap();
    build_btns.set_size_request(-1, -1);
    build_btns.add(&run_button);
    build_btns.add(&build_button);
    build_btns.add(&test_button);
    build_btns.add(&clean_button);
    build_btns.add(&stop_button);

    let build_term = gtk::VteTerminal::new().unwrap();

    let mut build_pane = gtk::Box::new(gtk::Orientation::Vertical, 0).unwrap();
    build_pane.pack_start(&build_btns, false, true, 0);
    build_pane.pack_start(&build_term, true, true, 0);

    let mut content = gtk::Box::new(gtk::Orientation::Vertical, 0).unwrap();
    content.pack_start(&editor_pane, false, true, 0);
    content.pack_start(&build_pane, true, true, 0);

    let mut hbox = gtk::Box::new(gtk::Orientation::Horizontal, 0).unwrap();
    hbox.pack_start(&proj_pane, false, true, 0);
    hbox.pack_start(&content, true, true, 0);
    window.add(&hbox);

    // populate the project tree

    let mut state = ::utils::State{
        projects: HashSet::new(),
        expansions: HashSet::new(),
        selection: None,
        tree_model: &model,
        tree_store: &store,
        tree_selection: &selection,
        rename_button: &rename_button,
        remove_button: &remove_button,
    };

    ::utils::create_data_dir();
    ::utils::read_prefs(&mut state);
    ::ui::update_project_tree(&mut state, &mut proj_tree);

    // connect to the signals

    new_button.connect(gtk::signals::Clicked::new(|| {
        ::projects::new_project(&mut state, &mut proj_tree);
    }));
    import_button.connect(gtk::signals::Clicked::new(|| {
        ::projects::import_project(&mut state, &mut proj_tree);
    }));
    rename_button.connect(gtk::signals::Clicked::new(|| {
        ::projects::rename_file(&mut state);
    }));
    remove_button.connect(gtk::signals::Clicked::new(|| {
        ::projects::remove_item(&mut state);
    }));
    selection.connect(gtk::signals::Changed::new(|| {
        ::projects::update_selection(&mut state);
    }));
    proj_tree.connect(gtk::signals::RowCollapsed::new(|iter_raw, _| {
        let iter = gtk::TreeIter::wrap_pointer(iter_raw);
        ::projects::remove_expansion(&mut state, &iter);
    }));
    proj_tree.connect(gtk::signals::RowExpanded::new(|iter_raw, _| {
        let iter = gtk::TreeIter::wrap_pointer(iter_raw);
        ::projects::add_expansion(&mut state, &iter);
    }));

    // show the window

    window.show_all();
    gtk::main();
}

fn main() {
    let mut pty = gtk::VtePty::new().unwrap();

    let mut nvim_from_gui : [ffi::c_int, ..2] = [0, ..2];
    let mut gui_from_nvim : [ffi::c_int, ..2] = [0, ..2];
    unsafe {
        ffi::pipe(nvim_from_gui.as_mut_ptr());
        ffi::pipe(gui_from_nvim.as_mut_ptr());
    };

    let pid = unsafe { ffi::fork() };

    if pid > 0 {
        spawn(proc() {
            let mut buf : [ffi::c_uchar, ..100] = [0, ..100];
            unsafe {
                loop {
                    let buf_ptr = buf.as_mut_ptr() as *mut ffi::c_void;
                    let n = ffi::read(gui_from_nvim[0], buf_ptr, 100);
                    if n < 0 { break; }
                    let msg = str::from_utf8(buf.slice_to(n as uint)).unwrap();
                    println!("Received: {}", msg);
                }
            }
        });
        gui_main(&mut pty, pid, gui_from_nvim[0], gui_from_nvim[1]);
    } else {
        pty.child_setup();
        spawn(proc() {
            io::timer::sleep(time::Duration::seconds(1));
            unsafe {
                let id = ffi::channel_from_fds(nvim_from_gui[0], gui_from_nvim[1]);
                let cmd = "test";
                ffi::channel_subscribe(id, cmd.to_c_str().as_ptr());

                let msg = "Hello, world!";
                let msg_c = msg.to_c_str();
                let msg_ptr = msg_c.as_ptr() as *const ffi::c_void;
                ffi::write(gui_from_nvim[1], msg_ptr, msg_c.len() as ffi::size_t);
            }
        });
        nvim_main();
    }
}
<|endoftext|>


<commit_before>extern crate base64;
extern crate hyper;
extern crate hyper_native_tls;
extern crate pem;
#[macro_use]
extern crate prettytable;
extern crate ring;
extern crate serde_json;
extern crate url;

extern crate ct_submitter;

use std::{env, process};
use std::fs::File;
use std::io::Read;

use ring::digest;

use ct_submitter::{fetch_trusted_ct_logs, submit_cert_to_logs, AddChainRequest};

fn pems_to_chain(data: &str) -> Vec<Vec<u8>> {
    return pem::parse_many(data)
               .into_iter()
               .filter(|p| p.tag == "CERTIFICATE")
               .map(|p| p.contents)
               .collect::<Vec<_>>();
}

fn build_chain(http_client: &hyper::Client, cert: &[u8]) -> Vec<Vec<u8>> {
    let body = url::form_urlencoded::Serializer::new(String::new())
        .append_pair("b64cert", &base64::encode(cert))
        .finish();
    let body_bytes = body.as_bytes();
    let response = http_client
        .post("https://crt.sh/gen-add-chain")
        .header(hyper::header::ContentType::form_url_encoded())
        .body(hyper::client::Body::BufBody(body_bytes, body_bytes.len()))
        .send()
        .unwrap();

    let add_chain_request: AddChainRequest = serde_json::from_reader(response).unwrap();
    return add_chain_request
        .chain
        .iter()
        .map(|c| base64::decode(c).unwrap())
        .collect();
}

fn main() {
    if env::args().len() == 1 {
        println!("Usage: {} <cert-or-chain.pem ...>", env::args().next().unwrap());
        process::exit(1);
    }

    let http_client = hyper::Client::with_connector(
        hyper::net::HttpsConnector::new(hyper_native_tls::NativeTlsClient::new().unwrap())
    );
    // TODO: timeout on the http_client, but for submit_cert_to_logs only, not build_chain
    let logs = fetch_trusted_ct_logs(&http_client);

    for path in env::args().skip(1) {
        println!("Submitting {} ...", path);

        let mut contents = String::new();
        File::open(path)
            .unwrap()
            .read_to_string(&mut contents)
            .unwrap();

        let mut chain = pems_to_chain(&contents);
        if chain.len() == 1 {
            // TODO: There's got to be some way to do this ourselves, instead of using crt.sh as a
            // glorified AIA chaser.
            println!("Only one certificate in chain, using crt.sh to build a full chain ...");
            chain = build_chain(&http_client, &chain[0]);
        }
        let scts = submit_cert_to_logs(&http_client, &logs, &chain);

        let mut table = prettytable::Table::new();
        table.add_row(row!["Log", "SCT"]);
        for (log, sct) in scts {
            table.add_row(row![log.description, base64::encode(&sct.to_raw_bytes())]);
        }
        table.printstd();

        // TODO: is there a better way to do this hex-encoding?
        println!("Find the cert on crt.sh: https://crt.sh?q={}",
                 digest::digest(&digest::SHA256, &chain[0])
                     .as_ref()
                     .iter()
                     .map(|b| format!("{:02X}", b))
                     .collect::<String>());
         println!();
         println!();
    }

}
<commit_msg>cargo fmt<commit_after>extern crate base64;
extern crate hyper;
extern crate hyper_native_tls;
extern crate pem;
#[macro_use]
extern crate prettytable;
extern crate ring;
extern crate serde_json;
extern crate url;

extern crate ct_submitter;

use std::{env, process};
use std::fs::File;
use std::io::Read;

use ring::digest;

use ct_submitter::{fetch_trusted_ct_logs, submit_cert_to_logs, AddChainRequest};

fn pems_to_chain(data: &str) -> Vec<Vec<u8>> {
    return pem::parse_many(data)
               .into_iter()
               .filter(|p| p.tag == "CERTIFICATE")
               .map(|p| p.contents)
               .collect::<Vec<_>>();
}

fn build_chain(http_client: &hyper::Client, cert: &[u8]) -> Vec<Vec<u8>> {
    let body = url::form_urlencoded::Serializer::new(String::new())
        .append_pair("b64cert", &base64::encode(cert))
        .finish();
    let body_bytes = body.as_bytes();
    let response = http_client
        .post("https://crt.sh/gen-add-chain")
        .header(hyper::header::ContentType::form_url_encoded())
        .body(hyper::client::Body::BufBody(body_bytes, body_bytes.len()))
        .send()
        .unwrap();

    let add_chain_request: AddChainRequest = serde_json::from_reader(response).unwrap();
    return add_chain_request
               .chain
               .iter()
               .map(|c| base64::decode(c).unwrap())
               .collect();
}

fn main() {
    if env::args().len() == 1 {
        println!("Usage: {} <cert-or-chain.pem ...>",
                 env::args().next().unwrap());
        process::exit(1);
    }

    let http_client = hyper::Client::with_connector(
        hyper::net::HttpsConnector::new(hyper_native_tls::NativeTlsClient::new().unwrap())
    );
    // TODO: timeout on the http_client, but for submit_cert_to_logs only, not build_chain
    let logs = fetch_trusted_ct_logs(&http_client);

    for path in env::args().skip(1) {
        println!("Submitting {} ...", path);

        let mut contents = String::new();
        File::open(path)
            .unwrap()
            .read_to_string(&mut contents)
            .unwrap();

        let mut chain = pems_to_chain(&contents);
        if chain.len() == 1 {
            // TODO: There's got to be some way to do this ourselves, instead of using crt.sh as a
            // glorified AIA chaser.
            println!("Only one certificate in chain, using crt.sh to build a full chain ...");
            chain = build_chain(&http_client, &chain[0]);
        }
        let scts = submit_cert_to_logs(&http_client, &logs, &chain);

        let mut table = prettytable::Table::new();
        table.add_row(row!["Log", "SCT"]);
        for (log, sct) in scts {
            table.add_row(row![log.description, base64::encode(&sct.to_raw_bytes())]);
        }
        table.printstd();

        // TODO: is there a better way to do this hex-encoding?
        println!("Find the cert on crt.sh: https://crt.sh?q={}",
                 digest::digest(&digest::SHA256, &chain[0])
                     .as_ref()
                     .iter()
                     .map(|b| format!("{:02X}", b))
                     .collect::<String>());
        println!();
        println!();
    }

}
<|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"),
    ("SOCKS4 NOAUTH Request", r"^\x04\x01\x00\x50"),
    ("SOCKS5 NOAUTH Request", r"^\x05\x01\x00$"), // Tested ok-ish
    ("SOCKS5 USER/PASS Request", r"^\x05\x02\x00\x02$") // possibly broken
  ];
  let io_timeout = Duration::new(300, 0); // 5 minutes

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

  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");
      match Connection::open("portlurker.sqlite") {
        Ok(conn) => { 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(err) => {println!("Failed to open or create database: {} - Continuing without logging", err); app.sql_logging = false;},
      }
    }
  }
  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
            };

            match Connection::open("portlurker.sqlite2") {
              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(err) => {println!("Failed to open database: {} - Continuing without logging", err); 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>Improved SQLite error handling<commit_after>extern crate yaml_rust;
extern crate regex;
extern crate rusqlite;

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

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

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

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

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

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

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

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

  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");
      match Connection::open("portlurker.sqlite") {
        Ok(conn) => { 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(err) => {println!("Failed to open or create database: {} - Continuing without logging", err); app.sql_logging = false;},
      }
    }
  }
  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
            };

            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(err) => {println!("Failed to open database: {} - Continuing without logging", err); 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|>