Enable activating a set of roles (#1)

Author: demoray

Cargo.toml

@@ -1,5 +1,5 @@
 [package]
-authors = ["Brian Caswell <bcaswell@microsoft.com>"]
+authors = ["Brian Caswell <bcaswell@gmail.com>"]
 name = "azure-pim-cli"
 version = "0.0.1"
 edition = "2021"

README.md

@@ -6,8 +6,9 @@ CLI to list and enable Azure Privileged Identity Management roles
 Usage: az-pim <COMMAND>
 
 Commands:
-  list     List eligible assignments
-  elevate  Elevate to a specific role
+  list          List eligible assignments
+  activate      Activate a specific role
+  activate-set  Activate a set of roles
 
 Options:
   -h, --help
@@ -32,12 +33,12 @@ Options:
           Print version
 
 ```
-## az-pim elevate <ROLE> <SCOPE> <JUSTIFICATION>
+## az-pim activate <ROLE> <SCOPE> <JUSTIFICATION>
 
 ```
-Elevate to a specific role
+Activate a specific role
 
-Usage: elevate [OPTIONS] <ROLE> <SCOPE> <JUSTIFICATION>
+Usage: activate [OPTIONS] <ROLE> <SCOPE> <JUSTIFICATION>
 
 Arguments:
   <ROLE>
@@ -61,4 +62,40 @@ Options:
   -V, --version
           Print version
 
+```
+## az-pim activate-set <JUSTIFICATION>
+
+```
+Activate a set of roles
+
+This command can be used to activate multiple roles at once.  It can be used with a config file or by specifying roles on the command line.
+
+Usage: activate-set [OPTIONS] <JUSTIFICATION>
+
+Arguments:
+  <JUSTIFICATION>
+          Justification for the request
+
+Options:
+      --duration <DURATION>
+          Duration in minutes
+
+          [default: 480]
+
+      --config <CONFIG>
+          Path to a JSON config file containing a set of roles to elevate
+
+          Example config file: ` [ { "scope": "/subscriptions/00000000-0000-0000-0000-000000000000", "role": "Owner" }, { "scope": "/subscriptions/00000000-0000-0000-0000-000000000001", "role": "Owner" } ] `
+
+      --role <SCOPE=NAME>
+          Specify a role to elevate
+
+          Specify multiple times to include multiple key/value pairs
+
+  -h, --help
+          Print help (see a summary with '-h')
+
+  -V, --version
+          Print version
+
 ```

src/activate.rs

@@ -0,0 +1,68 @@
+use crate::az_cli::get_userid;
+use anyhow::{bail, Result};
+use reqwest::{blocking::Client, StatusCode};
+use serde_json::Value;
+use tracing::{debug, info};
+use uuid::Uuid;
+
+// NOTE: serde_json doesn't panic on failed index slicing, it returns a Value
+// that allows further nested nulls
+#[allow(clippy::indexing_slicing)]
+fn check_error_response(body: &Value) -> Result<()> {
+    if body["error"]["code"].as_str() == Some("RoleAssignmentExists") {
+        info!("role already assigned");
+        return Ok(());
+    }
+    if body["error"]["code"].as_str() == Some("RoleAssignmentRequestExists") {
+        info!("role assignment request already exists");
+        return Ok(());
+    }
+    bail!("unable to elevate: {body:#?}");
+}
+
+pub fn activate_role(
+    token: &str,
+    scope: &str,
+    role_definition_id: &str,
+    justification: &str,
+    duration: u32,
+) -> Result<()> {
+    let principal_id = get_userid()?;
+
+    let request_id = Uuid::new_v4();
+    let url = format!("https://management.azure.com{scope}/providers/Microsoft.Authorization/roleAssignmentScheduleRequests/{request_id}");
+    let body = serde_json::json!({
+        "properties": {
+            "principalId": principal_id,
+            "roleDefinitionId": role_definition_id,
+            "requestType": "SelfActivate",
+            "justification": justification,
+            "scheduleInfo": {
+                "expiration": {
+                    "duration": format!("PT{}M", duration),
+                    "type": "AfterDuration",
+                }
+            }
+        }
+    });
+
+    let response = Client::new()
+        .put(url)
+        .query(&[("api-version", "2020-10-01")])
+        .bearer_auth(token)
+        .json(&body)
+        .send()?;
+
+    let status = response.status();
+
+    if status == StatusCode::BAD_REQUEST {
+        return check_error_response(&response.json()?);
+    }
+
+    let body: Value = response.error_for_status()?.json()?;
+
+    debug!("body: {status:#?} - {body:#?}");
+    info!("submitted request: {request_id}");
+
+    Ok(())
+}

src/bin/az-pim.rs

@@ -1,11 +1,12 @@
-use anyhow::{Context, Result};
-use azure_pim_cli::{
-    az_cli::get_token,
-    elevate::{elevate_role, ElevateConfig},
-    roles::list,
-};
+use anyhow::{bail, Context, Result};
+use azure_pim_cli::{activate::activate_role, az_cli::get_token, roles::list_roles};
 use clap::{Command, CommandFactory, Parser, Subcommand};
-use std::{cmp::min, io::stdout};
+use serde::Deserialize;
+use std::{
+    cmp::min, collections::BTreeSet, error::Error, fs::File, io::stdout, path::PathBuf,
+    str::FromStr,
+};
+use tracing::{error, info};
 
 #[derive(Parser)]
 #[command(
@@ -25,14 +26,76 @@ enum SubCommand {
     /// List eligible assignments
     List,
 
-    /// Elevate to a specific role
-    Elevate(ElevateConfig),
+    /// Activate a specific role
+    Activate {
+        /// Name of the role to elevate
+        role: String,
+        /// Scope to elevate
+        scope: String,
+        /// Justification for the request
+        justification: String,
+        /// Duration in minutes
+        #[clap(long, default_value_t = 480)]
+        duration: u32,
+    },
+
+    /// Activate a set of roles
+    ///
+    /// This command can be used to activate multiple roles at once.  It can be
+    /// used with a config file or by specifying roles on the command line.
+    ActivateSet {
+        /// Justification for the request
+        justification: String,
+        #[clap(long, default_value_t = 480)]
+        /// Duration in minutes
+        duration: u32,
+        #[clap(long)]
+        /// Path to a JSON config file containing a set of roles to elevate
+        ///
+        /// Example config file:
+        /// `
+        ///     [
+        ///         {
+        ///             "scope": "/subscriptions/00000000-0000-0000-0000-000000000000",
+        ///             "role": "Owner"
+        ///         },
+        ///         {
+        ///             "scope": "/subscriptions/00000000-0000-0000-0000-000000000001",
+        ///             "role": "Owner"
+        ///         }
+        ///     ]
+        /// `
+        config: Option<PathBuf>,
+        #[clap(long, conflicts_with = "config", value_name = "SCOPE=NAME", value_parser = parse_key_val::<String, String>, action = clap::ArgAction::Append)]
+        /// Specify a role to elevate
+        ///
+        /// Specify multiple times to include multiple key/value pairs
+        role: Option<Vec<(String, String)>>,
+    },
 
     #[command(hide = true)]
     Readme,
 }
 
-fn build_readme(cmd: &mut Command, mut names: Vec<String>) -> String {
+/// Parse a single key-value pair of `X=Y` into a typed tuple of `(X, Y)`.
+///
+/// # Errors
+/// Returns an `Err` if any of the keys or values cannot be parsed or if no `=` is found.
+pub fn parse_key_val<T, U>(s: &str) -> Result<(T, U), Box<dyn Error + Send + Sync + 'static>>
+where
+    T: FromStr,
+    T::Err: Error + Send + Sync + 'static,
+    U: FromStr,
+    U::Err: Error + Send + Sync + 'static,
+{
+    if let Some((key, value)) = s.split_once('=') {
+        Ok((key.parse()?, value.parse()?))
+    } else {
+        Err(format!("invalid KEY=value: no `=` found in `{s}`").into())
+    }
+}
+
+fn build_readme_entry(cmd: &mut Command, mut names: Vec<String>) -> String {
     let mut readme = String::new();
     let base_name = cmd.get_name().to_owned();
 
@@ -60,11 +123,37 @@ fn build_readme(cmd: &mut Command, mut names: Vec<String>) -> String {
         if cmd.get_name() == "readme" {
             continue;
         }
-        readme.push_str(&build_readme(cmd, names.clone()));
+        readme.push_str(&build_readme_entry(cmd, names.clone()));
     }
     readme
 }
 
+fn build_readme() {
+    let mut cmd = Cmd::command();
+    let readme = build_readme_entry(&mut cmd, Vec::new())
+        .replace("azure-pim-cli", "az-pim")
+        .replacen(
+            "# az-pim",
+            &format!("# Azure PIM CLI\n\n{}", env!("CARGO_PKG_DESCRIPTION")),
+            1,
+        )
+        .lines()
+        .map(str::trim_end)
+        .collect::<Vec<_>>()
+        .join("\n")
+        .replace("\n\n\n", "\n");
+    print!("{readme}");
+}
+
+#[derive(Deserialize)]
+struct Role {
+    scope: String,
+    role: String,
+}
+
+#[derive(Deserialize)]
+struct Roles(Vec<Role>);
+
 fn main() -> Result<()> {
     tracing_subscriber::fmt()
         .with_env_filter(
@@ -80,31 +169,93 @@ fn main() -> Result<()> {
     match args.commands {
         SubCommand::List => {
             let token = get_token().context("unable to obtain access token")?;
-            let roles = list(&token).context("unable to list available roles in PIM")?;
+            let roles = list_roles(&token).context("unable to list available roles in PIM")?;
             serde_json::to_writer_pretty(stdout(), &roles)?;
             Ok(())
         }
-        SubCommand::Elevate(config) => {
+        SubCommand::Activate {
+            role,
+            scope,
+            justification,
+            duration,
+        } => {
             let token = get_token().context("unable to obtain access token")?;
-            let roles = list(&token).context("unable to list available roles in PIM")?;
-            elevate_role(&token, &config, &roles).context("unable to elevate to specified role")?;
+            let roles = list_roles(&token).context("unable to list available roles in PIM")?;
+            let entry = roles
+                .iter()
+                .find(|v| v.role == role && v.scope == scope)
+                .context("role not found")?;
+
+            info!("activating {role:?} in {}", entry.scope_name);
+
+            activate_role(
+                &token,
+                &entry.scope,
+                &entry.role_definition_id,
+                &justification,
+                duration,
+            )
+            .context("unable to elevate to specified role")?;
+            Ok(())
+        }
+        SubCommand::ActivateSet {
+            config,
+            role,
+            justification,
+            duration,
+        } => {
+            let mut desired_roles = role
+                .unwrap_or_default()
+                .into_iter()
+                .collect::<BTreeSet<_>>();
+
+            if let Some(path) = config {
+                let handle = File::open(path).context("unable to open activate-set config file")?;
+                let Roles(roles) =
+                    serde_json::from_reader(handle).context("unable to parse config file")?;
+                for entry in roles {
+                    desired_roles.insert((entry.scope, entry.role));
+                }
+            }
+
+            if desired_roles.is_empty() {
+                bail!("no roles specified");
+            }
+
+            let token = get_token().context("unable to obtain access token")?;
+            let available = list_roles(&token).context("unable to list available roles in PIM")?;
+
+            let mut to_add = BTreeSet::new();
+            for (scope, role) in &desired_roles {
+                let entry = &available
+                    .iter()
+                    .find(|v| &v.role == role && &v.scope == scope)
+                    .with_context(|| format!("role not found.  role:{role} scope:{scope}"))?;
+
+                to_add.insert((scope, role, &entry.role_definition_id, &entry.scope_name));
+            }
+
+            let mut success = true;
+            for (scope, role, role_definition_id, scope_name) in to_add {
+                info!("activating {role:?} in {scope_name}");
+                if let Err(error) =
+                    activate_role(&token, scope, role_definition_id, &justification, duration)
+                {
+                    error!(
+                        "scope: {scope} role_definition_id: {role_definition_id} error: {error:?}"
+                    );
+                    success = false;
+                }
+            }
+
+            if !success {
+                bail!("unable to elevate to all roles");
+            }
+
             Ok(())
         }
         SubCommand::Readme => {
-            let mut cmd = Cmd::command();
-            let readme = build_readme(&mut cmd, Vec::new())
-                .replace("azure-pim-cli", "az-pim")
-                .replacen(
-                    "# az-pim",
-                    &format!("# Azure PIM CLI\n\n{}", env!("CARGO_PKG_DESCRIPTION")),
-                    1,
-                )
-                .lines()
-                .map(str::trim_end)
-                .collect::<Vec<_>>()
-                .join("\n")
-                .replace("\n\n\n", "\n");
-            print!("{readme}");
+            build_readme();
             Ok(())
         }
     }