diff --git a/README.md b/README.md index 91d3b4c..2cfe000 100644 --- a/README.md +++ b/README.md @@ -241,6 +241,33 @@ Quantity: 2 Size: standard-2X ``` +**~block_ip** + +If you wish to block an IP address from accessing your application, you can do so with the ~block_ip command. + +``` +you: ~block_ip you_app_name ip_address_to_block +``` + +``` +you: ~block_ip testing-nell-bot 123.456.68 +crates-io-bot: @you IP address 123.456.68 +``` + +**~unblock_ip** + +If you wish to unblock an IP address that was previously +blocked for your application you can do so with the ~unblock_ip command: + +``` +you: ~unblock_ip you_app_name ip_address_to_unblock +``` + +``` +you: ~unblock_ip testing-nell-bot 123.456.68 +crates-io-bot: @you IP address 123.456.789 has been unblocked +``` + ## Setup To setup this Discord bot, you need: diff --git a/src/commands/heroku.rs b/src/commands/heroku.rs index 26ce77b..4f4fcd3 100644 --- a/src/commands/heroku.rs +++ b/src/commands/heroku.rs @@ -9,6 +9,9 @@ use serenity::model::prelude::*; use serenity::prelude::*; use std::collections::HashMap; +use std::collections::HashSet; + +use crate::utilities::*; #[derive(Debug, Deserialize)] struct HerokuApp { @@ -55,9 +58,7 @@ pub fn get_app(ctx: &mut Context, msg: &Message, mut args: Args) -> CommandResul Ok(()) } -// Config variables that can be updated through Discord -// Set only as "FOO" until we fill this in with the real config vars -// for our Heroku account that we want to allow to be updated +// App config variables that can be updated through Discord const AUTHORIZED_CONFIG_VARS: &[&str] = &["FOO"]; // Get app by name or id @@ -110,6 +111,154 @@ pub fn update_app_config(ctx: &mut Context, msg: &Message, mut args: Args) -> Co Ok(()) } +const BLOCKED_IPS_ENV_VAR: &str = "BLOCKED_IPS"; + +#[command] +#[num_args(2)] +pub fn block_ip(ctx: &mut Context, msg: &Message, mut args: Args) -> CommandResult { + let app_name = args + .single::() + .expect("You must include an app name"); + + let ip_addr = args + .single::() + .expect("You must include an IP address to block"); + + let current_config_vars = heroku_app_config_vars(&ctx, &app_name); + + // If the BLOCKED_IPS environmental variable does not + // currently exist, create it + if !blocked_ips_exist(¤t_config_vars) { + let response = heroku_client(&ctx).request(&config_vars::AppConfigVarUpdate { + app_id: &app_name, + params: empty_config_var(), + }); + + let response_err = response.is_err(); + + msg.reply( + &ctx, + match response { + Ok(_response) => format!("The {} environmental variable has been created for {}", BLOCKED_IPS_ENV_VAR, app_name), + Err(e) => format!( + "The {} environmental variable does not current exist for {}.\n There was an error when trying to create it: {}", + BLOCKED_IPS_ENV_VAR, app_name, e + ), + }, + )?; + + if response_err { + return Ok(()); + } + } + + let mut blocked_ips_set = current_blocked_ip_addresses(current_config_vars); + + if blocked_ips_set.contains(&ip_addr) { + msg.reply( + &ctx, + format!("{} is already blocked for {}", &ip_addr, app_name), + )?; + } else { + blocked_ips_set.insert(ip_addr.clone()); + + let updated_config_var = blocked_ips_config_var(blocked_ips_set); + + let response = heroku_client(ctx).request(&config_vars::AppConfigVarUpdate { + app_id: &app_name, + params: updated_config_var, + }); + + msg.reply( + ctx, + match response { + Ok(_response) => format!("IP address {} has been blocked", ip_addr.clone()), + Err(e) => format!( + "An error occurred when trying to block the IP address: {}\n{}", + ip_addr, e + ), + }, + )?; + }; + + Ok(()) +} + +#[command] +#[num_args(2)] +pub fn unblock_ip(ctx: &mut Context, msg: &Message, mut args: Args) -> CommandResult { + let app_name = args + .single::() + .expect("You must include an app name"); + + let ip_addr = args + .single::() + .expect("You must include an IP address to unblock"); + + let current_config_vars = heroku_app_config_vars(&ctx, &app_name); + + if !blocked_ips_exist(¤t_config_vars) { + msg.reply( + &ctx, + format!("No IP addresses are currently blocked for {}", &app_name), + )?; + + return Ok(()); + } + + let mut blocked_ips_set = current_blocked_ip_addresses(current_config_vars); + + if !blocked_ips_set.contains(&ip_addr) { + msg.reply( + &ctx, + format!("{} is not currently blocked for {}", &ip_addr, app_name), + )?; + } else { + blocked_ips_set.remove(&ip_addr); + + // Removes config variable from the Heroku application + // if there are no more blocked ip addresses + if blocked_ips_set.is_empty() { + let response = heroku_client(ctx).request(&config_vars::AppConfigVarDelete { + app_id: &app_name, + params: null_blocked_ips_config_var(), + }); + + msg.reply( + ctx, + match response { + Ok(_response) => format!( + "IP address {} has been unblocked, there are now no unblocked IP addresses", + ip_addr.clone() + ), + Err(e) => format!( + "An error occurred when trying to unblock the IP address: {}\n{}", + ip_addr, e + ), + }, + )?; + } else { + let response = heroku_client(ctx).request(&config_vars::AppConfigVarUpdate { + app_id: &app_name, + params: blocked_ips_config_var(blocked_ips_set), + }); + + msg.reply( + ctx, + match response { + Ok(_response) => format!("IP address {} has been unblocked", ip_addr.clone()), + Err(e) => format!( + "An error occurred when trying to unblock the IP address: {}\n{}", + ip_addr, e + ), + }, + )?; + }; + } + + Ok(()) +} + #[command] #[num_args(4)] pub fn scale_app(ctx: &mut Context, msg: &Message, mut args: Args) -> CommandResult { @@ -322,3 +471,55 @@ fn heroku_client(ctx: &Context) -> std::sync::Arc HashMap> { + let config_var_list = heroku_client(ctx) + .request(&config_vars::AppConfigVarDetails { app_id: &app_name }) + .unwrap(); + config_var_list +} + +fn block_ips_value(config_vars: HashMap>) -> String { + config_vars + .get(&BLOCKED_IPS_ENV_VAR.to_string()) + .unwrap() + .as_ref() + .unwrap() + .to_string() +} + +fn blocked_ips_config_var(blocked_ips_set: HashSet) -> HashMap { + let blocked_ips_set_string = parse_config_value_string(blocked_ips_set); + let blocked_ips_config_var = config_var(blocked_ips_set_string); + blocked_ips_config_var +} + +fn config_var(updated_blocked_ips_value: String) -> HashMap { + let mut config_var = HashMap::new(); + config_var.insert(BLOCKED_IPS_ENV_VAR.to_string(), updated_blocked_ips_value); + config_var +} + +fn empty_config_var() -> HashMap { + let mut config_var = HashMap::new(); + config_var.insert(BLOCKED_IPS_ENV_VAR.to_string(), "".to_string()); + config_var +} + +fn current_blocked_ip_addresses(config_vars: HashMap>) -> HashSet { + let blocked_ips_value = block_ips_value(config_vars); + + let blocked_ips_set = parse_config_value_set(blocked_ips_value); + blocked_ips_set +} + +fn null_blocked_ips_config_var() -> HashMap> { + let mut config_var = HashMap::new(); + config_var.insert(BLOCKED_IPS_ENV_VAR.to_string(), None); + config_var +} + +fn blocked_ips_exist(config_vars: &HashMap>) -> bool { + let exists = config_vars.get(&BLOCKED_IPS_ENV_VAR.to_string()).is_some(); + exists +} diff --git a/src/config.rs b/src/config.rs index c59c815..96ea331 100644 --- a/src/config.rs +++ b/src/config.rs @@ -1,3 +1,4 @@ +use crate::utilities::*; use serenity::prelude::TypeMapKey; use std::collections::HashSet; use std::sync::Arc; @@ -22,30 +23,3 @@ impl Config { impl TypeMapKey for Config { type Value = Arc; } - -fn parse_config_value_set(config_value: String) -> HashSet { - let mut value_set = HashSet::new(); - - let split_string = config_value.split(','); - - for string in split_string { - value_set.insert(String::from(string)); - } - - value_set -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn create_authorized_users_hashset() { - let test_string = String::from("123,456,789"); - let users_set = parse_config_value_set(test_string); - - assert!(users_set.contains("123")); - assert!(users_set.contains("456")); - assert!(users_set.contains("789")); - } -} diff --git a/src/lib.rs b/src/lib.rs index 1bb92ff..3bff33e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -15,6 +15,8 @@ mod authorizations; pub mod config; +pub mod utilities; + use crate::config::Config; use crate::authorizations::users::*; @@ -30,7 +32,9 @@ use crate::authorizations::users::*; scale_app, update_app_config, get_app_releases, - rollback_app + rollback_app, + block_ip, + unblock_ip )] struct General; diff --git a/src/utilities.rs b/src/utilities.rs new file mode 100644 index 0000000..85fce9b --- /dev/null +++ b/src/utilities.rs @@ -0,0 +1,38 @@ +use std::collections::HashSet; + +pub fn parse_config_value_set(config_value: String) -> HashSet { + config_value.split(',').map(String::from).collect() +} + +pub fn parse_config_value_string(config_value: HashSet) -> String { + let non_empty: Vec = config_value.into_iter().filter(|s| !s.is_empty()).collect(); + non_empty.join(",") +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn create_authorized_users_hashset() { + let test_string = String::from("123,456,789"); + let users_set = parse_config_value_set(test_string); + + assert!(users_set.contains("123")); + assert!(users_set.contains("456")); + assert!(users_set.contains("789")); + } + + #[test] + fn create_authorized_users_string() { + let mut users_hash_set = HashSet::new(); + users_hash_set.insert("123".to_string()); + users_hash_set.insert("456".to_string()); + users_hash_set.insert("789".to_string()); + + let users_string = parse_config_value_string(users_hash_set); + assert!(users_string.contains("123")); + assert!(users_string.contains("456")); + assert!(users_string.contains("789")); + } +}