<commit_before>//!Router asigns handlers to paths and resolves them per request

#[cfg(test)]
use http::method;
use http::method::Method;
use regex::Regex;
use std::collections::hashmap::HashMap;
use request::Request;
use response::Response;

/// A Route is the basic data structure that stores both the path
/// and the handler that gets executed for the route.
/// The path can contain variable pattern such as `user/:userid/invoices`
pub struct Route {
    pub path: String,
    pub method: Method,
    pub handler: fn(request: &Request, response: &mut Response),
    pub variables: HashMap<String, uint>,
    matcher: Regex
}

impl Clone for Route {
    fn clone(&self) -> Route {
        Route {
            path: self.path.clone(),
            method: self.method.clone(),
            handler: self.handler,
            matcher: self.matcher.clone(),
            variables: self.variables.clone()
        }
    }
}

/// A RouteResult is what the router returns when `match_route` is called.
/// It contains the matched `route` and also a `params` property holding
/// a HashMap with the keys being the variable names and the value being the
/// evaluated string
pub struct RouteResult<'a> {
    pub route: &'a Route,
    pub params: HashMap<String, String>
}

/// The PathUtils collects some small helper methods that operate on the path
struct PathUtils;

// matches named variables (e.g. :userid)
static REGEX_VAR_SEQ: Regex                 = regex!(r":([a-zA-Z0-9_-]*)");
static VAR_SEQ:&'static str                 = "[a-zA-Z0-9_-]*";
static VAR_SEQ_WITH_SLASH:&'static str      = "[/a-zA-Z0-9_-]*";
static VAR_SEQ_WITH_CAPTURE:&'static str    = "([a-zA-Z0-9_-]*)";
// matches request params (e.g. ?foo=true&bar=false)
static REGEX_PARAM_SEQ:&'static str               = "[a-zA-Z0-9_=&?-]*";
static REGEX_START:&'static str             = "^";
static REGEX_END:&'static str               = "$";

impl PathUtils {
    fn create_regex (route_path: &str) -> Regex {

        let updated_path = route_path.to_string()
                                     // first mark all double wildcards for replacement. We can't directly replace them
                                     // since the replacement does contain the * symbol as well, which would get overwritten
                                     // with the next replace call
                                     .replace("**", "___DOUBLE_WILDCARD___")
                                     // then replace the regular wildcard symbols (*) with the appropriate regex
                                     .replace("*", VAR_SEQ)
                                     // now replace the previously marked double wild cards (**)
                                     .replace("___DOUBLE_WILDCARD___", VAR_SEQ_WITH_SLASH);

        // then replace the variable symbols (:variable) with the appropriate regex
        let result = REGEX_START.to_string()
                                .append(REGEX_VAR_SEQ.replace_all(updated_path.as_slice(), VAR_SEQ_WITH_CAPTURE).as_slice())
                                .append(REGEX_PARAM_SEQ)
                                .append(REGEX_END);

        match Regex::new(result.as_slice()) {
            Ok(re) => re,
            Err(err) => fail!("{}", err)
        }
    }

    fn get_variable_info (route_path: &str) -> HashMap<String, uint> {
        REGEX_VAR_SEQ.captures_iter(route_path)
             .enumerate()
             .map(|(i, matched)| (matched.at(1).to_string(), i))
             .collect()
    }
}

/// The Router's job is it to hold routes and to resolve them later against
/// concrete URLs

#[deriving(Clone)]
pub struct Router{
    pub routes: Vec<Route>,
}

impl Router {
    pub fn new () -> Router {
        Router {
            routes: Vec::new()
        }
    }

    pub fn add_route (&mut self, method: Method, path: String, handler: fn(request: &Request, response: &mut Response)) -> () {
        let matcher = PathUtils::create_regex(path.as_slice());
        let variable_infos = PathUtils::get_variable_info(path.as_slice());
        let route = Route {
            path: path,
            method: method,
            matcher: matcher,
            handler: handler,
            variables: variable_infos
        };
        self.routes.push(route);
    }

    pub fn match_route(&self, method: Method, path: String) -> Option<RouteResult> {
        self.routes.iter().find(|item| item.method == method && item.matcher.is_match(path.as_slice()))
            .and_then(|route| {
                match route.matcher.captures(path.as_slice()) {
                    Some(captures) => {
                        let mut map = HashMap::new();
                        for (name, pos) in route.variables.iter() {
                            map.insert(name.to_string(), captures.at(pos + 1).to_string());
                        }
                        Some(RouteResult {
                            route: route,
                            params: map
                        })
                    },
                    None => Some(RouteResult{
                        route: route,
                        params: HashMap::new()
                    })
                }
            })
    }
}

#[test]
fn creates_map_with_var_variable_infos () {
    let map = PathUtils::get_variable_info("foo/:uid/bar/:groupid");

    assert_eq!(map.len(), 2);
    assert_eq!(map.get(&"uid".to_string()), &0);
    assert_eq!(map.get(&"groupid".to_string()), &1);
}

#[test]
fn creates_regex_with_captures () {
    let regex = PathUtils::create_regex("foo/:uid/bar/:groupid");
    let caps = regex.captures("foo/4711/bar/5490").unwrap();

    assert_eq!(caps.at(1), "4711");
    assert_eq!(caps.at(2), "5490");

    let regex = PathUtils::create_regex("foo/*/:uid/bar/:groupid");
    let caps = regex.captures("foo/test/4711/bar/5490").unwrap();

    assert_eq!(caps.at(1), "4711");
    assert_eq!(caps.at(2), "5490");

    let regex = PathUtils::create_regex("foo/**/:uid/bar/:groupid");
    let caps = regex.captures("foo/test/another/4711/bar/5490").unwrap();

    assert_eq!(caps.at(1), "4711");
    assert_eq!(caps.at(2), "5490");
}

#[test]
fn creates_valid_regex_for_routes () {
    let regex1 = PathUtils::create_regex("foo/:uid/bar/:groupid");
    let regex2 = PathUtils::create_regex("foo/*/bar");
    let regex3 = PathUtils::create_regex("foo/**/bar");

    assert_eq!(regex1.is_match("foo/4711/bar/5490"), true);
    assert_eq!(regex1.is_match("foo/4711/bar/5490?foo=true&bar=false"), true);
    assert_eq!(regex1.is_match("foo/4711/bar"), false);
    assert_eq!(regex1.is_match("foo/4711/bar?foo=true&bar=false"), false);

    assert_eq!(regex2.is_match("foo/4711/bar"), true);
    assert_eq!(regex2.is_match("foo/4711/bar?foo=true&bar=false"), true);
    assert_eq!(regex2.is_match("foo/4711/4712/bar"), false);
    assert_eq!(regex2.is_match("foo/4711/4712/bar?foo=true&bar=false"), false);

    assert_eq!(regex3.is_match("foo/4711/bar"), true);
    assert_eq!(regex3.is_match("foo/4711/bar?foo=true&bar=false"), true);
    assert_eq!(regex3.is_match("foo/4711/4712/bar"), true);
    assert_eq!(regex3.is_match("foo/4711/4712/bar?foo=true&bar=false"), true);
}

#[test]
fn can_match_var_routes () {
    let route_store = &mut Router::new();

    fn handler (_request: &Request, response: &mut Response) -> () {
        let _ = response.origin.write("hello from foo".as_bytes());
    };

    route_store.add_route(method::Get, "/foo/:userid".to_string(), handler);
    route_store.add_route(method::Get, "/bar".to_string(), handler);

    let route_result = route_store.match_route(method::Get, "/foo/4711".to_string()).unwrap();
    let route = route_result.route;

    assert_eq!(route_result.params.get(&"userid".to_string()), &"4711".to_string());

    //assert the route has identified the variable
    assert_eq!(route.variables.len(), 1);
    assert_eq!(route.variables.get(&"userid".to_string()), &0);

    let route_result = route_store.match_route(method::Get, "/bar/4711".to_string());

    let result = match route_result {
        Some(_res) => true,
        None => false
    };

    assert_eq!(result, false);

    let route_result = route_store.match_route(method::Get, "/foo".to_string());

    let result = match route_result{
        Some(_res) => true,
        None => false
    };

    assert_eq!(result, false);
}
<commit_msg>fix(router): fixes screwed up indentation<commit_after>//!Router asigns handlers to paths and resolves them per request

#[cfg(test)]
use http::method;
use http::method::Method;
use regex::Regex;
use std::collections::hashmap::HashMap;
use request::Request;
use response::Response;

/// A Route is the basic data structure that stores both the path
/// and the handler that gets executed for the route.
/// The path can contain variable pattern such as `user/:userid/invoices`
pub struct Route {
    pub path: String,
    pub method: Method,
    pub handler: fn(request: &Request, response: &mut Response),
    pub variables: HashMap<String, uint>,
    matcher: Regex
}

impl Clone for Route {
    fn clone(&self) -> Route {
        Route {
            path: self.path.clone(),
            method: self.method.clone(),
            handler: self.handler,
            matcher: self.matcher.clone(),
            variables: self.variables.clone()
        }
    }
}

/// A RouteResult is what the router returns when `match_route` is called.
/// It contains the matched `route` and also a `params` property holding
/// a HashMap with the keys being the variable names and the value being the
/// evaluated string
pub struct RouteResult<'a> {
    pub route: &'a Route,
    pub params: HashMap<String, String>
}

/// The PathUtils collects some small helper methods that operate on the path
struct PathUtils;

// matches named variables (e.g. :userid)
static REGEX_VAR_SEQ: Regex                 = regex!(r":([a-zA-Z0-9_-]*)");
static VAR_SEQ:&'static str                 = "[a-zA-Z0-9_-]*";
static VAR_SEQ_WITH_SLASH:&'static str      = "[/a-zA-Z0-9_-]*";
static VAR_SEQ_WITH_CAPTURE:&'static str    = "([a-zA-Z0-9_-]*)";
// matches request params (e.g. ?foo=true&bar=false)
static REGEX_PARAM_SEQ:&'static str         = "[a-zA-Z0-9_=&?-]*";
static REGEX_START:&'static str             = "^";
static REGEX_END:&'static str               = "$";

impl PathUtils {
    fn create_regex (route_path: &str) -> Regex {

        let updated_path = route_path.to_string()
                                     // first mark all double wildcards for replacement. We can't directly replace them
                                     // since the replacement does contain the * symbol as well, which would get overwritten
                                     // with the next replace call
                                     .replace("**", "___DOUBLE_WILDCARD___")
                                     // then replace the regular wildcard symbols (*) with the appropriate regex
                                     .replace("*", VAR_SEQ)
                                     // now replace the previously marked double wild cards (**)
                                     .replace("___DOUBLE_WILDCARD___", VAR_SEQ_WITH_SLASH);

        // then replace the variable symbols (:variable) with the appropriate regex
        let result = REGEX_START.to_string()
                                .append(REGEX_VAR_SEQ.replace_all(updated_path.as_slice(), VAR_SEQ_WITH_CAPTURE).as_slice())
                                .append(REGEX_PARAM_SEQ)
                                .append(REGEX_END);

        match Regex::new(result.as_slice()) {
            Ok(re) => re,
            Err(err) => fail!("{}", err)
        }
    }

    fn get_variable_info (route_path: &str) -> HashMap<String, uint> {
        REGEX_VAR_SEQ.captures_iter(route_path)
             .enumerate()
             .map(|(i, matched)| (matched.at(1).to_string(), i))
             .collect()
    }
}

/// The Router's job is it to hold routes and to resolve them later against
/// concrete URLs

#[deriving(Clone)]
pub struct Router{
    pub routes: Vec<Route>,
}

impl Router {
    pub fn new () -> Router {
        Router {
            routes: Vec::new()
        }
    }

    pub fn add_route (&mut self, method: Method, path: String, handler: fn(request: &Request, response: &mut Response)) -> () {
        let matcher = PathUtils::create_regex(path.as_slice());
        let variable_infos = PathUtils::get_variable_info(path.as_slice());
        let route = Route {
            path: path,
            method: method,
            matcher: matcher,
            handler: handler,
            variables: variable_infos
        };
        self.routes.push(route);
    }

    pub fn match_route(&self, method: Method, path: String) -> Option<RouteResult> {
        self.routes.iter().find(|item| item.method == method && item.matcher.is_match(path.as_slice()))
            .and_then(|route| {
                match route.matcher.captures(path.as_slice()) {
                    Some(captures) => {
                        let mut map = HashMap::new();
                        for (name, pos) in route.variables.iter() {
                            map.insert(name.to_string(), captures.at(pos + 1).to_string());
                        }
                        Some(RouteResult {
                            route: route,
                            params: map
                        })
                    },
                    None => Some(RouteResult{
                        route: route,
                        params: HashMap::new()
                    })
                }
            })
    }
}

#[test]
fn creates_map_with_var_variable_infos () {
    let map = PathUtils::get_variable_info("foo/:uid/bar/:groupid");

    assert_eq!(map.len(), 2);
    assert_eq!(map.get(&"uid".to_string()), &0);
    assert_eq!(map.get(&"groupid".to_string()), &1);
}

#[test]
fn creates_regex_with_captures () {
    let regex = PathUtils::create_regex("foo/:uid/bar/:groupid");
    let caps = regex.captures("foo/4711/bar/5490").unwrap();

    assert_eq!(caps.at(1), "4711");
    assert_eq!(caps.at(2), "5490");

    let regex = PathUtils::create_regex("foo/*/:uid/bar/:groupid");
    let caps = regex.captures("foo/test/4711/bar/5490").unwrap();

    assert_eq!(caps.at(1), "4711");
    assert_eq!(caps.at(2), "5490");

    let regex = PathUtils::create_regex("foo/**/:uid/bar/:groupid");
    let caps = regex.captures("foo/test/another/4711/bar/5490").unwrap();

    assert_eq!(caps.at(1), "4711");
    assert_eq!(caps.at(2), "5490");
}

#[test]
fn creates_valid_regex_for_routes () {
    let regex1 = PathUtils::create_regex("foo/:uid/bar/:groupid");
    let regex2 = PathUtils::create_regex("foo/*/bar");
    let regex3 = PathUtils::create_regex("foo/**/bar");

    assert_eq!(regex1.is_match("foo/4711/bar/5490"), true);
    assert_eq!(regex1.is_match("foo/4711/bar/5490?foo=true&bar=false"), true);
    assert_eq!(regex1.is_match("foo/4711/bar"), false);
    assert_eq!(regex1.is_match("foo/4711/bar?foo=true&bar=false"), false);

    assert_eq!(regex2.is_match("foo/4711/bar"), true);
    assert_eq!(regex2.is_match("foo/4711/bar?foo=true&bar=false"), true);
    assert_eq!(regex2.is_match("foo/4711/4712/bar"), false);
    assert_eq!(regex2.is_match("foo/4711/4712/bar?foo=true&bar=false"), false);

    assert_eq!(regex3.is_match("foo/4711/bar"), true);
    assert_eq!(regex3.is_match("foo/4711/bar?foo=true&bar=false"), true);
    assert_eq!(regex3.is_match("foo/4711/4712/bar"), true);
    assert_eq!(regex3.is_match("foo/4711/4712/bar?foo=true&bar=false"), true);
}

#[test]
fn can_match_var_routes () {
    let route_store = &mut Router::new();

    fn handler (_request: &Request, response: &mut Response) -> () {
        let _ = response.origin.write("hello from foo".as_bytes());
    };

    route_store.add_route(method::Get, "/foo/:userid".to_string(), handler);
    route_store.add_route(method::Get, "/bar".to_string(), handler);

    let route_result = route_store.match_route(method::Get, "/foo/4711".to_string()).unwrap();
    let route = route_result.route;

    assert_eq!(route_result.params.get(&"userid".to_string()), &"4711".to_string());

    //assert the route has identified the variable
    assert_eq!(route.variables.len(), 1);
    assert_eq!(route.variables.get(&"userid".to_string()), &0);

    let route_result = route_store.match_route(method::Get, "/bar/4711".to_string());

    let result = match route_result {
        Some(_res) => true,
        None => false
    };

    assert_eq!(result, false);

    let route_result = route_store.match_route(method::Get, "/foo".to_string());

    let result = match route_result{
        Some(_res) => true,
        None => false
    };

    assert_eq!(result, false);
}
<|endoftext|>


<commit_before>use std::collections::{BTreeMap, HashMap};
use std::env;
use std::fs::File;
use std::process;

use applescript::{run as run_applescript};

use plist::Plist;

use tera::{Context, Tera};

use urls;

macro_rules! error(
  ($($arg:tt)*) => { { return Err(format!($($arg)*)) } }
);

/// Returns true/false if Safari is running.
pub fn is_safari_running() -> bool {
  let cmd_result = process::Command::new("ps")
    .arg("-eaf")
    .output()
    .expect("Unable to test if Safari is running.");
  for line in String::from_utf8(cmd_result.stdout).unwrap().lines() {
    if line.contains("Safari.app/Contents/MacOS/Safari") {
      return true
    }
  }
  false
}

/// Return a URL from a Safari window.
///
/// Given a (window, tab) pair, this function looks up the URL of the tab.
/// Note that it doesn't do any error handling, so will throw an execution
/// error if it fails.
///
/// * `window` - Window index.  1 is frontmost.  If None, assumes the
///              frontmost window.
/// * `tab` - Tab index.  1 is leftmost.  If None, assumes the frontmost tab.
///
pub fn get_url(window: Option<i32>, tab: Option<i32>) -> Result<String, String> {
  // If a tab isn't specified, assume the user wants the frontmost tab.
  let command = match window {
    Some(w_idx) => {
      match tab {
        Some(t_idx) => format!("tell application \"Safari\" to get tab {} of window {}", t_idx, w_idx),
        None => format!("tell application \"Safari\" to get URL of document {}", w_idx),
      }
    },
    None => "tell application \"Safari\" to get URL of document 1".to_string(),
  };
  let output = run_applescript(&command);

  if output.status.success() {
    Ok(urls::tidy_url(output.stdout.trim()))
  } else {
    if output.stderr.contains("Invalid index") {
      error!("Invalid index: no such window or tab.")
    } else {
      error!("Unexpected error from osascript: {:?}", output.stderr)
    }
  }
}

/// Return a list of URLs from Safari.
///
/// This returns a list of URLs, one for every tab in Safari.  Iteration
/// order depends on AppleScript, which I don't think is guaranteed to be
/// stable (in particular, I think it depends on which window is frontmost).
///
pub fn get_all_urls() -> Result<Vec<String>, String> {
  let script = include_str!("scripts/list-open-tabs.scpt");
  let output = run_applescript(&script);
  if output.status.success() {
    Ok(parse_list_open_tabs_output(&output.stdout))
  } else {
    error!("Unexpected error from osascript: {:?}", output.stderr);
  }
}

/// Given AppleScript output from list-open-tabs.scpt, return a list of URLs.
fn parse_list_open_tabs_output(stdout: &str) -> Vec<String> {
  stdout
    .trim()
    .split(", ")
    .map(|url| urls::tidy_url(url))
    .filter(|url| url != "favorites://")
    .filter(|url| url != "://missing value")
    .collect()
}

macro_rules! parse_list_open_tabs_output_tests {
  ($($name:ident: $value:expr,)*) => {
    $(
      #[test]
      fn $name() {
        let (input, expected) = $value;
        assert_eq!(expected, parse_list_open_tabs_output(input));
      }
    )*
  }
}

parse_list_open_tabs_output_tests! {

  // I don't have a good way to put an empty vector on the right-hand side of
  // this comparison, so no test checks for nothing (yet).

  favorites_with_url: (
    "favorites://, http://example.org",
    vec!["http://example.org"]
  ),

  url_with_missing_value: (
    "://missing value, https://www.example.net",
    vec!["https://www.example.net"]
  ),

  single_url: (
    "http://foo_bar.com",
    vec!["http://foo_bar.com"]
  ),

  multiple_urls: (
    "http://example.org, https://www.example.net, http://test.co.uk",
    vec!["http://example.org", "https://www.example.net", "http://test.co.uk"]
  ),

  with_extra_whitespace: (
    "    http://space.org, https://www.nasa.gov   ",
    vec!["http://space.org", "https://www.nasa.gov"]
  ),

  applies_tidy_url_transform: (
    "https://mobile.twitter.com",
    vec!["https://twitter.com"]
  ),

  colon_slash_slash_with_url: (
    "://, http://example.org",
    vec!["http://example.org"]
  ),
}

/// Convert URL patterns into AppleScript conditions.
///
/// These conditions can be used in an `if` statement in AppleScript to
/// decide whether a tab should be closed.  Three patterns are supported,
/// a limited regex syntax:
///
///     example.com             matches anywhere in the URL
///     ^http://examples.com    matches at the start of the URL
///     example.com/$           matches at the end of the URL
///
fn parse_conditions(url_patterns: Vec<&str>) -> Vec<String> {
  url_patterns.iter().map(|p|
    if p.starts_with("^") {
      format!("starts with \"{}\"", p.replace("^", ""))
    } else if p.ends_with("$") {
      format!("ends with \"{}\"", p.replace("$", ""))
    } else {
      format!("contains \"{}\"", p)
    }
  ).collect()
}

/// Tests for parse_conditions().
#[cfg(test)]
mod tests {
  use safari::parse_conditions;

  #[test]
  fn test_parse_conditions() {
    let patterns = vec!["github.com", "^facebook.com", "twitter.com$"];
    let expected = vec!["contains \"github.com\"", "starts with \"facebook.com\"", "ends with \"twitter.com\""];
    let actual = parse_conditions(patterns);
    assert_eq!(actual, expected);
  }
}

/// Close tabs in Safari that match URL patterns.
///
/// Takes a list of URL patterns, and tries to close any matching Safari tabs.
/// Bugginess in AppleScript means this isn't always perfect, but we can
/// have a go.  In general it will fail to close tabs, rather than close
/// the wrong tabs.
///
pub fn close_tabs(url_patterns: Vec<&str>) {
  let conditions = parse_conditions(url_patterns);

  let clean_tabs_template = include_str!("scripts/clean-tabs.scpt");
  let mut context = Context::new();
  context.add("conditions", &conditions);
  let script = Tera::one_off(&clean_tabs_template, context, false).unwrap();

  // Run it twice to get around weird AppleScript bugs.
  run_applescript(&script);
  run_applescript(&script);
}

/// Get the Bookmarks.plist dict for a given title
fn read_bookmarks_plist(title: &str) -> Result<Plist, String> {

  // All Safari data lives at ~/Library/Safari/Bookmarks.plist
  // TODO: There's probably a more idiomatic Rust-like way to write this.
  let mut plist_path = match env::home_dir() {
    Some(v) => v,
    None => error!("Unable to get home directory?"),
  };
  plist_path.push("Library/Safari/Bookmarks.plist");

  let file = match File::open(plist_path) {
    Ok(v) => v,
    Err(e) => error!("Unable to open ~/Library/Safari/Bookmarks.plist: {:?}", e),
  };

  let plist = match Plist::read(file) {
    Ok(v) => v,
    Err(e) => error!("Unable to read ~/Library/Safari/Bookmarks.plist: {:?}", e),
  };

  let data = match plist.as_dictionary() {
    Some(v) => v,
    None => error!("Unable to parse ~/Library/Safari/Bookmarks.plist as dictionary?"),
  };

  // The structure of Bookmarks.plist is as follows:
  //
  //     <dict>
  //       <key>Children</key>
  //       <array>
  //         <dict>
  //           <key>Title</key><string>History</string>
  //           ... dict data ...
  //         </dict>
  //         <dict>
  //           <key>Title</key><string>com.apple.ReadingList</string>
  //           ... dict data ...
  //         </dict>
  //         ... other array items ...
  //       </array>
  //     </dict>
  //
  let children = match data.get("Children") {
    Some(child_key) => match child_key.as_array() {
      Some(v) => v,
      None => error!("Top-level children key in ~/Library/Safari/Bookmarks.plist isn't an array?"),
    },
    None => error!("Unable to find top-level Children key in ~/Library/Safari/Bookmarks.plist"),
  };

  let mut matching_children = children.iter().filter(|d|
    match d.as_dictionary() {
      Some(dict) => match dict.get("Title") {
        Some(title_elem) => match title_elem.as_string() {
          Some(v) => (v == title),
          None => false,
        },
        None => false,
      },
      None => false,
    }
  );

  // Check we got one, and only one result.
  let result = match matching_children.next() {
    Some(v) => v,
    None => error!("Unable to find key {} in Bookmarks.plist", title),
  };

  match matching_children.next() {
    Some(_) => error!("Got more than one result for {} in Bookmarks.plist", title),
    None => {},
  };

  Ok(result.to_owned())
}

/// Return a list of URLs from Reading List.
///
/// Iteration order depends on the order in which they're stored in
/// Bookmarks.plist, which is usually (but not guaranteed to be) newest first.
///
pub fn get_reading_list_urls() -> Result<Vec<String>, String> {
  let plist = match read_bookmarks_plist("com.apple.ReadingList") {
    Ok(v) => v,
    Err(e) => return Err(e),
  };

  // TODO: All these unwrap() calls should probably be handled better
  let children = plist
    .as_dictionary().unwrap()
    .get("Children").unwrap()
    .as_array().unwrap();

  Ok(children
    .iter()
    .map(|child|
      child
        .as_dictionary().unwrap()
        .get("URLString").unwrap()
        .as_string().unwrap()
    )
    .map(|url| urls::tidy_url(url))
    .collect()
  )
}

/// Get the com.apple.Safari.plist preferences file
fn read_safari_plist() -> Result<BTreeMap<String, Plist>, String> {

  // All Safari data lives at ~/Library/Safari/Bookmarks.plist
  // TODO: There's probably a more idiomatic Rust-like way to write this.
  let mut plist_path = match env::home_dir() {
    Some(v) => v,
    None => error!("Unable to get home directory?"),
  };
  plist_path.push("Library/SyncedPreferences/com.apple.Safari.plist");

  let file = match File::open(plist_path) {
    Ok(v) => v,
    Err(e) => error!("Unable to open com.apple.Safari.plist: {:?}", e),
  };

  let plist = match Plist::read(file) {
    Ok(v) => v,
    Err(e) => error!("Unable to read com.apple.Safari.plist: {:?}", e),
  };

  let data = match plist.as_dictionary() {
    Some(v) => v,
    None => error!("Unable to parse com.apple.Safari.plist as dictionary?"),
  };

  // The structure of com.apple.Safari.plist is as follows:
  //
  //      <dict>
  //        <key>values</key>
  //        <dict>
  //          <dict>
  //            <key>[[ UUID ]]</key>
  //            <dict>
  //              ... dict data ...
  //              <key>value</key>
  //              <dict>
  //                ... dict data ...
  //                <key>DeviceName</key><string>[[ Device name ]]</string>
  //                <key>Tabs</key>
  //                <array>
  //                  <dict>
  //                    <key>Title</key><string>[[ Title ]]</string>
  //                    <key>URL</key><string>[[ URL ]]</string>
  //                  </dict>
  //                  ... array items for other tabs ...
  //                </array>
  //              </dict>
  //            </dict>
  //            ... dicts for other devices ...
  //          </dict>
  //        </dict>
  //      </dict>
  //
  match data.get("values") {
    Some(child_key) => match child_key.as_dictionary() {
      Some(v) => Ok(v.to_owned()),
      None => error!("Top-level values key in com.apple.Safari.plist isn't an dictionary?"),
    },
    None => error!("Unable to find top-level values key in com.apple.Safari.plist"),
  }
}

/// Return a list of devices in iCloud Tabs.
pub fn list_icloud_tabs_devices() -> Result<Vec<String>, String> {
  let plist = match read_safari_plist() {
    Ok(plist) => plist,
    Err(e) => return Err(e),
  };

  Ok(plist
    .values()
    .map(|value| value
      .as_dictionary().unwrap()
      .get("value").unwrap()
      .as_dictionary().unwrap()
      .get("DeviceName").unwrap()
      .as_string().unwrap()
      .to_owned()
    )
    .collect()
  )
}

/// Return a list of URLs from iCloud Tabs.
pub fn get_icloud_tabs_urls() -> Result<HashMap<String, Vec<String>>, String> {
  let plist = match read_safari_plist() {
    Ok(plist) => plist,
    Err(e) => return Err(e),
  };

  // Within the `values` dictionary, each device is as follows:
  //
  //    <key>[[ device UUID ]]</key>
  //    <dict>
  //      [[ device data ]]
  //      <dict>
  //        <key>Tabs</key>
  //        <key>DeviceName</key><string>[[ device name ]]</key>
  //        <key>Tabs</key>
  //        <array>
  //          <dict>
  //            <key>URL</key><string>[[ URL ]]</string>
  //            [[ other tab data ]]
  //          </dict>
  //          [[ other tabs ]]
  //
  let mut result: HashMap<String, Vec<String>> = HashMap::new();
  let devices_with_tabs = plist
    .values()
    .map(|value| value
      .as_dictionary().unwrap()
      .get("value").unwrap()
      .as_dictionary().unwrap()
    );
  for data in devices_with_tabs {
    let name: String = data
      .get("DeviceName").unwrap()
      .as_string().unwrap()
      .to_owned();

    // If a device is registered with iCloud but Safari isn't running or
    // there aren't any tabs open, there isn't a Tabs field.
    if !data.contains_key("Tabs") {
      continue;
    }

    let urls: Vec<String> = data
      .get("Tabs").unwrap()
      .as_array().unwrap()
      .iter()
      .map(|tab| tab
        .as_dictionary().unwrap()
        .get("URL").unwrap()
        .as_string().unwrap()
      )
      .map(|url| urls::tidy_url(url))
      .collect();
    result.insert(name, urls);
  }
  Ok(result)
}
<commit_msg>Exclude :// from results in 'urls-all'<commit_after>use std::collections::{BTreeMap, HashMap};
use std::env;
use std::fs::File;
use std::process;

use applescript::{run as run_applescript};

use plist::Plist;

use tera::{Context, Tera};

use urls;

macro_rules! error(
  ($($arg:tt)*) => { { return Err(format!($($arg)*)) } }
);

/// Returns true/false if Safari is running.
pub fn is_safari_running() -> bool {
  let cmd_result = process::Command::new("ps")
    .arg("-eaf")
    .output()
    .expect("Unable to test if Safari is running.");
  for line in String::from_utf8(cmd_result.stdout).unwrap().lines() {
    if line.contains("Safari.app/Contents/MacOS/Safari") {
      return true
    }
  }
  false
}

/// Return a URL from a Safari window.
///
/// Given a (window, tab) pair, this function looks up the URL of the tab.
/// Note that it doesn't do any error handling, so will throw an execution
/// error if it fails.
///
/// * `window` - Window index.  1 is frontmost.  If None, assumes the
///              frontmost window.
/// * `tab` - Tab index.  1 is leftmost.  If None, assumes the frontmost tab.
///
pub fn get_url(window: Option<i32>, tab: Option<i32>) -> Result<String, String> {
  // If a tab isn't specified, assume the user wants the frontmost tab.
  let command = match window {
    Some(w_idx) => {
      match tab {
        Some(t_idx) => format!("tell application \"Safari\" to get tab {} of window {}", t_idx, w_idx),
        None => format!("tell application \"Safari\" to get URL of document {}", w_idx),
      }
    },
    None => "tell application \"Safari\" to get URL of document 1".to_string(),
  };
  let output = run_applescript(&command);

  if output.status.success() {
    Ok(urls::tidy_url(output.stdout.trim()))
  } else {
    if output.stderr.contains("Invalid index") {
      error!("Invalid index: no such window or tab.")
    } else {
      error!("Unexpected error from osascript: {:?}", output.stderr)
    }
  }
}

/// Return a list of URLs from Safari.
///
/// This returns a list of URLs, one for every tab in Safari.  Iteration
/// order depends on AppleScript, which I don't think is guaranteed to be
/// stable (in particular, I think it depends on which window is frontmost).
///
pub fn get_all_urls() -> Result<Vec<String>, String> {
  let script = include_str!("scripts/list-open-tabs.scpt");
  let output = run_applescript(&script);
  if output.status.success() {
    Ok(parse_list_open_tabs_output(&output.stdout))
  } else {
    error!("Unexpected error from osascript: {:?}", output.stderr);
  }
}

/// Given AppleScript output from list-open-tabs.scpt, return a list of URLs.
fn parse_list_open_tabs_output(stdout: &str) -> Vec<String> {
  stdout
    .trim()
    .split(", ")
    .map(|url| urls::tidy_url(url))
    .filter(|url| url != "favorites://")
    .filter(|url| url != "://missing value")
    .filter(|url| url != "://")
    .collect()
}

macro_rules! parse_list_open_tabs_output_tests {
  ($($name:ident: $value:expr,)*) => {
    $(
      #[test]
      fn $name() {
        let (input, expected) = $value;
        assert_eq!(expected, parse_list_open_tabs_output(input));
      }
    )*
  }
}

parse_list_open_tabs_output_tests! {

  // I don't have a good way to put an empty vector on the right-hand side of
  // this comparison, so no test checks for nothing (yet).

  favorites_with_url: (
    "favorites://, http://example.org",
    vec!["http://example.org"]
  ),

  url_with_missing_value: (
    "://missing value, https://www.example.net",
    vec!["https://www.example.net"]
  ),

  single_url: (
    "http://foo_bar.com",
    vec!["http://foo_bar.com"]
  ),

  multiple_urls: (
    "http://example.org, https://www.example.net, http://test.co.uk",
    vec!["http://example.org", "https://www.example.net", "http://test.co.uk"]
  ),

  with_extra_whitespace: (
    "    http://space.org, https://www.nasa.gov   ",
    vec!["http://space.org", "https://www.nasa.gov"]
  ),

  applies_tidy_url_transform: (
    "https://mobile.twitter.com",
    vec!["https://twitter.com"]
  ),

  colon_slash_slash_with_url: (
    "://, http://example.org",
    vec!["http://example.org"]
  ),
}

/// Convert URL patterns into AppleScript conditions.
///
/// These conditions can be used in an `if` statement in AppleScript to
/// decide whether a tab should be closed.  Three patterns are supported,
/// a limited regex syntax:
///
///     example.com             matches anywhere in the URL
///     ^http://examples.com    matches at the start of the URL
///     example.com/$           matches at the end of the URL
///
fn parse_conditions(url_patterns: Vec<&str>) -> Vec<String> {
  url_patterns.iter().map(|p|
    if p.starts_with("^") {
      format!("starts with \"{}\"", p.replace("^", ""))
    } else if p.ends_with("$") {
      format!("ends with \"{}\"", p.replace("$", ""))
    } else {
      format!("contains \"{}\"", p)
    }
  ).collect()
}

/// Tests for parse_conditions().
#[cfg(test)]
mod tests {
  use safari::parse_conditions;

  #[test]
  fn test_parse_conditions() {
    let patterns = vec!["github.com", "^facebook.com", "twitter.com$"];
    let expected = vec!["contains \"github.com\"", "starts with \"facebook.com\"", "ends with \"twitter.com\""];
    let actual = parse_conditions(patterns);
    assert_eq!(actual, expected);
  }
}

/// Close tabs in Safari that match URL patterns.
///
/// Takes a list of URL patterns, and tries to close any matching Safari tabs.
/// Bugginess in AppleScript means this isn't always perfect, but we can
/// have a go.  In general it will fail to close tabs, rather than close
/// the wrong tabs.
///
pub fn close_tabs(url_patterns: Vec<&str>) {
  let conditions = parse_conditions(url_patterns);

  let clean_tabs_template = include_str!("scripts/clean-tabs.scpt");
  let mut context = Context::new();
  context.add("conditions", &conditions);
  let script = Tera::one_off(&clean_tabs_template, context, false).unwrap();

  // Run it twice to get around weird AppleScript bugs.
  run_applescript(&script);
  run_applescript(&script);
}

/// Get the Bookmarks.plist dict for a given title
fn read_bookmarks_plist(title: &str) -> Result<Plist, String> {

  // All Safari data lives at ~/Library/Safari/Bookmarks.plist
  // TODO: There's probably a more idiomatic Rust-like way to write this.
  let mut plist_path = match env::home_dir() {
    Some(v) => v,
    None => error!("Unable to get home directory?"),
  };
  plist_path.push("Library/Safari/Bookmarks.plist");

  let file = match File::open(plist_path) {
    Ok(v) => v,
    Err(e) => error!("Unable to open ~/Library/Safari/Bookmarks.plist: {:?}", e),
  };

  let plist = match Plist::read(file) {
    Ok(v) => v,
    Err(e) => error!("Unable to read ~/Library/Safari/Bookmarks.plist: {:?}", e),
  };

  let data = match plist.as_dictionary() {
    Some(v) => v,
    None => error!("Unable to parse ~/Library/Safari/Bookmarks.plist as dictionary?"),
  };

  // The structure of Bookmarks.plist is as follows:
  //
  //     <dict>
  //       <key>Children</key>
  //       <array>
  //         <dict>
  //           <key>Title</key><string>History</string>
  //           ... dict data ...
  //         </dict>
  //         <dict>
  //           <key>Title</key><string>com.apple.ReadingList</string>
  //           ... dict data ...
  //         </dict>
  //         ... other array items ...
  //       </array>
  //     </dict>
  //
  let children = match data.get("Children") {
    Some(child_key) => match child_key.as_array() {
      Some(v) => v,
      None => error!("Top-level children key in ~/Library/Safari/Bookmarks.plist isn't an array?"),
    },
    None => error!("Unable to find top-level Children key in ~/Library/Safari/Bookmarks.plist"),
  };

  let mut matching_children = children.iter().filter(|d|
    match d.as_dictionary() {
      Some(dict) => match dict.get("Title") {
        Some(title_elem) => match title_elem.as_string() {
          Some(v) => (v == title),
          None => false,
        },
        None => false,
      },
      None => false,
    }
  );

  // Check we got one, and only one result.
  let result = match matching_children.next() {
    Some(v) => v,
    None => error!("Unable to find key {} in Bookmarks.plist", title),
  };

  match matching_children.next() {
    Some(_) => error!("Got more than one result for {} in Bookmarks.plist", title),
    None => {},
  };

  Ok(result.to_owned())
}

/// Return a list of URLs from Reading List.
///
/// Iteration order depends on the order in which they're stored in
/// Bookmarks.plist, which is usually (but not guaranteed to be) newest first.
///
pub fn get_reading_list_urls() -> Result<Vec<String>, String> {
  let plist = match read_bookmarks_plist("com.apple.ReadingList") {
    Ok(v) => v,
    Err(e) => return Err(e),
  };

  // TODO: All these unwrap() calls should probably be handled better
  let children = plist
    .as_dictionary().unwrap()
    .get("Children").unwrap()
    .as_array().unwrap();

  Ok(children
    .iter()
    .map(|child|
      child
        .as_dictionary().unwrap()
        .get("URLString").unwrap()
        .as_string().unwrap()
    )
    .map(|url| urls::tidy_url(url))
    .collect()
  )
}

/// Get the com.apple.Safari.plist preferences file
fn read_safari_plist() -> Result<BTreeMap<String, Plist>, String> {

  // All Safari data lives at ~/Library/Safari/Bookmarks.plist
  // TODO: There's probably a more idiomatic Rust-like way to write this.
  let mut plist_path = match env::home_dir() {
    Some(v) => v,
    None => error!("Unable to get home directory?"),
  };
  plist_path.push("Library/SyncedPreferences/com.apple.Safari.plist");

  let file = match File::open(plist_path) {
    Ok(v) => v,
    Err(e) => error!("Unable to open com.apple.Safari.plist: {:?}", e),
  };

  let plist = match Plist::read(file) {
    Ok(v) => v,
    Err(e) => error!("Unable to read com.apple.Safari.plist: {:?}", e),
  };

  let data = match plist.as_dictionary() {
    Some(v) => v,
    None => error!("Unable to parse com.apple.Safari.plist as dictionary?"),
  };

  // The structure of com.apple.Safari.plist is as follows:
  //
  //      <dict>
  //        <key>values</key>
  //        <dict>
  //          <dict>
  //            <key>[[ UUID ]]</key>
  //            <dict>
  //              ... dict data ...
  //              <key>value</key>
  //              <dict>
  //                ... dict data ...
  //                <key>DeviceName</key><string>[[ Device name ]]</string>
  //                <key>Tabs</key>
  //                <array>
  //                  <dict>
  //                    <key>Title</key><string>[[ Title ]]</string>
  //                    <key>URL</key><string>[[ URL ]]</string>
  //                  </dict>
  //                  ... array items for other tabs ...
  //                </array>
  //              </dict>
  //            </dict>
  //            ... dicts for other devices ...
  //          </dict>
  //        </dict>
  //      </dict>
  //
  match data.get("values") {
    Some(child_key) => match child_key.as_dictionary() {
      Some(v) => Ok(v.to_owned()),
      None => error!("Top-level values key in com.apple.Safari.plist isn't an dictionary?"),
    },
    None => error!("Unable to find top-level values key in com.apple.Safari.plist"),
  }
}

/// Return a list of devices in iCloud Tabs.
pub fn list_icloud_tabs_devices() -> Result<Vec<String>, String> {
  let plist = match read_safari_plist() {
    Ok(plist) => plist,
    Err(e) => return Err(e),
  };

  Ok(plist
    .values()
    .map(|value| value
      .as_dictionary().unwrap()
      .get("value").unwrap()
      .as_dictionary().unwrap()
      .get("DeviceName").unwrap()
      .as_string().unwrap()
      .to_owned()
    )
    .collect()
  )
}

/// Return a list of URLs from iCloud Tabs.
pub fn get_icloud_tabs_urls() -> Result<HashMap<String, Vec<String>>, String> {
  let plist = match read_safari_plist() {
    Ok(plist) => plist,
    Err(e) => return Err(e),
  };

  // Within the `values` dictionary, each device is as follows:
  //
  //    <key>[[ device UUID ]]</key>
  //    <dict>
  //      [[ device data ]]
  //      <dict>
  //        <key>Tabs</key>
  //        <key>DeviceName</key><string>[[ device name ]]</key>
  //        <key>Tabs</key>
  //        <array>
  //          <dict>
  //            <key>URL</key><string>[[ URL ]]</string>
  //            [[ other tab data ]]
  //          </dict>
  //          [[ other tabs ]]
  //
  let mut result: HashMap<String, Vec<String>> = HashMap::new();
  let devices_with_tabs = plist
    .values()
    .map(|value| value
      .as_dictionary().unwrap()
      .get("value").unwrap()
      .as_dictionary().unwrap()
    );
  for data in devices_with_tabs {
    let name: String = data
      .get("DeviceName").unwrap()
      .as_string().unwrap()
      .to_owned();

    // If a device is registered with iCloud but Safari isn't running or
    // there aren't any tabs open, there isn't a Tabs field.
    if !data.contains_key("Tabs") {
      continue;
    }

    let urls: Vec<String> = data
      .get("Tabs").unwrap()
      .as_array().unwrap()
      .iter()
      .map(|tab| tab
        .as_dictionary().unwrap()
        .get("URL").unwrap()
        .as_string().unwrap()
      )
      .map(|url| urls::tidy_url(url))
      .collect();
    result.insert(name, urls);
  }
  Ok(result)
}
<|endoftext|>


<commit_before>use std::io;
use std::net::{Ipv4Addr, SocketAddrV4};
use std::net::UdpSocket;
use std::str;

use curl::http;

use xml;

use gateway::Gateway;

// Content of the request.
const SEARCH_REQUEST: &'static str =
"M-SEARCH * HTTP/1.1\r
Host:239.255.255.250:1900\r
ST:urn:schemas-upnp-org:device:InternetGatewayDevice:1\r
Man:\"ssdp:discover\"\r
MX:3\r\n\r\n";

// Error type this module emits.
#[derive(Debug)]
pub enum SearchError {
    IoError(io::Error),
    InvalidResponse,
}

impl From<io::Error> for SearchError {
    fn from(err: io::Error) -> SearchError {
        SearchError::IoError(err)
    }
}

// Try to find the gateway on the local network.
// Bind to the given interface.
pub fn search_gateway_from(ip: Ipv4Addr) -> Result<Gateway, SearchError> {
    let addr = SocketAddrV4::new(ip, 0);
    let socket = try!(UdpSocket::bind(addr));

    // send the request on the broadcast address
    try!(socket.send_to(SEARCH_REQUEST.as_bytes(), "239.255.255.250:1900"));
    let mut buf = [0u8; 1024];
    let (read, _) = try!(socket.recv_from(&mut buf));
    let text = str::from_utf8(&buf[..read]).unwrap();
    match parse_result(text) {
        None => Err(SearchError::InvalidResponse),
        Some(location) => {
            Ok(Gateway::new(location.0, match get_control_url(&location) {
                Some(u) => u,
                None => return Err(SearchError::InvalidResponse),
            }))
        },
    }
}

// Try to find the gateway on the local network.
// Bind to all interfaces.
pub fn search_gateway() -> Result<Gateway, SearchError> {
    search_gateway_from(Ipv4Addr::new(0, 0, 0, 0))
}

// Parse the result.
fn parse_result(text: &str) -> Option<(SocketAddrV4, String)> {
    let re = regex!(r"(?i:Location):\s*http://(\d+\.\d+\.\d+\.\d+):(\d+)(/.*)\r");
    for line in text.lines() {
        match re.captures(line) {
            None => continue,
            Some(cap) => {
                // these shouldn't fail if the regex matched.
                let addr = cap.at(1).unwrap();
                let port = cap.at(2).unwrap();
                return Some(
                    (SocketAddrV4::new(
                        addr.parse::<Ipv4Addr>().unwrap(),
                        port.parse::<u16>().unwrap()),
                     cap.at(3).unwrap().to_string()));
            },
        }
    }
    None
}

fn get_control_url(location: &(SocketAddrV4, String))
                   -> Option<String> {
    let resp = match http::handle()
        .get(format!("http://{}{}", location.0, location.1))
        .exec() {
            Ok(r) => r,
            Err(_) => return None,
        };
    let text = match str::from_utf8(resp.get_body()) {
        Ok(t) => t,
        Err(_) => return None,
    };
    let text = io::Cursor::new(text.as_bytes());
    let mut parser = xml::EventReader::new(text);
    let mut chain = Vec::<String>::with_capacity(4);

    struct Service {
        service_type: String,
        control_url: String,
    }

    let mut service = Service{
        service_type: "".to_string(),
        control_url: "".to_string(),
    };

    for e in parser.events() {
        match e {
            xml::reader::events::XmlEvent::StartElement { name, .. } => {
                chain.push(name.to_repr());
                let tail = if chain.len() >= 3 {
                    chain.iter().skip(chain.len() - 3)
                } else {
                    continue
                };

                if vec!["device", "serviceList", "service"].iter().zip(tail)
                    .all(|(l, r)| l == r) {
                    service.service_type.clear();
                    service.control_url.clear();
                }
            },
            xml::reader::events::XmlEvent::EndElement { .. } => {
                let top = chain.pop();
                let tail = if top == Some("service".to_string())
                    && chain.len() >= 2 {
                    chain.iter().skip(chain.len() - 2)
                } else {
                    continue
                };

                if vec!["device", "serviceList"].iter().zip(tail)
                    .all(|(l, r)| l == r) {
                    if "urn:schemas-upnp-org:service:WANIPConnection:1"
                        == service.service_type
                        && service.control_url.len() != 0 {
                        return Some(service.control_url);
                    }
                }
            },
            xml::reader::events::XmlEvent::Characters(text) => {
                let tail = if chain.len() >= 4 {
                    chain.iter().skip(chain.len() - 4)
                } else {
                    continue
                };

                if vec!["device", "serviceList", "service", "serviceType"]
                    .iter().zip(tail.clone()).all(|(l, r)| l == r) {
                    service.service_type.push_str(&text);
                }
                if vec!["device", "serviceList", "service", "controlURL"]
                    .iter().zip(tail).all(|(l, r)| l == r) {
                    service.control_url.push_str(&text);
                }
            },
            xml::reader::events::XmlEvent::Error(_) =>  return None,
            _ => (),
        }
    }
    None
}

#[test]
fn test_parse_result_case_insensitivity() {
    assert!(parse_result("location:http://0.0.0.0:0/").is_some());
    assert!(parse_result("LOCATION:http://0.0.0.0:0/").is_some());
}
<commit_msg>Repair the case insensitivity test and add new test for parse_result<commit_after>use std::io;
use std::net::{Ipv4Addr, SocketAddrV4};
use std::net::UdpSocket;
use std::str;

use curl::http;

use xml;

use gateway::Gateway;

// Content of the request.
const SEARCH_REQUEST: &'static str =
"M-SEARCH * HTTP/1.1\r
Host:239.255.255.250:1900\r
ST:urn:schemas-upnp-org:device:InternetGatewayDevice:1\r
Man:\"ssdp:discover\"\r
MX:3\r\n\r\n";

// Error type this module emits.
#[derive(Debug)]
pub enum SearchError {
    IoError(io::Error),
    InvalidResponse,
}

impl From<io::Error> for SearchError {
    fn from(err: io::Error) -> SearchError {
        SearchError::IoError(err)
    }
}

// Try to find the gateway on the local network.
// Bind to the given interface.
pub fn search_gateway_from(ip: Ipv4Addr) -> Result<Gateway, SearchError> {
    let addr = SocketAddrV4::new(ip, 0);
    let socket = try!(UdpSocket::bind(addr));

    // send the request on the broadcast address
    try!(socket.send_to(SEARCH_REQUEST.as_bytes(), "239.255.255.250:1900"));
    let mut buf = [0u8; 1024];
    let (read, _) = try!(socket.recv_from(&mut buf));
    let text = str::from_utf8(&buf[..read]).unwrap();
    match parse_result(text) {
        None => Err(SearchError::InvalidResponse),
        Some(location) => {
            Ok(Gateway::new(location.0, match get_control_url(&location) {
                Some(u) => u,
                None => return Err(SearchError::InvalidResponse),
            }))
        },
    }
}

// Try to find the gateway on the local network.
// Bind to all interfaces.
pub fn search_gateway() -> Result<Gateway, SearchError> {
    search_gateway_from(Ipv4Addr::new(0, 0, 0, 0))
}

// Parse the result.
fn parse_result(text: &str) -> Option<(SocketAddrV4, String)> {
    let re = regex!(r"(?i:Location):\s*http://(\d+\.\d+\.\d+\.\d+):(\d+)(/[^\r]*)");
    for line in text.lines() {
        match re.captures(line) {
            None => continue,
            Some(cap) => {
                // these shouldn't fail if the regex matched.
                let addr = cap.at(1).unwrap();
                let port = cap.at(2).unwrap();
                return Some(
                    (SocketAddrV4::new(
                        addr.parse::<Ipv4Addr>().unwrap(),
                        port.parse::<u16>().unwrap()),
                     cap.at(3).unwrap().to_string()));
            },
        }
    }
    None
}

fn get_control_url(location: &(SocketAddrV4, String))
                   -> Option<String> {
    let resp = match http::handle()
        .get(format!("http://{}{}", location.0, location.1))
        .exec() {
            Ok(r) => r,
            Err(_) => return None,
        };
    let text = match str::from_utf8(resp.get_body()) {
        Ok(t) => t,
        Err(_) => return None,
    };
    let text = io::Cursor::new(text.as_bytes());
    let mut parser = xml::EventReader::new(text);
    let mut chain = Vec::<String>::with_capacity(4);

    struct Service {
        service_type: String,
        control_url: String,
    }

    let mut service = Service{
        service_type: "".to_string(),
        control_url: "".to_string(),
    };

    for e in parser.events() {
        match e {
            xml::reader::events::XmlEvent::StartElement { name, .. } => {
                chain.push(name.to_repr());
                let tail = if chain.len() >= 3 {
                    chain.iter().skip(chain.len() - 3)
                } else {
                    continue
                };

                if vec!["device", "serviceList", "service"].iter().zip(tail)
                    .all(|(l, r)| l == r) {
                    service.service_type.clear();
                    service.control_url.clear();
                }
            },
            xml::reader::events::XmlEvent::EndElement { .. } => {
                let top = chain.pop();
                let tail = if top == Some("service".to_string())
                    && chain.len() >= 2 {
                    chain.iter().skip(chain.len() - 2)
                } else {
                    continue
                };

                if vec!["device", "serviceList"].iter().zip(tail)
                    .all(|(l, r)| l == r) {
                    if "urn:schemas-upnp-org:service:WANIPConnection:1"
                        == service.service_type
                        && service.control_url.len() != 0 {
                        return Some(service.control_url);
                    }
                }
            },
            xml::reader::events::XmlEvent::Characters(text) => {
                let tail = if chain.len() >= 4 {
                    chain.iter().skip(chain.len() - 4)
                } else {
                    continue
                };

                if vec!["device", "serviceList", "service", "serviceType"]
                    .iter().zip(tail.clone()).all(|(l, r)| l == r) {
                    service.service_type.push_str(&text);
                }
                if vec!["device", "serviceList", "service", "controlURL"]
                    .iter().zip(tail).all(|(l, r)| l == r) {
                    service.control_url.push_str(&text);
                }
            },
            xml::reader::events::XmlEvent::Error(_) =>  return None,
            _ => (),
        }
    }
    None
}

#[test]
fn test_parse_result_case_insensitivity() {
    assert!(parse_result("location:http://0.0.0.0:0/control_url").is_some());
    assert!(parse_result("LOCATION:http://0.0.0.0:0/control_url").is_some());
}

#[test]
fn test_parse_result() {
    let result = parse_result("location:http://0.0.0.0:0/control_url").unwrap();
    assert_eq!(result.0.ip(), &Ipv4Addr::new(0,0,0,0));
    assert_eq!(result.0.port(), 0);
    assert_eq!(&result.1[..], "/control_url");
}
<|endoftext|>


<commit_before>use super::command::{Command, Query};
use super::command::Command::Select;
use std::collections::HashMap;
use util;
use command_query::CommandQuery;
use queryable::Queryable;
use commandable::Commandable;
use command_line::CommandLine;
use selectable::fragmentable::Fragmentable;
use selectable::fragmentable::{OrderedFragment, QueryFragment};

#[derive(Clone, PartialEq, Eq, Debug)]
pub struct SelectCommand {
    command: Command,
    table: String,
    arguments: HashMap<String, String>,
}

impl Default for SelectCommand {
    fn default() -> SelectCommand {
        SelectCommand {
            command: Select,
            table: "".to_string(),
            arguments: HashMap::new(),
        }
    }
}

impl SelectCommand {
    pub fn new(table: String) -> SelectCommand {
        let default: SelectCommand = Default::default();
        SelectCommand {
            command: default.command,
            table: table,
            arguments: default.arguments,
        }
    }

    pub fn filter(mut self, filter: String) -> SelectCommand {
        let encoded = format!("'{}'", filter);
        self.arguments.insert("filter".to_string(), encoded.clone());
        self
    }

    pub fn query(mut self, query: String) -> SelectCommand {
        let encoded = format!("'{}'", query);
        self.arguments.insert("query".to_string(), encoded.clone());
        self
    }

    pub fn match_columns(mut self, columns: Vec<String>) -> SelectCommand {
        let string = util::split_values_vec(columns);
        self.arguments.insert("match_columns".to_string(), string.clone());
        self
    }

    pub fn output_columns(mut self, columns: Vec<String>) -> SelectCommand {
        let string = util::split_values_vec(columns);
        self.arguments.insert("output_columns".to_string(), string.clone());
        self
    }

    pub fn offset(mut self, offset: u64) -> SelectCommand {
        let offset = format!("{}", offset);
        self.arguments.insert("offset".to_string(), offset.clone());
        self
    }

    pub fn limit(mut self, limit: u64) -> SelectCommand {
        let limit = format!("{}", limit);
        self.arguments.insert("limit".to_string(), limit.clone());
        self
    }

    pub fn cache(mut self, cache: bool) -> SelectCommand {
        let flag = match cache {
            true => "yes",
            false => "no",
        };
        self.arguments.insert("cache".to_string(), flag.to_string());
        self
    }

    pub fn build(self) -> (Command, Query) {
        let mut query: Query = vec![("table".to_string(), self.table)];
        for (key, value) in self.arguments.iter() {
            query.push((key.to_owned(), value.to_owned()));
        }
        (Select, query)
    }
}

impl Queryable for SelectCommand {
    fn to_query(self) -> String {
        let (command, query) = self.build();
        let mut command = CommandQuery::new(command, query);
        command.encode()
    }
}

impl Commandable for SelectCommand {
    fn to_command(self) -> String {
        let (command, query) = self.build();
        let mut command = CommandLine::new(command, query);
        command.encode()
    }
}

impl Fragmentable for SelectCommand {
    fn to_fragment(self) -> (OrderedFragment, QueryFragment) {
        let mut select_fragment = HashMap::new();
        let ordered_fragment = vec![("table".to_string(), self.table)];
        for (key, value) in self.arguments.clone() {
            select_fragment.insert(key, value);
        }
        (ordered_fragment, select_fragment)
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use command::Query;
    use command::Command::Select;
    use std::collections::HashMap;
    use queryable::Queryable;
    use commandable::Commandable;

    #[test]
    fn test_new() {
        let vanilla_select = SelectCommand::new("test".to_string());
        let expected = SelectCommand {
            command: Select,
            table: "test".to_string(),
            arguments: HashMap::new(),
        };
        assert_eq!(expected, vanilla_select);
    }

    #[test]
    fn test_filter() {
        let select = SelectCommand::new("test".to_string())
                         .filter("output_column @ \"type_safe\"".to_string());
        let mut arg: HashMap<String, String> = HashMap::new();
        arg.insert("filter".to_string(),
                   "'output_column @ \"type_safe\"'".to_string());
        let expected = SelectCommand {
            command: Select,
            table: "test".to_string(),
            arguments: arg,
        };
        assert_eq!(expected, select);
    }

    #[test]
    fn test_query() {
        let select = SelectCommand::new("test".to_string())
                         .query("_key:\"http://example.org/\"".to_string());
        let mut arg: HashMap<String, String> = HashMap::new();
        arg.insert("query".to_string(),
                   "\'_key:\"http://example.org/\"\'".to_string());
        let expected = SelectCommand {
            command: Select,
            table: "test".to_string(),
            arguments: arg,
        };
        assert_eq!(expected, select);
    }

    #[test]
    fn test_match_columns() {
        let select = SelectCommand::new("test".to_string())
                         .match_columns(vec!["test".to_string(), "piyo".to_string()]);
        let mut arg: HashMap<String, String> = HashMap::new();
        arg.insert("match_columns".to_string(), "test,piyo".to_string());
        let expected = SelectCommand {
            command: Select,
            table: "test".to_string(),
            arguments: arg,
        };
        assert_eq!(expected, select);
    }

    #[test]
    fn test_output_columns() {
        let select = SelectCommand::new("test".to_string())
                         .output_columns(vec!["test".to_string(), "piyo".to_string()]);
        let mut arg: HashMap<String, String> = HashMap::new();
        arg.insert("output_columns".to_string(), "test,piyo".to_string());
        let expected = SelectCommand {
            command: Select,
            table: "test".to_string(),
            arguments: arg,
        };
        assert_eq!(expected, select);
    }

    #[test]
    fn test_offset() {
        let select = SelectCommand::new("test".to_string()).offset(100);
        let mut arg: HashMap<String, String> = HashMap::new();
        arg.insert("offset".to_string(), "100".to_string());
        let expected = SelectCommand {
            command: Select,
            arguments: arg,
            table: "test".to_string(),
        };
        assert_eq!(expected, select);
    }

    #[test]
    fn test_limit() {
        let select = SelectCommand::new("test".to_string()).limit(50);
        let mut arg: HashMap<String, String> = HashMap::new();
        arg.insert("limit".to_string(), "50".to_string());
        let expected = SelectCommand {
            command: Select,
            table: "test".to_string(),
            arguments: arg,
        };
        assert_eq!(expected, select);
    }

    #[test]
    fn test_cache() {
        let select_yes = SelectCommand::new("test".to_string()).cache(true);
        let mut arg_yes: HashMap<String, String> = HashMap::new();
        arg_yes.insert("cache".to_string(), "yes".to_string());
        let expected_yes = SelectCommand {
            command: Select,
            table: "test".to_string(),
            arguments: arg_yes,
        };
        assert_eq!(expected_yes, select_yes);
        let select_no = SelectCommand::new("test".to_string()).cache(false);
        let mut arg_no: HashMap<String, String> = HashMap::new();
        arg_no.insert("cache".to_string(), "no".to_string());
        let expected_no = SelectCommand {
            command: Select,
            table: "test".to_string(),
            arguments: arg_no,
        };
        assert_eq!(expected_no, select_no);
    }

    #[test]
    fn test_build() {
        let actual = SelectCommand::new("test".to_string())
                         .filter("output_column @ \"type_safe\"".to_string())
                         .build();
        let expected_query: Query = vec![("table".to_string(), "test".to_string()),
                                         ("filter".to_string(),
                                          "'output_column @ \"type_safe\"'".to_string())];
        let expected = (Select, expected_query);
        assert_eq!(expected, actual);
    }

    #[test]
    fn test_queryable() {
        let query = SelectCommand::new("Test".to_string())
                        .filter("output_column @ \"type_safe\"".to_string())
                        .to_query();
        let url_encoded = "/d/select?table=Test&filter=%27output_column+%40+%22type_safe%22%27";
        assert_eq!(url_encoded.to_string(), query);
    }

    #[test]
    fn test_commandable() {
        let query = SelectCommand::new("Test".to_string())
                        .filter("output_column @ \"type_safe\"".to_string())
                        .to_command();
        let url_encoded = "select --table Test --filter \'output_column @ \"type_safe\"\'";
        assert_eq!(url_encoded.to_string(), query);
    }
}
<commit_msg>select: Add Add operator overload to concat with drilldown type<commit_after>use super::command::{Command, Query};
use super::command::Command::Select;
use std::collections::HashMap;
use util;
use command_query::CommandQuery;
use queryable::Queryable;
use commandable::Commandable;
use command_line::CommandLine;
use selectable::fragmentable::Fragmentable;
use selectable::fragmentable::{OrderedFragment, QueryFragment};
use selectable::drilldown::Drilldown;
use selectable::drilldown_builder::DrilldownBuilder;
use std::ops::Add;

#[derive(Clone, PartialEq, Eq, Debug)]
pub struct SelectCommand {
    command: Command,
    table: String,
    arguments: HashMap<String, String>,
}

impl Default for SelectCommand {
    fn default() -> SelectCommand {
        SelectCommand {
            command: Select,
            table: "".to_string(),
            arguments: HashMap::new(),
        }
    }
}

impl SelectCommand {
    pub fn new(table: String) -> SelectCommand {
        let default: SelectCommand = Default::default();
        SelectCommand {
            command: default.command,
            table: table,
            arguments: default.arguments,
        }
    }

    pub fn filter(mut self, filter: String) -> SelectCommand {
        let encoded = format!("'{}'", filter);
        self.arguments.insert("filter".to_string(), encoded.clone());
        self
    }

    pub fn query(mut self, query: String) -> SelectCommand {
        let encoded = format!("'{}'", query);
        self.arguments.insert("query".to_string(), encoded.clone());
        self
    }

    pub fn match_columns(mut self, columns: Vec<String>) -> SelectCommand {
        let string = util::split_values_vec(columns);
        self.arguments.insert("match_columns".to_string(), string.clone());
        self
    }

    pub fn output_columns(mut self, columns: Vec<String>) -> SelectCommand {
        let string = util::split_values_vec(columns);
        self.arguments.insert("output_columns".to_string(), string.clone());
        self
    }

    pub fn offset(mut self, offset: u64) -> SelectCommand {
        let offset = format!("{}", offset);
        self.arguments.insert("offset".to_string(), offset.clone());
        self
    }

    pub fn limit(mut self, limit: u64) -> SelectCommand {
        let limit = format!("{}", limit);
        self.arguments.insert("limit".to_string(), limit.clone());
        self
    }

    pub fn cache(mut self, cache: bool) -> SelectCommand {
        let flag = match cache {
            true => "yes",
            false => "no",
        };
        self.arguments.insert("cache".to_string(), flag.to_string());
        self
    }

    pub fn build(self) -> (Command, Query) {
        let mut query: Query = vec![("table".to_string(), self.table)];
        for (key, value) in self.arguments.iter() {
            query.push((key.to_owned(), value.to_owned()));
        }
        (Select, query)
    }
}

impl Queryable for SelectCommand {
    fn to_query(self) -> String {
        let (command, query) = self.build();
        let mut command = CommandQuery::new(command, query);
        command.encode()
    }
}

impl Commandable for SelectCommand {
    fn to_command(self) -> String {
        let (command, query) = self.build();
        let mut command = CommandLine::new(command, query);
        command.encode()
    }
}

impl Fragmentable for SelectCommand {
    fn to_fragment(self) -> (OrderedFragment, QueryFragment) {
        let mut select_fragment = HashMap::new();
        let ordered_fragment = vec![("table".to_string(), self.table)];
        for (key, value) in self.arguments.clone() {
            select_fragment.insert(key, value);
        }
        (ordered_fragment, select_fragment)
    }
}

impl Add<Drilldown> for SelectCommand {
    type Output = DrilldownBuilder;

    fn add(self, rhs: Drilldown) -> DrilldownBuilder {
        let drilldown_builder = DrilldownBuilder::new(self, rhs);
        drilldown_builder
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use command::Query;
    use command::Command::Select;
    use std::collections::HashMap;
    use queryable::Queryable;
    use commandable::Commandable;
    use selectable::drilldown::Drilldown;
    use selectable::drilldown_builder::DrilldownBuilder;

    #[test]
    fn test_new() {
        let vanilla_select = SelectCommand::new("test".to_string());
        let expected = SelectCommand {
            command: Select,
            table: "test".to_string(),
            arguments: HashMap::new(),
        };
        assert_eq!(expected, vanilla_select);
    }

    #[test]
    fn test_filter() {
        let select = SelectCommand::new("test".to_string())
                         .filter("output_column @ \"type_safe\"".to_string());
        let mut arg: HashMap<String, String> = HashMap::new();
        arg.insert("filter".to_string(),
                   "'output_column @ \"type_safe\"'".to_string());
        let expected = SelectCommand {
            command: Select,
            table: "test".to_string(),
            arguments: arg,
        };
        assert_eq!(expected, select);
    }

    #[test]
    fn test_query() {
        let select = SelectCommand::new("test".to_string())
                         .query("_key:\"http://example.org/\"".to_string());
        let mut arg: HashMap<String, String> = HashMap::new();
        arg.insert("query".to_string(),
                   "\'_key:\"http://example.org/\"\'".to_string());
        let expected = SelectCommand {
            command: Select,
            table: "test".to_string(),
            arguments: arg,
        };
        assert_eq!(expected, select);
    }

    #[test]
    fn test_match_columns() {
        let select = SelectCommand::new("test".to_string())
                         .match_columns(vec!["test".to_string(), "piyo".to_string()]);
        let mut arg: HashMap<String, String> = HashMap::new();
        arg.insert("match_columns".to_string(), "test,piyo".to_string());
        let expected = SelectCommand {
            command: Select,
            table: "test".to_string(),
            arguments: arg,
        };
        assert_eq!(expected, select);
    }

    #[test]
    fn test_output_columns() {
        let select = SelectCommand::new("test".to_string())
                         .output_columns(vec!["test".to_string(), "piyo".to_string()]);
        let mut arg: HashMap<String, String> = HashMap::new();
        arg.insert("output_columns".to_string(), "test,piyo".to_string());
        let expected = SelectCommand {
            command: Select,
            table: "test".to_string(),
            arguments: arg,
        };
        assert_eq!(expected, select);
    }

    #[test]
    fn test_offset() {
        let select = SelectCommand::new("test".to_string()).offset(100);
        let mut arg: HashMap<String, String> = HashMap::new();
        arg.insert("offset".to_string(), "100".to_string());
        let expected = SelectCommand {
            command: Select,
            arguments: arg,
            table: "test".to_string(),
        };
        assert_eq!(expected, select);
    }

    #[test]
    fn test_limit() {
        let select = SelectCommand::new("test".to_string()).limit(50);
        let mut arg: HashMap<String, String> = HashMap::new();
        arg.insert("limit".to_string(), "50".to_string());
        let expected = SelectCommand {
            command: Select,
            table: "test".to_string(),
            arguments: arg,
        };
        assert_eq!(expected, select);
    }

    #[test]
    fn test_cache() {
        let select_yes = SelectCommand::new("test".to_string()).cache(true);
        let mut arg_yes: HashMap<String, String> = HashMap::new();
        arg_yes.insert("cache".to_string(), "yes".to_string());
        let expected_yes = SelectCommand {
            command: Select,
            table: "test".to_string(),
            arguments: arg_yes,
        };
        assert_eq!(expected_yes, select_yes);
        let select_no = SelectCommand::new("test".to_string()).cache(false);
        let mut arg_no: HashMap<String, String> = HashMap::new();
        arg_no.insert("cache".to_string(), "no".to_string());
        let expected_no = SelectCommand {
            command: Select,
            table: "test".to_string(),
            arguments: arg_no,
        };
        assert_eq!(expected_no, select_no);
    }

    #[test]
    fn test_build() {
        let actual = SelectCommand::new("test".to_string())
                         .filter("output_column @ \"type_safe\"".to_string())
                         .build();
        let expected_query: Query = vec![("table".to_string(), "test".to_string()),
                                         ("filter".to_string(),
                                          "'output_column @ \"type_safe\"'".to_string())];
        let expected = (Select, expected_query);
        assert_eq!(expected, actual);
    }

    #[test]
    fn test_queryable() {
        let query = SelectCommand::new("Test".to_string())
                        .filter("output_column @ \"type_safe\"".to_string())
                        .to_query();
        let url_encoded = "/d/select?table=Test&filter=%27output_column+%40+%22type_safe%22%27";
        assert_eq!(url_encoded.to_string(), query);
    }

    #[test]
    fn test_commandable() {
        let query = SelectCommand::new("Test".to_string())
                        .filter("output_column @ \"type_safe\"".to_string())
                        .to_command();
        let url_encoded = "select --table Test --filter \'output_column @ \"type_safe\"\'";
        assert_eq!(url_encoded.to_string(), query);
    }

    #[test]
    fn test_add_ops_drilldown() {
        let select = SelectCommand::new("Entries".to_string())
                         .filter("content @ \"fast\"".to_string());
        let drilldown = Drilldown::new().drilldown("tag".to_string());
        let ops_builder = (select.clone() + drilldown.clone()).build();
        let drilldown_builder = DrilldownBuilder::new(select.clone(), drilldown.clone()).build();
        assert_eq!(ops_builder, drilldown_builder);
    }
}
<|endoftext|>


<commit_before>#![feature(phase, globs, slicing_syntax, if_let, unboxed_closures)]

extern crate getopts;
extern crate serialize;
extern crate iron;
extern crate persistent;
extern crate error;
extern crate regex;
extern crate "conduit-mime-types" as conduit_mime;
extern crate mime;
extern crate mustache;

#[phase(plugin)]
extern crate lazy_static;

use std::{ str, os };
use std::str::from_str;
use std::path::{ Path, GenericPath };
use std::io::{ fs, Reader };
use std::io::fs::{ File, PathExtensions };
use std::default::Default;
use std::sync::{ Arc };

use regex::Regex;

use conduit_mime::Types;

use getopts::{ optopt, optflag, getopts, usage, OptGroup };

use serialize::json;
use serialize::json::Json;

use iron::prelude::*;
use iron::response::modifiers::*;
use iron::status;
use iron::mime::*;
use iron::middleware::ChainBuilder;
use iron::typemap::Assoc;

use persistent::Read;

use mustache::{ Template, VecBuilder, MapBuilder };

pub mod constants;

#[deriving(Send, Clone, Default, Encodable, Decodable)]
struct Options {
    host: Option<String>,
    port: Option<u16>,
    root: Option<Path>,
    filter: Option<String>,
    max_size: Option<u64>,
    template: Option<String>
}

struct OptCarrier;
impl Assoc<Arc<Options>> for OptCarrier {}

#[deriving(Send, Clone)]
struct State {
    template: Template
}

struct StateCarrier;
impl Assoc<Arc<State>> for StateCarrier {}

static UNITS: &'static [&'static str] = &["B", "kB", "MB", "GB", "TB"];
const BRIEF: &'static str = "A minimal directory server, written in Rust with Iron.";

const KEY_TITLE: &'static str = "title";
const KEY_CONTENT: &'static str = "content";
const KEY_URL: &'static str = "url";
const KEY_SIZE: &'static str = "size";
const KEY_NAME: &'static str = "name";

lazy_static! {
    static ref MIME: Types = Types::new().ok().unwrap();
}

fn size_with_unit(mut size: u64) -> String {
    let mut frac = 0;
    let mut index = 0;

    while size > 1000 && index + 1 < UNITS.len() {
        frac = size % 1000;
        size /= 1000;
        index += 1;
    }

    format!("{}.{} {}", size, frac, UNITS[index])
}

fn render(template: Template, root: Path, dir: Path, files: Vec<Path>, filter: Option<Regex>) -> String {
    let data = MapBuilder::new()
        .insert_str(KEY_TITLE, dir.display().as_cow().into_owned())
        .insert_vec(KEY_CONTENT, |mut vec: VecBuilder| {
            let item = |map: MapBuilder, url: &Path, size: u64, name: String| {
                map.insert(KEY_URL, &format!("{}", url.display())[]).unwrap()
                   .insert(KEY_SIZE, &size_with_unit(size)[]).unwrap()
                   .insert_str(KEY_NAME, name)
            };
            let mut up = dir.clone();
            up.pop();
            if root.is_ancestor_of(&up) {
                vec = vec.push_map(|map: MapBuilder| item(map, &up.path_relative_from(&root).unwrap(), 0, "..".into_string()));
            }

            for file in files.iter() {
                let relative = file.path_relative_from(&root).unwrap();
                let stat = file.stat().unwrap();
                let filename = file.filename_display().as_cow().into_owned();
                if filter.as_ref().map_or(true, |f| f.is_match(filename[])) {
                    vec = vec.push_map(|map| item(map, &relative, stat.size, filename.clone()));
                }
            }
            vec
        }).build();

    let mut out = Vec::new(); // with_capacity(template.len())
    template.render_data(&mut out, &data);
    unsafe { String::from_utf8_unchecked(out) }
}

fn plain<B: Bodyable>(content: B) -> IronResult<Response> {
    Ok(Response::new()
       .set(Status(status::Ok))
       .set(Body(content)))
}

fn html<B: Bodyable>(content: B) -> IronResult<Response> {
    plain(content).map(|r| r.set(ContentType(Mime(Text, Html, Vec::new()))))
}

fn serve(req: &mut Request) -> IronResult<Response> {
    let (root, filter_str, max_size) = {
        let o = req.get::<Read<OptCarrier, Arc<Options>>>().unwrap();
        (o.root.clone().unwrap_or_else(|| os::getcwd().ok().unwrap()),
         o.filter.clone(),
         o.max_size)
    };

    let template = {
        let s = req.get::<Read<StateCarrier, Arc<State>>>().unwrap();
        s.template.clone()
    };

    let mut path = root.clone();
    for part in req.url.path.iter() { path.push(part[]) }
    if !path.exists() { return html("Well, no... We don't have that today."); }

    let filter = filter_str.and_then(|s| Regex::new(s[]).ok());

    if path.is_file() && root.is_ancestor_of(&path) {
        let stat = path.stat();
        if stat.as_ref().ok().is_some() && max_size.is_some() && stat.ok().unwrap().size > max_size.unwrap() {
            return html("I'm afraid, I'm too lazy to serve the requested file. It's pretty big...")
        }
        let content = match File::open(&path).read_to_end() {
            Ok(s) => s,
            Err(e) => return html(e.desc)
        };

        if filter.as_ref().map_or(false, |f| !f.is_match(path.filename_str().unwrap())) {
            return html("I don't think you're allowed to do this.");
        }
        let mime: Option<iron::mime::Mime> = path.extension_str()
            .map_or(None, |e| MIME.get_mime_type(e))
            .map_or(None, |m| from_str(m));
        if mime.as_ref().is_some() {
            plain(content[]).map(|r| r.set(ContentType((*mime.as_ref().unwrap()).clone())))
        } else {
            plain(content[])
        }
    } else {
        let mut content = match fs::readdir(&path) {
            Ok(s) => s,
            Err(e) => return html(e.desc)
        };
        content.sort_by(|a, b| a.filename_str().unwrap().cmp(b.filename_str().unwrap()));
        html(render(template, root, path, content, filter)[])
    }
}

fn print_usage(program: &str, opts: &[OptGroup]) {
    println!("Usage: {} [options]\n", program);
    println!("{}", usage(BRIEF, opts));
}

fn main() {
    let args: Vec<String> = os::args();
    let program = args[0].clone();
    let opts = &[
        optopt("c", "config", "set config file name", "NAME"),
        optopt("a", "address", "the address to bind to", "HOST"),
        optopt("p", "port", "the port to serve", "PORT"),
        optopt("r", "root", "the uppermost directory to serve", "ROOT"),
        optopt("f", "filter", "a regular expression to filter the filenames", "REGEX"),
        optopt("s", "size", "the maximum size of a file that will be served", "BYTES"),
        optopt("t", "template", "a mustache template to use for rendering", "TEMPLATE"),
        optflag("h", "help", "print this help menu")
    ];
    let matches = match getopts(args.tail(), opts) {
        Ok(m) => { m }
        Err(f) => {
            println!("{}", f.to_string());
            return;
        }
    };

    if matches.opt_present("h") {
        print_usage(program[], opts);
        return;
    }

    let mut options: Options = Default::default();

    matches.opt_str("c").map(|conf_file| {
        let conf_file = File::open(&Path::new(conf_file));
        conf_file.as_ref().map_err::<()>(|e| panic!("{}", e.desc)).unwrap();
        let json = match json::from_reader(&mut conf_file.ok().unwrap()) {
            Ok(Json::Object(o)) => o,
            _ => panic!("Invalid configuration file. Doesn't contain top-level object.")
        };
        options.host = match json.get("address") {
            Some(&Json::String(ref s)) => Some((*s).clone()),
            None => None,
            _ => panic!("Invalid configuration file. `address` field must be a string.")
        };
        options.port = match json.get("port") {
            Some(&Json::U64(u)) => Some(u as u16),
            None => None,
            _ => panic!("Invalid configuration file. `port` field must be an unsigned integer.")
        };
        options.root = match json.get("root") {
            Some(&Json::String(ref s)) => Some(Path::new((*s).clone())),
            None => None,
            _ => panic!("Invalid configuration file. `root` field must be a string.")
        };
        options.filter = match json.get("filter") {
            Some(&Json::String(ref s)) => Some((*s).clone()),
            None => None,
            _ => panic!("Invalid configuration file. `filter` field must be a string.")
        };
        options.max_size = match json.get("size") {
            Some(&Json::U64(u)) => Some(u),
            None => None,
            _ => panic!("Invalid configuration file. `size` field must be an unsigned integer.")
        };
        options.template = match json.get("template") {
            Some(&Json::String(ref s)) => Some((*s).clone()),
            None => None,
            _ => panic!("Invalid configuration file. `template` field must be a string.")
        };
    });

    let (host, port) = {
        options.host = options.host.clone().or(matches.opt_str("a"));
        options.port = options.port.or(matches.opt_str("p").and_then(|p| str::from_str(p[])));
        options.root = options.root.clone().or(matches.opt_str("r").and_then(|p| Path::new_opt(p)));
        options.filter = options.filter.clone().or(matches.opt_str("f"));
        options.max_size = options.max_size.or(matches.opt_str("s").and_then(|s| str::from_str(s[])));
        options.template = options.template.clone().or(matches.opt_str("t"));
        (options.host.clone().unwrap_or("0.0.0.0".into_string()),
         options.port.clone().unwrap_or(8080))
    };

    let template = mustache::compile_str(options.template.clone().unwrap_or(constants::OPT_TEMPLATE.into_string())[]);
    let state = State {
        template: template
    };

    let mut chain = ChainBuilder::new(serve);
    chain.link(Read::<OptCarrier, Arc<Options>>::both(Arc::new(options)));
    chain.link(Read::<StateCarrier, Arc<State>>::both(Arc::new(state)));
    match Iron::new(chain).listen((host[], port)) {
        Ok(_) => (),
        Err(e) => println!("I'm sorry, I failed you. {}", e)
    }
}
<commit_msg>Fixed configuration overriding<commit_after>#![feature(phase, globs, slicing_syntax, if_let, unboxed_closures)]

extern crate getopts;
extern crate serialize;
extern crate iron;
extern crate persistent;
extern crate error;
extern crate regex;
extern crate "conduit-mime-types" as conduit_mime;
extern crate mime;
extern crate mustache;

#[phase(plugin)]
extern crate lazy_static;

use std::{ str, os };
use std::str::from_str;
use std::path::{ Path, GenericPath };
use std::io::{ fs, Reader };
use std::io::fs::{ File, PathExtensions };
use std::default::Default;
use std::sync::{ Arc };

use regex::Regex;

use conduit_mime::Types;

use getopts::{ optopt, optflag, getopts, usage, OptGroup };

use serialize::json;
use serialize::json::Json;

use iron::prelude::*;
use iron::response::modifiers::*;
use iron::status;
use iron::mime::*;
use iron::middleware::ChainBuilder;
use iron::typemap::Assoc;

use persistent::Read;

use mustache::{ Template, VecBuilder, MapBuilder };

pub mod constants;

#[deriving(Send, Clone, Default, Encodable, Decodable)]
struct Options {
    host: Option<String>,
    port: Option<u16>,
    root: Option<Path>,
    filter: Option<String>,
    max_size: Option<u64>,
    template: Option<String>
}

struct OptCarrier;
impl Assoc<Arc<Options>> for OptCarrier {}

#[deriving(Send, Clone)]
struct State {
    template: Template
}

struct StateCarrier;
impl Assoc<Arc<State>> for StateCarrier {}

static UNITS: &'static [&'static str] = &["B", "kB", "MB", "GB", "TB"];
const BRIEF: &'static str = "A minimal directory server, written in Rust with Iron.";

const KEY_TITLE: &'static str = "title";
const KEY_CONTENT: &'static str = "content";
const KEY_URL: &'static str = "url";
const KEY_SIZE: &'static str = "size";
const KEY_NAME: &'static str = "name";

lazy_static! {
    static ref MIME: Types = Types::new().ok().unwrap();
}

fn size_with_unit(mut size: u64) -> String {
    let mut frac = 0;
    let mut index = 0;

    while size > 1000 && index + 1 < UNITS.len() {
        frac = size % 1000;
        size /= 1000;
        index += 1;
    }

    format!("{}.{} {}", size, frac, UNITS[index])
}

fn render(template: Template, root: Path, dir: Path, files: Vec<Path>, filter: Option<Regex>) -> String {
    let data = MapBuilder::new()
        .insert_str(KEY_TITLE, dir.display().as_cow().into_owned())
        .insert_vec(KEY_CONTENT, |mut vec: VecBuilder| {
            let item = |map: MapBuilder, url: &Path, size: u64, name: String| {
                map.insert(KEY_URL, &format!("{}", url.display())[]).unwrap()
                   .insert(KEY_SIZE, &size_with_unit(size)[]).unwrap()
                   .insert_str(KEY_NAME, name)
            };
            let mut up = dir.clone();
            up.pop();
            if root.is_ancestor_of(&up) {
                vec = vec.push_map(|map: MapBuilder| item(map, &up.path_relative_from(&root).unwrap(), 0, "..".into_string()));
            }

            for file in files.iter() {
                let relative = file.path_relative_from(&root).unwrap();
                let stat = file.stat().unwrap();
                let filename = file.filename_display().as_cow().into_owned();
                if filter.as_ref().map_or(true, |f| f.is_match(filename[])) {
                    vec = vec.push_map(|map| item(map, &relative, stat.size, filename.clone()));
                }
            }
            unsafe { std::mem::transmute(vec) }
        }).build();

    let mut out = Vec::new(); // with_capacity(template.len())
    template.render_data(&mut out, &data);
    unsafe { String::from_utf8_unchecked(out) }
}

fn plain<B: Bodyable>(content: B) -> IronResult<Response> {
    Ok(Response::new()
       .set(Status(status::Ok))
       .set(Body(content)))
}

fn html<B: Bodyable>(content: B) -> IronResult<Response> {
    plain(content).map(|r| r.set(ContentType(Mime(Text, Html, Vec::new()))))
}

fn serve(req: &mut Request) -> IronResult<Response> {
    let (root, filter_str, max_size) = {
        let o = req.get::<Read<OptCarrier, Arc<Options>>>().unwrap();
        (o.root.clone().unwrap_or_else(|| os::getcwd().ok().unwrap()),
         o.filter.clone(),
         o.max_size)
    };

    let template = {
        let s = req.get::<Read<StateCarrier, Arc<State>>>().unwrap();
        s.template.clone()
    };

    let mut path = root.clone();
    for part in req.url.path.iter() { path.push(part[]) }
    if !path.exists() { return html("Well, no... We don't have that today."); }

    let filter = filter_str.and_then(|s| Regex::new(s[]).ok());

    if path.is_file() && root.is_ancestor_of(&path) {
        let stat = path.stat();
        if stat.as_ref().ok().is_some() && max_size.is_some() && stat.ok().unwrap().size > max_size.unwrap() {
            return html("I'm afraid, I'm too lazy to serve the requested file. It's pretty big...")
        }
        let content = match File::open(&path).read_to_end() {
            Ok(s) => s,
            Err(e) => return html(e.desc)
        };

        if filter.as_ref().map_or(false, |f| !f.is_match(path.filename_str().unwrap())) {
            return html("I don't think you're allowed to do this.");
        }
        let mime: Option<iron::mime::Mime> = path.extension_str()
            .map_or(None, |e| MIME.get_mime_type(e))
            .map_or(None, |m| from_str(m));
        if mime.as_ref().is_some() {
            plain(content[]).map(|r| r.set(ContentType((*mime.as_ref().unwrap()).clone())))
        } else {
            plain(content[])
        }
    } else {
        let mut content = match fs::readdir(&path) {
            Ok(s) => s,
            Err(e) => return html(e.desc)
        };
        content.sort_by(|a, b| a.filename_str().unwrap().cmp(b.filename_str().unwrap()));
        html(render(template, root, path, content, filter)[])
    }
}

fn print_usage(program: &str, opts: &[OptGroup]) {
    println!("Usage: {} [options]\n", program);
    println!("{}", usage(BRIEF, opts));
}

fn main() {
    let args: Vec<String> = os::args();
    let program = args[0].clone();
    let opts = &[
        optopt("c", "config", "set config file name", "NAME"),
        optopt("a", "address", "the address to bind to", "HOST"),
        optopt("p", "port", "the port to serve", "PORT"),
        optopt("r", "root", "the uppermost directory to serve", "ROOT"),
        optopt("f", "filter", "a regular expression to filter the filenames", "REGEX"),
        optopt("s", "size", "the maximum size of a file that will be served", "BYTES"),
        optopt("t", "template", "a mustache template to use for rendering", "TEMPLATE"),
        optflag("h", "help", "print this help menu")
    ];
    let matches = match getopts(args.tail(), opts) {
        Ok(m) => { m }
        Err(f) => {
            println!("{}", f.to_string());
            return;
        }
    };

    if matches.opt_present("h") {
        print_usage(program[], opts);
        return;
    }

    let mut options: Options = Default::default();

    matches.opt_str("c").map(|conf_file| {
        let conf_file = File::open(&Path::new(conf_file));
        conf_file.as_ref().map_err::<()>(|e| panic!("{}", e.desc)).unwrap();
        let json = match json::from_reader(&mut conf_file.ok().unwrap()) {
            Ok(Json::Object(o)) => o,
            _ => panic!("Invalid configuration file. Doesn't contain top-level object.")
        };
        options.host = match json.get("address") {
            Some(&Json::String(ref s)) => Some((*s).clone()),
            None => None,
            _ => panic!("Invalid configuration file. `address` field must be a string.")
        };
        options.port = match json.get("port") {
            Some(&Json::U64(u)) => Some(u as u16),
            None => None,
            _ => panic!("Invalid configuration file. `port` field must be an unsigned integer.")
        };
        options.root = match json.get("root") {
            Some(&Json::String(ref s)) => Some(Path::new((*s).clone())),
            None => None,
            _ => panic!("Invalid configuration file. `root` field must be a string.")
        };
        options.filter = match json.get("filter") {
            Some(&Json::String(ref s)) => Some((*s).clone()),
            None => None,
            _ => panic!("Invalid configuration file. `filter` field must be a string.")
        };
        options.max_size = match json.get("size") {
            Some(&Json::U64(u)) => Some(u),
            None => None,
            _ => panic!("Invalid configuration file. `size` field must be an unsigned integer.")
        };
        options.template = match json.get("template") {
            Some(&Json::String(ref s)) => Some((*s).clone()),
            None => None,
            _ => panic!("Invalid configuration file. `template` field must be a string.")
        };
    });

    let (host, port) = {
        options.host = matches.opt_str("a").or(options.host);
        options.port = matches.opt_str("p").and_then(|p| str::from_str(p[])).or(options.port);
        options.root = matches.opt_str("r").and_then(|p| Path::new_opt(p)).or(options.root);
        options.filter = matches.opt_str("f").or(options.filter);
        options.max_size = matches.opt_str("s").and_then(|s| str::from_str(s[])).or(options.max_size);
        options.template = matches.opt_str("t").or(options.template);
        (options.host.clone().unwrap_or("0.0.0.0".into_string()),
         options.port.clone().unwrap_or(8080))
    };

    let template = mustache::compile_str(options.template.clone().unwrap_or(constants::OPT_TEMPLATE.into_string())[]);
    let state = State {
        template: template
    };

    let mut chain = ChainBuilder::new(serve);
    chain.link(Read::<OptCarrier, Arc<Options>>::both(Arc::new(options)));
    chain.link(Read::<StateCarrier, Arc<State>>::both(Arc::new(state)));
    match Iron::new(chain).listen((host[], port)) {
        Ok(_) => (),
        Err(e) => println!("I'm sorry, I failed you. {}", e)
    }
}
<|endoftext|>