Allow ip command

src/cli/allow_ip.rs

@@ -0,0 +1,160 @@
+use crate::cli::client_error_to_shell_error;
+use crate::cli::util::{
+    cluster_from_conn_str, cluster_identifiers_from, find_org_id, find_project_id,
+    get_active_cluster,
+};
+use crate::state::State;
+use log::{debug, info};
+use nu_engine::CallExt;
+use nu_protocol::ast::Call;
+use nu_protocol::engine::{Command, EngineState, Stack};
+use nu_protocol::{Category, PipelineData, ShellError, Signature, SyntaxShape, Value};
+use std::ops::Add;
+use std::sync::{Arc, Mutex};
+use tokio::time::Instant;
+
+#[derive(Clone)]
+pub struct AllowIP {
+    state: Arc<Mutex<State>>,
+}
+
+impl crate::cli::AllowIP {
+    pub fn new(state: Arc<Mutex<State>>) -> Self {
+        Self { state }
+    }
+}
+
+impl Command for crate::cli::AllowIP {
+    fn name(&self) -> &str {
+        "allow ip"
+    }
+
+    fn signature(&self) -> Signature {
+        Signature::build("allow ip")
+            .category(Category::Custom("couchbase".to_string()))
+            .named(
+                "clusters",
+                SyntaxShape::String,
+                "the clusters which should be contacted",
+                None,
+            )
+            .optional(
+                "address",
+                SyntaxShape::String,
+                "ip address to allow access to the cluster",
+            )
+    }
+
+    fn usage(&self) -> &str {
+        "Adds IP address to allowlist on a Capella cluster"
+    }
+
+    fn run(
+        &self,
+        engine_state: &EngineState,
+        stack: &mut Stack,
+        call: &Call,
+        input: PipelineData,
+    ) -> Result<PipelineData, ShellError> {
+        allow_ip(self.state.clone(), engine_state, stack, call, input)
+    }
+}
+
+fn allow_ip(
+    state: Arc<Mutex<State>>,
+    engine_state: &EngineState,
+    stack: &mut Stack,
+    call: &Call,
+    input: PipelineData,
+) -> Result<PipelineData, ShellError> {
+    let span = call.head;
+    let ctrl_c = engine_state.ctrlc.as_ref().unwrap().clone();
+
+    let cluster_identifiers = cluster_identifiers_from(engine_state, stack, &state, call, true)?;
+    let guard = state.lock().unwrap();
+
+    debug!("Running allow_ip");
+
+    let ip_address = match input.into_value(span)? {
+        Value::String { val, .. } => format_ip_address(val),
+        Value::Nothing { .. } => {
+            if let Some(address) = call.opt(engine_state, stack, 0)? {
+                format_ip_address(address)
+            } else {
+                return Err(ShellError::GenericError {
+                    error: "No IP address provided".to_string(),
+                    msg: "".to_string(),
+                    span: None,
+                    help: Some("Provide IP as positional parameter or piped input".into()),
+                    inner: vec![],
+                });
+            }
+        }
+        _ => {
+            return Err(ShellError::GenericError {
+                error: "IP address must be a string".to_string(),
+                msg: "".to_string(),
+                span: None,
+                help: None,
+                inner: vec![],
+            })
+        }
+    };
+
+    for identifier in cluster_identifiers {
+        let cluster = get_active_cluster(identifier.clone(), &guard, span)?;
+
+        let org = if let Some(cluster_org) = cluster.capella_org() {
+            guard.get_capella_org(cluster_org)
+        } else {
+            guard.active_capella_org()
+        }?;
+
+        let client = org.client();
+        let deadline = Instant::now().add(org.timeout());
+
+        let org_id = find_org_id(ctrl_c.clone(), &client, deadline, span)?;
+
+        let project_id = find_project_id(
+            ctrl_c.clone(),
+            guard.active_project().unwrap(),
+            &client,
+            deadline,
+            span,
+            org_id.clone(),
+        )?;
+
+        let json_cluster = cluster_from_conn_str(
+            identifier,
+            ctrl_c.clone(),
+            cluster.hostnames().clone(),
+            &client,
+            deadline,
+            span,
+            org_id.clone(),
+            project_id.clone(),
+        )?;
+
+        client
+            .allow_ip_address(
+                org_id,
+                project_id,
+                json_cluster.id(),
+                ip_address.clone(),
+                deadline,
+                ctrl_c.clone(),
+            )
+            .map_err(|e| client_error_to_shell_error(e, span))?;
+    }
+
+    Ok(PipelineData::empty())
+}
+
+fn format_ip_address(ip_address: String) -> String {
+    if !ip_address.contains('/') {
+        info!("IP address supplied without a subnet mask, defaulting to '/32'");
+        format!("{}/32", ip_address)
+    } else {
+        ip_address
+    }
+}

src/cli/mod.rs

@@ -1,3 +1,4 @@
+mod allow_ip;
 mod analytics;
 mod analytics_buckets;
 mod analytics_datasets;
@@ -83,6 +84,7 @@ mod vector_enrich_text;
 mod vector_search;
 mod version;
 
+pub use allow_ip::AllowIP;
 pub use analytics::Analytics;
 pub use analytics_buckets::AnalyticsBuckets;
 pub use analytics_datasets::AnalyticsDatasets;

src/client/cloud.rs

@@ -540,6 +540,32 @@ impl CapellaClient {
             }),
         }
     }
+
+    pub fn allow_ip_address(
+        &self,
+        org_id: String,
+        project_id: String,
+        cluster_id: String,
+        address: String,
+        deadline: Instant,
+        ctrl_c: Arc<AtomicBool>,
+    ) -> Result<(), ClientError> {
+        let request = CapellaRequest::AllowIPAddress {
+            org_id,
+            project_id,
+            cluster_id,
+            payload: format!("{{\"cidr\": \"{}\"}}", address.clone()),
+        };
+        let response = self.capella_request(request, deadline, ctrl_c)?;
+
+        match response.status() {
+            201 => Ok(()),
+            _ => Err(ClientError::RequestFailed {
+                reason: Some(response.content().into()),
+                key: None,
+            }),
+        }
+    }
 }
 
 #[allow(dead_code)]
@@ -549,6 +575,12 @@ pub enum CapellaRequest {
         org_id: String,
         payload: String,
     },
+    AllowIPAddress {
+        org_id: String,
+        project_id: String,
+        cluster_id: String,
+        payload: String,
+    },
     ProjectDelete {
         org_id: String,
         project_id: String,
@@ -626,6 +658,17 @@ impl CapellaRequest {
             Self::ProjectDelete { org_id, project_id } => {
                 format!("/v4/organizations/{}/projects/{}", org_id, project_id)
             }
+            Self::AllowIPAddress {
+                org_id,
+                project_id,
+                cluster_id,
+                ..
+            } => {
+                format!(
+                    "/v4/organizations/{}/projects/{}/clusters/{}/allowedcidrs",
+                    org_id, project_id, cluster_id
+                )
+            }
             Self::ProjectCreate { org_id, .. } => {
                 format!("/v4/organizations/{}/projects", org_id)
             }
@@ -762,6 +805,7 @@ impl CapellaRequest {
             Self::BucketLoadSample { .. } => HttpVerb::Post,
             Self::BucketList { .. } => HttpVerb::Get,
             Self::BucketUpdate { .. } => HttpVerb::Put,
+            Self::AllowIPAddress { .. } => HttpVerb::Post,
             Self::CredentialsCreate { .. } => HttpVerb::Post,
         }
     }
@@ -773,6 +817,7 @@ impl CapellaRequest {
             Self::BucketCreate { payload, .. } => Some(payload.as_bytes().into()),
             Self::BucketLoadSample { payload, .. } => Some(payload.as_bytes().into()),
             Self::BucketUpdate { payload, .. } => Some(payload.as_bytes().into()),
+            Self::AllowIPAddress { payload, .. } => Some(payload.as_bytes().into()),
             Self::CredentialsCreate { payload, .. } => Some(payload.as_bytes().into()),
             _ => None,
         }

src/main.rs

@@ -704,6 +704,7 @@ fn make_state(
 fn merge_couchbase_delta(context: &mut EngineState, state: Arc<Mutex<State>>) {
     let delta = {
         let mut working_set = StateWorkingSet::new(context);
+        working_set.add_decl(Box::new(AllowIP::new(state.clone())));
         working_set.add_decl(Box::new(Analytics::new(state.clone())));
         working_set.add_decl(Box::new(AnalyticsBuckets::new(state.clone())));
         working_set.add_decl(Box::new(AnalyticsDatasets::new(state.clone())));