<commit_before>use crate::core::{
    prelude::*,
    util::{geo, validate},
};
use std::collections::HashMap;
use uuid::Uuid;

//TODO: move usecases into separate files

mod confirm_email;
mod create_new_entry;
mod create_new_event;
pub mod create_new_user;
mod delete_event;
mod find_duplicates;
mod login;
mod query_events;
mod rate_entry;
mod search;
#[cfg(test)]
pub mod tests;
mod update_entry;
mod update_event;

pub use self::{
    confirm_email::*, create_new_entry::*, create_new_event::*, create_new_user::*,
    delete_event::*, find_duplicates::*, login::*, query_events::*, rate_entry::*, search::*,
    update_entry::*, update_event::*,
};

// TODO: Remove this function after replacing Bbox with MapBbox
fn map_bbox(bbox: &Bbox) -> Option<geo::MapBbox> {
    let sw = MapPoint::try_from_lat_lng_deg(bbox.south_west.lat, bbox.south_west.lng);
    let ne = MapPoint::try_from_lat_lng_deg(bbox.north_east.lat, bbox.north_east.lng);
    if let (Some(sw), Some(ne)) = (sw, ne) {
        Some(geo::MapBbox::new(sw, ne))
    } else {
        warn!("Invalid Bbox: {:?}", bbox);
        None
    }
}

pub fn get_ratings<D: Db>(db: &D, ids: &[String]) -> Result<Vec<Rating>> {
    Ok(db
        .all_ratings()?
        .iter()
        .filter(|x| ids.iter().any(|id| *id == x.id))
        .cloned()
        .collect())
}

pub fn get_ratings_by_entry_ids<D: Db>(
    db: &D,
    ids: &[String],
) -> Result<HashMap<String, Vec<Rating>>> {
    let ratings = db.all_ratings()?;
    Ok(ids
        .iter()
        .map(|e_id| {
            (
                e_id.clone(),
                ratings
                    .iter()
                    .filter(|r| r.entry_id == **e_id)
                    .cloned()
                    .collect(),
            )
        })
        .collect())
}

pub fn get_comments_by_rating_ids<D: Db>(
    db: &D,
    ids: &[String],
) -> Result<HashMap<String, Vec<Comment>>> {
    let comments = db.all_comments()?;
    Ok(ids
        .iter()
        .map(|r_id| {
            (
                r_id.clone(),
                comments
                    .iter()
                    .filter_map(|comment| {
                        if comment.rating_id == *r_id {
                            Some(comment)
                        } else {
                            None
                        }
                    })
                    .cloned()
                    .collect(),
            )
        })
        .collect())
}

pub fn get_entries<D: Db>(db: &D, ids: &[String]) -> Result<Vec<Entry>> {
    let entries = db
        .all_entries()?
        .into_iter()
        .filter(|e| ids.iter().any(|id| *id == e.id))
        .collect();
    Ok(entries)
}

pub fn get_user<D: Db>(
    db: &mut D,
    logged_in_username: &str,
    requested_username: &str,
) -> Result<(String, String)> {
    let u: User = db.get_user(requested_username)?;
    if logged_in_username != requested_username {
        return Err(Error::Parameter(ParameterError::Forbidden));
    }
    Ok((u.username, u.email))
}

pub fn get_event<D: Db>(db: &D, id: &str) -> Result<Event> {
    let mut e: Event = db.get_event(id)?;
    if let Some(ref username) = e.created_by {
        let u = db.get_user(username)?;
        e.created_by = Some(u.email);
    }
    Ok(e)
}

pub fn delete_user(db: &mut Db, login_id: &str, u_id: &str) -> Result<()> {
    if login_id != u_id {
        return Err(Error::Parameter(ParameterError::Forbidden));
    }
    db.delete_user(login_id)?;
    Ok(())
}

pub fn subscribe_to_bbox(coordinates: &[Coordinate], username: &str, db: &mut Db) -> Result<()> {
    if coordinates.len() != 2 {
        return Err(Error::Parameter(ParameterError::Bbox));
    }
    let bbox = Bbox {
        south_west: coordinates[0].clone(),
        north_east: coordinates[1].clone(),
    };
    validate::bbox(&bbox)?;

    // TODO: support multiple subscriptions in KVM (frontend)
    // In the meanwile we just replace existing subscriptions
    // with a new one.
    unsubscribe_all_bboxes_by_username(db, username)?;

    let id = Uuid::new_v4().to_simple_ref().to_string();
    db.create_bbox_subscription(&BboxSubscription {
        id,
        bbox,
        username: username.into(),
    })?;
    Ok(())
}

pub fn get_bbox_subscriptions(username: &str, db: &Db) -> Result<Vec<BboxSubscription>> {
    Ok(db
        .all_bbox_subscriptions()?
        .into_iter()
        .filter(|s| s.username == username)
        .collect())
}

pub fn unsubscribe_all_bboxes_by_username(db: &mut Db, username: &str) -> Result<()> {
    let user_subscriptions: Vec<_> = db
        .all_bbox_subscriptions()?
        .into_iter()
        .filter(|s| s.username == username)
        .map(|s| s.id)
        .collect();
    for s_id in user_subscriptions {
        db.delete_bbox_subscription(&s_id)?;
    }
    Ok(())
}

pub fn bbox_subscriptions_by_coordinate(
    db: &mut Db,
    x: &Coordinate,
) -> Result<Vec<BboxSubscription>> {
    Ok(db
        .all_bbox_subscriptions()?
        .into_iter()
        .filter(|s| {
            if let Some(bbox) = map_bbox(&s.bbox) {
                let pt = geo::MapPoint::from_lat_lng_deg(x.lat, x.lng);
                bbox.contains_point(&pt)
            } else {
                false
            }
        })
        .collect())
}

pub fn email_addresses_from_subscriptions(
    db: &mut Db,
    subs: &[BboxSubscription],
) -> Result<Vec<String>> {
    let usernames: Vec<_> = subs.iter().map(|s| &s.username).collect();

    let mut addresses: Vec<_> = db
        .all_users()?
        .into_iter()
        .filter(|u| usernames.iter().any(|x| **x == u.username))
        .map(|u| u.email)
        .collect();
    addresses.dedup();
    Ok(addresses)
}

pub fn email_addresses_by_coordinate(db: &mut Db, lat: &f64, lng: &f64) -> Result<Vec<String>> {
    let subs = bbox_subscriptions_by_coordinate(
        db,
        &Coordinate {
            lat: *lat,
            lng: *lng,
        },
    )?;
    let addresses = email_addresses_from_subscriptions(db, &subs)?;
    Ok(addresses)
}

pub fn prepare_tag_list(tags: Vec<String>) -> Vec<String> {
    let mut tags: Vec<_> = tags
        .into_iter()
        .map(|t| t.trim().to_owned())
        .filter(|t| !t.is_empty())
        .map(|t| t.split(" ").map(|x| x.to_owned()).collect::<Vec<_>>())
        .flatten()
        .map(|t| t.replace("#", ""))
        .collect();
    tags.dedup();
    tags
}

pub fn check_for_owned_tags<D: Db>(
    db: &D,
    tags: &[String],
    org: &Option<Organization>,
) -> Result<()> {
    let owned_tags = db.get_all_tags_owned_by_orgs()?;
    for t in tags {
        if owned_tags.iter().any(|id| id == t) {
            match org {
                Some(ref o) => {
                    if !o.owned_tags.iter().any(|x| x == t) {
                        return Err(ParameterError::OwnedTag.into());
                    }
                }
                None => {
                    return Err(ParameterError::OwnedTag.into());
                }
            }
        }
    }
    Ok(())
}
<commit_msg>Optimize special cases when fetching entries<commit_after>use crate::core::{
    prelude::*,
    util::{geo, validate},
};
use std::collections::HashMap;
use uuid::Uuid;

//TODO: move usecases into separate files

mod confirm_email;
mod create_new_entry;
mod create_new_event;
pub mod create_new_user;
mod delete_event;
mod find_duplicates;
mod login;
mod query_events;
mod rate_entry;
mod search;
#[cfg(test)]
pub mod tests;
mod update_entry;
mod update_event;

pub use self::{
    confirm_email::*, create_new_entry::*, create_new_event::*, create_new_user::*,
    delete_event::*, find_duplicates::*, login::*, query_events::*, rate_entry::*, search::*,
    update_entry::*, update_event::*,
};

// TODO: Remove this function after replacing Bbox with MapBbox
fn map_bbox(bbox: &Bbox) -> Option<geo::MapBbox> {
    let sw = MapPoint::try_from_lat_lng_deg(bbox.south_west.lat, bbox.south_west.lng);
    let ne = MapPoint::try_from_lat_lng_deg(bbox.north_east.lat, bbox.north_east.lng);
    if let (Some(sw), Some(ne)) = (sw, ne) {
        Some(geo::MapBbox::new(sw, ne))
    } else {
        warn!("Invalid Bbox: {:?}", bbox);
        None
    }
}

pub fn get_ratings<D: Db>(db: &D, ids: &[String]) -> Result<Vec<Rating>> {
    Ok(db
        .all_ratings()?
        .iter()
        .filter(|x| ids.iter().any(|id| *id == x.id))
        .cloned()
        .collect())
}

pub fn get_ratings_by_entry_ids<D: Db>(
    db: &D,
    ids: &[String],
) -> Result<HashMap<String, Vec<Rating>>> {
    let ratings = db.all_ratings()?;
    Ok(ids
        .iter()
        .map(|e_id| {
            (
                e_id.clone(),
                ratings
                    .iter()
                    .filter(|r| r.entry_id == **e_id)
                    .cloned()
                    .collect(),
            )
        })
        .collect())
}

pub fn get_comments_by_rating_ids<D: Db>(
    db: &D,
    ids: &[String],
) -> Result<HashMap<String, Vec<Comment>>> {
    let comments = db.all_comments()?;
    Ok(ids
        .iter()
        .map(|r_id| {
            (
                r_id.clone(),
                comments
                    .iter()
                    .filter_map(|comment| {
                        if comment.rating_id == *r_id {
                            Some(comment)
                        } else {
                            None
                        }
                    })
                    .cloned()
                    .collect(),
            )
        })
        .collect())
}

pub fn get_entries<D: Db>(db: &D, ids: &[String]) -> Result<Vec<Entry>> {
    let entries = match ids.len() {
        0 => vec![],
        1 => vec![db.get_entry(&ids[0])?],
        // TODO: Retrieve multiple entries in batches!!!
        _ => db
            .all_entries()?
            .into_iter()
            .filter(|e| ids.iter().any(|id| *id == e.id))
            .collect(),
    };
    Ok(entries)
}

pub fn get_user<D: Db>(
    db: &mut D,
    logged_in_username: &str,
    requested_username: &str,
) -> Result<(String, String)> {
    let u: User = db.get_user(requested_username)?;
    if logged_in_username != requested_username {
        return Err(Error::Parameter(ParameterError::Forbidden));
    }
    Ok((u.username, u.email))
}

pub fn get_event<D: Db>(db: &D, id: &str) -> Result<Event> {
    let mut e: Event = db.get_event(id)?;
    if let Some(ref username) = e.created_by {
        let u = db.get_user(username)?;
        e.created_by = Some(u.email);
    }
    Ok(e)
}

pub fn delete_user(db: &mut Db, login_id: &str, u_id: &str) -> Result<()> {
    if login_id != u_id {
        return Err(Error::Parameter(ParameterError::Forbidden));
    }
    db.delete_user(login_id)?;
    Ok(())
}

pub fn subscribe_to_bbox(coordinates: &[Coordinate], username: &str, db: &mut Db) -> Result<()> {
    if coordinates.len() != 2 {
        return Err(Error::Parameter(ParameterError::Bbox));
    }
    let bbox = Bbox {
        south_west: coordinates[0].clone(),
        north_east: coordinates[1].clone(),
    };
    validate::bbox(&bbox)?;

    // TODO: support multiple subscriptions in KVM (frontend)
    // In the meanwile we just replace existing subscriptions
    // with a new one.
    unsubscribe_all_bboxes_by_username(db, username)?;

    let id = Uuid::new_v4().to_simple_ref().to_string();
    db.create_bbox_subscription(&BboxSubscription {
        id,
        bbox,
        username: username.into(),
    })?;
    Ok(())
}

pub fn get_bbox_subscriptions(username: &str, db: &Db) -> Result<Vec<BboxSubscription>> {
    Ok(db
        .all_bbox_subscriptions()?
        .into_iter()
        .filter(|s| s.username == username)
        .collect())
}

pub fn unsubscribe_all_bboxes_by_username(db: &mut Db, username: &str) -> Result<()> {
    let user_subscriptions: Vec<_> = db
        .all_bbox_subscriptions()?
        .into_iter()
        .filter(|s| s.username == username)
        .map(|s| s.id)
        .collect();
    for s_id in user_subscriptions {
        db.delete_bbox_subscription(&s_id)?;
    }
    Ok(())
}

pub fn bbox_subscriptions_by_coordinate(
    db: &mut Db,
    x: &Coordinate,
) -> Result<Vec<BboxSubscription>> {
    Ok(db
        .all_bbox_subscriptions()?
        .into_iter()
        .filter(|s| {
            if let Some(bbox) = map_bbox(&s.bbox) {
                let pt = geo::MapPoint::from_lat_lng_deg(x.lat, x.lng);
                bbox.contains_point(&pt)
            } else {
                false
            }
        })
        .collect())
}

pub fn email_addresses_from_subscriptions(
    db: &mut Db,
    subs: &[BboxSubscription],
) -> Result<Vec<String>> {
    let usernames: Vec<_> = subs.iter().map(|s| &s.username).collect();

    let mut addresses: Vec<_> = db
        .all_users()?
        .into_iter()
        .filter(|u| usernames.iter().any(|x| **x == u.username))
        .map(|u| u.email)
        .collect();
    addresses.dedup();
    Ok(addresses)
}

pub fn email_addresses_by_coordinate(db: &mut Db, lat: &f64, lng: &f64) -> Result<Vec<String>> {
    let subs = bbox_subscriptions_by_coordinate(
        db,
        &Coordinate {
            lat: *lat,
            lng: *lng,
        },
    )?;
    let addresses = email_addresses_from_subscriptions(db, &subs)?;
    Ok(addresses)
}

pub fn prepare_tag_list(tags: Vec<String>) -> Vec<String> {
    let mut tags: Vec<_> = tags
        .into_iter()
        .map(|t| t.trim().to_owned())
        .filter(|t| !t.is_empty())
        .map(|t| t.split(" ").map(|x| x.to_owned()).collect::<Vec<_>>())
        .flatten()
        .map(|t| t.replace("#", ""))
        .collect();
    tags.dedup();
    tags
}

pub fn check_for_owned_tags<D: Db>(
    db: &D,
    tags: &[String],
    org: &Option<Organization>,
) -> Result<()> {
    let owned_tags = db.get_all_tags_owned_by_orgs()?;
    for t in tags {
        if owned_tags.iter().any(|id| id == t) {
            match org {
                Some(ref o) => {
                    if !o.owned_tags.iter().any(|x| x == t) {
                        return Err(ParameterError::OwnedTag.into());
                    }
                }
                None => {
                    return Err(ParameterError::OwnedTag.into());
                }
            }
        }
    }
    Ok(())
}
<|endoftext|>


<commit_before>use runtime::Runtime;
use storage::backend::{StorageBackendError, StorageBackend};

use module::Module;
use module::ModuleError;
use module::CommandResult;
use module::CommandEnv;

use module::bm::header::build_header;
use module::bm::header::get_tags_from_header;
use storage::json::parser::JsonHeaderParser;
use storage::parser::{Parser, FileHeaderParser};
use storage::file::File;
use ui::file::{FilePrinter, TablePrinter};
use std::vec::IntoIter;

use clap::ArgMatches;
use regex::Regex;

pub fn add_command(module: &Module, env: CommandEnv) -> CommandResult {
    let url = env.matches.value_of("url").unwrap();
    let tags = get_tags(env.rt, env.matches);
    info!("Adding url '{}' with tags '{:?}'", url, tags);

    let header  = build_header(&String::from(url), &tags);
    let file    = File::new_with_header(module, header);
    let parser  = Parser::new(JsonHeaderParser::new(None));
    let putres  = env.bk.put_file(file, &parser);

    putres.map_err(|sberr| {
        let mut err = ModuleError::new("Storage Backend Error");
        err.caused_by = Some(Box::new(sberr));
        err
    })
}

pub fn list_command(module: &Module, env: CommandEnv) -> CommandResult {
    let printer = TablePrinter::new(env.rt.is_verbose(), env.rt.is_debugging());
    let files   = get_filtered_files_from_backend(module, &env);

    debug!("Printing files now");
    files.map(|f| printer.print_files(f));

    Ok(())
}

pub fn remove_command(module: &Module, env: CommandEnv) -> CommandResult {
    let checked : bool = run_removal_checking(&env);
    debug!("Checked mode: {}", checked);
    if let Some(id) = get_id(env.rt, env.matches) {
        debug!("Remove by id: {}", id);

        let parser = Parser::new(JsonHeaderParser::new(None));
        let file   = env.bk.get_file_by_id(module, &id.into(), &parser).unwrap();
        debug!("Remove file  : {:?}", file);

        if let Err(e) = env.bk.remove_file(module, file, checked) {
            debug!("Remove failed");
            let mut err = ModuleError::new("Removing file failed");
            err.caused_by = Some(Box::new(e));
            Err(err)
        } else {
            debug!("Remove worked");
            Ok(())
        }
    } else {
        debug!("Remove more than one file");

        let files = get_filtered_files_from_backend(module, &env);
        let nfiles = files.len();
        info!("Removing {} Files", nfiles);

        let errs = files.map(|file| {
                debug!("Remove file: {:?}", file);
                env.bk.remove_file(module, file, checked)
            })
            .filter(|e| e.is_err())
            .map(|e| {
                let err = e.err().unwrap();
                warn!("Error occured in Filesystem operation: {}", err);
                err
            })
            .collect::<Vec<StorageBackendError>>();

        let nerrs = errs.len();

        if nerrs != 0 {
            warn!("{} Errors occured while removing {} files", nerrs, nfiles);
            let moderr = ModuleError::new("File removal failed");

            // TODO : Collect StorageBackendErrors

            Err(moderr)
        } else {
            Ok(())
        }
    }
}

/*
 *
 * Private helpers
 *
 */

fn get_filtered_files_from_backend<'a>(module: &'a Module,
                                       env: &CommandEnv) -> IntoIter<File<'a>>
{
    let parser = Parser::new(JsonHeaderParser::new(None));
    let tags = get_tags(env.rt, env.matches);
    debug!("Tags: {:?}", tags);
    env.bk.iter_files(module, &parser).and_then(|files| {
        let f = files.filter(|file| {
            if tags.len() != 0 {
                debug!("Checking tags of: {:?}", file.id());
                get_tags_from_header(&file.header()).iter()
                    .any(|t| tags.contains(t))
            } else {
                true
            }
        }).filter(|file| {
            debug!("Checking matches of: {:?}", file.id());
            get_matcher(env.rt, env.matches)
                .and_then(|r| Some(file.matches_with(&r)))
                .unwrap_or(true)
        }).collect::<Vec<File>>();
        Some(f)
    }).unwrap_or(Vec::<File>::new()).into_iter()
}

fn get_tags<'a>(rt: &Runtime, sub: &ArgMatches<'a, 'a>) -> Vec<String> {
    debug!("Fetching tags from commandline");
    sub.value_of("tags").and_then(|tags|
                                  Some(tags.split(",")
                                       .into_iter()
                                       .map(|s| s.to_string())
                                       .filter(|e|
                                            if e.contains(" ") {
                                                warn!("Tag contains spaces: '{}'", e);
                                                false
                                            } else {
                                                true
                                            }).collect()
                                      )
                                 ).or(Some(vec![])).unwrap()

}

fn get_matcher<'a>(rt: &Runtime, sub: &ArgMatches<'a, 'a>) -> Option<Regex> {
    debug!("Fetching matcher from commandline");
    if let Some(s) = sub.value_of("match") {
        if let Ok(r) = Regex::new(s) {
            return Some(r)
        } else {
            error!("Regex error, continuing without regex");
        }
    }
    None

}

fn get_id<'a>(rt: &Runtime, sub: &ArgMatches<'a, 'a>) -> Option<String> {
    debug!("Fetching id from commandline");
    sub.value_of("id").and_then(|s| Some(String::from(s)))
}

/*
 * Checks whether the commandline call was set to run the removal "checked",
 * so if another entry from the store refers to this ID, do not remove the file.
 */
fn run_removal_checking(env: &CommandEnv) -> bool {
    env.matches.is_present("check")
}
<commit_msg>Rewrite BMModule::get_filtered_files_from_backend() with the backend speaking Result<> now instead of Option<><commit_after>use runtime::Runtime;
use storage::backend::{StorageBackendError, StorageBackend};

use module::Module;
use module::ModuleError;
use module::CommandResult;
use module::CommandEnv;

use module::bm::header::build_header;
use module::bm::header::get_tags_from_header;
use storage::json::parser::JsonHeaderParser;
use storage::parser::{Parser, FileHeaderParser};
use storage::file::File;
use ui::file::{FilePrinter, TablePrinter};
use std::vec::IntoIter;

use clap::ArgMatches;
use regex::Regex;

pub fn add_command(module: &Module, env: CommandEnv) -> CommandResult {
    let url = env.matches.value_of("url").unwrap();
    let tags = get_tags(env.rt, env.matches);
    info!("Adding url '{}' with tags '{:?}'", url, tags);

    let header  = build_header(&String::from(url), &tags);
    let file    = File::new_with_header(module, header);
    let parser  = Parser::new(JsonHeaderParser::new(None));
    let putres  = env.bk.put_file(file, &parser);

    putres.map_err(|sberr| {
        let mut err = ModuleError::new("Storage Backend Error");
        err.caused_by = Some(Box::new(sberr));
        err
    })
}

pub fn list_command(module: &Module, env: CommandEnv) -> CommandResult {
    let printer = TablePrinter::new(env.rt.is_verbose(), env.rt.is_debugging());
    let files   = get_filtered_files_from_backend(module, &env);

    debug!("Printing files now");
    files.map(|f| printer.print_files(f));

    Ok(())
}

pub fn remove_command(module: &Module, env: CommandEnv) -> CommandResult {
    let checked : bool = run_removal_checking(&env);
    debug!("Checked mode: {}", checked);
    if let Some(id) = get_id(env.rt, env.matches) {
        debug!("Remove by id: {}", id);

        let parser = Parser::new(JsonHeaderParser::new(None));
        let file   = env.bk.get_file_by_id(module, &id.into(), &parser).unwrap();
        debug!("Remove file  : {:?}", file);

        if let Err(e) = env.bk.remove_file(module, file, checked) {
            debug!("Remove failed");
            let mut err = ModuleError::new("Removing file failed");
            err.caused_by = Some(Box::new(e));
            Err(err)
        } else {
            debug!("Remove worked");
            Ok(())
        }
    } else {
        debug!("Remove more than one file");

        get_filtered_files_from_backend(module, &env).and_then(|files| {
            let nfiles = files.len();
            info!("Removing {} Files", nfiles);

            let errs = files.map(|file| {
                    debug!("Remove file: {:?}", file);
                    env.bk.remove_file(module, file, checked)
                })
                .filter(|e| e.is_err())
                .map(|e| {
                    let err = e.err().unwrap();
                    warn!("Error occured in Filesystem operation: {}", err);
                    err
                })
                .collect::<Vec<StorageBackendError>>();

            let nerrs = errs.len();

            if nerrs != 0 {
                warn!("{} Errors occured while removing {} files", nerrs, nfiles);
                let moderr = ModuleError::new("File removal failed");

                // TODO : Collect StorageBackendErrors

                Err(moderr)
            } else {
                Ok(())
            }
        })
    }
}

/*
 *
 * Private helpers
 *
 */

fn get_filtered_files_from_backend<'a>(module: &'a Module,
                                       env: &CommandEnv)
    -> Result<IntoIter<File<'a>>, ModuleError>
{
    let parser = Parser::new(JsonHeaderParser::new(None));
    let tags = get_tags(env.rt, env.matches);
    debug!("Tags: {:?}", tags);
    env.bk.iter_files(module, &parser)
        .map(|files| {
            let f = files.filter(|file| {
                if tags.len() != 0 {
                    debug!("Checking tags of: {:?}", file.id());
                    get_tags_from_header(&file.header()).iter()
                        .any(|t| tags.contains(t))
                } else {
                    true
                }
            }).filter(|file| {
                debug!("Checking matches of: {:?}", file.id());
                get_matcher(env.rt, env.matches)
                    .and_then(|r| Some(file.matches_with(&r)))
                    .unwrap_or(true)
            }).collect::<Vec<File>>();
            f.into_iter()
        }).map_err(|e| {
            let mut merr = ModuleError::new("Could not filter files");
            merr.caused_by = Some(Box::new(e));
            merr
        })
}

fn get_tags<'a>(rt: &Runtime, sub: &ArgMatches<'a, 'a>) -> Vec<String> {
    debug!("Fetching tags from commandline");
    sub.value_of("tags").and_then(|tags|
                                  Some(tags.split(",")
                                       .into_iter()
                                       .map(|s| s.to_string())
                                       .filter(|e|
                                            if e.contains(" ") {
                                                warn!("Tag contains spaces: '{}'", e);
                                                false
                                            } else {
                                                true
                                            }).collect()
                                      )
                                 ).or(Some(vec![])).unwrap()

}

fn get_matcher<'a>(rt: &Runtime, sub: &ArgMatches<'a, 'a>) -> Option<Regex> {
    debug!("Fetching matcher from commandline");
    if let Some(s) = sub.value_of("match") {
        if let Ok(r) = Regex::new(s) {
            return Some(r)
        } else {
            error!("Regex error, continuing without regex");
        }
    }
    None

}

fn get_id<'a>(rt: &Runtime, sub: &ArgMatches<'a, 'a>) -> Option<String> {
    debug!("Fetching id from commandline");
    sub.value_of("id").and_then(|s| Some(String::from(s)))
}

/*
 * Checks whether the commandline call was set to run the removal "checked",
 * so if another entry from the store refers to this ID, do not remove the file.
 */
fn run_removal_checking(env: &CommandEnv) -> bool {
    env.matches.is_present("check")
}
<|endoftext|>


<commit_before>use std::fmt::Write;

use anyhow::Result;

use crate::{analysis::dis::zydis, arch::Arch, workspace::PEWorkspace, VA};

#[derive(Default, Clone)]
struct OriginalHooks {
    print_address_abs: Option<zydis::Hook>,
    pre_instruction:   Option<zydis::Hook>,
}

struct UserData<'a> {
    ws:   &'a PEWorkspace,
    orig: OriginalHooks,
}

pub struct FormatterOptions {
    colors: bool,
}

pub struct FormatterBuilder {
    options: FormatterOptions,
}

impl FormatterBuilder {
    pub fn build(self) -> Formatter {
        Formatter::from_options(self.options)
    }

    pub fn with_colors(mut self, colors: bool) -> FormatterBuilder {
        self.options.colors = colors;
        self
    }
}

pub struct Formatter {
    options: FormatterOptions,
    inner:   zydis::Formatter,
    orig:    OriginalHooks,
}

pub const TOKEN_USER_SYMBOLNAME: zydis::Token = zydis::Token(zydis::TOKEN_USER.0 + 1);

impl Formatter {
    const COLOR_ADDRESS_ABS: ansi_term::Color = ansi_term::Color::Blue;
    const COLOR_ADDRESS_REL: ansi_term::Color = ansi_term::Color::Blue;
    // grey
    const COLOR_DECORATOR: ansi_term::Color = ansi_term::Color::Fixed(242);
    // grey
    const COLOR_DELIMITER: ansi_term::Color = ansi_term::Color::Fixed(242);
    const COLOR_DISPLACEMENT: ansi_term::Color = ansi_term::Color::Blue;
    const COLOR_IMMEDIATE: ansi_term::Color = ansi_term::Color::Blue;
    // default theme
    // TODO: move this to a struct that can be configured
    const COLOR_INVALID: ansi_term::Color = ansi_term::Color::Red;
    const COLOR_MNEMONIC: ansi_term::Color = ansi_term::Color::Green;
    // grey
    const COLOR_PARENTHESIS_CLOSE: ansi_term::Color = ansi_term::Color::Fixed(242);
    // grey
    const COLOR_PARENTHESIS_OPEN: ansi_term::Color = ansi_term::Color::Fixed(242);
    // grey
    const COLOR_PREFIX: ansi_term::Color = ansi_term::Color::Fixed(242);
    const COLOR_REGISTER: ansi_term::Color = ansi_term::Color::Yellow;
    // grey
    const COLOR_SYMBOL: ansi_term::Color = ansi_term::Color::Fixed(242);
    const COLOR_SYMBOLNAME: ansi_term::Color = ansi_term::Color::Purple;
    // grey
    const COLOR_TYPECAST: ansi_term::Color = ansi_term::Color::Fixed(242);
    const COLOR_USER: ansi_term::Color = ansi_term::Color::Fixed(242);
    const COLOR_WHITESPACE: ansi_term::Color = ansi_term::Color::Black;

    pub fn new() -> FormatterBuilder {
        FormatterBuilder {
            options: FormatterOptions { colors: true },
        }
    }

    pub fn from_options(options: FormatterOptions) -> Formatter {
        let mut inner = zydis::Formatter::new(zydis::FormatterStyle::INTEL).unwrap();

        let mut orig: OriginalHooks = Default::default();

        // TODO: align mnemonic

        let f = inner
            .set_pre_instruction(Box::new(
                |_formatter: &zydis::Formatter,
                 buf: &mut zydis::FormatterBuffer,
                 ctx: &mut zydis::FormatterContext,
                 userdata: Option<&mut dyn core::any::Any>|
                 -> zydis::Result<()> {
                    // programming error: userdata must be provided. this is guaranteed within
                    // Formatter.
                    let userdata = userdata.expect("no userdata");

                    // programming error: userdata must be a Box<UserData>. this is guaranteed
                    // within Formatter.
                    let userdata = userdata.downcast_ref::<UserData>().expect("incorrect userdata");

                    let va = ctx.runtime_address;

                    if let Some(sec) = userdata
                        .ws
                        .pe
                        .module
                        .sections
                        .iter()
                        .find(|&sec| sec.virtual_range.contains(&va))
                    {
                        buf.append(TOKEN_USER_SYMBOLNAME)?;
                        buf.get_string()?.append(&sec.name)?;
                    } else {
                        buf.append(zydis::TOKEN_INVALID)?;
                        buf.get_string()?.append("???")?;
                    }
                    buf.append(zydis::TOKEN_DELIMITER)?;
                    buf.get_string()?.append(":")?;

                    // TODO: insn bytes

                    buf.append(zydis::TOKEN_ADDRESS_ABS)?;
                    match userdata.ws.pe.module.arch {
                        Arch::X32 => {
                            buf.get_string()?.append(&format!("{:08x}", va))?;
                        }
                        Arch::X64 => {
                            buf.get_string()?.append(&format!("{:016x}", va))?;
                        }
                    }

                    buf.append(zydis::TOKEN_WHITESPACE)?;
                    buf.get_string()?.append("  ")?;

                    Ok(())
                },
            ))
            .unwrap();
        orig.pre_instruction = Some(f);

        let f = inner
            .set_print_address_abs(Box::new(
                |formatter: &zydis::Formatter,
                 buf: &mut zydis::FormatterBuffer,
                 ctx: &mut zydis::FormatterContext,
                 userdata: Option<&mut dyn core::any::Any>|
                 -> zydis::Result<()> {
                    // programming error: userdata must be provided. this is guaranteed within
                    // Formatter.
                    let userdata = userdata.expect("no userdata");

                    // programming error: userdata must be a Box<UserData>. this is guaranteed
                    // within Formatter.
                    let userdata = userdata.downcast_ref::<UserData>().expect("incorrect userdata");

                    let absolute_address = unsafe {
                        // safety: the insn and operands come from zydis, so we assume they contain
                        // valid data.
                        let insn: &zydis::DecodedInstruction = &*ctx.instruction;
                        let op: &zydis::DecodedOperand = &*ctx.operand;
                        insn.calc_absolute_address(ctx.runtime_address, op)
                            .expect("failed to calculate absolute address")
                    };

                    if let Some(name) = userdata.ws.analysis.names.names_by_address.get(&absolute_address) {
                        // name is found in map, use that.
                        buf.append(TOKEN_USER_SYMBOLNAME)?;
                        return buf.get_string()?.append(name);
                    } else {
                        // name is not found, use original formatter.

                        // programming error: the original hook must be recorded. this is guaranteed
                        // within Formatter.
                        let orig = userdata.orig.print_address_abs.as_ref().expect("no original hook");

                        if let zydis::Hook::PrintAddressAbs(Some(f)) = orig {
                            // safety: zydis::Formatter <-> zydis::ffi::ZydisFormatter is safe according to
                            // here: https://docs.rs/zydis/3.1.2/src/zydis/formatter.rs.html#306
                            let status =
                                unsafe { f(formatter as *const _ as *const zydis::ffi::ZydisFormatter, buf, ctx) };
                            if status.is_error() {
                                return Err(status);
                            } else {
                                return Ok(());
                            }
                        } else {
                            // I'm not sure how this could ever be the case, as zydis initializes the hook
                            // with a default. I suppose if you explicitly set
                            // the callback to NULL/None? Which we don't do here.
                            panic!("unexpected original hook");
                        }
                    }
                },
            ))
            .unwrap();
        orig.print_address_abs = Some(f);

        Formatter { options, inner, orig }
    }

    // grey

    fn get_token_color(token: zydis::Token) -> ansi_term::Color {
        match token {
            zydis::TOKEN_INVALID => Formatter::COLOR_INVALID,
            zydis::TOKEN_WHITESPACE => Formatter::COLOR_WHITESPACE,
            zydis::TOKEN_DELIMITER => Formatter::COLOR_DELIMITER,
            zydis::TOKEN_PARENTHESIS_OPEN => Formatter::COLOR_PARENTHESIS_OPEN,
            zydis::TOKEN_PARENTHESIS_CLOSE => Formatter::COLOR_PARENTHESIS_CLOSE,
            zydis::TOKEN_PREFIX => Formatter::COLOR_PREFIX,
            zydis::TOKEN_MNEMONIC => Formatter::COLOR_MNEMONIC,
            zydis::TOKEN_REGISTER => Formatter::COLOR_REGISTER,
            zydis::TOKEN_ADDRESS_ABS => Formatter::COLOR_ADDRESS_ABS,
            zydis::TOKEN_ADDRESS_REL => Formatter::COLOR_ADDRESS_REL,
            zydis::TOKEN_DISPLACEMENT => Formatter::COLOR_DISPLACEMENT,
            zydis::TOKEN_IMMEDIATE => Formatter::COLOR_IMMEDIATE,
            zydis::TOKEN_TYPECAST => Formatter::COLOR_TYPECAST,
            zydis::TOKEN_DECORATOR => Formatter::COLOR_DECORATOR,
            zydis::TOKEN_SYMBOL => Formatter::COLOR_SYMBOL,
            zydis::TOKEN_USER => Formatter::COLOR_USER,
            TOKEN_USER_SYMBOLNAME => Formatter::COLOR_SYMBOLNAME,
            _ => unimplemented!("token: {}", token),
        }
    }

    fn render_token<T: Write>(&self, o: &mut T, token: zydis::Token, s: &str) -> Result<()> {
        if self.options.colors {
            o.write_str(&Formatter::get_token_color(token).paint(s).to_string())?;
        } else {
            o.write_str(s)?;
        }

        Ok(())
    }

    pub fn format_instruction(&self, ws: &PEWorkspace, insn: &zydis::DecodedInstruction, va: VA) -> Result<String> {
        let mut buffer = [0u8; 200];

        // we pass our userdata to ZydisFormatterFormatInstruction.
        // but to make it work, we have to play games with the lifetimes:
        // we need to convince the compiler that the userdata pointer lives long enough
        // to by used by the callbacks.
        //
        // we do this by extending the CFG lifetime from '_ to 'static.
        //
        // userdata is passed into ZydisFormatterFormatInstruction,
        // which passes userdata to each of the formatter callbacks.
        // those read strictly from insn/ctx/userdata and write strictly to output
        // buffer. there is no state maintained within these routines.
        // the callbacks won't be invoked beyond the call into FormatInstruction.
        //
        // therefore, i believe its safe to extend the lifetime here to work with zydis.
        let x = unsafe { std::mem::transmute::<&'_ PEWorkspace, &'static PEWorkspace>(ws) };

        let mut ud = UserData {
            orig: self.orig.clone(),
            ws:   x,
        };

        let mut out = String::new();
        for (token, s) in self
            .inner
            .tokenize_instruction(&insn, &mut buffer, Some(va), Some(&mut ud))?
        {
            self.render_token(&mut out, token, s)?;
        }

        Ok(out)
    }
}

#[cfg(test)]
mod tests {
    use super::{super::*, *};
    use crate::rsrc::*;
    use anyhow::Result;

    #[test]
    fn with_colors() -> Result<()> {
        let buf = get_buf(Rsrc::NOP);
        let pe = crate::loader::pe::PE::from_bytes(&buf)?;

        let ws = PEWorkspace::from_pe(pe)?;

        let fmt = Formatter::new().with_colors(true).build();

        // ```
        //     .text:00401C4E 000 68 F4 61 40 00          push    offset ModuleName ; "mscoree.dll"
        //     .text:00401C53 004 FF 15 00 60 40 00       call    ds:GetModuleHandleA
        //     .text:00401C59 000 85 C0                   test    eax, eax
        // ```
        let insn = crate::test::read_insn(&ws.pe.module, 0x401C53);
        let s = fmt.format_instruction(&ws, &insn, 0x401C53)?;
        assert!(s.contains("\u{1b}"));
        assert!(s.contains("call"));
        assert!(s.contains("GetModuleHandleA"));

        Ok(())
    }

    #[test]
    fn no_colors() -> Result<()> {
        let buf = get_buf(Rsrc::NOP);
        let pe = crate::loader::pe::PE::from_bytes(&buf)?;

        let ws = PEWorkspace::from_pe(pe)?;

        let fmt = Formatter::new().with_colors(false).build();

        // ```
        //     .text:00401C4E 000 68 F4 61 40 00          push    offset ModuleName ; "mscoree.dll"
        //     .text:00401C53 004 FF 15 00 60 40 00       call    ds:GetModuleHandleA
        //     .text:00401C59 000 85 C0                   test    eax, eax
        // ```
        let insn = crate::test::read_insn(&ws.pe.module, 0x401C53);
        assert_eq!(
            fmt.format_instruction(&ws, &insn, 0x401C53)?,
            "call [kernel32.dll!GetModuleHandleA]"
        );

        Ok(())
    }
}
<commit_msg>formatter: documentation<commit_after>use std::fmt::Write;

use anyhow::Result;

use crate::{analysis::dis::zydis, arch::Arch, workspace::PEWorkspace, VA};

#[derive(Default, Clone)]
struct OriginalHooks {
    print_address_abs: Option<zydis::Hook>,
    pre_instruction:   Option<zydis::Hook>,
}

struct UserData<'a> {
    ws:   &'a PEWorkspace,
    orig: OriginalHooks,
}

pub struct FormatterOptions {
    colors: bool,
}

pub struct FormatterBuilder {
    options: FormatterOptions,
}

impl FormatterBuilder {
    pub fn build(self) -> Formatter {
        Formatter::from_options(self.options)
    }

    pub fn with_colors(mut self, colors: bool) -> FormatterBuilder {
        self.options.colors = colors;
        self
    }
}

pub struct Formatter {
    options: FormatterOptions,
    inner:   zydis::Formatter,
    orig:    OriginalHooks,
}

pub const TOKEN_USER_SYMBOLNAME: zydis::Token = zydis::Token(zydis::TOKEN_USER.0 + 1);

impl Formatter {
    // default theme
    // TODO: move this to a struct that can be configured
    const COLOR_ADDRESS_ABS: ansi_term::Color = ansi_term::Color::Blue;
    const COLOR_ADDRESS_REL: ansi_term::Color = ansi_term::Color::Blue;
    const COLOR_DECORATOR: ansi_term::Color = Formatter::GREY;
    const COLOR_DELIMITER: ansi_term::Color = Formatter::GREY;
    const COLOR_DISPLACEMENT: ansi_term::Color = ansi_term::Color::Blue;
    const COLOR_IMMEDIATE: ansi_term::Color = ansi_term::Color::Blue;
    const COLOR_INVALID: ansi_term::Color = ansi_term::Color::Red;
    const COLOR_MNEMONIC: ansi_term::Color = ansi_term::Color::Green;
    const COLOR_PARENTHESIS_CLOSE: ansi_term::Color = Formatter::GREY;
    const COLOR_PARENTHESIS_OPEN: ansi_term::Color = Formatter::GREY;
    const COLOR_PREFIX: ansi_term::Color = Formatter::GREY;
    const COLOR_REGISTER: ansi_term::Color = ansi_term::Color::Yellow;
    const COLOR_SYMBOL: ansi_term::Color = Formatter::GREY;
    const COLOR_SYMBOLNAME: ansi_term::Color = ansi_term::Color::Purple;
    const COLOR_TYPECAST: ansi_term::Color = Formatter::GREY;
    const COLOR_USER: ansi_term::Color = Formatter::GREY;
    const COLOR_WHITESPACE: ansi_term::Color = ansi_term::Color::Black;
    const GREY: ansi_term::Color = ansi_term::Color::Fixed(242);

    pub fn new() -> FormatterBuilder {
        FormatterBuilder {
            options: FormatterOptions { colors: true },
        }
    }

    pub fn from_options(options: FormatterOptions) -> Formatter {
        let mut inner = zydis::Formatter::new(zydis::FormatterStyle::INTEL).unwrap();

        let mut orig: OriginalHooks = Default::default();

        // TODO: align mnemonic

        let f = inner
            .set_pre_instruction(Box::new(
                |_formatter: &zydis::Formatter,
                 buf: &mut zydis::FormatterBuffer,
                 ctx: &mut zydis::FormatterContext,
                 userdata: Option<&mut dyn core::any::Any>|
                 -> zydis::Result<()> {
                    // programming error: userdata must be provided. this is guaranteed within
                    // Formatter.
                    let userdata = userdata.expect("no userdata");

                    // programming error: userdata must be a Box<UserData>. this is guaranteed
                    // within Formatter.
                    let userdata = userdata.downcast_ref::<UserData>().expect("incorrect userdata");

                    let va = ctx.runtime_address;

                    if let Some(sec) = userdata
                        .ws
                        .pe
                        .module
                        .sections
                        .iter()
                        .find(|&sec| sec.virtual_range.contains(&va))
                    {
                        buf.append(TOKEN_USER_SYMBOLNAME)?;
                        buf.get_string()?.append(&sec.name)?;
                    } else {
                        buf.append(zydis::TOKEN_INVALID)?;
                        buf.get_string()?.append("???")?;
                    }
                    buf.append(zydis::TOKEN_DELIMITER)?;
                    buf.get_string()?.append(":")?;

                    // TODO: insn bytes

                    buf.append(zydis::TOKEN_ADDRESS_ABS)?;
                    match userdata.ws.pe.module.arch {
                        Arch::X32 => {
                            buf.get_string()?.append(&format!("{:08x}", va))?;
                        }
                        Arch::X64 => {
                            buf.get_string()?.append(&format!("{:016x}", va))?;
                        }
                    }

                    buf.append(zydis::TOKEN_WHITESPACE)?;
                    buf.get_string()?.append("  ")?;

                    Ok(())
                },
            ))
            .unwrap();
        orig.pre_instruction = Some(f);

        let f = inner
            .set_print_address_abs(Box::new(
                |formatter: &zydis::Formatter,
                 buf: &mut zydis::FormatterBuffer,
                 ctx: &mut zydis::FormatterContext,
                 userdata: Option<&mut dyn core::any::Any>|
                 -> zydis::Result<()> {
                    // programming error: userdata must be provided. this is guaranteed within
                    // Formatter.
                    let userdata = userdata.expect("no userdata");

                    // programming error: userdata must be a Box<UserData>. this is guaranteed
                    // within Formatter.
                    let userdata = userdata.downcast_ref::<UserData>().expect("incorrect userdata");

                    let absolute_address = unsafe {
                        // safety: the insn and operands come from zydis, so we assume they contain
                        // valid data.
                        let insn: &zydis::DecodedInstruction = &*ctx.instruction;
                        let op: &zydis::DecodedOperand = &*ctx.operand;
                        insn.calc_absolute_address(ctx.runtime_address, op)
                            .expect("failed to calculate absolute address")
                    };

                    if let Some(name) = userdata.ws.analysis.names.names_by_address.get(&absolute_address) {
                        // name is found in map, use that.
                        buf.append(TOKEN_USER_SYMBOLNAME)?;
                        return buf.get_string()?.append(name);
                    } else {
                        // name is not found, use original formatter.

                        // programming error: the original hook must be recorded. this is guaranteed
                        // within Formatter.
                        let orig = userdata.orig.print_address_abs.as_ref().expect("no original hook");

                        if let zydis::Hook::PrintAddressAbs(Some(f)) = orig {
                            // safety: zydis::Formatter <-> zydis::ffi::ZydisFormatter is safe according to
                            // here: https://docs.rs/zydis/3.1.2/src/zydis/formatter.rs.html#306
                            let status =
                                unsafe { f(formatter as *const _ as *const zydis::ffi::ZydisFormatter, buf, ctx) };
                            if status.is_error() {
                                return Err(status);
                            } else {
                                return Ok(());
                            }
                        } else {
                            // I'm not sure how this could ever be the case, as zydis initializes the hook
                            // with a default. I suppose if you explicitly set
                            // the callback to NULL/None? Which we don't do here.
                            panic!("unexpected original hook");
                        }
                    }
                },
            ))
            .unwrap();
        orig.print_address_abs = Some(f);

        Formatter { options, inner, orig }
    }

    // grey

    fn get_token_color(token: zydis::Token) -> ansi_term::Color {
        match token {
            zydis::TOKEN_INVALID => Formatter::COLOR_INVALID,
            zydis::TOKEN_WHITESPACE => Formatter::COLOR_WHITESPACE,
            zydis::TOKEN_DELIMITER => Formatter::COLOR_DELIMITER,
            zydis::TOKEN_PARENTHESIS_OPEN => Formatter::COLOR_PARENTHESIS_OPEN,
            zydis::TOKEN_PARENTHESIS_CLOSE => Formatter::COLOR_PARENTHESIS_CLOSE,
            zydis::TOKEN_PREFIX => Formatter::COLOR_PREFIX,
            zydis::TOKEN_MNEMONIC => Formatter::COLOR_MNEMONIC,
            zydis::TOKEN_REGISTER => Formatter::COLOR_REGISTER,
            zydis::TOKEN_ADDRESS_ABS => Formatter::COLOR_ADDRESS_ABS,
            zydis::TOKEN_ADDRESS_REL => Formatter::COLOR_ADDRESS_REL,
            zydis::TOKEN_DISPLACEMENT => Formatter::COLOR_DISPLACEMENT,
            zydis::TOKEN_IMMEDIATE => Formatter::COLOR_IMMEDIATE,
            zydis::TOKEN_TYPECAST => Formatter::COLOR_TYPECAST,
            zydis::TOKEN_DECORATOR => Formatter::COLOR_DECORATOR,
            zydis::TOKEN_SYMBOL => Formatter::COLOR_SYMBOL,
            zydis::TOKEN_USER => Formatter::COLOR_USER,
            TOKEN_USER_SYMBOLNAME => Formatter::COLOR_SYMBOLNAME,
            _ => unimplemented!("token: {}", token),
        }
    }

    fn render_token<T: Write>(&self, o: &mut T, token: zydis::Token, s: &str) -> Result<()> {
        if self.options.colors {
            o.write_str(&Formatter::get_token_color(token).paint(s).to_string())?;
        } else {
            o.write_str(s)?;
        }

        Ok(())
    }

    pub fn format_instruction(&self, ws: &PEWorkspace, insn: &zydis::DecodedInstruction, va: VA) -> Result<String> {
        let mut buffer = [0u8; 200];

        // we pass our userdata to ZydisFormatterFormatInstruction.
        // but to make it work, we have to play games with the lifetimes:
        // we need to convince the compiler that the userdata pointer lives long enough
        // to by used by the callbacks.
        //
        // we do this by extending the CFG lifetime from '_ to 'static.
        //
        // userdata is passed into ZydisFormatterFormatInstruction,
        // which passes userdata to each of the formatter callbacks.
        // those read strictly from insn/ctx/userdata and write strictly to output
        // buffer. there is no state maintained within these routines.
        // the callbacks won't be invoked beyond the call into FormatInstruction.
        //
        // therefore, i believe its safe to extend the lifetime here to work with zydis.
        let x = unsafe { std::mem::transmute::<&'_ PEWorkspace, &'static PEWorkspace>(ws) };

        let mut ud = UserData {
            orig: self.orig.clone(),
            ws:   x,
        };

        let mut out = String::new();
        for (token, s) in self
            .inner
            .tokenize_instruction(&insn, &mut buffer, Some(va), Some(&mut ud))?
        {
            self.render_token(&mut out, token, s)?;
        }

        Ok(out)
    }
}

#[cfg(test)]
mod tests {
    use super::{super::*, *};
    use crate::rsrc::*;
    use anyhow::Result;

    #[test]
    fn with_colors() -> Result<()> {
        let buf = get_buf(Rsrc::NOP);
        let pe = crate::loader::pe::PE::from_bytes(&buf)?;

        let ws = PEWorkspace::from_pe(pe)?;

        let fmt = Formatter::new().with_colors(true).build();

        // ```
        //     .text:00401C4E 000 68 F4 61 40 00          push    offset ModuleName ; "mscoree.dll"
        //     .text:00401C53 004 FF 15 00 60 40 00       call    ds:GetModuleHandleA
        //     .text:00401C59 000 85 C0                   test    eax, eax
        // ```
        let insn = crate::test::read_insn(&ws.pe.module, 0x401C53);
        let s = fmt.format_instruction(&ws, &insn, 0x401C53)?;
        assert!(s.contains("\u{1b}"));
        assert!(s.contains("call"));
        assert!(s.contains("GetModuleHandleA"));

        Ok(())
    }

    #[test]
    fn no_colors() -> Result<()> {
        let buf = get_buf(Rsrc::NOP);
        let pe = crate::loader::pe::PE::from_bytes(&buf)?;

        let ws = PEWorkspace::from_pe(pe)?;

        let fmt = Formatter::new().with_colors(false).build();

        // ```
        //     .text:00401C4E 000 68 F4 61 40 00          push    offset ModuleName ; "mscoree.dll"
        //     .text:00401C53 004 FF 15 00 60 40 00       call    ds:GetModuleHandleA
        //     .text:00401C59 000 85 C0                   test    eax, eax
        // ```
        let insn = crate::test::read_insn(&ws.pe.module, 0x401C53);
        assert_eq!(
            fmt.format_instruction(&ws, &insn, 0x401C53)?,
            "call [kernel32.dll!GetModuleHandleA]"
        );

        Ok(())
    }
}
<|endoftext|>


<commit_before>use sdl::{render, pixels, video, keycode};
use sdl::rect::Point;
use std::default::Default;

use self::entity::{Entity,Asteroid};
mod entity;

type Renderer = render::Renderer<video::Window>;

static TURN_SPEED: f32 = 0.001;
static ACCELERATION: f32 = 0.1;

pub struct Loop {
    renderer: Renderer,
    state: State,
}

impl Loop {
    pub fn new(renderer: Renderer) -> Loop {
        let output_size = renderer.get_output_size().unwrap();
        Loop {
            renderer: renderer,
            state: State::new(output_size)
        }
    }

    pub fn handlekey(&mut self, keyDown: bool, key: keycode::KeyCode) {
        // I'm feeling a bit iffy about this. Ideally, no dummy variable
        // should be needed.
        let mut dummy = false;
        let keys = &mut self.state.keys;
        let oldkeys = *keys;
        *match key {
            keycode::UpKey => &mut keys.up,
            keycode::DownKey => &mut keys.down,
            keycode::LeftKey => &mut keys.left,
            keycode::RightKey => &mut keys.right,
            keycode::SpaceKey => &mut keys.space,
            _ => &mut dummy // Ignore any other keys
        } = keyDown;
        if *keys != oldkeys {
            println!("{}", *keys);
        }
    }

    pub fn update(&mut self) {
        {
            let keys = self.state.keys;
            let player = &mut self.state.player1;
            if keys.up { player.accelerate() }
            if keys.down { player.decelerate() }
            if keys.left { player.turn_left() }
            if keys.right { player.turn_right() }
        }
        self.state.player1.update();
        for entity in self.state.entities.iter() {
            match entity.update(self.state.entities.as_slice(), &mut self.state.delete_list, &mut self.state.add_list) {
                Some(e) => self.state.updated_entities.push(e),
                None => {}
            }
        }
        ::std::mem::swap(&mut self.state.entities, &mut self.state.updated_entities);
        self.state.updated_entities.truncate(0);
        self.draw();
    }

    pub fn draw(&self) {
        let renderer = &self.renderer;
        renderer.set_draw_color(pixels::RGB(0, 0, 0)).unwrap();
        renderer.clear().unwrap();
        self.state.player1.draw(renderer);
        let screen_centre = self.state.player1.position();
        for entity in self.state.entities.iter() {
            entity.draw(renderer, screen_centre);
        }
        renderer.present();
    }
}

struct State {
    keys: KeyState,
    player1: PlayerShip,
    entities: Vec<Entity>,
    updated_entities: Vec<Entity>,
    delete_list: Vec<uint>,
    add_list: Vec<Entity>
}

impl State {
    fn new(dims: (int, int)) -> State {
        State {
            keys: Default::default(),
            player1: PlayerShip::new(dims),
            entities: vec![Asteroid.new_entity((10.0, 10.0), 4.0)],
            updated_entities: vec![],
            delete_list: vec![],
            add_list: vec![],
        }
    }
}

#[deriving(Default, PartialEq, Eq, Show)]
struct KeyState {
    down: bool,
    left: bool,
    right: bool,
    up: bool,
    space: bool,
}

// For smoother movement, everything is stored to 64th of a pixel accuracy
struct PlayerShip {
    pos: Point,
    vel: Point,
    angle: f32, // Zero points upwards.
    shape: Vec<(f32,f32)>,
    midpoint: Point,
}

impl PlayerShip {
    fn new((xmax, ymax): (int, int)) -> PlayerShip {
        let midpoint = {
            Point::new((xmax/2) as i32,(ymax/2) as i32)
        };

        let mut shape = Vec::new();
        shape.push_all([(0.0,-10.0),(-5.0,10.0),(0.0,3.0),(5.0,10.0),(0.0,-10.0)]);

        PlayerShip {
            pos: Point::new(0, 0),
            vel: Point::new(0, 0),
            angle: 0.0,
            shape: shape,
            midpoint: midpoint
        }
    }

    fn accelerate(&mut self) {
    }
    fn decelerate(&mut self) {
    }
    fn turn_left(&mut self) {
        self.angle = self.angle - TURN_SPEED;
    }
    fn turn_right(&mut self) {
        self.angle = self.angle + TURN_SPEED;
    }

    fn update(&mut self) {
    }

    fn draw(&self, renderer: &Renderer) {
        let rotated_shape = self.shape
            .iter()
            .map(|&point| rotate(point, self.angle))
            .map(|(x,y)| translate(Point::new(x as i32, y as i32),self.midpoint))
            .collect::<Vec<Point>>();
        renderer.set_draw_color(pixels::RGB(255, 255, 255)).unwrap();
        renderer.draw_lines(rotated_shape.as_slice());
    }

    fn position(&self) -> Point {
        self.pos
    }
}

fn rotate((x,y): (f32,f32), angle: f32) -> (f32,f32) {
    (x * angle.cos() + y * angle.sin(), x * angle.sin() - y * angle.cos())
}

fn translate(point: Point, by: Point) -> Point {
    Point::new(point.x + by.x, point.y + by.y)
}
<commit_msg>Fix ship rotation.<commit_after>use sdl::{render, pixels, video, keycode};
use sdl::rect::Point;
use std::default::Default;

use self::entity::{Entity,Asteroid};
mod entity;

type Renderer = render::Renderer<video::Window>;

static TURN_SPEED: f32 = 0.001;
static ACCELERATION: f32 = 0.1;

pub struct Loop {
    renderer: Renderer,
    state: State,
}

impl Loop {
    pub fn new(renderer: Renderer) -> Loop {
        let output_size = renderer.get_output_size().unwrap();
        Loop {
            renderer: renderer,
            state: State::new(output_size)
        }
    }

    pub fn handlekey(&mut self, keyDown: bool, key: keycode::KeyCode) {
        // I'm feeling a bit iffy about this. Ideally, no dummy variable
        // should be needed.
        let mut dummy = false;
        let keys = &mut self.state.keys;
        let oldkeys = *keys;
        *match key {
            keycode::UpKey => &mut keys.up,
            keycode::DownKey => &mut keys.down,
            keycode::LeftKey => &mut keys.left,
            keycode::RightKey => &mut keys.right,
            keycode::SpaceKey => &mut keys.space,
            _ => &mut dummy // Ignore any other keys
        } = keyDown;
        if *keys != oldkeys {
            println!("{}", *keys);
        }
    }

    pub fn update(&mut self) {
        {
            let keys = self.state.keys;
            let player = &mut self.state.player1;
            if keys.up { player.accelerate() }
            if keys.down { player.decelerate() }
            if keys.left { player.turn_left() }
            if keys.right { player.turn_right() }
        }
        self.state.player1.update();
        for entity in self.state.entities.iter() {
            match entity.update(self.state.entities.as_slice(), &mut self.state.delete_list, &mut self.state.add_list) {
                Some(e) => self.state.updated_entities.push(e),
                None => {}
            }
        }
        ::std::mem::swap(&mut self.state.entities, &mut self.state.updated_entities);
        self.state.updated_entities.truncate(0);
        self.draw();
    }

    pub fn draw(&self) {
        let renderer = &self.renderer;
        renderer.set_draw_color(pixels::RGB(0, 0, 0)).unwrap();
        renderer.clear().unwrap();
        self.state.player1.draw(renderer);
        let screen_centre = self.state.player1.position();
        for entity in self.state.entities.iter() {
            entity.draw(renderer, screen_centre);
        }
        renderer.present();
    }
}

struct State {
    keys: KeyState,
    player1: PlayerShip,
    entities: Vec<Entity>,
    updated_entities: Vec<Entity>,
    delete_list: Vec<uint>,
    add_list: Vec<Entity>
}

impl State {
    fn new(dims: (int, int)) -> State {
        State {
            keys: Default::default(),
            player1: PlayerShip::new(dims),
            entities: vec![Asteroid.new_entity((10.0, 10.0), 4.0)],
            updated_entities: vec![],
            delete_list: vec![],
            add_list: vec![],
        }
    }
}

#[deriving(Default, PartialEq, Eq, Show)]
struct KeyState {
    down: bool,
    left: bool,
    right: bool,
    up: bool,
    space: bool,
}

// For smoother movement, everything is stored to 64th of a pixel accuracy
struct PlayerShip {
    pos: Point,
    vel: Point,
    angle: f32, // Zero points upwards.
    shape: Vec<(f32,f32)>,
    midpoint: Point,
}

impl PlayerShip {
    fn new((xmax, ymax): (int, int)) -> PlayerShip {
        let midpoint = {
            Point::new((xmax/2) as i32,(ymax/2) as i32)
        };

        let mut shape = Vec::new();
        shape.push_all([(0.0,-10.0),(-5.0,10.0),(0.0,3.0),(5.0,10.0),(0.0,-10.0)]);

        PlayerShip {
            pos: Point::new(0, 0),
            vel: Point::new(0, 0),
            angle: 0.0,
            shape: shape,
            midpoint: midpoint
        }
    }

    fn accelerate(&mut self) {
    }
    fn decelerate(&mut self) {
    }
    fn turn_left(&mut self) {
        self.angle = self.angle + TURN_SPEED;
    }
    fn turn_right(&mut self) {
        self.angle = self.angle - TURN_SPEED;
    }

    fn update(&mut self) {
    }

    fn draw(&self, renderer: &Renderer) {
        let rotated_shape = self.shape
            .iter()
            .map(|&point| rotate(point, self.angle))
            .map(|(x,y)| translate(Point::new(x as i32, y as i32),self.midpoint))
            .collect::<Vec<Point>>();
        renderer.set_draw_color(pixels::RGB(255, 255, 255)).unwrap();
        renderer.draw_lines(rotated_shape.as_slice());
    }

    fn position(&self) -> Point {
        self.pos
    }
}

fn rotate((x,y): (f32,f32), angle: f32) -> (f32,f32) {
    (x * angle.cos() + y * angle.sin(), - x * angle.sin() + y * angle.cos())
}

fn translate(point: Point, by: Point) -> Point {
    Point::new(point.x + by.x, point.y + by.y)
}
<|endoftext|>


<commit_before><commit_msg>Adds stub definition to the ETL exercise<commit_after>use std::collections::BTreeMap;

pub fn transform(h: &BTreeMap<i32, Vec<char>>) -> BTreeMap<char, i32> {
    unimplemented!("How will you transform the tree {:?}?", h)
}
<|endoftext|>


<commit_before>//! Implements `PartialEq` for vector types.
#![allow(unused)]

macro_rules! impl_partial_eq {
    ($id:ident) => {
        impl ::cmp::PartialEq<$id> for $id {
            #[inline]
            fn eq(&self, other: &Self) -> bool {
                $id::eq(*self, *other).all()
            }
            #[inline]
            fn ne(&self, other: &Self) -> bool {
                $id::ne(*self, *other).all()
            }
        }
    };
}

#[cfg(test)]
macro_rules! test_partial_eq {
    ($id:ident, $true:expr, $false:expr) => {
        #[test]
        fn partial_eq() {
            use coresimd::simd::*;

            let a = $id::splat($false);
            let b = $id::splat($true);

            assert!(a != b);
            assert!(!(a == b));
            assert!(a == a);
            assert!(!(a != a));
        }
    };
}
<commit_msg>Fix PartialEq issue #511<commit_after>//! Implements `PartialEq` for vector types.
#![allow(unused)]

macro_rules! impl_partial_eq {
    ($id:ident) => {
        impl ::cmp::PartialEq<$id> for $id {
            #[inline]
            fn eq(&self, other: &Self) -> bool {
                $id::eq(*self, *other).all()
            }
            #[inline]
            fn ne(&self, other: &Self) -> bool {
                $id::ne(*self, *other).any()
            }
        }
    };
}

#[cfg(test)]
macro_rules! test_partial_eq {
    ($id:ident, $true:expr, $false:expr) => {
        #[test]
        fn partial_eq() {
            use coresimd::simd::*;

            let a = $id::splat($false);
            let b = $id::splat($true);

            assert!(a != b);
            assert!(!(a == b));
            assert!(a == a);
            assert!(!(a != a));

            // Test further to make sure comparisons work with non-splatted
            // values.
            // This is to test the fix for #511

            let a = $id::splat($false).replace(0, $true);
            let b = $id::splat($true);

            assert!(a != b);
            assert!(!(a == b));
            assert!(a == a);
            assert!(!(a != a));
        }
    };
}
<|endoftext|>


<commit_before>use reqwest;
use serde_json;
use std;
use std::collections::HashMap;
use std::fs::File;
use std::io::Read;
use toml;

#[derive(Deserialize, Debug, Clone)]
pub struct Config {
    pub token: String,
    pub org: String,
    pub team: Vec<TeamConfig>,
    #[serde(default = "default_endpoint")]
    pub endpoint: String,
    #[serde(default = "default_home")]
    pub home: String,
    #[serde(default = "default_sh")]
    pub sh: String,
    #[serde(default = "default_cache_duration")]
    pub cache_duration: u64,
    #[serde(default = "default_cert_path")]
    pub cert_path: String,
    #[serde(default = "default_user_conf_path")]
    pub user_conf_path: String,
    pub proxy_url: Option<String>,
}

fn default_endpoint() -> String { String::from("https://api.github.com") }
fn default_home() -> String { String::from("/home/{}") }
fn default_sh() -> String { String::from("/bin/bash") }
fn default_cache_duration() -> u64 { 3600 }
fn default_cert_path() -> String { String::from("/etc/ssl/certs/ca-certificates.crt") }
fn default_user_conf_path() -> String { String::from(".config/ghteam-auth.toml") }

impl Config {
    pub fn new(configpath: &std::path::Path) -> Result<Self, CliError> {
        let mut file = File::open(configpath)?;
        let mut contents = String::new();
        file.read_to_string(&mut contents)?;
        Ok(toml::from_str::<Config>(&contents)?)
    }
}

#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct UserConfig {
    pub sh: Option<String>,
    pub pass: Option<String>,
}

impl UserConfig {
    #[allow(dead_code)]
    pub fn new(configpath: &std::path::Path) -> Result<Self, CliError> {
        let mut file = File::open(configpath)?;
        let mut contents = String::new();
        file.read_to_string(&mut contents)?;
        Ok(toml::from_str::<UserConfig>(&contents)?)
    }
}

#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct Team {
    pub id: u64,
    pub name: String,
}

#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct TeamConfig {
    pub name: String,
    pub gid: Option<u64>,
    pub group: Option<String>,
}

#[derive(Debug, Clone)]
pub struct TeamGroup {
    pub team: Team,
    pub gid: Option<u64>,
    pub group: Option<String>,
    pub members: HashMap<String, Member>,
}

impl TeamGroup {
    #[allow(dead_code)]
    pub fn get_gid(&self) -> u64 { self.gid.unwrap_or(self.team.id) }
    #[allow(dead_code)]
    pub fn get_group(&self) -> String { self.group.clone().unwrap_or(self.team.name.clone()) }
}

#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct Member {
    pub id: u64,
    pub login: String,
}

#[derive(Serialize, Deserialize, Debug)]
pub struct PublicKey {
    pub id: u64,
    pub key: String,
}

#[derive(Debug)]
pub enum CliError {
    Serde(serde_json::Error),
    Reqwest(reqwest::Error),
    Io(std::io::Error),
    Toml(toml::de::Error),
}

impl From<serde_json::Error> for CliError {
    fn from(err: serde_json::Error) -> CliError { CliError::Serde(err) }
}
impl From<reqwest::Error> for CliError {
    fn from(err: reqwest::Error) -> CliError { CliError::Reqwest(err) }
}
impl From<std::io::Error> for CliError {
    fn from(err: std::io::Error) -> CliError { CliError::Io(err) }
}
impl From<toml::de::Error> for CliError {
    fn from(err: toml::de::Error) -> CliError { CliError::Toml(err) }
}
<commit_msg>Add repo config to the struct<commit_after>use reqwest;
use serde_json;
use std;
use std::collections::HashMap;
use std::fs::File;
use std::io::Read;
use toml;

#[derive(Deserialize, Debug, Clone)]
pub struct Config {
    pub token: String,
    pub org: String,
    pub team: Vec<TeamConfig>,
    pub repo: Option<Vec<RepoConfig>>,
    #[serde(default = "default_endpoint")]
    pub endpoint: String,
    #[serde(default = "default_home")]
    pub home: String,
    #[serde(default = "default_sh")]
    pub sh: String,
    #[serde(default = "default_cache_duration")]
    pub cache_duration: u64,
    #[serde(default = "default_cert_path")]
    pub cert_path: String,
    #[serde(default = "default_user_conf_path")]
    pub user_conf_path: String,
    pub proxy_url: Option<String>,
}

fn default_endpoint() -> String { String::from("https://api.github.com") }
fn default_home() -> String { String::from("/home/{}") }
fn default_sh() -> String { String::from("/bin/bash") }
fn default_cache_duration() -> u64 { 3600 }
fn default_cert_path() -> String { String::from("/etc/ssl/certs/ca-certificates.crt") }
fn default_user_conf_path() -> String { String::from(".config/ghteam-auth.toml") }

impl Config {
    pub fn new(configpath: &std::path::Path) -> Result<Self, CliError> {
        let mut file = File::open(configpath)?;
        let mut contents = String::new();
        file.read_to_string(&mut contents)?;
        Ok(toml::from_str::<Config>(&contents)?)
    }
}

#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct UserConfig {
    pub sh: Option<String>,
    pub pass: Option<String>,
}

impl UserConfig {
    #[allow(dead_code)]
    pub fn new(configpath: &std::path::Path) -> Result<Self, CliError> {
        let mut file = File::open(configpath)?;
        let mut contents = String::new();
        file.read_to_string(&mut contents)?;
        Ok(toml::from_str::<UserConfig>(&contents)?)
    }
}

#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct Team {
    pub id: u64,
    pub name: String,
}

#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct TeamConfig {
    pub name: String,
    pub gid: Option<u64>,
    pub group: Option<String>,
}

#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct RepoConfig {
    pub name: String,
    pub gid: Option<u64>,
    pub group: Option<String>,
}

#[derive(Debug, Clone)]
pub struct TeamGroup {
    pub team: Team,
    pub gid: Option<u64>,
    pub group: Option<String>,
    pub members: HashMap<String, Member>,
}

impl TeamGroup {
    #[allow(dead_code)]
    pub fn get_gid(&self) -> u64 { self.gid.unwrap_or(self.team.id) }
    #[allow(dead_code)]
    pub fn get_group(&self) -> String { self.group.clone().unwrap_or(self.team.name.clone()) }
}

#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct Member {
    pub id: u64,
    pub login: String,
}

#[derive(Serialize, Deserialize, Debug)]
pub struct PublicKey {
    pub id: u64,
    pub key: String,
}

#[derive(Debug)]
pub enum CliError {
    Serde(serde_json::Error),
    Reqwest(reqwest::Error),
    Io(std::io::Error),
    Toml(toml::de::Error),
}

impl From<serde_json::Error> for CliError {
    fn from(err: serde_json::Error) -> CliError { CliError::Serde(err) }
}
impl From<reqwest::Error> for CliError {
    fn from(err: reqwest::Error) -> CliError { CliError::Reqwest(err) }
}
impl From<std::io::Error> for CliError {
    fn from(err: std::io::Error) -> CliError { CliError::Io(err) }
}
impl From<toml::de::Error> for CliError {
    fn from(err: toml::de::Error) -> CliError { CliError::Toml(err) }
}
<|endoftext|>


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

//! Implementation of ElasticSearch Mapping operation

use std::collections::HashMap;

use ::Client;
use ::error::EsError;

pub type DocType<'a>    = HashMap<&'a str, HashMap<&'a str, &'a str>>;
pub type DocTypes<'a>   = HashMap<&'a str, DocType<'a>>;
pub type Mapping<'a>    = HashMap<&'a str, DocTypes<'a>>;

/// An indexing operation
pub struct MappingOperation<'a, 'b> {
    /// The HTTP client that this operation will use
    client:    &'a mut Client,

    /// The index that will be created and eventually mapped
    index:     &'b str,

    /// A map containing the doc types and their type mapping
    doc_types: &'b DocTypes<'b>
}

impl<'a, 'b> MappingOperation<'a, 'b> {
    pub fn new(client: &'a mut Client, index: &'b str, doc_types: &'b DocTypes) -> MappingOperation<'a, 'b> {
        MappingOperation {
            client:    client,
            index:     index,
            doc_types: doc_types
        }
    }

    pub fn send(&'b mut self) -> Result<MappingResult, EsError> {
        let mut mappings: Mapping = HashMap::new();
        for (type_name, mapping) in self.doc_types.into_iter() {
            let doc_type = hashmap! { "properties" => mapping.clone() };
            mappings.insert(type_name, doc_type.to_owned());
        }

        let url    = format!("{}", self.index);
        let (_, _) = try!(self.client.put_body_op(&url, &mappings));
        Ok(MappingResult)
    }
}

/// The result of a mapping operation
#[derive(Debug)]
pub struct MappingResult;

#[cfg(test)]
pub mod tests {
    extern crate env_logger;

    use std::collections::HashMap;

    use super::MappingOperation;

    #[test]
    fn test_mapping() {
        let index_name = "tests_test_mapping";
        let mut client = ::tests::make_client();

        client.delete_index(index_name);

        let mapping = hashmap! { // DocTypes
            "post" => hashmap! { // DocType
                "created_at" => hashmap! {
                    "type" => "date",
                    "format" => "date_time"
                },

                "title" => hashmap! {
                    "type" => "string",
                    "index" => "not_analyzed"
                }
            },

            "author" => hashmap! { // DocType
                "name" => hashmap! {
                    "type" => "string",
                }
            },
        };

        let result = MappingOperation::new(&mut client, index_name, &mapping).send();
        assert!(result.is_ok());
    }
}
<commit_msg>Fix query for mapping the index<commit_after>/*
 * Copyright 2016 Ben Ashford
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

//! Implementation of ElasticSearch Mapping operation

use std::collections::HashMap;

use ::Client;
use ::error::EsError;

pub type DocType<'a>    = HashMap<&'a str, HashMap<&'a str, &'a str>>;
pub type DocTypes<'a>   = HashMap<&'a str, DocType<'a>>;
pub type Mapping<'a>    = HashMap<&'a str, DocTypes<'a>>;

/// An indexing operation
pub struct MappingOperation<'a, 'b> {
    /// The HTTP client that this operation will use
    client:    &'a mut Client,

    /// The index that will be created and eventually mapped
    index:     &'b str,

    /// A map containing the doc types and their type mapping
    doc_types: &'b DocTypes<'b>
}

impl<'a, 'b> MappingOperation<'a, 'b> {
    pub fn new(client: &'a mut Client, index: &'b str, doc_types: &'b DocTypes) -> MappingOperation<'a, 'b> {
        MappingOperation {
            client:    client,
            index:     index,
            doc_types: doc_types
        }
    }

    pub fn send(&'b mut self) -> Result<MappingResult, EsError> {
        let mut mappings: Mapping = HashMap::new();
        for (type_name, mapping) in self.doc_types.into_iter() {
            let doc_type = hashmap! { "properties" => mapping.clone() };
            mappings.insert(type_name, doc_type.to_owned());
        }

        let body = hashmap! { "mappings" => mappings };

        let url    = format!("{}", self.index);
        let (_, _) = try!(self.client.put_body_op(&url, &body));
        Ok(MappingResult)
    }
}

/// The result of a mapping operation
#[derive(Debug)]
pub struct MappingResult;

#[cfg(test)]
pub mod tests {
    extern crate env_logger;

    use std::collections::HashMap;

    use super::MappingOperation;

    #[derive(Debug, RustcDecodable, RustcEncodable)]
    pub struct Author {
        pub name: String
    }

    #[test]
    fn test_mapping() {
        let index_name = "tests_test_mapping";
        let mut client = ::tests::make_client();

        client.delete_index(index_name);

        let mapping = hashmap! { // DocTypes
            "post" => hashmap! { // DocType
                "created_at" => hashmap! {
                    "type" => "date",
                    "format" => "date_time"
                },

                "title" => hashmap! {
                    "type" => "string",
                    "index" => "not_analyzed"
                }
            },

            "author" => hashmap! { // DocType
                "name" => hashmap! {
                    "type" => "string",
                }
            },
        };

        let result = MappingOperation::new(&mut client, index_name, &mapping).send();
        assert!(result.is_ok());

        {
            let result_wrapped = client
                .index(index_name, "post")
                .with_doc(&Author { name: "Homu".to_owned() })
                .send();

            assert!(result_wrapped.is_ok());

            let result = result_wrapped.unwrap();
            assert!(result.created);
        }
    }
}
<|endoftext|>