<commit_before>extern crate lancelot;
extern crate log;

use failure::{Error, Fail};
use log::{error, info, trace};
use std::env;
use std::process;

use lancelot::arch::*;
use lancelot::workspace::Workspace;

// TODO: removeme
use goblin::{Object};

#[derive(Debug, Fail)]
pub enum MainError {
    #[fail(display = "foo")]
    Foo,
}

pub struct Config {
    pub mode: u8,
    pub filename: String,
}

impl Config {
    pub fn from_args(args: env::Args) -> Result<Config, &'static str> {
        let args: Vec<String> = args.collect();

        if args.len() < 3 {
            return Err("not enough arguments: provide `lancelot.exe 32|64 /path/to/input`");
        }

        let mode = match args[1].as_ref() {
            "32" => 32,
            "64" => 64,
            _ => return Err("invalid mode, pick one of `32` or `64`"),
        };

        let filename = args[2].clone();
        trace!("config: parsed filename: {:?}", filename);

        Ok(Config { mode, filename })
    }
}

pub fn setup_logging(_args: &Config) {
    simplelog::TermLogger::init(simplelog::LevelFilter::Info, simplelog::Config::default())
        .expect("failed to setup logging");
}

pub fn run(args: &Config) -> Result<(), Error> {
    info!("filename: {:?}", args.filename);

    if args.mode == 32 {
        let mut ws = Workspace::<Arch32>::from_file(&args.filename)?.load()?;
    } else if args.mode == 64 {
        let mut ws = Workspace::<Arch64>::from_file(&args.filename)?.load()?;

        if ws.loader.get_name() == "Windows/64/PE" {
            if let Ok(Object::PE(pe)) = Object::parse(&ws.buf) {
                let entry = pe.entry;
                let exports: Vec<(usize, Option<String>)> = pe.exports.iter()
                    // re-exports are simply strings that point to a `DLL.export_name` ASCII string.
                    // therefore, they're not functions/code.
                    .filter(|exp| exp.reexport.is_none())
                    .map(|exp| {
                        (exp.rva, exp.name.map(|n| n.to_string()))
                    }).collect();

                // TODO: need to add symbols for re-exports

                info!("PE entry: {:#x}", entry);
                ws.make_symbol(entry as i64, "entry");
                ws.make_insn(entry as i64)?;
                ws.analyze()?;
                for (rva, name) in exports.iter() {
                    info!("export: {:#x}", rva);
                    ws.make_insn(*rva as i64)?;
                    if let Some(name) = name {
                        ws.make_symbol(*rva as i64, &name)?;
                    }
                    ws.analyze()?;
                }
            }
        }
    }

    Ok(())
}

fn main() {
    let args = Config::from_args(env::args()).unwrap_or_else(|err| {
        eprintln!("error parsing arguments: {}", err);
        process::exit(1);
    });

    setup_logging(&args);

    if let Err(e) = run(&args) {
        error!("{:?}", e)
    }
}
<commit_msg>main: make functions, not insns<commit_after>extern crate lancelot;
extern crate log;

use failure::{Error, Fail};
use log::{error, info, trace};
use std::env;
use std::process;

use lancelot::arch::*;
use lancelot::workspace::Workspace;

// TODO: removeme
use goblin::{Object};

#[derive(Debug, Fail)]
pub enum MainError {
    #[fail(display = "foo")]
    Foo,
}

pub struct Config {
    pub mode: u8,
    pub filename: String,
}

impl Config {
    pub fn from_args(args: env::Args) -> Result<Config, &'static str> {
        let args: Vec<String> = args.collect();

        if args.len() < 3 {
            return Err("not enough arguments: provide `lancelot.exe 32|64 /path/to/input`");
        }

        let mode = match args[1].as_ref() {
            "32" => 32,
            "64" => 64,
            _ => return Err("invalid mode, pick one of `32` or `64`"),
        };

        let filename = args[2].clone();
        trace!("config: parsed filename: {:?}", filename);

        Ok(Config { mode, filename })
    }
}

pub fn setup_logging(_args: &Config) {
    simplelog::TermLogger::init(simplelog::LevelFilter::Info, simplelog::Config::default())
        .expect("failed to setup logging");
}

pub fn run(args: &Config) -> Result<(), Error> {
    info!("filename: {:?}", args.filename);

    if args.mode == 32 {
        let mut ws = Workspace::<Arch32>::from_file(&args.filename)?.load()?;
    } else if args.mode == 64 {
        let mut ws = Workspace::<Arch64>::from_file(&args.filename)?.load()?;

        if ws.loader.get_name() == "Windows/64/PE" {
            if let Ok(Object::PE(pe)) = Object::parse(&ws.buf) {
                let entry = pe.entry;
                let exports: Vec<(usize, Option<String>)> = pe.exports.iter()
                    // re-exports are simply strings that point to a `DLL.export_name` ASCII string.
                    // therefore, they're not functions/code.
                    .filter(|exp| exp.reexport.is_none())
                    .map(|exp| {
                        (exp.rva, exp.name.map(|n| n.to_string()))
                    }).collect();

                // TODO: need to add symbols for re-exports

                info!("PE entry: {:#x}", entry);
                ws.make_symbol(entry as i64, "entry")?;
                ws.make_function(entry as i64)?;
                ws.analyze()?;
                for (rva, name) in exports.iter() {
                    info!("export: {:#x}", rva);
                    if let Some(name) = name {
                        ws.make_symbol(*rva as i64, &name)?;
                    }
                    ws.make_function(*rva as i64)?;
                    ws.analyze()?;
                }
            }
        }
    }

    Ok(())
}

fn main() {
    let args = Config::from_args(env::args()).unwrap_or_else(|err| {
        eprintln!("error parsing arguments: {}", err);
        process::exit(1);
    });

    setup_logging(&args);

    if let Err(e) = run(&args) {
        error!("{:?}", e)
    }
}
<|endoftext|>


<commit_before>#![feature(proc_macro, slice_patterns)]

extern crate hyper;
extern crate iron;
extern crate regex;
extern crate router;
extern crate regex;
extern crate serde;
#[macro_use]
extern crate serde_derive;
extern crate serde_json;
extern crate scoped_threadpool;
extern crate tera;

use std::collections::{BTreeMap, HashMap};
use std::env;
use std::sync::{Arc, RwLock};
use std::thread::sleep;
use std::time::{Instant, Duration, SystemTime, UNIX_EPOCH};

use iron::headers::ContentType;
use iron::prelude::*;
use iron::status;
use router::Router;
use scoped_threadpool::Pool;
use tera::{Context, Tera};

pub mod battle_net_api_client;
pub mod thread_throttler;

use battle_net_api_client::{AuctionListing, BattleNetApiClient, Region};

/// Represents a single option available for sale from the blood vendor.
#[derive(Debug, Deserialize)]
struct BloodVendorItem {
    name: String,
    quantity: u64,
    id: u64,
}

/// Value of an item on a realm.
#[derive(Debug)]
struct ItemValue {
    id: u64,
    value: u64,
}

/// The calculated values for items on a particular realm.
#[derive(Debug)]
struct CurrentRealmValues {
    last_update: u64,  // The last time we got this info, as reported by the Blizzard API.
    auction_values: Arc<Vec<ItemValue>>,  // Should be sorted by value.
}

/// All the data in a single row in our price list for a realm.
#[derive(Debug, Serialize)]
struct PriceRow {
    name: String,
    quantity: u64,
    icon: String,
    value_ratio: u64,
    gold: u64,
    silver: u64,
    copper: u64,
}

/// Number of threads to use when fetching auction house results.
const NUM_AUCTION_DATA_THREADS: u32 = 5;

/// Number of seconds to wait between fetching new auction results.
const RESULT_FETCH_PERIOD: u64 = 60 * 30;

/// Given a vec of auction listings for a realm and a map of the items we care about,
/// returns a vec of (item_id, value) sorted by decreasing value, where value is
/// based on the 5th percentile buyout price.
fn calculate_auction_values(listings: &Vec<AuctionListing>, items: &HashMap<u64, BloodVendorItem>) -> Vec<ItemValue> {
    // Calculate 5th percentiles for the items we care about.
    let mut price_points: BTreeMap<u64, Vec<(u64, u64)>> = BTreeMap::new();
    for listing in listings {
        if items.contains_key(&listing.item) && listing.buyout > 0 {
            price_points.entry(listing.item).or_insert(Vec::new()).push((listing.quantity, listing.buyout / listing.quantity));
        }
    }
    for quantities_and_buyouts in price_points.values_mut() {
        quantities_and_buyouts.sort_by_key(|a| a.1);  // Sort by buyout price.
    }
    let total_item_quantities: BTreeMap<u64, u64> =
        price_points.iter().map(|(k, v)| {
            (*k, v.iter().fold(0, |sum, quantity_and_buyout| sum + quantity_and_buyout.0))
        }).collect();
    let fifth_percentile_price_points: BTreeMap<u64, u64> =
        price_points.iter().map(|(item_id, ref item_listings)| {
            let fifth_percentile_quantity = total_item_quantities.get(item_id).unwrap() / 20;
            let mut running_sum: u64 = 0;
            let fifth_percentile_listing = item_listings.iter().find(|&&(quantity, _)| {
                running_sum += quantity;
                running_sum >= fifth_percentile_quantity
            }).unwrap();
            (*item_id, fifth_percentile_listing.1)
        }).collect();
    let mut item_values: Vec<ItemValue> = items.values().map(|item|
        ItemValue {
            id: item.id,
            value: fifth_percentile_price_points.get(&item.id).unwrap_or(&0u64) * item.quantity,
        }
    ).collect();
    item_values.sort_by_key(|item_value| !item_value.value);
    item_values
}

/// Given a battle net Region, return the first path of the app's URL.
fn app_url_for_region(region: &Region) -> &'static str {
    match region {
        &Region::US => "blood-money",
        &Region::EU => "blood-money-eu",
    }
}

fn main() {
    let token = match env::args().nth(1) {
        Some(token) => token,
        None => {
            println!("Usage: bloodmoney <api token> (us|eu)");
            return;
        }
    };
    let locale = match env::args().nth(2) {
        Some(ref s) if s == "us" => Region::US,
        Some(ref s) if s == "eu" => Region::EU,
        _ => {
            println!("Usage: bloodmoney <api token> (us|eu)");
            return;
        }
    };
    let client = Arc::new(BattleNetApiClient::new(&token, locale));

    // Process our item options and grab their icon names.
    let items: Vec<BloodVendorItem> = serde_json::from_str(include_str!("../catalog/items.json"))
        .expect("Error reading items.");
    let item_id_map: Arc<HashMap<u64, BloodVendorItem>> = Arc::new(items.into_iter().map(|x| (x.id, x)).collect());
    let item_icons: Arc<HashMap<u64, String>> = Arc::new(item_id_map.keys().map(|&id| (id, client.get_item_info(id).icon)).collect());

    // Get the list of realms and create an empty price map so we can render pages while
    // waiting for the auction results to be retrieved.
    let realms = Arc::new(client.get_realms());
    let connected_realms = BattleNetApiClient::process_connected_realms(&realms);
    let price_map: Arc<BTreeMap<String, RwLock<CurrentRealmValues>>> =
        Arc::new(realms.iter().map(|realm| (realm.slug.clone(), RwLock::new(CurrentRealmValues {
            last_update: 0,
            auction_values: Arc::new(item_id_map.keys().map(|id| ItemValue{id: *id, value: 0u64}).collect()),
        }))).collect());

    // Set up our web-app.
    let tera = Arc::new(Tera::new("templates/**/*"));
    let mut router = Router::new();
    {
        let realms = realms.clone();
        let tera = tera.clone();
        router.get(format!("/{}", app_url_for_region(&locale)), move |_: &mut Request| {
            let mut context = Context::new();
            context.add("realms", &realms);
            context.add("is_eu", &(locale == Region::EU));
            Ok(Response::with((ContentType::html().0, status::Ok, tera.render("index.html", context).unwrap())))
        }, "index");
    }
    {
        let price_map = price_map.clone();
        let item_id_map = item_id_map.clone();
        let realms = realms.clone();
        let tera = tera.clone();
        router.get(format!("/{}/:realm", app_url_for_region(&locale)), move |req : &mut Request| {
            let realm = req.extensions.get::<Router>().unwrap().find("realm").unwrap();
            if let Some(realm_prices_lock) = price_map.get(realm) {
                let mut context = Context::new();
                let realm_prices = realm_prices_lock.read().unwrap();
                // Build up a list of entries.
                let highest_value = realm_prices.auction_values.get(0).unwrap().value;
                let price_rows: Vec<PriceRow> = realm_prices.auction_values.iter().map(|&ItemValue{id, value}| {
                    let item_info = item_id_map.get(&id).unwrap();
                    let gold = value / (10_000);
                    let silver = (value - gold * 10_000) / 100;
                    let copper = value - gold * 10_000 - silver * 100;
                    let value_ratio = match highest_value {
                        0u64 => 0u64,
                        _ => value*100/highest_value,  // Percentile!
                    };
                    PriceRow {
                        name: item_info.name.clone(),
                        quantity: item_info.quantity,
                        icon: item_icons.get(&id).unwrap().clone(),
                        value_ratio: value_ratio,
                        gold: gold,
                        silver: silver,
                        copper: copper,
                    }
                }).collect();
                context.add("realm_name", &realms.iter().find(|&realm_info| &realm_info.slug == realm).unwrap().name);
                context.add("price_rows", &price_rows);
                // TODO: Change this to something more human readable.
                if realm_prices.last_update == 0 {
                    context.add("update_age", &-1);
                } else {
                    context.add("update_age",
                        &((SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs() - realm_prices.last_update / 1000) / 60));
                }
                context.add("is_eu", &(locale == Region::EU));
                Ok(Response::with((ContentType::html().0, status::Ok, tera.render("prices.html", context).unwrap())))
            } else {
                return Ok(Response::with(status::NotFound));
            }
        }, "realm-prices");
    }
    let http_result = Iron::new(router).http(format!("localhost:{}", match locale {
        Region::US => 3000,
        Region::EU => 3001,
    }).as_str());
    println!("Ready for web traffic.");

    // Now that the webserver is up, periodically fetch
    // new auction house data.
    let mut pool = Pool::new(NUM_AUCTION_DATA_THREADS);
    loop {
        let download_start = Instant::now();
        let next_download_time = download_start + Duration::from_secs(RESULT_FETCH_PERIOD);
        println!("Starting download of auction data.");
        pool.scoped(|scope| {
            for realm_list in &connected_realms {
                // We have to move realm_list into the closure.
                // Clone other values.
                let client = client.clone();
                let price_map = price_map.clone();
                let item_id_map = item_id_map.clone();
                scope.execute(move || {
                    let lead_realm = realm_list.get(0).unwrap();
                    let update_time: u64;
                    let auction_listings: Vec<AuctionListing>;
                    println!("Downloading {}", lead_realm);
                    {
                        let current_realm_values =
                            price_map.get(lead_realm).unwrap().read().unwrap();
                        match client.get_auction_listings(lead_realm, current_realm_values.last_update) {
                            Some((ts, al)) => {
                                update_time = ts;
                                auction_listings = al;
                            },
                            None => return,
                        }
                    }
                    let auction_values = Arc::new(calculate_auction_values(&auction_listings, &item_id_map));
                    for realm in realm_list {
                        println!("Updating {}", realm);
                        let mut current_realm_values =
                            price_map.get(realm).unwrap().write().unwrap();
                        current_realm_values.auction_values = auction_values.clone();
                        current_realm_values.last_update = update_time;
                    }
                })
            }
            scope.join_all();
        });
        let download_end_time = Instant::now();
        println!("Downloading all realms took {} seconds.", download_end_time.duration_since(download_start).as_secs());
        if download_end_time < next_download_time {
            println!("Sleeping for {}", next_download_time.duration_since(download_end_time).as_secs());
            sleep(next_download_time.duration_since(download_end_time));
        }
    }
}
<commit_msg>Fix a bad import<commit_after>#![feature(proc_macro, slice_patterns)]

extern crate hyper;
extern crate iron;
extern crate regex;
extern crate router;
extern crate serde;
#[macro_use]
extern crate serde_derive;
extern crate serde_json;
extern crate scoped_threadpool;
extern crate tera;

use std::collections::{BTreeMap, HashMap};
use std::env;
use std::sync::{Arc, RwLock};
use std::thread::sleep;
use std::time::{Instant, Duration, SystemTime, UNIX_EPOCH};

use iron::headers::ContentType;
use iron::prelude::*;
use iron::status;
use router::Router;
use scoped_threadpool::Pool;
use tera::{Context, Tera};

pub mod battle_net_api_client;
pub mod thread_throttler;

use battle_net_api_client::{AuctionListing, BattleNetApiClient, Region};

/// Represents a single option available for sale from the blood vendor.
#[derive(Debug, Deserialize)]
struct BloodVendorItem {
    name: String,
    quantity: u64,
    id: u64,
}

/// Value of an item on a realm.
#[derive(Debug)]
struct ItemValue {
    id: u64,
    value: u64,
}

/// The calculated values for items on a particular realm.
#[derive(Debug)]
struct CurrentRealmValues {
    last_update: u64,  // The last time we got this info, as reported by the Blizzard API.
    auction_values: Arc<Vec<ItemValue>>,  // Should be sorted by value.
}

/// All the data in a single row in our price list for a realm.
#[derive(Debug, Serialize)]
struct PriceRow {
    name: String,
    quantity: u64,
    icon: String,
    value_ratio: u64,
    gold: u64,
    silver: u64,
    copper: u64,
}

/// Number of threads to use when fetching auction house results.
const NUM_AUCTION_DATA_THREADS: u32 = 5;

/// Number of seconds to wait between fetching new auction results.
const RESULT_FETCH_PERIOD: u64 = 60 * 30;

/// Given a vec of auction listings for a realm and a map of the items we care about,
/// returns a vec of (item_id, value) sorted by decreasing value, where value is
/// based on the 5th percentile buyout price.
fn calculate_auction_values(listings: &Vec<AuctionListing>, items: &HashMap<u64, BloodVendorItem>) -> Vec<ItemValue> {
    // Calculate 5th percentiles for the items we care about.
    let mut price_points: BTreeMap<u64, Vec<(u64, u64)>> = BTreeMap::new();
    for listing in listings {
        if items.contains_key(&listing.item) && listing.buyout > 0 {
            price_points.entry(listing.item).or_insert(Vec::new()).push((listing.quantity, listing.buyout / listing.quantity));
        }
    }
    for quantities_and_buyouts in price_points.values_mut() {
        quantities_and_buyouts.sort_by_key(|a| a.1);  // Sort by buyout price.
    }
    let total_item_quantities: BTreeMap<u64, u64> =
        price_points.iter().map(|(k, v)| {
            (*k, v.iter().fold(0, |sum, quantity_and_buyout| sum + quantity_and_buyout.0))
        }).collect();
    let fifth_percentile_price_points: BTreeMap<u64, u64> =
        price_points.iter().map(|(item_id, ref item_listings)| {
            let fifth_percentile_quantity = total_item_quantities.get(item_id).unwrap() / 20;
            let mut running_sum: u64 = 0;
            let fifth_percentile_listing = item_listings.iter().find(|&&(quantity, _)| {
                running_sum += quantity;
                running_sum >= fifth_percentile_quantity
            }).unwrap();
            (*item_id, fifth_percentile_listing.1)
        }).collect();
    let mut item_values: Vec<ItemValue> = items.values().map(|item|
        ItemValue {
            id: item.id,
            value: fifth_percentile_price_points.get(&item.id).unwrap_or(&0u64) * item.quantity,
        }
    ).collect();
    item_values.sort_by_key(|item_value| !item_value.value);
    item_values
}

/// Given a battle net Region, return the first path of the app's URL.
fn app_url_for_region(region: &Region) -> &'static str {
    match region {
        &Region::US => "blood-money",
        &Region::EU => "blood-money-eu",
    }
}

fn main() {
    let token = match env::args().nth(1) {
        Some(token) => token,
        None => {
            println!("Usage: bloodmoney <api token> (us|eu)");
            return;
        }
    };
    let locale = match env::args().nth(2) {
        Some(ref s) if s == "us" => Region::US,
        Some(ref s) if s == "eu" => Region::EU,
        _ => {
            println!("Usage: bloodmoney <api token> (us|eu)");
            return;
        }
    };
    let client = Arc::new(BattleNetApiClient::new(&token, locale));

    // Process our item options and grab their icon names.
    let items: Vec<BloodVendorItem> = serde_json::from_str(include_str!("../catalog/items.json"))
        .expect("Error reading items.");
    let item_id_map: Arc<HashMap<u64, BloodVendorItem>> = Arc::new(items.into_iter().map(|x| (x.id, x)).collect());
    let item_icons: Arc<HashMap<u64, String>> = Arc::new(item_id_map.keys().map(|&id| (id, client.get_item_info(id).icon)).collect());

    // Get the list of realms and create an empty price map so we can render pages while
    // waiting for the auction results to be retrieved.
    let realms = Arc::new(client.get_realms());
    let connected_realms = BattleNetApiClient::process_connected_realms(&realms);
    let price_map: Arc<BTreeMap<String, RwLock<CurrentRealmValues>>> =
        Arc::new(realms.iter().map(|realm| (realm.slug.clone(), RwLock::new(CurrentRealmValues {
            last_update: 0,
            auction_values: Arc::new(item_id_map.keys().map(|id| ItemValue{id: *id, value: 0u64}).collect()),
        }))).collect());

    // Set up our web-app.
    let tera = Arc::new(Tera::new("templates/**/*"));
    let mut router = Router::new();
    {
        let realms = realms.clone();
        let tera = tera.clone();
        router.get(format!("/{}", app_url_for_region(&locale)), move |_: &mut Request| {
            let mut context = Context::new();
            context.add("realms", &realms);
            context.add("is_eu", &(locale == Region::EU));
            Ok(Response::with((ContentType::html().0, status::Ok, tera.render("index.html", context).unwrap())))
        }, "index");
    }
    {
        let price_map = price_map.clone();
        let item_id_map = item_id_map.clone();
        let realms = realms.clone();
        let tera = tera.clone();
        router.get(format!("/{}/:realm", app_url_for_region(&locale)), move |req : &mut Request| {
            let realm = req.extensions.get::<Router>().unwrap().find("realm").unwrap();
            if let Some(realm_prices_lock) = price_map.get(realm) {
                let mut context = Context::new();
                let realm_prices = realm_prices_lock.read().unwrap();
                // Build up a list of entries.
                let highest_value = realm_prices.auction_values.get(0).unwrap().value;
                let price_rows: Vec<PriceRow> = realm_prices.auction_values.iter().map(|&ItemValue{id, value}| {
                    let item_info = item_id_map.get(&id).unwrap();
                    let gold = value / (10_000);
                    let silver = (value - gold * 10_000) / 100;
                    let copper = value - gold * 10_000 - silver * 100;
                    let value_ratio = match highest_value {
                        0u64 => 0u64,
                        _ => value*100/highest_value,  // Percentile!
                    };
                    PriceRow {
                        name: item_info.name.clone(),
                        quantity: item_info.quantity,
                        icon: item_icons.get(&id).unwrap().clone(),
                        value_ratio: value_ratio,
                        gold: gold,
                        silver: silver,
                        copper: copper,
                    }
                }).collect();
                context.add("realm_name", &realms.iter().find(|&realm_info| &realm_info.slug == realm).unwrap().name);
                context.add("price_rows", &price_rows);
                // TODO: Change this to something more human readable.
                if realm_prices.last_update == 0 {
                    context.add("update_age", &-1);
                } else {
                    context.add("update_age",
                        &((SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs() - realm_prices.last_update / 1000) / 60));
                }
                context.add("is_eu", &(locale == Region::EU));
                Ok(Response::with((ContentType::html().0, status::Ok, tera.render("prices.html", context).unwrap())))
            } else {
                return Ok(Response::with(status::NotFound));
            }
        }, "realm-prices");
    }
    let http_result = Iron::new(router).http(format!("localhost:{}", match locale {
        Region::US => 3000,
        Region::EU => 3001,
    }).as_str());
    println!("Ready for web traffic.");

    // Now that the webserver is up, periodically fetch
    // new auction house data.
    let mut pool = Pool::new(NUM_AUCTION_DATA_THREADS);
    loop {
        let download_start = Instant::now();
        let next_download_time = download_start + Duration::from_secs(RESULT_FETCH_PERIOD);
        println!("Starting download of auction data.");
        pool.scoped(|scope| {
            for realm_list in &connected_realms {
                // We have to move realm_list into the closure.
                // Clone other values.
                let client = client.clone();
                let price_map = price_map.clone();
                let item_id_map = item_id_map.clone();
                scope.execute(move || {
                    let lead_realm = realm_list.get(0).unwrap();
                    let update_time: u64;
                    let auction_listings: Vec<AuctionListing>;
                    println!("Downloading {}", lead_realm);
                    {
                        let current_realm_values =
                            price_map.get(lead_realm).unwrap().read().unwrap();
                        match client.get_auction_listings(lead_realm, current_realm_values.last_update) {
                            Some((ts, al)) => {
                                update_time = ts;
                                auction_listings = al;
                            },
                            None => return,
                        }
                    }
                    let auction_values = Arc::new(calculate_auction_values(&auction_listings, &item_id_map));
                    for realm in realm_list {
                        println!("Updating {}", realm);
                        let mut current_realm_values =
                            price_map.get(realm).unwrap().write().unwrap();
                        current_realm_values.auction_values = auction_values.clone();
                        current_realm_values.last_update = update_time;
                    }
                })
            }
            scope.join_all();
        });
        let download_end_time = Instant::now();
        println!("Downloading all realms took {} seconds.", download_end_time.duration_since(download_start).as_secs());
        if download_end_time < next_download_time {
            println!("Sleeping for {}", next_download_time.duration_since(download_end_time).as_secs());
            sleep(next_download_time.duration_since(download_end_time));
        }
    }
}
<|endoftext|>


<commit_before>#![feature(slice_patterns)]

use std::collections::BTreeMap;
use std::fs;
use std::io;
use std::io::{Read, Write};
use std::thread;

#[derive(Debug)]
#[derive(PartialEq)]
#[derive(Eq)]
enum Op {
    Add(u8),
    Mov(isize),
    In,
    Out,
    Loop(OpStream),

    // extra optimized ops
    Transfer(u8, Vec<(isize, u8)>),
}

use Op::*;

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

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

                    if let &mut Loop(ref mut stream) = &mut self.ops[i] {
                        stream.optimize();
                        maybe_new_op = stream.find_alternative();
                    } else {
                        unreachable!()
                    }

                    if let Some(new_op) = maybe_new_op {
                        self.ops[i] = new_op;
                    }

                    i += 1
                }
                _ => i += 1,
            }
        }
    }

    fn find_alternative(&self) -> Option<Op> {
        let mut map: BTreeMap<isize, u8> = BTreeMap::new();
        let mut rel_index = 0;

        for op in &self.ops[..] {
            match op {
                &Add(x) => {
                    let new_val = map.get(&rel_index).unwrap_or(&0).wrapping_add(x);
                    map.insert(rel_index, new_val);
                }
                &Mov(x) => {
                    rel_index += x;
                }
                _ => {
                    return None;
                }
            }
        }

        if rel_index != 0 {
            return None;
        }

        Some(Transfer(map.remove(&0).unwrap_or(0), map.into_iter().collect()))
    }

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

static ZERO: u8 = 0;

struct State {
    index: usize,
    memory: Vec<u8>,
}

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

    fn rel_index(&self, relative: isize) -> usize {
        (self.index as isize + relative) as usize
    }

    fn step(&mut self, op: &Op) -> bool {
        match *op {
            Add(i) => {
                self[0] = self[0].wrapping_add(i);
            }
            Mov(n) => {
                self.index = self.rel_index(n);
            }
            In => {
                let mut c = [0u8];
                if io::stdin().read(&mut c).unwrap() == 0 {
                    return false;
                }
                self[0] = c[0];
            }
            Out => {
                io::stdout().write(&[self[0]]).unwrap();
            }
            Loop(ref ops) => {
                while self[0] != 0 {
                    if !self.run(ops.get()) {
                        return false;
                    }
                }
            }
            Transfer(d, ref map) => {
                if self[0] == 0 {
                    return true;
                }

                let mut _v = self[0];
                let mut _n = 0;

                while _v != 0 {
                    _v = _v.wrapping_add(d);
                    if _v == self[0] {
                        // stalled: the current transfer will never complete
                        loop {
                            thread::park()
                        }
                    }
                    _n += 1
                }

                self[0] = 0;
                for &(k, v) in &map[..] {
                    self[k] = self[k].wrapping_add(v.wrapping_mul(_n));
                }
            }
        }
        return true;
    }

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

impl std::ops::Index<isize> for State {
    type Output = u8;
    fn index(&self, _index: isize) -> &u8 {
        let idx = self.rel_index(_index);
        if idx >= self.memory.len() {
            &ZERO
        } else {
            &self.memory[idx]
        }
    }
}

impl std::ops::IndexMut<isize> for State {
    fn index_mut(&mut self, _index: isize) -> &mut u8 {
        let idx = self.rel_index(_index);
        if idx >= self.memory.len() {
            self.memory.resize(idx * 2 + 1, 0);
        }
        &mut self.memory[idx]
    }
}

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

    for filename in &filenames[1..] {
        match fs::File::open(filename)
                  .map(|f| io::BufReader::new(f))
                  .and_then(|mut reader| {
                      let mut buffer = Vec::new();
                      reader.read_to_end(&mut buffer).map(|_| buffer)
                  })
                  .map_err(|e| format!("{}", e))
                  .and_then(|buffer| parse(&buffer[..])) {
            Ok(ops) => {
                let mut opstream = OpStream { ops: ops };
                opstream.optimize();
                State::new().run(opstream.get());
            }
            Err(e) => {
                writeln!(&mut io::stderr(),
                         "Error while processing {}: {}",
                         filename,
                         e)
                    .unwrap();
            }
        }
    }
}

fn parse(text: &[u8]) -> Result<Vec<Op>, String> {
    let mut stack = vec![];
    let mut current = vec![];
    let mut line = 1;
    let mut column = 1;

    for c in text {
        match *c as char {
            '+' => current.push(Add(0x01)),
            '-' => current.push(Add(0xff)),
            '>' => current.push(Mov(1)),
            '<' => current.push(Mov(-1)),
            '.' => current.push(Out),
            ',' => current.push(In),
            '[' => {
                stack.push(current);
                current = vec![];
            }
            ']' => {
                let opstream = OpStream { ops: current };
                current = match stack.pop() {
                    Some(v) => v,
                    None => return Err(format!("Stray ] in line {}, column {}", line, column)),
                };
                current.push(Loop(opstream));
            }
            _ => {}
        }
        if *c as char == '\n' {
            line += 1;
            column = 1;
        } else {
            column += 1;
        }
    }
    if !stack.is_empty() {
        return Err(format!("Missing ] in line {}, column {}", line, column));
    }
    Ok(current)
}

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

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

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

    #[test]
    fn test_state_step_add() {
        let mut state = State::new();
        state.step(&Add(23));
        assert_eq!(23, state[0]);
        state.step(&Add(42));
        assert_eq!(65, state[0]);
        state.step(&Add(190));
        assert_eq!(255, state[0]);
    }

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

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

    #[test]
    fn test_state_step_transfer() {
        let mut state = State::new();
        state[0] = 15;
        state[1] = 7;
        state.step(&Transfer(5, vec![(1, 2)]));
        assert_eq!(0, state[0]);
        assert_eq!(1, state[1]);
    }

    #[test]
    fn test_opstream_optimize() {
        let mut opstream = OpStream {
            ops: vec![Mov(1),
                      Mov(1),
                      Add(0x01),
                      Add(0x0ff),
                      Add(0x0ff),
                      Mov(1),
                      Mov(-1),
                      Loop(OpStream { ops: vec![Mov(2), Mov(3)] })],
        };
        opstream.optimize();

        assert_eq!(opstream,
                   OpStream { ops: vec![Mov(2), Add(0xff), Loop(OpStream { ops: vec![Mov(5)] })] });
    }

    #[test]
    fn test_opstream_optimize_transfer() {
        let mut opstream = OpStream {
            ops: vec![Loop(OpStream { ops: vec![Add(0x01), Mov(3), Add(0xff), Mov(-3)] })],
        };
        opstream.optimize();

        assert_eq!(opstream,
                   OpStream { ops: vec![Transfer(1, vec![(3, 255)])] });
    }

    #[test]
    fn test_parse() {
        let input = b"+-[+.,]+";
        assert_eq!(parse(&input[..]),
                   Ok(vec![Add(0x01),
                           Add(0xff),
                           Loop(OpStream { ops: vec![Add(1), Out, In] }),
                           Add(0x01)]));
    }

    #[test]
    fn test_parse_stray() {
        let input = include_bytes!("../test_cases/stray.bf");
        assert_eq!(parse(&input[..]),
                   Err("Stray ] in line 3, column 3".to_string()));
    }

    #[test]
    fn test_parse_incomplete() {
        let input = include_bytes!("../test_cases/incomplete.bf");
        assert_eq!(parse(&input[..]),
                   Err("Missing ] in line 4, column 1".to_string()));
    }
}
<commit_msg>refactor: clean up character types<commit_after>#![feature(slice_patterns)]

use std::collections::BTreeMap;
use std::fs;
use std::io;
use std::io::{Read, Write};
use std::thread;

#[derive(Debug)]
#[derive(PartialEq)]
#[derive(Eq)]
enum Op {
    Add(u8),
    Mov(isize),
    In,
    Out,
    Loop(OpStream),

    // extra optimized ops
    Transfer(u8, Vec<(isize, u8)>),
}

use Op::*;

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

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

                    if let &mut Loop(ref mut stream) = &mut self.ops[i] {
                        stream.optimize();
                        maybe_new_op = stream.find_alternative();
                    } else {
                        unreachable!()
                    }

                    if let Some(new_op) = maybe_new_op {
                        self.ops[i] = new_op;
                    }

                    i += 1
                }
                _ => i += 1,
            }
        }
    }

    fn find_alternative(&self) -> Option<Op> {
        let mut map: BTreeMap<isize, u8> = BTreeMap::new();
        let mut rel_index = 0;

        for op in &self.ops[..] {
            match op {
                &Add(x) => {
                    let new_val = map.get(&rel_index).unwrap_or(&0).wrapping_add(x);
                    map.insert(rel_index, new_val);
                }
                &Mov(x) => {
                    rel_index += x;
                }
                _ => {
                    return None;
                }
            }
        }

        if rel_index != 0 {
            return None;
        }

        Some(Transfer(map.remove(&0).unwrap_or(0), map.into_iter().collect()))
    }

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

static ZERO: u8 = 0;

struct State {
    index: usize,
    memory: Vec<u8>,
}

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

    fn rel_index(&self, relative: isize) -> usize {
        (self.index as isize + relative) as usize
    }

    fn step(&mut self, op: &Op) -> bool {
        match *op {
            Add(i) => {
                self[0] = self[0].wrapping_add(i);
            }
            Mov(n) => {
                self.index = self.rel_index(n);
            }
            In => {
                let mut c = [0u8];
                if io::stdin().read(&mut c).unwrap() == 0 {
                    return false;
                }
                self[0] = c[0];
            }
            Out => {
                io::stdout().write(&[self[0]]).unwrap();
            }
            Loop(ref ops) => {
                while self[0] != 0 {
                    if !self.run(ops.get()) {
                        return false;
                    }
                }
            }
            Transfer(d, ref map) => {
                if self[0] == 0 {
                    return true;
                }

                let mut _v = self[0];
                let mut _n = 0;

                while _v != 0 {
                    _v = _v.wrapping_add(d);
                    if _v == self[0] {
                        // stalled: the current transfer will never complete
                        loop {
                            thread::park()
                        }
                    }
                    _n += 1
                }

                self[0] = 0;
                for &(k, v) in &map[..] {
                    self[k] = self[k].wrapping_add(v.wrapping_mul(_n));
                }
            }
        }
        return true;
    }

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

impl std::ops::Index<isize> for State {
    type Output = u8;
    fn index(&self, _index: isize) -> &u8 {
        let idx = self.rel_index(_index);
        if idx >= self.memory.len() {
            &ZERO
        } else {
            &self.memory[idx]
        }
    }
}

impl std::ops::IndexMut<isize> for State {
    fn index_mut(&mut self, _index: isize) -> &mut u8 {
        let idx = self.rel_index(_index);
        if idx >= self.memory.len() {
            self.memory.resize(idx * 2 + 1, 0);
        }
        &mut self.memory[idx]
    }
}

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

    for filename in &filenames[1..] {
        match fs::File::open(filename)
                  .map(|f| io::BufReader::new(f))
                  .and_then(|mut reader| {
                      let mut buffer = Vec::new();
                      reader.read_to_end(&mut buffer).map(|_| buffer)
                  })
                  .map_err(|e| format!("{}", e))
                  .and_then(|buffer| parse(&buffer[..])) {
            Ok(ops) => {
                let mut opstream = OpStream { ops: ops };
                opstream.optimize();
                State::new().run(opstream.get());
            }
            Err(e) => {
                writeln!(&mut io::stderr(),
                         "Error while processing {}: {}",
                         filename,
                         e)
                    .unwrap();
            }
        }
    }
}

fn parse(text: &[u8]) -> Result<Vec<Op>, String> {
    let mut stack = vec![];
    let mut current = vec![];
    let mut line = 1;
    let mut column = 1;

    for c in text {
        match *c {
            b'+' => current.push(Add(0x01)),
            b'-' => current.push(Add(0xff)),
            b'>' => current.push(Mov(1)),
            b'<' => current.push(Mov(-1)),
            b'.' => current.push(Out),
            b',' => current.push(In),
            b'[' => {
                stack.push(current);
                current = vec![];
            }
            b']' => {
                let opstream = OpStream { ops: current };
                current = match stack.pop() {
                    Some(v) => v,
                    None => return Err(format!("Stray ] in line {}, column {}", line, column)),
                };
                current.push(Loop(opstream));
            }
            _ => {}
        }
        if *c == b'\n' {
            line += 1;
            column = 1;
        } else {
            column += 1;
        }
    }
    if !stack.is_empty() {
        return Err(format!("Missing ] in line {}, column {}", line, column));
    }
    Ok(current)
}

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

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

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

    #[test]
    fn test_state_step_add() {
        let mut state = State::new();
        state.step(&Add(23));
        assert_eq!(23, state[0]);
        state.step(&Add(42));
        assert_eq!(65, state[0]);
        state.step(&Add(190));
        assert_eq!(255, state[0]);
    }

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

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

    #[test]
    fn test_state_step_transfer() {
        let mut state = State::new();
        state[0] = 15;
        state[1] = 7;
        state.step(&Transfer(5, vec![(1, 2)]));
        assert_eq!(0, state[0]);
        assert_eq!(1, state[1]);
    }

    #[test]
    fn test_opstream_optimize() {
        let mut opstream = OpStream {
            ops: vec![Mov(1),
                      Mov(1),
                      Add(0x01),
                      Add(0x0ff),
                      Add(0x0ff),
                      Mov(1),
                      Mov(-1),
                      Loop(OpStream { ops: vec![Mov(2), Mov(3)] })],
        };
        opstream.optimize();

        assert_eq!(opstream,
                   OpStream { ops: vec![Mov(2), Add(0xff), Loop(OpStream { ops: vec![Mov(5)] })] });
    }

    #[test]
    fn test_opstream_optimize_transfer() {
        let mut opstream = OpStream {
            ops: vec![Loop(OpStream { ops: vec![Add(0x01), Mov(3), Add(0xff), Mov(-3)] })],
        };
        opstream.optimize();

        assert_eq!(opstream,
                   OpStream { ops: vec![Transfer(1, vec![(3, 255)])] });
    }

    #[test]
    fn test_parse() {
        let input = b"+-[+.,]+";
        assert_eq!(parse(&input[..]),
                   Ok(vec![Add(0x01),
                           Add(0xff),
                           Loop(OpStream { ops: vec![Add(1), Out, In] }),
                           Add(0x01)]));
    }

    #[test]
    fn test_parse_stray() {
        let input = include_bytes!("../test_cases/stray.bf");
        assert_eq!(parse(&input[..]),
                   Err("Stray ] in line 3, column 3".to_string()));
    }

    #[test]
    fn test_parse_incomplete() {
        let input = include_bytes!("../test_cases/incomplete.bf");
        assert_eq!(parse(&input[..]),
                   Err("Missing ] in line 4, column 1".to_string()));
    }
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate clap;

use clap::{App, ArgGroup, ArgMatches, SubCommand};

arg_enum!{
    #[derive(Debug)]
    enum DomainRec {
        A,
        AAAA,
        CNAME,
        MX,
        NS,
        SRV,
        TXT
    }
}

fn parse_args<'a, 'b>() -> ArgMatches<'a, 'b> {
    let dns_args = "[name]     'The name to use'
                    [data]     'The user data'
                    [priority] 'The priority to set'
                    [port]     'The port to use'
                    [weight]   'The weight value'";
    App::new("docli")
        .version(&format!("v{}", crate_version!()))
        .about("A utility for managing DigitalOcean infrastructure")
        .author("Kevin K. <kbknapp@gmail.com>")
        .args_from_usage("-d --debug         'Displays JSON being sent to server'
                          -n --nosend        'Does NOT execute network send of JSON'
                          -t --token [token] 'Digital Ocean Auth Token (Defaults to contents \
                                              of DO_AUTH_TOKEN env var if omitted)'")
        .subcommand(SubCommand::new("list")
            .about("Commands for displaying available information")
            .subcommand(SubCommand::new("regions")
                .about("Displays available regions"))
            .subcommand(SubCommand::new("sizes")
                .about("Displays available droplet sizes"))
            .subcommand(SubCommand::new("images")
                .args_from_usage("--distrobutions   'Displays all distrobution images'
                                  --applications    'Displays all application images'
                                  --private         'Displays all private user images'
                                  --available       'Displays all available images (Default)'")
                .about("Displays avaiable droplet images")
                .arg_group(ArgGroup::with_name("images").add_all(vec!["distrobutions",
                                                                      "applications",
                                                                      "private",
                                                                      "available"])))
            .subcommand(SubCommand::new("ssh-keys")
                .about("Displays available SSH keys"))
            .subcommand(SubCommand::new("droplets")
                .about("Displays available droplets"))
            .subcommand(SubCommand::new("domains")
                .about("Displays available domains"))
            .subcommand(SubCommand::new("account-info")
                .about("Displays all account information"))
            .subcommand(SubCommand::new("account-actions")
                .about("Displays all current and previous account actions")))
        .subcommand(SubCommand::new("account")
            .about("Commands related to a single account")
            .subcommand(SubCommand::new("retrieve-actions")
                .about("Gets information about a particular account action")
                .arg_from_usage("<action_id> 'The action id to retrieve")))
        .subcommand(SubCommand::new("domains")
            .about("Commands for managing domains")
            .subcommand(SubCommand::new("create")
                .about("Creates a new domain")
                .args_from_usage("<name> 'The name for the domain'
                                  <ip>   'The IP address of the domain'"))
            .subcommand(SubCommand::new("retrieve")
                .about("Gets information on a particular domain")
                .arg_from_usage("<name> 'The name of the domain to get'"))
            .subcommand(SubCommand::new("delete")
                .about("Deletes a domain")
                .arg_from_usage("<name> 'The domain to delete'")))
        .subcommand(SubCommand::new("domain")
            .about("Commands for managing a specific domain")
            .arg_from_usage("<name> 'The domain name to use'")
            .subcommand(SubCommand::new("create-record")
                .about("Creates a new DNS record for a domain")
                .arg_from_usage("<type> 'The type of record to create (i.e. A, AAAA, CNAME, \
                                         MX, NS, SRV, TXT)'")
                .args_from_usage(dns_args))
            .subcommand(SubCommand::new("retrieve-record")
                .about("Gets information on a specific DNS record")
                .arg_from_usage("<id>   'The domain ID to retrieve info on'"))
            .subcommand(SubCommand::new("update-record")
                .about("Updates a DNS record")
                .arg_from_usage("<type> 'The type of record to create (i.e. A, AAAA, CNAME, \
                                         MX, NS, SRV, TXT)'")
                .args_from_usage(dns_args))
            .subcommand(SubCommand::new("delete-record")
                .about("Deletes a DNS record")
                .arg_from_usage("<id>   'The domain ID to delete'")))
        .subcommand(SubCommand::new("droplets")
            .about("Commands for managing droplets")
            .subcommand(SubCommand::new("list-all-neighbors")
                .about("Displays all droplets running on the same physical hardware"))
            .subcommand(SubCommand::new("list-upgrades")
                .about("Displays all droplets with pending upgrades"))
            .subcommand(SubCommand::new("create")
                .about("Creates a new droplet")
                .args_from_usage("<name>                      'The name of the droplet'
                                  -r --region <region>        'The region of the droplet'
                                  -z --size <size>            'The size of the droplet'
                                  -i --image <image>          'The image to use'
                                  -s --ssh-keys [ssh_keys]... 'Any ssh keys to add'
                                  --backups                   'Allow backups'
                                  --ipv6                      'Use IPv6'
                                  --private-networking        'Use private networking'
                                  -u --user-data [data]       'User data'")))
        .subcommand(SubCommand::new("droplet")
            .about("Commands for managing a single droplet")
            .arg_from_usage("<id> 'The droplet ID to use'")
            .subcommand(SubCommand::new("retrieve")
                .about("Display the information on the droplet"))
            .subcommand(SubCommand::new("list-kernels")
                .about("Display all available kernels"))
            .subcommand(SubCommand::new("list-snapshots")
                .about("Display all snapshots"))
            .subcommand(SubCommand::new("list-backups")
                .about("Display all backups"))
            .subcommand(SubCommand::new("list-actions")
                .about("Display all current and previous actions"))
            .subcommand(SubCommand::new("list-neighbors")
                .about("Display all droplets running on the same physical hardware"))
            .subcommand(SubCommand::new("delete")
                .about("Deletes a droplet"))
            .subcommand(SubCommand::new("disable-backups")
                .about("Disables backups for a droplet"))
            .subcommand(SubCommand::new("reboot")
                .about("Reboots a droplet"))
            .subcommand(SubCommand::new("power-cycle")
                .about("Performs a power cycle on a droplet"))
            .subcommand(SubCommand::new("shutdown")
                .about("Shutsdown a droplet"))
            .subcommand(SubCommand::new("power-off")
                .about("Powers a droplet off"))
            .subcommand(SubCommand::new("power-on")
                .about("Turns on a droplet"))
            .subcommand(SubCommand::new("restore")
                .about("Restores a droplet from an image")
                .arg_from_usage("<image> 'The image to restore to'"))
            .subcommand(SubCommand::new("password-reset")
                .about("Resets the root password for a droplet"))
            .subcommand(SubCommand::new("resize")
                .about("Resizes a droplet")
                .args_from_usage("--disk 'Resizes the disk'
                                  <size> 'The new size to use'"))
            .subcommand(SubCommand::new("rebuild")
                .about("Rebuilds a droplet from an image")
                .arg_from_usage("<image> 'The image to use'"))
            .subcommand(SubCommand::new("rename")
                .about("Renames a droplet")
                .arg_from_usage("<name> 'The new name of the droplet'"))
            .subcommand(SubCommand::new("change-kernel")
                .about("Changes the kernel of a droplet")
                .arg_from_usage("<kernel_id> 'The kernel ID to use'"))
            .subcommand(SubCommand::new("enable-ipv6")
                .about("Enables IPv6 addresses"))
            .subcommand(SubCommand::new("enable-private-networking")
                .about("Enables private networking"))
            .subcommand(SubCommand::new("snapshot")
                .about("Creates a snapshot of a droplet"))
            .subcommand(SubCommand::new("retrieve-action")
                .about("Displays a specific action for a droplet")
                .arg_from_usage("<action_id> 'The action ID to display'"))
            .subcommand(SubCommand::new("upgrade")
                .about("Performs pending upgrades")))
        .subcommand(SubCommand::new("image")
            .about("Gets or sets information on a particular image")
            .arg_from_usage("<id> 'The image ID to use'")
            .subcommand(SubCommand::new("list-actions")
                .about("Lists all previous and current actions for an image"))
            .subcommand(SubCommand::new("retrieve")
                .about("Displays a particular image")
                .arg_from_usage("--slug 'The <id> provided to \'docli image\' is a slug and \
                                         NOT an image ID'"))
            .subcommand(SubCommand::new("update")
                .about("Performs pending updates"))
            .subcommand(SubCommand::new("delete")
                .about("Deletes an image"))
            .subcommand(SubCommand::new("transfer")
                .about("Transfers an image to a new region")
                .arg_from_usage("<region> 'The region to transfer to'"))
            .subcommand(SubCommand::new("convert")
                .about("Converts a an image (i.e. from a snapshot to a backup)"))
            .subcommand(SubCommand::new("retrieve-action")
                .about("Displays a particular action of an image")
                .arg_from_usage("<action_id> 'The action ID to display'")))
        .subcommand(SubCommand::new("ssh-keys")
            .about("Gets or sets information on SSH keys")
            .subcommand(SubCommand::new("create")
                .about("Creatse a new SSH key")
                .arg_from_usage("<name>       'The name of the SSH key'
                                 <public_key> 'The public key of the SSH key'"))
            .subcommand(SubCommand::new("retrieve")
                .about("Displays information on a particular key")
                .args_from_usage("<key_id>       'The key ID of the key to display'
                                  <finger_print> 'The fingerprint of the key to display'"))
            .subcommand(SubCommand::new("update")
                .about("Updates a particular SSH key")
                .args_from_usage("<name>         'The new name to use'
                                  <key_id>       'The key ID to update'
                                  <finger_print> 'The fingerprint of the key to update'")
                .arg_group(ArgGroup::with_name("sshkeys")
                    .add_all(vec!["key_id",
                                  "finger_print"])
                    .required(true)))
            .subcommand(SubCommand::new("destroy")
                .about("Destroys a particular SSH key")
                .args_from_usage("<key_id>       'The key ID to update'
                                  <finger_print> 'The fingerprint of the key to update'")
                .arg_group(ArgGroup::with_name("sshkeys")
                    .add_all(vec!["key_id",
                                  "finger_print"])
                    .required(true))))
        .subcommand(SubCommand::new("list-regions")
            .about("Display all available regions"))
        .subcommand(SubCommand::new("list-sizes")
            .about("Display all available sizes"))
        .get_matches()
}

fn get_auth_token(m: &ArgMatches) -> String {
    let tok = if let Some(auth_tok) = m.value_of("token") {
            auth_tok.to_owned()
    } else {
        std::env::vars().filter(|&(ref k, _)| k == "DO_AUTH_TOKEN")
                        .map(|(_, v)| v.clone() )
                        .next().unwrap_or("".to_owned())
    };
    if tok.len() != 64 {
        println!("No DigitalOcean Auth Token found.\n\n\
        Use `docli --token <token>` or set the DO_AUTH_TOKEN environment variable and try again");
        std::process::exit(1);
    }
    tok
}

fn main() {
    let m = parse_args();

    let DEBUG = m.is_present("debug");
    let SEND = m.is_present("nosend");
    let AUTH_TOKEN = get_auth_token(&m);

    match m.subcommand() {
        ("account", Some(m))  =>(),
        ("domains", Some(m))  =>(),
        ("domain", Some(m))   =>(),
        ("droplets", Some(m)) =>(),
        ("droplet", Some(m))  =>(),
        ("image", Some(m))    =>(),
        ("ssh-keys", Some(m)) =>(),
        ("list", Some(m))     =>(),
        _                     => println!("No command was provided\n\n{}",m.usage())
    }

    println!("Done");
}
<commit_msg>fixup! refactor(subcommands): re-order and change subcommands<commit_after>#[macro_use]
extern crate clap;

use clap::{App, ArgGroup, ArgMatches, SubCommand};

arg_enum!{
    #[derive(Debug)]
    enum DomainRec {
        A,
        AAAA,
        CNAME,
        MX,
        NS,
        SRV,
        TXT
    }
}

fn parse_args<'a, 'b>() -> ArgMatches<'a, 'b> {
    let dns_args = "[name]     'The name to use'
                    [data]     'The user data'
                    [priority] 'The priority to set'
                    [port]     'The port to use'
                    [weight]   'The weight value'";
    App::new("docli")
        .version(&format!("v{}", crate_version!()))
        .about("A utility for managing DigitalOcean infrastructure")
        .author("Kevin K. <kbknapp@gmail.com>")
        .args_from_usage("-d --debug         'Displays JSON being sent to server'
                          -n --nosend        'Does NOT execute network send of JSON'
                          -t --token [token] 'Digital Ocean Auth Token (Defaults to contents \
                                              of DO_AUTH_TOKEN env var if omitted)'")
        .subcommand(SubCommand::new("list")
            .about("Commands for displaying available information")
            .subcommand(SubCommand::new("regions")
                .about("Displays available regions"))
            .subcommand(SubCommand::new("sizes")
                .about("Displays available droplet sizes"))
            .subcommand(SubCommand::new("images")
                .args_from_usage("--distrobutions   'Displays all distrobution images'
                                  --applications    'Displays all application images'
                                  --private         'Displays all private user images'
                                  --available       'Displays all available images (Default)'")
                .about("Displays avaiable droplet images")
                .arg_group(ArgGroup::with_name("images").add_all(vec!["distrobutions",
                                                                      "applications",
                                                                      "private",
                                                                      "available"])))
            .subcommand(SubCommand::new("ssh-keys")
                .about("Displays available SSH keys"))
            .subcommand(SubCommand::new("droplets")
                .about("Displays available droplets"))
            .subcommand(SubCommand::new("domains")
                .about("Displays available domains"))
            .subcommand(SubCommand::new("account-info")
                .about("Displays all account information"))
            .subcommand(SubCommand::new("account-actions")
                .about("Displays all current and previous account actions")))
        .subcommand(SubCommand::new("account")
            .about("Commands related to a single account")
            .subcommand(SubCommand::new("retrieve-actions")
                .about("Gets information about a particular account action")
                .arg_from_usage("<action_id> 'The action id to retrieve")))
        .subcommand(SubCommand::new("domains")
            .about("Commands for managing domains")
            .subcommand(SubCommand::new("create")
                .about("Creates a new domain")
                .args_from_usage("<name> 'The name for the domain'
                                  <ip>   'The IP address of the domain'"))
            .subcommand(SubCommand::new("retrieve")
                .about("Gets information on a particular domain")
                .arg_from_usage("<name> 'The name of the domain to get'"))
            .subcommand(SubCommand::new("delete")
                .about("Deletes a domain")
                .arg_from_usage("<name> 'The domain to delete'")))
        .subcommand(SubCommand::new("domain")
            .about("Commands for managing a specific domain")
            .arg_from_usage("<name> 'The domain name to use'")
            .subcommand(SubCommand::new("create-record")
                .about("Creates a new DNS record for a domain")
                .arg_from_usage("<type> 'The type of record to create (i.e. A, AAAA, CNAME, \
                                         MX, NS, SRV, TXT)'")
                .args_from_usage(dns_args))
            .subcommand(SubCommand::new("retrieve-record")
                .about("Gets information on a specific DNS record")
                .arg_from_usage("<id>   'The domain ID to retrieve info on'"))
            .subcommand(SubCommand::new("update-record")
                .about("Updates a DNS record")
                .arg_from_usage("<type> 'The type of record to create (i.e. A, AAAA, CNAME, \
                                         MX, NS, SRV, TXT)'")
                .args_from_usage(dns_args))
            .subcommand(SubCommand::new("delete-record")
                .about("Deletes a DNS record")
                .arg_from_usage("<id>   'The domain ID to delete'")))
        .subcommand(SubCommand::new("droplets")
            .about("Commands for managing droplets")
            .subcommand(SubCommand::new("list-all-neighbors")
                .about("Displays all droplets running on the same physical hardware"))
            .subcommand(SubCommand::new("list-upgrades")
                .about("Displays all droplets with pending upgrades"))
            .subcommand(SubCommand::new("create")
                .about("Creates a new droplet")
                .args_from_usage("<name>                      'The name of the droplet'
                                  -r --region <region>        'The region of the droplet'
                                  -z --size <size>            'The size of the droplet'
                                  -i --image <image>          'The image to use'
                                  -s --ssh-keys [ssh_keys]... 'Any ssh keys to add'
                                  --backups                   'Allow backups'
                                  --ipv6                      'Use IPv6'
                                  --private-networking        'Use private networking'
                                  -u --user-data [data]       'User data'")))
        .subcommand(SubCommand::new("droplet")
            .about("Commands for managing a single droplet")
            .arg_from_usage("<id> 'The droplet ID to use'")
            .subcommand(SubCommand::new("retrieve")
                .about("Display the information on the droplet"))
            .subcommand(SubCommand::new("list-kernels")
                .about("Display all available kernels"))
            .subcommand(SubCommand::new("list-snapshots")
                .about("Display all snapshots"))
            .subcommand(SubCommand::new("list-backups")
                .about("Display all backups"))
            .subcommand(SubCommand::new("list-actions")
                .about("Display all current and previous actions"))
            .subcommand(SubCommand::new("list-neighbors")
                .about("Display all droplets running on the same physical hardware"))
            .subcommand(SubCommand::new("delete")
                .about("Deletes a droplet"))
            .subcommand(SubCommand::new("disable-backups")
                .about("Disables backups for a droplet"))
            .subcommand(SubCommand::new("reboot")
                .about("Reboots a droplet"))
            .subcommand(SubCommand::new("power-cycle")
                .about("Performs a power cycle on a droplet"))
            .subcommand(SubCommand::new("shutdown")
                .about("Shutsdown a droplet"))
            .subcommand(SubCommand::new("power-off")
                .about("Powers a droplet off"))
            .subcommand(SubCommand::new("power-on")
                .about("Turns on a droplet"))
            .subcommand(SubCommand::new("restore")
                .about("Restores a droplet from an image")
                .arg_from_usage("<image> 'The image to restore to'"))
            .subcommand(SubCommand::new("password-reset")
                .about("Resets the root password for a droplet"))
            .subcommand(SubCommand::new("resize")
                .about("Resizes a droplet")
                .args_from_usage("--disk 'Resizes the disk'
                                  <size> 'The new size to use'"))
            .subcommand(SubCommand::new("rebuild")
                .about("Rebuilds a droplet from an image")
                .arg_from_usage("<image> 'The image to use'"))
            .subcommand(SubCommand::new("rename")
                .about("Renames a droplet")
                .arg_from_usage("<name> 'The new name of the droplet'"))
            .subcommand(SubCommand::new("change-kernel")
                .about("Changes the kernel of a droplet")
                .arg_from_usage("<kernel_id> 'The kernel ID to use'"))
            .subcommand(SubCommand::new("enable-ipv6")
                .about("Enables IPv6 addresses"))
            .subcommand(SubCommand::new("enable-private-networking")
                .about("Enables private networking"))
            .subcommand(SubCommand::new("snapshot")
                .about("Creates a snapshot of a droplet"))
            .subcommand(SubCommand::new("retrieve-action")
                .about("Displays a specific action for a droplet")
                .arg_from_usage("<action_id> 'The action ID to display'"))
            .subcommand(SubCommand::new("upgrade")
                .about("Performs pending upgrades")))
        .subcommand(SubCommand::new("image")
            .about("Gets or sets information on a particular image")
            .arg_from_usage("<id> 'The image ID to use'")
            .subcommand(SubCommand::new("list-actions")
                .about("Lists all previous and current actions for an image"))
            .subcommand(SubCommand::new("retrieve")
                .about("Displays a particular image")
                .arg_from_usage("--slug 'The <id> provided to \'docli image\' is a slug and \
                                         NOT an image ID'"))
            .subcommand(SubCommand::new("update")
                .about("Performs pending updates"))
            .subcommand(SubCommand::new("delete")
                .about("Deletes an image"))
            .subcommand(SubCommand::new("transfer")
                .about("Transfers an image to a new region")
                .arg_from_usage("<region> 'The region to transfer to'"))
            .subcommand(SubCommand::new("convert")
                .about("Converts a an image (i.e. from a snapshot to a backup)"))
            .subcommand(SubCommand::new("retrieve-action")
                .about("Displays a particular action of an image")
                .arg_from_usage("<action_id> 'The action ID to display'")))
        .subcommand(SubCommand::new("ssh-keys")
            .about("Gets or sets information on SSH keys")
            .subcommand(SubCommand::new("create")
                .about("Creatse a new SSH key")
                .arg_from_usage("<name>       'The name of the SSH key'
                                 <public_key> 'The public key of the SSH key'"))
            .subcommand(SubCommand::new("retrieve")
                .about("Displays information on a particular key")
                .args_from_usage("<key_id>       'The key ID of the key to display'
                                  <finger_print> 'The fingerprint of the key to display'"))
            .subcommand(SubCommand::new("update")
                .about("Updates a particular SSH key")
                .args_from_usage("<name>         'The new name to use'
                                  <key_id>       'The key ID to update'
                                  <finger_print> 'The fingerprint of the key to update'")
                .arg_group(ArgGroup::with_name("sshkeys")
                    .add_all(vec!["key_id",
                                  "finger_print"])
                    .required(true)))
            .subcommand(SubCommand::new("destroy")
                .about("Destroys a particular SSH key")
                .args_from_usage("<key_id>       'The key ID to update'
                                  <finger_print> 'The fingerprint of the key to update'")
                .arg_group(ArgGroup::with_name("sshkeys")
                    .add_all(vec!["key_id",
                                  "finger_print"])
                    .required(true))))
        .subcommand(SubCommand::new("list-regions")
            .about("Display all available regions"))
        .subcommand(SubCommand::new("list-sizes")
            .about("Display all available sizes"))
        .get_matches()
}

fn get_auth_token(m: &ArgMatches) -> String {
    let tok = if let Some(auth_tok) = m.value_of("token") {
            auth_tok.to_owned()
    } else {
        std::env::vars().filter(|&(ref k, _)| k == "DO_AUTH_TOKEN")
                        .map(|(_, v)| v.clone() )
                        .next().unwrap_or("".to_owned())
    };
    if tok.len() != 64 {
        println!("No DigitalOcean Auth Token found.\n\n\
        Use `docli --token <token>` or set the DO_AUTH_TOKEN environment variable and try again");
        std::process::exit(1);
    }
    tok
}

fn main() {
    let m = parse_args();

    let DEBUG = m.is_present("debug");
    let SEND = m.is_present("nosend");
    let AUTH_TOKEN = get_auth_token(&m);

    match m.subcommand() {
        ("account", Some(m))  =>(),
        ("domains", Some(m))  =>(),
        ("domain", Some(m))   =>(),
        ("droplets", Some(m)) =>(),
        ("droplet", Some(m))  =>(),
        ("image", Some(m))    =>(),
        ("ssh-keys", Some(m)) =>(),
        ("list", Some(m))     =>(),
        _                     => println!("No command was provided\n\n{}\n\n\
                                           For more information try --help",m.usage())
    }
}
<|endoftext|>