<commit_before>#![feature(slice_splits, core)]
#![cfg_attr(test, feature(test))]
extern crate core;
extern crate itertools;
extern crate rand;

use core::ptr;
use itertools::partition;
use std::cmp::Ordering::{self, Less};
use std::mem;

#[derive(Debug, Clone)]
pub struct LazySort<T> {
    inner: LazySortInternal<T>
}

impl<T: Ord> Iterator for LazySort<T> {
    type Item = T;

    fn next(&mut self) -> Option<T> {
        self.inner.next()
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        self.inner.size_hint()
    }
}

#[derive(Debug, Clone)]
enum LazySortInternal<T> {
    Base(Vec<T>),
    Recursive(Recursive<T>),
}

impl<T: Ord> LazySortInternal<T> {
    fn new(mut v: Vec<T>) -> LazySortInternal<T> {
        if v.len() <= 32 {
            insertion_sort(&mut v, |a, b| b.cmp(a));
            LazySortInternal::Base(v)
        } else {
            LazySortInternal::Recursive(Recursive::new(v))
        }
    }
}

impl<T: Ord> Iterator for LazySortInternal<T> {
    type Item = T;

    fn next(&mut self) -> Option<T> {
        match *self {
            LazySortInternal::Base(ref mut v) => v.pop(),
            LazySortInternal::Recursive(ref mut r) => r.next()
        }
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        match *self {
            LazySortInternal::Base(ref v) => (v.len(), Some(v.len())),
            LazySortInternal::Recursive(ref r) => r.size_hint()
        }
    }
}

#[derive(Clone, Debug)]
struct Recursive<T> {
    greater: Vec<T>,
    less: Option<Box<LazySortInternal<T>>>,
}

impl<T: Ord> Recursive<T> {
    fn new(v: Vec<T>) -> Recursive<T> {
        Recursive { greater: v, less: None }
    }

    fn split_greater(&mut self) -> Option<T> {
        match self.greater.len() {
            0 => None,
            1 => self.greater.pop(),
            _ => {
                let pivot_idx = self.greater.len() - 1;
                let split_idx = {
                    let mid_idx = self.greater.len() / 2;
                    // I've chosen the element in the middle of the vec as the pivot.
                    // However, we first swap the pivot with the last element so that there is
                    // a contiguous space in memory to be partitioned.
                    self.greater.swap(pivot_idx, mid_idx);
                    let (pivot, rest) = self.greater.split_last_mut().unwrap();
                    // partition all but the last element, which is the pivot. This makes the vec
                    // look like [greater, greater, ..., greater, less, less, ..., less, pivot]
                    partition(rest, |el| el > pivot)
                };
                // Swapping the pivot with the first less element allows us to split off
                // vec[split_idx + 1..] to create a new vec with all the elements less than pivot.
                self.greater.swap(pivot_idx, split_idx);
                let split_off_idx = split_idx + 1;
                if split_off_idx < self.greater.len() {
                    let mut less = Box::new(LazySortInternal::new(self.greater.split_off(split_off_idx)));
                    // Recursively compute the next element from the LazySortInternal struct containing
                    // the elements less than the pivot.
                    let next = less.next();
                    self.less = Some(less);
                    next
                } else {
                    // If there were no elements less than the pivot, then return the pivot.
                    self.greater.pop()
                }
            }
        }
    }
}

impl<T: Ord> Iterator for Recursive<T> {
    type Item = T;

    fn next(&mut self) -> Option<T> {
        let next = if let Some(ref mut less) = self.less {
            less.next()
        } else {
            return self.split_greater();
        };
        if next.is_some() {
            next
        } else {
            self.less = None;
            // The pivot is always the last element in the vec, and it's the first element
            // to be returned once all of the elements less than it have been returned.
            self.greater.pop()
        }
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        match self.less {
            None => (self.greater.len(), Some(self.greater.len())),
            Some(ref less) => {
                let (lower, upper) = less.size_hint();
                (lower + self.greater.len(),
                 upper.map(|upper| upper + self.greater.len()))
            }
        }
    }
}

pub trait LazySortIterator: Iterator
    where Self: Sized,
          Self::Item: Ord
{
    fn lazy_sort(self) -> LazySort<Self::Item> {
        LazySort { inner: LazySortInternal::new(self.collect()) }
    }
}

impl<T> LazySortIterator for T
    where T: Iterator,
          T::Item: Ord { }

#[test]
fn test_sort() {
    let mut v = vec![2, 4, 2, 5, 8, 4, 3, 4, 6];
    let v2: Vec<_> = v.iter().cloned().lazy_sort().collect();
    v.sort();
    assert_eq!(v, v2);
}

#[test]
fn test_empty() {
    let v: Vec<u64> = vec![];
    let v2: Vec<_> = v.iter().cloned().lazy_sort().collect();
    assert_eq!(v, v2);
}

#[test]
fn test_size_hint() {
    let v = vec![2, 4, 2, 5, 8, 4, 3, 4, 6];
    let mut sort_iter = v.iter().cloned().lazy_sort();
    for i in 0..v.len() {
        let (lower, upper) = sort_iter.size_hint();
        assert_eq!(v.len() - i, lower);
        assert_eq!(Some(v.len() - i), upper);
        sort_iter.next();
    }
}

// This is copied from libcollections/slice.rs
fn insertion_sort<T, F>(v: &mut [T], mut compare: F) where F: FnMut(&T, &T) -> Ordering {
    let len = v.len() as isize;
    let buf_v = v.as_mut_ptr();

    // 1 <= i < len;
    for i in 1..len {
        // j satisfies: 0 <= j <= i;
        let mut j = i;
        unsafe {
            // `i` is in bounds.
            let read_ptr = buf_v.offset(i) as *const T;

            // find where to insert, we need to do strict <,
            // rather than <=, to maintain stability.

            // 0 <= j - 1 < len, so .offset(j - 1) is in bounds.
            while j > 0 &&
                    compare(&*read_ptr, &*buf_v.offset(j - 1)) == Less {
                j -= 1;
            }

            // shift everything to the right, to make space to
            // insert this value.

            // j + 1 could be `len` (for the last `i`), but in
            // that case, `i == j` so we don't copy. The
            // `.offset(j)` is always in bounds.

            if i != j {
                let tmp = ptr::read(read_ptr);
                ptr::copy(&*buf_v.offset(j),
                          buf_v.offset(j + 1),
                          (i - j) as usize);
                ptr::copy_nonoverlapping(&tmp, buf_v.offset(j), 1);
                mem::forget(tmp);
            }
        }
    }
}

#[cfg(test)]
mod bench {
    extern crate test;

    use self::test::Bencher;
    use rand::{thread_rng, Rng};
    use super::LazySortIterator;

    #[bench]
    fn take_1000_lazy(b: &mut Bencher) {
        let mut rng = thread_rng();
        let v: Vec<u32> = rng.gen_iter().take(50_000).collect();
        b.iter(|| v.iter().cloned().lazy_sort().take(1000).collect::<Vec<_>>());
    }

    #[bench]
    fn take_1000_eager(b: &mut Bencher) {
        let mut rng = thread_rng();
        let v: Vec<u32> = rng.gen_iter().take(50_000).collect();
        b.iter(|| {
            let mut v = v.clone();
            v.sort();
            v.iter().cloned().take(1000).collect::<Vec<_>>();
        });
    }

    #[bench]
    fn take_10_000_lazy(b: &mut Bencher) {
        let mut rng = thread_rng();
        let v: Vec<u32> = rng.gen_iter().take(50_000).collect();
        b.iter(|| v.iter().cloned().lazy_sort().take(10_000).collect::<Vec<_>>());
    }

    #[bench]
    fn take_10_000_eager(b: &mut Bencher) {
        let mut rng = thread_rng();
        let v: Vec<u32> = rng.gen_iter().take(50_000).collect();
        b.iter(|| {
            let mut v = v.clone();
            v.sort();
            v.iter().cloned().take(10_000).collect::<Vec<_>>();
        });
    }

    #[bench]
    fn take_50_000_lazy(b: &mut Bencher) {
        let mut rng = thread_rng();
        let v: Vec<u32> = rng.gen_iter().take(50_000).collect();
        b.iter(|| v.iter().cloned().lazy_sort().take(50_000).collect::<Vec<_>>());
    }

    #[bench]
    fn take_50_000_eager(b: &mut Bencher) {
        let mut rng = thread_rng();
        let v: Vec<u32> = rng.gen_iter().take(50_000).collect();
        b.iter(|| {
            let mut v = v.clone();
            v.sort();
            v.iter().cloned().take(50_000).collect::<Vec<_>>();
        });
    }
}
<commit_msg>run rustfmt<commit_after>#![feature(slice_splits, core)]
#![cfg_attr(test, feature(test))]
extern crate core;
extern crate itertools;
extern crate rand;

use core::ptr;
use itertools::partition;
use std::cmp::Ordering::{self, Less};
use std::mem;

#[derive(Debug, Clone)]
pub struct LazySort<T> {
    inner: LazySortInternal<T>,
}

impl<T: Ord> Iterator for LazySort<T> {
    type Item = T;

    fn next(&mut self) -> Option<T> {
        self.inner.next()
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        self.inner.size_hint()
    }
}

#[derive(Debug, Clone)]
enum LazySortInternal<T> {
    Base(Vec<T>),
    Recursive(Recursive<T>),
}

impl<T: Ord> LazySortInternal<T> {
    fn new(mut v: Vec<T>) -> LazySortInternal<T> {
        if v.len() <= 32 {
            insertion_sort(&mut v, |a, b| b.cmp(a));
            LazySortInternal::Base(v)
        } else {
            LazySortInternal::Recursive(Recursive::new(v))
        }
    }
}

impl<T: Ord> Iterator for LazySortInternal<T> {
    type Item = T;

    fn next(&mut self) -> Option<T> {
        match *self {
            LazySortInternal::Base(ref mut v) => v.pop(),
            LazySortInternal::Recursive(ref mut r) => r.next(),
        }
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        match *self {
            LazySortInternal::Base(ref v) => (v.len(), Some(v.len())),
            LazySortInternal::Recursive(ref r) => r.size_hint(),
        }
    }
}

#[derive(Clone, Debug)]
struct Recursive<T> {
    greater: Vec<T>,
    less: Option<Box<LazySortInternal<T>>>,
}

impl<T: Ord> Recursive<T> {
    fn new(v: Vec<T>) -> Recursive<T> {
        Recursive {
            greater: v,
            less: None,
        }
    }

    fn split_greater(&mut self) -> Option<T> {
        match self.greater.len() {
            0 => None,
            1 => self.greater.pop(),
            _ => {
                let pivot_idx = self.greater.len() - 1;
                let split_idx = {
                    let mid_idx = self.greater.len() / 2;
                    // I've chosen the element in the middle of the vec as the pivot.
                    // However, we first swap the pivot with the last element so that there is
                    // a contiguous space in memory to be partitioned.
                    self.greater.swap(pivot_idx, mid_idx);
                    let (pivot, rest) = self.greater.split_last_mut().unwrap();
                    // partition all but the last element, which is the pivot. This makes the vec
                    // look like [greater, greater, ..., greater, less, less, ..., less, pivot]
                    partition(rest, |el| el > pivot)
                };
                // Swapping the pivot with the first less element allows us to split off
                // vec[split_idx + 1..] to create a new vec with all the elements less than pivot.
                self.greater.swap(pivot_idx, split_idx);
                let split_off_idx = split_idx + 1;
                if split_off_idx < self.greater.len() {
                    let mut less = Box::new(LazySortInternal::new(self.greater
                                                                      .split_off(split_off_idx)));
                    // Recursively compute the next element from the LazySortInternal struct containing
                    // the elements less than the pivot.
                    let next = less.next();
                    self.less = Some(less);
                    next
                } else {
                    // If there were no elements less than the pivot, then return the pivot.
                    self.greater.pop()
                }
            }
        }
    }
}

impl<T: Ord> Iterator for Recursive<T> {
    type Item = T;

    fn next(&mut self) -> Option<T> {
        let next = if let Some(ref mut less) = self.less {
            less.next()
        } else {
            return self.split_greater();
        };
        if next.is_some() {
            next
        } else {
            self.less = None;
            // The pivot is always the last element in the vec, and it's the first element
            // to be returned once all of the elements less than it have been returned.
            self.greater.pop()
        }
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        match self.less {
            None => (self.greater.len(), Some(self.greater.len())),
            Some(ref less) => {
                let (lower, upper) = less.size_hint();
                (lower + self.greater.len(),
                 upper.map(|upper| upper + self.greater.len()))
            }
        }
    }
}

pub trait LazySortIterator: Iterator
    where Self: Sized,
          Self::Item: Ord
{
    fn lazy_sort(self) -> LazySort<Self::Item> {
        LazySort { inner: LazySortInternal::new(self.collect()) }
    }
}

impl<T> LazySortIterator for T
    where T: Iterator,
          T::Item: Ord { }

#[test]
fn test_sort() {
    let mut v = vec![2, 4, 2, 5, 8, 4, 3, 4, 6];
    let v2: Vec<_> = v.iter().cloned().lazy_sort().collect();
    v.sort();
    assert_eq!(v, v2);
}

#[test]
fn test_empty() {
    let v: Vec<u64> = vec![];
    let v2: Vec<_> = v.iter().cloned().lazy_sort().collect();
    assert_eq!(v, v2);
}

#[test]
fn test_size_hint() {
    let v = vec![2, 4, 2, 5, 8, 4, 3, 4, 6];
    let mut sort_iter = v.iter().cloned().lazy_sort();
    for i in 0..v.len() {
        let (lower, upper) = sort_iter.size_hint();
        assert_eq!(v.len() - i, lower);
        assert_eq!(Some(v.len() - i), upper);
        sort_iter.next();
    }
}

// This is copied from libcollections/slice.rs
fn insertion_sort<T, F>(v: &mut [T], mut compare: F)
    where F: FnMut(&T, &T) -> Ordering
{
    let len = v.len() as isize;
    let buf_v = v.as_mut_ptr();

    // 1 <= i < len;
    for i in 1..len {
        // j satisfies: 0 <= j <= i;
        let mut j = i;
        unsafe {
            // `i` is in bounds.
            let read_ptr = buf_v.offset(i) as *const T;

            // find where to insert, we need to do strict <,
            // rather than <=, to maintain stability.

            // 0 <= j - 1 < len, so .offset(j - 1) is in bounds.
            while j > 0 && compare(&*read_ptr, &*buf_v.offset(j - 1)) == Less {
                j -= 1;
            }

            // shift everything to the right, to make space to
            // insert this value.

            // j + 1 could be `len` (for the last `i`), but in
            // that case, `i == j` so we don't copy. The
            // `.offset(j)` is always in bounds.

            if i != j {
                let tmp = ptr::read(read_ptr);
                ptr::copy(&*buf_v.offset(j), buf_v.offset(j + 1), (i - j) as usize);
                ptr::copy_nonoverlapping(&tmp, buf_v.offset(j), 1);
                mem::forget(tmp);
            }
        }
    }
}

#[cfg(test)]
mod bench {
    extern crate test;

    use self::test::Bencher;
    use rand::{thread_rng, Rng};
    use super::LazySortIterator;

    #[bench]
    fn take_1000_lazy(b: &mut Bencher) {
        let mut rng = thread_rng();
        let v: Vec<u32> = rng.gen_iter().take(50_000).collect();
        b.iter(|| v.iter().cloned().lazy_sort().take(1000).collect::<Vec<_>>());
    }

    #[bench]
    fn take_1000_eager(b: &mut Bencher) {
        let mut rng = thread_rng();
        let v: Vec<u32> = rng.gen_iter().take(50_000).collect();
        b.iter(|| {
            let mut v = v.clone();
            v.sort();
            v.iter().cloned().take(1000).collect::<Vec<_>>();
        });
    }

    #[bench]
    fn take_10_000_lazy(b: &mut Bencher) {
        let mut rng = thread_rng();
        let v: Vec<u32> = rng.gen_iter().take(50_000).collect();
        b.iter(|| v.iter().cloned().lazy_sort().take(10_000).collect::<Vec<_>>());
    }

    #[bench]
    fn take_10_000_eager(b: &mut Bencher) {
        let mut rng = thread_rng();
        let v: Vec<u32> = rng.gen_iter().take(50_000).collect();
        b.iter(|| {
            let mut v = v.clone();
            v.sort();
            v.iter().cloned().take(10_000).collect::<Vec<_>>();
        });
    }

    #[bench]
    fn take_50_000_lazy(b: &mut Bencher) {
        let mut rng = thread_rng();
        let v: Vec<u32> = rng.gen_iter().take(50_000).collect();
        b.iter(|| v.iter().cloned().lazy_sort().take(50_000).collect::<Vec<_>>());
    }

    #[bench]
    fn take_50_000_eager(b: &mut Bencher) {
        let mut rng = thread_rng();
        let v: Vec<u32> = rng.gen_iter().take(50_000).collect();
        b.iter(|| {
            let mut v = v.clone();
            v.sort();
            v.iter().cloned().take(50_000).collect::<Vec<_>>();
        });
    }
}
<|endoftext|>


<commit_before>// Copyright 2017 Daniel P. Clark & other digits Developers
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.
//
#![deny(missing_docs,trivial_casts,trivial_numeric_casts,
        missing_debug_implementations, missing_copy_implementations,
        unsafe_code,unstable_features,unused_import_braces,unused_qualifications)
]
//! # digits
//!
//! The digits crate is a linked list implementation of a score card flipper.  But
//! in this case it's with any characters you want and you can enumerate through
//! possibilities beyond the numeric limits intrinsic in basic numerc types like `u64`.
//!
//! Primary use case would be brute forcing character sequences.
extern crate base_custom;
use base_custom::BaseCustom;
use std::fmt;

/// # struct Digits
///
/// This struct acts similar to a full number with a custom numeric character base.
/// But the underlying implementation is a linked list where all the methods recurse
/// as far as need to to implement the operations.
#[derive(Clone)]
pub struct Digits<'a> {
  mapping: &'a BaseCustom<char>,
  digit: u64,
  left: Option<Box<Digits<'a>>>,
}

impl<'a> Digits<'a> {
  /// `new` creates a new Digits instance with the provided character set and value.
  ///
  /// The first parameter must be a BaseCustom object which defines and maps all values.
  /// The second parameter is a string value with all valid characters from the BaseCustom set.
  ///
  /// You can view the numbers current value at any time with `to_s()` or `to_string()`.
  pub fn new<S>(mapping: &'a BaseCustom<char>, number: S) -> Digits<'a>
  where S: Into<String> {
    let number = number.into();
    if number.is_empty() { return Digits { mapping: mapping, digit: 0, left: None }; };
    let (last, rest) = {
      let mut n = number.chars().rev();
      (n.next().unwrap(), n.rev().collect::<String>())
    };

    let continuation = {
      if rest.is_empty() {
        None
      } else {
        Some(Box::new(Digits::new(&mapping, rest)))
      }
    };
    Digits {
      mapping: mapping,
      digit: mapping.decimal(last.to_string()),
      left: continuation,
    }
  }

  /// `replicate` — alias for clone (useful for unboxing)
  pub fn replicate(self) -> Self { self.clone() }

  /// `propagate` creates a new number from the same underlying numeric base
  pub fn propagate<S>(&self, number: S) -> Self
  where S: Into<String> {
    Digits::new(self.mapping, number)
  }

  /// Gives the full value of all digits within the linked list as a String.
  pub fn to_s(&self) -> String {
    let num = self.mapping.gen(self.digit);
    match &self.left {
      &None => format!("{}", num),
      &Some(ref bx) => format!("{}{}", bx.to_s(), num),
    }
  }

  /// Gives the full value of all digits within the linked list as a String.
  pub fn to_string(&self) -> String {
    self.to_s()
  }

  fn head_tail(self) -> (u64, Option<Box<Self>>) {
    match self.left {
      Some(bx) => (self.digit, Some(bx)),
      None => (self.digit, None),
    }
  }

  fn set_left(&mut self, d: Digits<'a>) {
    if d.is_end() {
      self.left = None;
    } else {
      self.left = Some(Box::new(d));
    }
  }

  /// `length` returns a `usize` of the total linked list length.
  pub fn length(&self) -> usize {
    match &self.left {
      &None => 1,
      &Some(ref l) => { l.length() + 1 }
    }
  }

  /// `zero` returns a Digits instance with value of zero and the current character mapping.
  pub fn zero(&self) -> Self {
    Digits { mapping: self.mapping, digit: 0, left: None }
  }

  /// `new_zero` returns a Digits instance with value of zero and the current character mapping.
  pub fn new_zero(mapping: &'a BaseCustom<char>) -> Self {
    Digits { mapping: mapping, digit: 0, left: None }
  }

  /// `is_zero` returns bool value of if the number is zero
  pub fn is_zero(&self) -> bool {
    if self.digit != 0 { return false }
    match &self.left {
      &None => { true },
      &Some(ref bx) => { bx.is_zero() },
    }
  }

  /// `pinky` is the smallest digit.
  /// a.k.a. current digit in the linked list.
  /// a.k.a. the right most digit.
  /// This will be a char value for that digit.
  pub fn pinky(&self) -> char {
    self.mapping.char(self.digit as usize).unwrap()
  }

  /// `one` returns a Digits instance with value of one and the current character mapping.
  pub fn one(&self) -> Self {
    Digits { mapping: self.mapping, digit: 1, left: None }
  }

  /// `new_one` returns a Digits instance with value of one and the current character mapping.
  pub fn new_one(mapping: &'a BaseCustom<char>) -> Self {
    Digits { mapping: mapping, digit: 1, left: None }
  }

  fn is_end(&self) -> bool {
    self.digit == 0 && match self.left { None => true, _ => false }
  }

  /// Add two Digits instances together.  The one the `add` method is called on
  /// must be mutable and modifies itself.  The other is consumed.
  ///
  /// Returns a clone of the updated `Self` as well.
  pub fn add(&mut self, other: Self) -> Self {
    assert!(self.mapping == other.mapping);
    if other.is_end() { return self.clone(); };
    let (last, rest) = other.head_tail();

    // sums current single digit
    let added = self.propagate(self.mapping.gen(last + self.digit));
    let (l, r) = added.head_tail();
    self.digit = l;

    // sums for left
    let mut intermediate = Digits::new_zero(self.mapping);
    if let Some(dg) = r { intermediate.add(dg.replicate()); }
    if let Some(dg) = rest { intermediate.add(dg.replicate()); }

    let current_left = match self.left.clone() {
      None => { Box::new(self.zero()) },
      Some(bx) => { bx },
    }.as_ref().clone();

    self.set_left( intermediate.add(current_left).clone() );
    self.clone()
  }
}

/// NODOC
pub trait Into<String> {
  /// NODOC
  fn into(self) -> String;
}

impl<'a> From<Digits<'a>> for String {
  fn from(d: Digits) -> String {
    d.to_s()
  }
}

impl<'a> Into<String> for Digits<'a> {
  fn into(self) -> String {
    self.to_s()
  }
}

impl Into<String> for String {
  fn into(self) -> String {
    self.clone()
  }
}

impl<'a> fmt::Display for Digits<'a> {
  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
   write!(f, "Digits<'a> — (Character: '{}', Decimal Value: {}{})",
     self.mapping.gen(self.digit), self.digit, {
       match self.left {
         None => "".to_string(),
         Some(ref l) => format!(", With Preceeding: '{}'", l.to_s()),
       }
     }
     )
  }
}

impl<'a> fmt::Debug for Digits<'a> {
  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    write!(f, "{} with base: {:?}", self, self.mapping)
  }
}
<commit_msg>minor refactor<commit_after>// Copyright 2017 Daniel P. Clark & other digits Developers
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.
//
#![deny(missing_docs,trivial_casts,trivial_numeric_casts,
        missing_debug_implementations, missing_copy_implementations,
        unsafe_code,unstable_features,unused_import_braces,unused_qualifications)
]
//! # digits
//!
//! The digits crate is a linked list implementation of a score card flipper.  But
//! in this case it's with any characters you want and you can enumerate through
//! possibilities beyond the numeric limits intrinsic in basic numerc types like `u64`.
//!
//! Primary use case would be brute forcing character sequences.
extern crate base_custom;
use base_custom::BaseCustom;
use std::fmt;

/// # struct Digits
///
/// This struct acts similar to a full number with a custom numeric character base.
/// But the underlying implementation is a linked list where all the methods recurse
/// as far as need to to implement the operations.
#[derive(Clone)]
pub struct Digits<'a> {
  mapping: &'a BaseCustom<char>,
  digit: u64,
  left: Option<Box<Digits<'a>>>,
}

impl<'a> Digits<'a> {
  /// `new` creates a new Digits instance with the provided character set and value.
  ///
  /// The first parameter must be a BaseCustom object which defines and maps all values.
  /// The second parameter is a string value with all valid characters from the BaseCustom set.
  ///
  /// You can view the numbers current value at any time with `to_s()` or `to_string()`.
  pub fn new<S>(mapping: &'a BaseCustom<char>, number: S) -> Digits<'a>
  where S: Into<String> {
    let number = number.into();
    if number.is_empty() { return Digits { mapping: mapping, digit: 0, left: None }; };
    let (last, rest) = {
      let mut n = number.chars().rev();
      (n.next().unwrap(), n.rev().collect::<String>())
    };

    let continuation = {
      if rest.is_empty() {
        None
      } else {
        Some(Box::new(Digits::new(&mapping, rest)))
      }
    };
    Digits {
      mapping: mapping,
      digit: mapping.decimal(last.to_string()),
      left: continuation,
    }
  }

  /// `replicate` — alias for clone (useful for unboxing)
  pub fn replicate(self) -> Self { self.clone() }

  /// `propagate` creates a new number from the same underlying numeric base
  pub fn propagate<S>(&self, number: S) -> Self
  where S: Into<String> {
    Digits::new(self.mapping, number)
  }

  /// Gives the full value of all digits within the linked list as a String.
  pub fn to_s(&self) -> String {
    let num = self.mapping.gen(self.digit);
    match &self.left {
      &None => format!("{}", num),
      &Some(ref bx) => format!("{}{}", bx.to_s(), num),
    }
  }

  /// Gives the full value of all digits within the linked list as a String.
  pub fn to_string(&self) -> String {
    self.to_s()
  }

  // the way to recurse and process Digits
  fn head_tail(self) -> (u64, Option<Box<Self>>) {
    match self.left {
      Some(bx) => (self.digit, Some(bx)),
      None => (self.digit, None),
    }
  }

  // logic for setting left linked list continuation
  fn set_left(&mut self, d: Digits<'a>) {
    if d.is_end() {
      self.left = None;
    } else {
      self.left = Some(Box::new(d));
    }
  }

  /// `length` returns a `usize` of the total linked list length.
  pub fn length(&self) -> usize {
    match &self.left {
      &None => 1,
      &Some(ref l) => { l.length() + 1 }
    }
  }

  /// `zero` returns a Digits instance with value of zero and the current character mapping.
  pub fn zero(&self) -> Self {
    Digits::new_zero(self.mapping)
  }

  /// `new_zero` returns a Digits instance with value of zero and the current character mapping.
  pub fn new_zero(mapping: &'a BaseCustom<char>) -> Self {
    Digits { mapping: mapping, digit: 0, left: None }
  }

  /// `is_zero` returns bool value of if the number is zero
  pub fn is_zero(&self) -> bool {
    if self.digit != 0 { return false }
    match &self.left {
      &None => { true },
      &Some(ref bx) => { bx.is_zero() },
    }
  }

  /// `pinky` is the smallest digit.
  /// a.k.a. current digit in the linked list.
  /// a.k.a. the right most digit.
  /// This will be a char value for that digit.
  pub fn pinky(&self) -> char {
    self.mapping.char(self.digit as usize).unwrap()
  }

  /// `one` returns a Digits instance with value of one and the current character mapping.
  pub fn one(&self) -> Self {
    Digits::new_one(self.mapping)
  }

  /// `new_one` returns a Digits instance with value of one and the current character mapping.
  pub fn new_one(mapping: &'a BaseCustom<char>) -> Self {
    Digits { mapping: mapping, digit: 1, left: None }
  }

  // A non-consuming quick end check.
  // More efficient than calling `is_zero` when this applies.
  fn is_end(&self) -> bool {
    self.digit == 0 && match self.left { None => true, _ => false }
  }

  /// Add two Digits instances together.  The one the `add` method is called on
  /// must be mutable and modifies itself.  The other is consumed.
  ///
  /// Returns a clone of the updated `Self` as well.
  pub fn add(&mut self, other: Self) -> Self {
    assert!(self.mapping == other.mapping);
    if other.is_end() { return self.clone(); };
    let (last, rest) = other.head_tail();

    // sums current single digit
    let added = self.propagate(self.mapping.gen(last + self.digit));
    let (l, r) = added.head_tail();
    self.digit = l;

    // sums for left
    let mut intermediate = Digits::new_zero(self.mapping);
    if let Some(dg) = r { intermediate.add(dg.replicate()); }
    if let Some(dg) = rest { intermediate.add(dg.replicate()); }

    let current_left = match self.left.clone() {
      None => { Box::new(self.zero()) },
      Some(bx) => { bx },
    }.as_ref().clone();

    self.set_left( intermediate.add(current_left).clone() );
    self.clone()
  }
}

/// NODOC
pub trait Into<String> {
  /// NODOC
  fn into(self) -> String;
}

impl<'a> From<Digits<'a>> for String {
  fn from(d: Digits) -> String {
    d.to_s()
  }
}

impl<'a> Into<String> for Digits<'a> {
  fn into(self) -> String {
    self.to_s()
  }
}

impl Into<String> for String {
  fn into(self) -> String {
    self.clone()
  }
}

impl<'a> fmt::Display for Digits<'a> {
  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
   write!(f, "Digits<'a> — (Character: '{}', Decimal Value: {}{})",
     self.mapping.gen(self.digit), self.digit, {
       match self.left {
         None => "".to_string(),
         Some(ref l) => format!(", With Preceeding: '{}'", l.to_s()),
       }
     }
     )
  }
}

impl<'a> fmt::Debug for Digits<'a> {
  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    write!(f, "{} with base: {:?}", self, self.mapping)
  }
}
<|endoftext|>


<commit_before>#![feature(core, fs_walk, path_ext)]

//! # Handlebars for Iron
//!
//! This library combines [Handlebars templating library](https://github.com/sunng87/handlebars-rust) and [Iron web framework](http://ironframework.io) together. It gives you a `HandlebarsEngine` as Iron `AfterMiddleware`, so you can render your data with specified handlebars template.
//!
//! ## Setup
//!
//! Given the template root directory (prefix) and template file extension (suffix), you can create `HandlebarsEngine` with `Handlebars::new("/prefix/path", ".hbs")` function. HandlebarsEngine will scan the directory and its sub-directories (with Unix glob **/*), and register these templates with `path/name` as name.
//!
//! ## Usage
//!
//! From any of your handler, you can set template name and data into our `Template` struct. Remember you need to make your data `ToJson`-able, which is required by handlebars-rust.
//!
//! We have implemented Modifier for `Template` on `Response`, so you can just use `response.set` to put set template into response and let it processed by our middleware.
//!
//! Also we made `Response` plugin for `Template` via `HandlebarsEngine`. So you can test your handler from a test case, and retrieve the `Template` you set into it by `response.get::<HandlebarsEngine>`.
//!

extern crate iron;
extern crate "rustc-serialize" as serialize;
extern crate handlebars;
extern crate modifier;
extern crate plugin;
extern crate hyper;

pub use self::middleware::Template;
pub use self::middleware::HandlebarsEngine;

mod middleware;
<commit_msg>(fix) fix dash in crate name<commit_after>#![feature(core, fs_walk, path_ext)]

//! # Handlebars for Iron
//!
//! This library combines [Handlebars templating library](https://github.com/sunng87/handlebars-rust) and [Iron web framework](http://ironframework.io) together. It gives you a `HandlebarsEngine` as Iron `AfterMiddleware`, so you can render your data with specified handlebars template.
//!
//! ## Setup
//!
//! Given the template root directory (prefix) and template file extension (suffix), you can create `HandlebarsEngine` with `Handlebars::new("/prefix/path", ".hbs")` function. HandlebarsEngine will scan the directory and its sub-directories (with Unix glob **/*), and register these templates with `path/name` as name.
//!
//! ## Usage
//!
//! From any of your handler, you can set template name and data into our `Template` struct. Remember you need to make your data `ToJson`-able, which is required by handlebars-rust.
//!
//! We have implemented Modifier for `Template` on `Response`, so you can just use `response.set` to put set template into response and let it processed by our middleware.
//!
//! Also we made `Response` plugin for `Template` via `HandlebarsEngine`. So you can test your handler from a test case, and retrieve the `Template` you set into it by `response.get::<HandlebarsEngine>`.
//!

extern crate iron;
extern crate rustc_serialize as serialize;
extern crate handlebars;
extern crate modifier;
extern crate plugin;
extern crate hyper;

pub use self::middleware::Template;
pub use self::middleware::HandlebarsEngine;

mod middleware;
<|endoftext|>


<commit_before>extern crate regex;

extern crate rusoto_core;
extern crate rusoto_sts;
extern crate rusoto_dynamodb;

pub mod SessionHandler {
    use std::fs::{File, OpenOptions};
    use std::io::BufReader;
    use std::io;
    use std::io::prelude::*;
    use std::io::{stdin, stdout};
    use std::env;
    use rusoto_core::{ProfileProvider, Region, default_tls_client};
    use rusoto_sts::{Sts, StsClient, GetSessionTokenRequest, GetSessionTokenResponse,
        GetSessionTokenError, AssumeRoleRequest, AssumeRoleResponse, AssumeRoleError, Credentials};

    use std::collections::HashMap;

    struct AWSConfig {
        profiles: HashMap<String, AWSProfile>,
    }

    #[derive(Debug)]
    struct AWSProfile {
        source_profile: Option<String>,
        region: Option<String>,
        output: Option<String>,
        role_arn: Option<String>,
        mfa_serial: Option<String>,
        ot_session_name: Option<String>,
        ot_expiration: Option<String>,
        ot_source_profile: Option<String>,
        ot_role_arn: Option<String>,
    }

    impl AWSConfig {
        fn save(&self) -> Result<(), io::Error> {
            println!("Saving config to disk");
            let mut bluenine_config_path = env::home_dir().unwrap().display().to_string();
            bluenine_config_path.push_str("/.aws/config");
            let mut file = OpenOptions::new()
                        .write(true)
                        .truncate(true)
                        .open(bluenine_config_path)
                        .unwrap();

            let mut prof = String::new();
            for (name, profile) in &self.profiles {
                prof.push_str(&format!("[profile {}]\n", name));

                let source_profile = &profile.source_profile;
                match source_profile {
                    &Some(ref source_profile) => {
                        prof.push_str(&format!("source_profile = {}\n", source_profile));
                    },
                    &None => {},
                };

                let region = &profile.region;
                match region {
                    &Some(ref region) => {
                        prof.push_str(&format!("region = {}\n", region));
                    },
                    &None => {},
                };

                let output = &profile.output;
                match output {
                    &Some(ref output) => {
                        prof.push_str(&format!("output = {}\n", output));
                    },
                    &None => {},
                };

                let role_arn = &profile.role_arn;
                match role_arn {
                    &Some(ref role_arn) => {
                        prof.push_str(&format!("role_arn = {}\n", role_arn));
                    },
                    &None => {},
                };

                let mfa_serial = &profile.mfa_serial;
                match mfa_serial {
                    &Some(ref mfa_serial) => {
                        prof.push_str(&format!("mfa_serial = {}\n", mfa_serial));
                    },
                    &None => {},
                };

                let ot_session_name = &profile.ot_session_name;
                match ot_session_name {
                    &Some(ref ot_session_name) => {
                        prof.push_str(&format!("ot_session_name = {}\n", ot_session_name));
                    },
                    &None => {},
                };

                let ot_expiration = &profile.ot_expiration;
                match ot_expiration {
                    &Some(ref ot_expiration) => {
                        prof.push_str(&format!("ot_expiration = {}\n", ot_expiration));
                    },
                    &None => {},
                };

                let ot_expiration = &profile.ot_expiration;
                match ot_expiration {
                    &Some(ref ot_expiration) => {
                        prof.push_str(&format!("ot_expiration = {}\n", ot_expiration));
                    },
                    &None => {},
                };

                let ot_source_profile = &profile.ot_source_profile;
                match ot_source_profile {
                    &Some(ref ot_source_profile) => {
                        prof.push_str(&format!("ot_source_profile = {}\n", ot_source_profile));
                    },
                    &None => {},
                };

                let ot_role_arn = &profile.ot_role_arn;
                match ot_role_arn {
                    &Some(ref ot_role_arn) => {
                        prof.push_str(&format!("ot_role_arn = {}\n", ot_role_arn));
                    },
                    &None => {},
                };

                prof.push_str("\n");
            }

            try!(file.write_all(prof.as_bytes()));
            Ok(())
        }

        fn get_profile(&self, name: &str) -> &AWSProfile {
            if (&self.profiles).contains_key(name) {
                &self.profiles[name]
            }
            else {
                panic!("Could not find profile {}", name);
            }
        }

        fn contains_profile(&self, name: &str) -> bool {
            if (&self.profiles).contains_key(name) {
                true
            }
            else {
                false
            }
        }
    }

    pub fn create(profile_name: &str) {
        println!("Loading config file");
        let aws_config = load_config();

        let mut session_name = profile_name.clone().to_string();
        session_name.push_str("-session");
        if aws_config.contains_profile(&session_name) {
            println!("Session {} for profile {} already exists", session_name, profile_name);
        }

        println!("Creating session for profile \"{}\"...", profile_name);
        let aws_profile = aws_config.get_profile(profile_name);

        // Check if this is a child profile
        let source_profile = &aws_profile.source_profile;
        match source_profile {
            &Some(ref source_profile) => {
                println!("This is a child profile");
                let mut session_name = source_profile.clone().to_string();
                session_name.push_str("-session");
                if aws_config.contains_profile(&session_name) {
                    println!("Session {} for profile {} already exists", session_name, source_profile);
                    let mut provider = ProfileProvider::new().unwrap();
                    provider.set_profile(session_name);
                    let client = StsClient::new(default_tls_client().unwrap(), provider, Region::EuWest1);

                    let mut role_arn = &aws_profile.role_arn;
                    let role_arn_string = role_arn.as_ref().unwrap();
                    let v: Vec<&str> = role_arn_string.split('/').collect();

                    let request = AssumeRoleRequest {
                        duration_seconds: None,
                        external_id: None,
                        policy: None,
                        role_arn: role_arn_string.to_owned(),
                        role_session_name: v[1].to_owned(),
                        serial_number: None,
                        token_code: None,
                    };
                    let response = client.assume_role(&request);
                    match response {
                        Ok(response) => {
                            match save_profile(profile_name, &aws_profile) {
                                Ok(_) => println!("Saved profile to config file."),
                                Err(err) => println!("Error saving profile config to file: {:?}", err)
                            };
                            let credentials = response.credentials.unwrap();
                            match save_credentials(profile_name, credentials) {
                                Ok(_) => println!("Saved to credentials file."),
                                Err(err) => println!("Error saving credentials to file: {:?}", err)
                            };
                        },
                        Err(err) => panic!("Failed to get session token for profile {}: {:?}", profile_name, err),
                    };
                }
                else {
                    println!("Parent profile does not exist. Will have to create it");
                    return;
                }
            },
            &None => {
                let mut provider = ProfileProvider::new().unwrap();
                provider.set_profile(profile_name);
                let client = StsClient::new(default_tls_client().unwrap(), provider, Region::EuWest1);

                let mfa_serial = &aws_profile.mfa_serial;
                if mfa_serial.is_some() {
                    print!("Enter AWS MFA code for profile [{}]: ", profile_name);
                    stdout().flush();
                    let mut token_code = String::new();
                    stdin().read_line(&mut token_code)
                        .ok()
                        .expect("Couldn't read line");
                    token_code.pop(); // Remove newline

                    let request = GetSessionTokenRequest {
                        duration_seconds: None,
                        serial_number: Some(mfa_serial.as_ref().unwrap().to_owned()),
                        token_code: Some(token_code),
                    };
                    let response = client.get_session_token(&request);
                    match response {
                        Ok(response) => {
                            match save_profile(profile_name, &aws_profile) {
                                Ok(_) => println!("Saved profile to config file."),
                                Err(err) => println!("Error saving profile config to file: {:?}", err)
                            };
                            let credentials = response.credentials.unwrap();
                            match save_credentials(profile_name, credentials) {
                                Ok(_) => println!("Saved to credentials file."),
                                Err(err) => println!("Error saving credentials to file: {:?}", err)
                            };
                        },
                        Err(err) => panic!("Failed to get session token for profile {}: {:?}", profile_name, err),
                    };
                }
            },
        };
    }
    pub fn show() {
        println!("Showing bluenine sessions...");
        println!("Loading config file");
        let aws_config = load_config();
        println!("Active sessions:\n");

        for (name, aws_profile) in aws_config.profiles {
            if name.contains("-session") {
                println!("{}", name);
            }
        }
    }

    pub fn clean(profile_name: &str) {
        println!("Loading config file");
        let mut aws_config = load_config();
        let mut session_name = profile_name.clone().to_string();
        session_name.push_str("-session");

        if aws_config.contains_profile(&session_name) {
            println!("Cleaning session {}", session_name);
            aws_config.profiles.remove(&session_name);
            aws_config.save();
            // Remove credentials also
            remove_credentials(&session_name);
        }
        else {
            println!("No profiles to clean");
        }
    }

    fn load_config() -> AWSConfig {
        let aws_config_file = read_aws_config_file();
        let profiles = split_config_file(aws_config_file);

        // Create an instance of the AWSConfig struct
        let mut aws_config = AWSConfig {
            profiles: HashMap::new()
        };

        // Iterate over profile chunks
        for profile in profiles {
            // Create a vector of string that holds each line
            let split = profile.split("\n");
            let lines: Vec<String> = split.map(|s| s.to_string()).collect();

            // Get the profile name from the first line
            let profile_line = lines[0].to_owned();

            // Lets ignore the default for now
            if profile_line == "[default]" {
                continue;
            }

            let split2 = profile_line.split(" ");
            let mut words: Vec<String> = split2.map(|s| s.to_string()).collect();
            let mut profile_name = String::new();
            if words.len() > 1 {
                words[1].pop();
                profile_name = words[1].trim_right().to_owned();
            }

            let mut aws_profile = AWSProfile{
                source_profile: None,
                region: None,
                output: None,
                role_arn: None,
                mfa_serial: None,
                ot_expiration: None,
                ot_session_name: None,
                ot_source_profile: None,
                ot_role_arn: None,
            };

            for i in 1..lines.len() {
                let split = lines[i].split(" = ");
                let config: Vec<String> = split.map(|s| s.to_string()).collect();
                let key = &config[0];
                let value = config[1].to_owned();
                match key.as_ref() {
                    "mfa_serial" => aws_profile.mfa_serial = Some(value),
                    "output" => aws_profile.output = Some(value),
                    "region" => aws_profile.region = Some(value),
                    "role_arn" => aws_profile.role_arn = Some(value),
                    "source_profile" => aws_profile.source_profile = Some(value),
                    "ot_session_name" => aws_profile.ot_session_name = Some(value),
                    "ot_expiration" => aws_profile.ot_expiration = Some(value),
                    "ot_source_profile" => aws_profile.ot_source_profile = Some(value),
                    "ot_role_arn" => aws_profile.ot_role_arn = Some(value),
                    _ => (),
                }
            }

            if profile_name != "" {
                aws_config.profiles.insert(profile_name.clone(), aws_profile);
            }
        }

        aws_config
    }

    fn read_aws_config_file() -> String {
        let path = match env::home_dir() {
            Some(path) => path,
            None => panic!("Could not retrieve user's home directory."),
        };
        let config_file_path = format!("{}/.aws/config", path.display());

        let f = File::open(config_file_path).expect("Could not find AWS config file.");
        let mut buf_reader = BufReader::new(f);
        let mut contents = String::new();
        buf_reader.read_to_string(&mut contents).expect("Found config file but could not read it.");

        contents
    }

    fn read_aws_credentials_file() -> String {
        let path = match env::home_dir() {
            Some(path) => path,
            None => panic!("Could not retrieve user's home directory."),
        };
        let config_file_path = format!("{}/.aws/credentials", path.display());

        let f = File::open(config_file_path).expect("Could not find AWS credentials file.");
        let mut buf_reader = BufReader::new(f);
        let mut contents = String::new();
        buf_reader.read_to_string(&mut contents).expect("Found credentials file but could not read it.");

        contents
    }

    fn split_config_file(aws_config: String) -> Vec<String> {
        let split = aws_config.split("\n\n");
        let mut profiles: Vec<String> = split.map(|s| s.to_string()).collect();;

        profiles.pop(); // Remove last element as it is always empty
        profiles
    }

    fn save_credentials(profile_name: &str, credentials: Credentials) -> Result<(), io::Error> {
        let mut aws_credentials_path = env::home_dir().unwrap().display().to_string();
        aws_credentials_path.push_str("/.aws/credentials");
        let mut file = OpenOptions::new()
                       .write(true)
                       .append(true)
                       .open(aws_credentials_path)
                       .unwrap();

        let mut creds = String::new();
        creds.push_str(&format!("[{}-session]\n", profile_name));
        creds.push_str(&format!("aws_access_key_id = {}\n", credentials.access_key_id));
        creds.push_str(&format!("expiration = {}\n", credentials.expiration));
        creds.push_str(&format!("aws_secret_access_key = {}\n", credentials.secret_access_key));
        creds.push_str(&format!("aws_session_token = {}\n\n", credentials.session_token));
        try!(file.write_all(creds.as_bytes()));
        Ok(())
    }

    fn remove_credentials(profile_name: &str) -> Result<(), io::Error> {
        println!("Saving credentials to disk");
        let aws_credentials_file = read_aws_credentials_file();
        let credentials = split_config_file(aws_credentials_file);
        let mut aws_credentials_path = env::home_dir().unwrap().display().to_string();
        aws_credentials_path.push_str("/.aws/credentials");
        let mut file = OpenOptions::new()
                .write(true)
                .truncate(true)
                .open(aws_credentials_path)
                .unwrap();

        let mut creds = String::new();
        for credential in credentials {
            let split = credential.split("\n");
            let lines: Vec<String> = split.map(|s| s.to_string()).collect();

            // Get the profile name from the first line
            let profile_line = lines[0].to_owned();
            if profile_line == format!("[{}]", profile_name) {
                continue;
            }
            creds.push_str(&credential);
            creds.push_str("\n\n");
        }
        try!(file.write_all(creds.as_bytes()));
        Ok(())
    }

    fn save_profile(profile_name: &str, aws_profile: &AWSProfile) -> Result<(), io::Error> {
        let mut aws_config_path = env::home_dir().unwrap().display().to_string();
        aws_config_path.push_str("/.aws/config");
        let mut file = OpenOptions::new()
                       .write(true)
                       .append(true)
                       .open(aws_config_path)
                       .unwrap();

        let region = &aws_profile.region;
        match region {
            &Some(ref region) => {
                let mut prof = String::new();
                prof.push_str(&format!("[profile {}-session]\n", profile_name));
                prof.push_str(&format!("region = {}\n\n", region));
                try!(file.write_all(prof.as_bytes()));
                Ok(())
            },
            &None => {
                let region = "us-east-1";
                let mut prof = String::new();
                prof.push_str(&format!("[profile {}-session]\n", profile_name));
                prof.push_str(&format!("region = {}\n\n", region));
                try!(file.write_all(prof.as_bytes()));
                Ok(())
            },
        }
    }
}
<commit_msg>Remove unusud profile parameters<commit_after>extern crate regex;

extern crate rusoto_core;
extern crate rusoto_sts;
extern crate rusoto_dynamodb;

pub mod SessionHandler {
    use std::fs::{File, OpenOptions};
    use std::io::BufReader;
    use std::io;
    use std::io::prelude::*;
    use std::io::{stdin, stdout};
    use std::env;
    use rusoto_core::{ProfileProvider, Region, default_tls_client};
    use rusoto_sts::{Sts, StsClient, GetSessionTokenRequest, GetSessionTokenResponse,
        GetSessionTokenError, AssumeRoleRequest, AssumeRoleResponse, AssumeRoleError, Credentials};

    use std::collections::HashMap;

    struct AWSConfig {
        profiles: HashMap<String, AWSProfile>,
    }

    #[derive(Debug)]
    struct AWSProfile {
        source_profile: Option<String>,
        region: Option<String>,
        output: Option<String>,
        role_arn: Option<String>,
        mfa_serial: Option<String>,
    }

    impl AWSConfig {
        fn save(&self) -> Result<(), io::Error> {
            println!("Saving config to disk");
            let mut bluenine_config_path = env::home_dir().unwrap().display().to_string();
            bluenine_config_path.push_str("/.aws/config");
            let mut file = OpenOptions::new()
                        .write(true)
                        .truncate(true)
                        .open(bluenine_config_path)
                        .unwrap();

            let mut prof = String::new();
            for (name, profile) in &self.profiles {
                prof.push_str(&format!("[profile {}]\n", name));

                let source_profile = &profile.source_profile;
                match source_profile {
                    &Some(ref source_profile) => {
                        prof.push_str(&format!("source_profile = {}\n", source_profile));
                    },
                    &None => {},
                };

                let region = &profile.region;
                match region {
                    &Some(ref region) => {
                        prof.push_str(&format!("region = {}\n", region));
                    },
                    &None => {},
                };

                let output = &profile.output;
                match output {
                    &Some(ref output) => {
                        prof.push_str(&format!("output = {}\n", output));
                    },
                    &None => {},
                };

                let role_arn = &profile.role_arn;
                match role_arn {
                    &Some(ref role_arn) => {
                        prof.push_str(&format!("role_arn = {}\n", role_arn));
                    },
                    &None => {},
                };

                let mfa_serial = &profile.mfa_serial;
                match mfa_serial {
                    &Some(ref mfa_serial) => {
                        prof.push_str(&format!("mfa_serial = {}\n", mfa_serial));
                    },
                    &None => {},
                };

                prof.push_str("\n");
            }

            try!(file.write_all(prof.as_bytes()));
            Ok(())
        }

        fn get_profile(&self, name: &str) -> &AWSProfile {
            if (&self.profiles).contains_key(name) {
                &self.profiles[name]
            }
            else {
                panic!("Could not find profile {}", name);
            }
        }

        fn contains_profile(&self, name: &str) -> bool {
            if (&self.profiles).contains_key(name) {
                true
            }
            else {
                false
            }
        }
    }

    pub fn create(profile_name: &str) {
        println!("Loading config file");
        let aws_config = load_config();

        let mut session_name = profile_name.clone().to_string();
        session_name.push_str("-session");
        if aws_config.contains_profile(&session_name) {
            println!("Session {} for profile {} already exists", session_name, profile_name);
        }

        println!("Creating session for profile \"{}\"...", profile_name);
        let aws_profile = aws_config.get_profile(profile_name);

        // Check if this is a child profile
        let source_profile = &aws_profile.source_profile;
        match source_profile {
            &Some(ref source_profile) => {
                println!("This is a child profile");
                let mut session_name = source_profile.clone().to_string();
                session_name.push_str("-session");
                if aws_config.contains_profile(&session_name) {
                    println!("Session {} for profile {} already exists", session_name, source_profile);
                    let mut provider = ProfileProvider::new().unwrap();
                    provider.set_profile(session_name);
                    let client = StsClient::new(default_tls_client().unwrap(), provider, Region::EuWest1);

                    let mut role_arn = &aws_profile.role_arn;
                    let role_arn_string = role_arn.as_ref().unwrap();
                    let v: Vec<&str> = role_arn_string.split('/').collect();

                    let request = AssumeRoleRequest {
                        duration_seconds: None,
                        external_id: None,
                        policy: None,
                        role_arn: role_arn_string.to_owned(),
                        role_session_name: v[1].to_owned(),
                        serial_number: None,
                        token_code: None,
                    };
                    let response = client.assume_role(&request);
                    match response {
                        Ok(response) => {
                            match save_profile(profile_name, &aws_profile) {
                                Ok(_) => println!("Saved profile to config file."),
                                Err(err) => println!("Error saving profile config to file: {:?}", err)
                            };
                            let credentials = response.credentials.unwrap();
                            match save_credentials(profile_name, credentials) {
                                Ok(_) => println!("Saved to credentials file."),
                                Err(err) => println!("Error saving credentials to file: {:?}", err)
                            };
                        },
                        Err(err) => panic!("Failed to get session token for profile {}: {:?}", profile_name, err),
                    };
                }
                else {
                    println!("Parent profile does not exist. Will have to create it");
                    return;
                }
            },
            &None => {
                let mut provider = ProfileProvider::new().unwrap();
                provider.set_profile(profile_name);
                let client = StsClient::new(default_tls_client().unwrap(), provider, Region::EuWest1);

                let mfa_serial = &aws_profile.mfa_serial;
                if mfa_serial.is_some() {
                    print!("Enter AWS MFA code for profile [{}]: ", profile_name);
                    stdout().flush();
                    let mut token_code = String::new();
                    stdin().read_line(&mut token_code)
                        .ok()
                        .expect("Couldn't read line");
                    token_code.pop(); // Remove newline

                    let request = GetSessionTokenRequest {
                        duration_seconds: None,
                        serial_number: Some(mfa_serial.as_ref().unwrap().to_owned()),
                        token_code: Some(token_code),
                    };
                    let response = client.get_session_token(&request);
                    match response {
                        Ok(response) => {
                            match save_profile(profile_name, &aws_profile) {
                                Ok(_) => println!("Saved profile to config file."),
                                Err(err) => println!("Error saving profile config to file: {:?}", err)
                            };
                            let credentials = response.credentials.unwrap();
                            match save_credentials(profile_name, credentials) {
                                Ok(_) => println!("Saved to credentials file."),
                                Err(err) => println!("Error saving credentials to file: {:?}", err)
                            };
                        },
                        Err(err) => panic!("Failed to get session token for profile {}: {:?}", profile_name, err),
                    };
                }
            },
        };
    }
    pub fn show() {
        println!("Showing bluenine sessions...");
        println!("Loading config file");
        let aws_config = load_config();
        println!("Active sessions:\n");

        for (name, aws_profile) in aws_config.profiles {
            if name.contains("-session") {
                println!("{}", name);
            }
        }
    }

    pub fn clean(profile_name: &str) {
        println!("Loading config file");
        let mut aws_config = load_config();
        let mut session_name = profile_name.clone().to_string();
        session_name.push_str("-session");

        if aws_config.contains_profile(&session_name) {
            println!("Cleaning session {}", session_name);
            aws_config.profiles.remove(&session_name);
            aws_config.save();
            // Remove credentials also
            remove_credentials(&session_name);
        }
        else {
            println!("No profiles to clean");
        }
    }

    fn load_config() -> AWSConfig {
        let aws_config_file = read_aws_config_file();
        let profiles = split_config_file(aws_config_file);

        // Create an instance of the AWSConfig struct
        let mut aws_config = AWSConfig {
            profiles: HashMap::new()
        };

        // Iterate over profile chunks
        for profile in profiles {
            // Create a vector of string that holds each line
            let split = profile.split("\n");
            let lines: Vec<String> = split.map(|s| s.to_string()).collect();

            // Get the profile name from the first line
            let profile_line = lines[0].to_owned();

            // Lets ignore the default for now
            if profile_line == "[default]" {
                continue;
            }

            let split2 = profile_line.split(" ");
            let mut words: Vec<String> = split2.map(|s| s.to_string()).collect();
            let mut profile_name = String::new();
            if words.len() > 1 {
                words[1].pop();
                profile_name = words[1].trim_right().to_owned();
            }

            let mut aws_profile = AWSProfile{
                source_profile: None,
                region: None,
                output: None,
                role_arn: None,
                mfa_serial: None,
            };

            for i in 1..lines.len() {
                let split = lines[i].split(" = ");
                let config: Vec<String> = split.map(|s| s.to_string()).collect();
                let key = &config[0];
                let value = config[1].to_owned();
                match key.as_ref() {
                    "mfa_serial" => aws_profile.mfa_serial = Some(value),
                    "output" => aws_profile.output = Some(value),
                    "region" => aws_profile.region = Some(value),
                    "role_arn" => aws_profile.role_arn = Some(value),
                    "source_profile" => aws_profile.source_profile = Some(value),
                    _ => (),
                }
            }

            if profile_name != "" {
                aws_config.profiles.insert(profile_name.clone(), aws_profile);
            }
        }

        aws_config
    }

    fn read_aws_config_file() -> String {
        let path = match env::home_dir() {
            Some(path) => path,
            None => panic!("Could not retrieve user's home directory."),
        };
        let config_file_path = format!("{}/.aws/config", path.display());

        let f = File::open(config_file_path).expect("Could not find AWS config file.");
        let mut buf_reader = BufReader::new(f);
        let mut contents = String::new();
        buf_reader.read_to_string(&mut contents).expect("Found config file but could not read it.");

        contents
    }

    fn read_aws_credentials_file() -> String {
        let path = match env::home_dir() {
            Some(path) => path,
            None => panic!("Could not retrieve user's home directory."),
        };
        let config_file_path = format!("{}/.aws/credentials", path.display());

        let f = File::open(config_file_path).expect("Could not find AWS credentials file.");
        let mut buf_reader = BufReader::new(f);
        let mut contents = String::new();
        buf_reader.read_to_string(&mut contents).expect("Found credentials file but could not read it.");

        contents
    }

    fn split_config_file(aws_config: String) -> Vec<String> {
        let split = aws_config.split("\n\n");
        let mut profiles: Vec<String> = split.map(|s| s.to_string()).collect();;

        profiles.pop(); // Remove last element as it is always empty
        profiles
    }

    fn save_credentials(profile_name: &str, credentials: Credentials) -> Result<(), io::Error> {
        let mut aws_credentials_path = env::home_dir().unwrap().display().to_string();
        aws_credentials_path.push_str("/.aws/credentials");
        let mut file = OpenOptions::new()
                       .write(true)
                       .append(true)
                       .open(aws_credentials_path)
                       .unwrap();

        let mut creds = String::new();
        creds.push_str(&format!("[{}-session]\n", profile_name));
        creds.push_str(&format!("aws_access_key_id = {}\n", credentials.access_key_id));
        creds.push_str(&format!("expiration = {}\n", credentials.expiration));
        creds.push_str(&format!("aws_secret_access_key = {}\n", credentials.secret_access_key));
        creds.push_str(&format!("aws_session_token = {}\n\n", credentials.session_token));
        try!(file.write_all(creds.as_bytes()));
        Ok(())
    }

    fn remove_credentials(profile_name: &str) -> Result<(), io::Error> {
        println!("Saving credentials to disk");
        let aws_credentials_file = read_aws_credentials_file();
        let credentials = split_config_file(aws_credentials_file);
        let mut aws_credentials_path = env::home_dir().unwrap().display().to_string();
        aws_credentials_path.push_str("/.aws/credentials");
        let mut file = OpenOptions::new()
                .write(true)
                .truncate(true)
                .open(aws_credentials_path)
                .unwrap();

        let mut creds = String::new();
        for credential in credentials {
            let split = credential.split("\n");
            let lines: Vec<String> = split.map(|s| s.to_string()).collect();

            // Get the profile name from the first line
            let profile_line = lines[0].to_owned();
            if profile_line == format!("[{}]", profile_name) {
                continue;
            }
            creds.push_str(&credential);
            creds.push_str("\n\n");
        }
        try!(file.write_all(creds.as_bytes()));
        Ok(())
    }

    fn save_profile(profile_name: &str, aws_profile: &AWSProfile) -> Result<(), io::Error> {
        let mut aws_config_path = env::home_dir().unwrap().display().to_string();
        aws_config_path.push_str("/.aws/config");
        let mut file = OpenOptions::new()
                       .write(true)
                       .append(true)
                       .open(aws_config_path)
                       .unwrap();

        let region = &aws_profile.region;
        match region {
            &Some(ref region) => {
                let mut prof = String::new();
                prof.push_str(&format!("[profile {}-session]\n", profile_name));
                prof.push_str(&format!("region = {}\n\n", region));
                try!(file.write_all(prof.as_bytes()));
                Ok(())
            },
            &None => {
                let region = "us-east-1";
                let mut prof = String::new();
                prof.push_str(&format!("[profile {}-session]\n", profile_name));
                prof.push_str(&format!("region = {}\n\n", region));
                try!(file.write_all(prof.as_bytes()));
                Ok(())
            },
        }
    }
}
<|endoftext|>


<commit_before>// Module for parsing ISO Base Media Format aka video/mp4 streams.

// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at https://mozilla.org/MPL/2.0/.

/// Basic ISO box structure.
pub struct BoxHeader {
    /// Four character box type
pub name: u32,
    /// Size of the box in bytes
pub size: u64,
    /// Offset to the start of the contained data (or header size).
pub offset: u64,
}

/// File type box 'ftyp'.
pub struct FileTypeBox {
    name: u32,
    size: u64,
    major_brand: u32,
    minor_version: u32,
    compatible_brands: Vec<u32>,
}

/// Movie header box 'mvhd'.
pub struct MovieHeaderBox {
    pub name: u32,
    pub size: u64,
    pub timescale: u32,
    pub duration: u64,
    // Ignore other fields.
}

pub struct TrackHeaderBox {
    pub name: u32,
    pub size: u64,
    pub track_id: u32,
    pub duration: u64,
    pub width: u32,
    pub height: u32,
}

extern crate byteorder;
use byteorder::{BigEndian, ReadBytesExt};
use std::io::{Read, Result, Seek, SeekFrom, Take};
use std::io::Cursor;

/// Parse a box out of a data buffer.
pub fn read_box_header<T: ReadBytesExt>(src: &mut T) -> Result<BoxHeader> {
    let tmp_size = try!(src.read_u32::<BigEndian>());
    let name = try!(src.read_u32::<BigEndian>());
    let size = match tmp_size {
        1 => try!(src.read_u64::<BigEndian>()),
        _ => tmp_size as u64,
    };
    assert!(size >= 8);
    if tmp_size == 1 {
        assert!(size >= 16);
    }
    let offset = match tmp_size {
        1 => 4 + 4 + 8,
        _ => 4 + 4,
    };
    assert!(offset <= size);
    Ok(BoxHeader{
      name: name,
      size: size,
      offset: offset,
    })
}

/// Parse the extra header fields for a full box.
fn read_fullbox_extra<T: ReadBytesExt>(src: &mut T) -> (u8, u32) {
    let version = src.read_u8().unwrap();
    let flags_a = src.read_u8().unwrap();
    let flags_b = src.read_u8().unwrap();
    let flags_c = src.read_u8().unwrap();
    (version, (flags_a as u32) << 16 |
              (flags_b as u32) <<  8 |
              (flags_c as u32))
}

/// Skip over the contents of a box.
pub fn skip_box_content<T: ReadBytesExt + Seek>
  (src: &mut T, header: &BoxHeader)
  -> std::io::Result<u64>
{
    src.seek(SeekFrom::Current((header.size - header.offset) as i64))
}

/// Helper to construct a Take over the contents of a box.
fn limit<'a, T: Read>(f: &'a mut T, h: &BoxHeader) -> Take<&'a mut T> {
    f.take(h.size - h.offset)
}

/// Helper to construct a Cursor over the contents of a box.
fn recurse<T: Read>(f: &mut T, h: &BoxHeader) {
    println!("{} -- recursing", h);
    // FIXME: I couldn't figure out how to do this without copying.
    // We use Seek on the Read we return in skip_box_content, but
    // that trait isn't implemented for a Take like our limit()
    // returns. Slurping the buffer and wrapping it in a Cursor
    // functions as a work around.
    let buf: Vec<u8> = limit(f, &h)
        .bytes()
        .map(|u| u.unwrap())
        .collect();
    let mut content = Cursor::new(buf);
    loop {
        read_box(&mut content).unwrap();
    }
    println!("{} -- end", h);
}

/// Read the contents of a box, including sub boxes.
/// Right now it just prints the box value rather than
/// returning anything.
pub fn read_box<T: Read + Seek>(f: &mut T) -> Result<()> {
    read_box_header(f).and_then(|h| {
        match &(fourcc_to_string(h.name))[..] {
            "ftyp" => {
                let mut content = limit(f, &h);
                let ftyp = read_ftyp(&mut content, &h).unwrap();
                println!("{}", ftyp);
            },
            "moov" => recurse(f, &h),
            "mvhd" => {
                let mut content = limit(f, &h);
                let mvhd = read_mvhd(&mut content, &h).unwrap();
                println!("  {}", mvhd);
            },
            "trak" => recurse(f, &h),
            "tkhd" => {
                let mut content = limit(f, &h);
                let tkhd = read_tkhd(&mut content, &h).unwrap();
                println!("  {}", tkhd);
            },
            _ => {
                // Skip the contents of unknown chunks.
                println!("{} (skipped)", h);
                skip_box_content(f, &h).unwrap();
            },
        };
        Ok(()) // and_then needs a Result to return.
    })
}

/// Entry point for C language callers.
/// Take a ptr, size and return a cursor over it.
fn cursor_from_cbuf<'a>(buffer: *const u8, size: usize) -> Cursor<&'a[u8]> {
    use std::slice;
    unsafe {
        let b = slice::from_raw_parts(buffer, size);
        Cursor::new(b)
    }
}

/// Take a buffer and call read_box() on it.
#[no_mangle]
pub unsafe extern fn read_box_from_buffer(buffer: *mut u8, size: usize) -> bool {
    use std::thread;

    // Validate arguments from C.
    if buffer.is_null() || size < 8 {
        return false;
    }

    // Parse in a subthread.
    let mut c = cursor_from_cbuf(buffer, size);
    let task = thread::spawn(move || {
        read_box(&mut c).unwrap();
    });
    // Catch any panics.
    task.join().is_ok()
}

/// Parse an ftype box.
pub fn read_ftyp<T: ReadBytesExt>(src: &mut T, head: &BoxHeader)
  -> Option<FileTypeBox> {
    let major = src.read_u32::<BigEndian>().unwrap();
    let minor = src.read_u32::<BigEndian>().unwrap();
    let brand_count = (head.size - 8 - 8) /4;
    let mut brands = Vec::new();
    for _ in 0..brand_count {
        brands.push(src.read_u32::<BigEndian>().unwrap());
    }
    Some(FileTypeBox{
        name: head.name,
        size: head.size,
        major_brand: major,
        minor_version: minor,
        compatible_brands: brands,
    })
}

/// Parse an mvhd box.
pub fn read_mvhd<T: ReadBytesExt>(src: &mut T, head: &BoxHeader)
  -> Option<MovieHeaderBox> {
    let (version, _) = read_fullbox_extra(src);
    match version {
        1 => {
            // 64 bit creation and modification times.
            let mut skip: Vec<u8> = vec![0; 16];
            let r = src.read(&mut skip).unwrap();
            assert!(r == skip.len());
        },
        0 => {
            // 32 bit creation and modification times.
            // 64 bit creation and modification times.
            let mut skip: Vec<u8> = vec![0; 8];
            let r = src.read(&mut skip).unwrap();
            assert!(r == skip.len());
        },
        _ => panic!("invalid mhdr version"),
    }
    let timescale = src.read_u32::<BigEndian>().unwrap();
    let duration = match version {
        1 => src.read_u64::<BigEndian>().unwrap(),
        0 => src.read_u32::<BigEndian>().unwrap() as u64,
        _ => panic!("invalid mhdr version"),
    };
    // Skip remaining fields.
    let mut skip: Vec<u8> = vec![0; 80];
    let r = src.read(&mut skip).unwrap();
    assert!(r == skip.len());
    Some(MovieHeaderBox {
        name: head.name,
        size: head.size,
        timescale: timescale,
        duration: duration,
    })
}

/// Parse a tkhd box.
pub fn read_tkhd<T: ReadBytesExt>(src: &mut T, head: &BoxHeader)
  -> Option<TrackHeaderBox> {
    let (version, flags) = read_fullbox_extra(src);
    if flags & 0x1u32 == 0 || flags & 0x2u32 == 0 {
        // Track is disabled.
        return None;
    }
    match version {
        1 => {
            // 64 bit creation and modification times.
            let mut skip: Vec<u8> = vec![0; 16];
            let r = src.read(&mut skip).unwrap();
            assert!(r == skip.len());
        },
        0 => {
            // 32 bit creation and modification times.
            // 64 bit creation and modification times.
            let mut skip: Vec<u8> = vec![0; 8];
            let r = src.read(&mut skip).unwrap();
            assert!(r == skip.len());
        },
        _ => panic!("invalid tkhd version"),
    }
    let track_id = src.read_u32::<BigEndian>().unwrap();
    let _reserved = src.read_u32::<BigEndian>().unwrap();
    assert!(_reserved == 0);
    let duration = match version {
        1 => {
            src.read_u64::<BigEndian>().unwrap()
        },
        0 => src.read_u32::<BigEndian>().unwrap() as u64,
        _ => panic!("invalid tkhd version"),
    };
    let _reserved = src.read_u32::<BigEndian>().unwrap();
    let _reserved = src.read_u32::<BigEndian>().unwrap();
    // Skip uninterested fields.
    let mut skip: Vec<u8> = vec![0; 44];
    let r = src.read(&mut skip).unwrap();
    assert!(r == skip.len());
    let width = src.read_u32::<BigEndian>().unwrap();
    let height = src.read_u32::<BigEndian>().unwrap();
    Some(TrackHeaderBox {
        name: head.name,
        size: head.size,
        track_id: track_id,
        duration: duration,
        width: width,
        height: height,
    })
}

/// Convert the iso box type or other 4-character value to a string.
fn fourcc_to_string(name: u32) -> String {
    let u32_to_vec = |u| {
        vec!((u >> 24 & 0xffu32) as u8,
             (u >> 16 & 0xffu32) as u8,
             (u >>  8 & 0xffu32) as u8,
             (u & 0xffu32) as u8)
    };
    let name_bytes = u32_to_vec(name);
    String::from_utf8_lossy(&name_bytes).into_owned()
}

use std::fmt;
impl fmt::Display for BoxHeader {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "'{}' {} bytes", fourcc_to_string(self.name), self.size)
    }
}

impl fmt::Display for FileTypeBox {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let name = fourcc_to_string(self.name);
        let brand = fourcc_to_string(self.major_brand);
        write!(f, "'{}' {} bytes '{}' v{}", name, self.size,
            brand, self.minor_version)
    }
}

impl fmt::Display for MovieHeaderBox {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let name = fourcc_to_string(self.name);
        write!(f, "'{}' {} bytes duration {}s", name, self.size,
            (self.duration as f64)/(self.timescale as f64))
    }
}

use std::u16;
impl fmt::Display for TrackHeaderBox {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let name = fourcc_to_string(self.name);
        // Dimensions are 16.16 fixed-point.
        let base = u16::MAX as f64 + 1.0;
        let width = (self.width as f64) / base;
        let height = (self.height as f64) / base;
        write!(f, "'{}' {} bytes duration {} id {} {}x{}",
            name, self.size, self.duration, self.track_id,
            width, height)
    }
}

#[test]
fn test_read_box_header() {
    use std::io::Cursor;
    use std::io::Write;
    let mut test: Vec<u8> = vec![0, 0, 0, 8];  // minimal box length
    write!(&mut test, "test").unwrap(); // box type
    let mut stream = Cursor::new(test);
    let parsed = read_box_header(&mut stream).unwrap();
    assert_eq!(parsed.name, 1952805748);
    assert_eq!(parsed.size, 8);
    println!("box {}", parsed);
}

#[test]
fn test_read_box_header_long() {
    use std::io::Cursor;
    let mut test: Vec<u8> = vec![0, 0, 0, 1]; // long box extension code
    test.extend("long".to_string().into_bytes()); // box type
    test.extend(vec![0, 0, 0, 0, 0, 0, 16, 0]); // 64 bit size
    // Skip generating box content.
    let mut stream = Cursor::new(test);
    let parsed = read_box_header(&mut stream).unwrap();
    assert_eq!(parsed.name, 1819242087);
    assert_eq!(parsed.size, 4096);
    println!("box {}", parsed);
}

#[test]
fn test_read_ftyp() {
    use std::io::Cursor;
    use std::io::Write;
    let mut test: Vec<u8> = vec![0, 0, 0, 24]; // size
    write!(&mut test, "ftyp").unwrap(); // type
    write!(&mut test, "mp42").unwrap(); // major brand
    test.extend(vec![0, 0, 0, 0]);      // minor version
    write!(&mut test, "isom").unwrap(); // compatible brands...
    write!(&mut test, "mp42").unwrap();
    assert_eq!(test.len(), 24);

    let mut stream = Cursor::new(test);
    let header = read_box_header(&mut stream).unwrap();
    let parsed = read_ftyp(&mut stream, &header).unwrap();
    assert_eq!(parsed.name, 1718909296);
    assert_eq!(parsed.size, 24);
    assert_eq!(parsed.major_brand, 1836069938);
    assert_eq!(parsed.minor_version, 0);
    assert_eq!(parsed.compatible_brands.len(), 2);
    println!("box {}", parsed);
}
<commit_msg>Move slice construction into the main C entry point.<commit_after>// Module for parsing ISO Base Media Format aka video/mp4 streams.

// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at https://mozilla.org/MPL/2.0/.

/// Basic ISO box structure.
pub struct BoxHeader {
    /// Four character box type
pub name: u32,
    /// Size of the box in bytes
pub size: u64,
    /// Offset to the start of the contained data (or header size).
pub offset: u64,
}

/// File type box 'ftyp'.
pub struct FileTypeBox {
    name: u32,
    size: u64,
    major_brand: u32,
    minor_version: u32,
    compatible_brands: Vec<u32>,
}

/// Movie header box 'mvhd'.
pub struct MovieHeaderBox {
    pub name: u32,
    pub size: u64,
    pub timescale: u32,
    pub duration: u64,
    // Ignore other fields.
}

pub struct TrackHeaderBox {
    pub name: u32,
    pub size: u64,
    pub track_id: u32,
    pub duration: u64,
    pub width: u32,
    pub height: u32,
}

extern crate byteorder;
use byteorder::{BigEndian, ReadBytesExt};
use std::io::{Read, Result, Seek, SeekFrom, Take};
use std::io::Cursor;

/// Parse a box out of a data buffer.
pub fn read_box_header<T: ReadBytesExt>(src: &mut T) -> Result<BoxHeader> {
    let tmp_size = try!(src.read_u32::<BigEndian>());
    let name = try!(src.read_u32::<BigEndian>());
    let size = match tmp_size {
        1 => try!(src.read_u64::<BigEndian>()),
        _ => tmp_size as u64,
    };
    assert!(size >= 8);
    if tmp_size == 1 {
        assert!(size >= 16);
    }
    let offset = match tmp_size {
        1 => 4 + 4 + 8,
        _ => 4 + 4,
    };
    assert!(offset <= size);
    Ok(BoxHeader{
      name: name,
      size: size,
      offset: offset,
    })
}

/// Parse the extra header fields for a full box.
fn read_fullbox_extra<T: ReadBytesExt>(src: &mut T) -> (u8, u32) {
    let version = src.read_u8().unwrap();
    let flags_a = src.read_u8().unwrap();
    let flags_b = src.read_u8().unwrap();
    let flags_c = src.read_u8().unwrap();
    (version, (flags_a as u32) << 16 |
              (flags_b as u32) <<  8 |
              (flags_c as u32))
}

/// Skip over the contents of a box.
pub fn skip_box_content<T: ReadBytesExt + Seek>
  (src: &mut T, header: &BoxHeader)
  -> std::io::Result<u64>
{
    src.seek(SeekFrom::Current((header.size - header.offset) as i64))
}

/// Helper to construct a Take over the contents of a box.
fn limit<'a, T: Read>(f: &'a mut T, h: &BoxHeader) -> Take<&'a mut T> {
    f.take(h.size - h.offset)
}

/// Helper to construct a Cursor over the contents of a box.
fn recurse<T: Read>(f: &mut T, h: &BoxHeader) {
    println!("{} -- recursing", h);
    // FIXME: I couldn't figure out how to do this without copying.
    // We use Seek on the Read we return in skip_box_content, but
    // that trait isn't implemented for a Take like our limit()
    // returns. Slurping the buffer and wrapping it in a Cursor
    // functions as a work around.
    let buf: Vec<u8> = limit(f, &h)
        .bytes()
        .map(|u| u.unwrap())
        .collect();
    let mut content = Cursor::new(buf);
    loop {
        read_box(&mut content).unwrap();
    }
    println!("{} -- end", h);
}

/// Read the contents of a box, including sub boxes.
/// Right now it just prints the box value rather than
/// returning anything.
pub fn read_box<T: Read + Seek>(f: &mut T) -> Result<()> {
    read_box_header(f).and_then(|h| {
        match &(fourcc_to_string(h.name))[..] {
            "ftyp" => {
                let mut content = limit(f, &h);
                let ftyp = read_ftyp(&mut content, &h).unwrap();
                println!("{}", ftyp);
            },
            "moov" => recurse(f, &h),
            "mvhd" => {
                let mut content = limit(f, &h);
                let mvhd = read_mvhd(&mut content, &h).unwrap();
                println!("  {}", mvhd);
            },
            "trak" => recurse(f, &h),
            "tkhd" => {
                let mut content = limit(f, &h);
                let tkhd = read_tkhd(&mut content, &h).unwrap();
                println!("  {}", tkhd);
            },
            _ => {
                // Skip the contents of unknown chunks.
                println!("{} (skipped)", h);
                skip_box_content(f, &h).unwrap();
            },
        };
        Ok(()) // and_then needs a Result to return.
    })
}

/// Entry point for C language callers.
/// Take a buffer and call read_box() on it.
#[no_mangle]
pub unsafe extern fn read_box_from_buffer(buffer: *const u8, size: usize) -> bool {
    use std::slice;
    use std::thread;

    // Validate arguments from C.
    if buffer.is_null() || size < 8 {
        return false;
    }

    // Wrap the buffer we've been give in a slice.
    let b = slice::from_raw_parts(buffer, size);
    let mut c = Cursor::new(b);

    // Parse in a subthread.
    let task = thread::spawn(move || {
        read_box(&mut c).unwrap();
    });
    // Catch any panics.
    task.join().is_ok()
}

/// Parse an ftype box.
pub fn read_ftyp<T: ReadBytesExt>(src: &mut T, head: &BoxHeader)
  -> Option<FileTypeBox> {
    let major = src.read_u32::<BigEndian>().unwrap();
    let minor = src.read_u32::<BigEndian>().unwrap();
    let brand_count = (head.size - 8 - 8) /4;
    let mut brands = Vec::new();
    for _ in 0..brand_count {
        brands.push(src.read_u32::<BigEndian>().unwrap());
    }
    Some(FileTypeBox{
        name: head.name,
        size: head.size,
        major_brand: major,
        minor_version: minor,
        compatible_brands: brands,
    })
}

/// Parse an mvhd box.
pub fn read_mvhd<T: ReadBytesExt>(src: &mut T, head: &BoxHeader)
  -> Option<MovieHeaderBox> {
    let (version, _) = read_fullbox_extra(src);
    match version {
        1 => {
            // 64 bit creation and modification times.
            let mut skip: Vec<u8> = vec![0; 16];
            let r = src.read(&mut skip).unwrap();
            assert!(r == skip.len());
        },
        0 => {
            // 32 bit creation and modification times.
            // 64 bit creation and modification times.
            let mut skip: Vec<u8> = vec![0; 8];
            let r = src.read(&mut skip).unwrap();
            assert!(r == skip.len());
        },
        _ => panic!("invalid mhdr version"),
    }
    let timescale = src.read_u32::<BigEndian>().unwrap();
    let duration = match version {
        1 => src.read_u64::<BigEndian>().unwrap(),
        0 => src.read_u32::<BigEndian>().unwrap() as u64,
        _ => panic!("invalid mhdr version"),
    };
    // Skip remaining fields.
    let mut skip: Vec<u8> = vec![0; 80];
    let r = src.read(&mut skip).unwrap();
    assert!(r == skip.len());
    Some(MovieHeaderBox {
        name: head.name,
        size: head.size,
        timescale: timescale,
        duration: duration,
    })
}

/// Parse a tkhd box.
pub fn read_tkhd<T: ReadBytesExt>(src: &mut T, head: &BoxHeader)
  -> Option<TrackHeaderBox> {
    let (version, flags) = read_fullbox_extra(src);
    if flags & 0x1u32 == 0 || flags & 0x2u32 == 0 {
        // Track is disabled.
        return None;
    }
    match version {
        1 => {
            // 64 bit creation and modification times.
            let mut skip: Vec<u8> = vec![0; 16];
            let r = src.read(&mut skip).unwrap();
            assert!(r == skip.len());
        },
        0 => {
            // 32 bit creation and modification times.
            // 64 bit creation and modification times.
            let mut skip: Vec<u8> = vec![0; 8];
            let r = src.read(&mut skip).unwrap();
            assert!(r == skip.len());
        },
        _ => panic!("invalid tkhd version"),
    }
    let track_id = src.read_u32::<BigEndian>().unwrap();
    let _reserved = src.read_u32::<BigEndian>().unwrap();
    assert!(_reserved == 0);
    let duration = match version {
        1 => {
            src.read_u64::<BigEndian>().unwrap()
        },
        0 => src.read_u32::<BigEndian>().unwrap() as u64,
        _ => panic!("invalid tkhd version"),
    };
    let _reserved = src.read_u32::<BigEndian>().unwrap();
    let _reserved = src.read_u32::<BigEndian>().unwrap();
    // Skip uninterested fields.
    let mut skip: Vec<u8> = vec![0; 44];
    let r = src.read(&mut skip).unwrap();
    assert!(r == skip.len());
    let width = src.read_u32::<BigEndian>().unwrap();
    let height = src.read_u32::<BigEndian>().unwrap();
    Some(TrackHeaderBox {
        name: head.name,
        size: head.size,
        track_id: track_id,
        duration: duration,
        width: width,
        height: height,
    })
}

/// Convert the iso box type or other 4-character value to a string.
fn fourcc_to_string(name: u32) -> String {
    let u32_to_vec = |u| {
        vec!((u >> 24 & 0xffu32) as u8,
             (u >> 16 & 0xffu32) as u8,
             (u >>  8 & 0xffu32) as u8,
             (u & 0xffu32) as u8)
    };
    let name_bytes = u32_to_vec(name);
    String::from_utf8_lossy(&name_bytes).into_owned()
}

use std::fmt;
impl fmt::Display for BoxHeader {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "'{}' {} bytes", fourcc_to_string(self.name), self.size)
    }
}

impl fmt::Display for FileTypeBox {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let name = fourcc_to_string(self.name);
        let brand = fourcc_to_string(self.major_brand);
        write!(f, "'{}' {} bytes '{}' v{}", name, self.size,
            brand, self.minor_version)
    }
}

impl fmt::Display for MovieHeaderBox {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let name = fourcc_to_string(self.name);
        write!(f, "'{}' {} bytes duration {}s", name, self.size,
            (self.duration as f64)/(self.timescale as f64))
    }
}

use std::u16;
impl fmt::Display for TrackHeaderBox {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let name = fourcc_to_string(self.name);
        // Dimensions are 16.16 fixed-point.
        let base = u16::MAX as f64 + 1.0;
        let width = (self.width as f64) / base;
        let height = (self.height as f64) / base;
        write!(f, "'{}' {} bytes duration {} id {} {}x{}",
            name, self.size, self.duration, self.track_id,
            width, height)
    }
}

#[test]
fn test_read_box_header() {
    use std::io::Cursor;
    use std::io::Write;
    let mut test: Vec<u8> = vec![0, 0, 0, 8];  // minimal box length
    write!(&mut test, "test").unwrap(); // box type
    let mut stream = Cursor::new(test);
    let parsed = read_box_header(&mut stream).unwrap();
    assert_eq!(parsed.name, 1952805748);
    assert_eq!(parsed.size, 8);
    println!("box {}", parsed);
}

#[test]
fn test_read_box_header_long() {
    use std::io::Cursor;
    let mut test: Vec<u8> = vec![0, 0, 0, 1]; // long box extension code
    test.extend("long".to_string().into_bytes()); // box type
    test.extend(vec![0, 0, 0, 0, 0, 0, 16, 0]); // 64 bit size
    // Skip generating box content.
    let mut stream = Cursor::new(test);
    let parsed = read_box_header(&mut stream).unwrap();
    assert_eq!(parsed.name, 1819242087);
    assert_eq!(parsed.size, 4096);
    println!("box {}", parsed);
}

#[test]
fn test_read_ftyp() {
    use std::io::Cursor;
    use std::io::Write;
    let mut test: Vec<u8> = vec![0, 0, 0, 24]; // size
    write!(&mut test, "ftyp").unwrap(); // type
    write!(&mut test, "mp42").unwrap(); // major brand
    test.extend(vec![0, 0, 0, 0]);      // minor version
    write!(&mut test, "isom").unwrap(); // compatible brands...
    write!(&mut test, "mp42").unwrap();
    assert_eq!(test.len(), 24);

    let mut stream = Cursor::new(test);
    let header = read_box_header(&mut stream).unwrap();
    let parsed = read_ftyp(&mut stream, &header).unwrap();
    assert_eq!(parsed.name, 1718909296);
    assert_eq!(parsed.size, 24);
    assert_eq!(parsed.major_brand, 1836069938);
    assert_eq!(parsed.minor_version, 0);
    assert_eq!(parsed.compatible_brands.len(), 2);
    println!("box {}", parsed);
}
<|endoftext|>