sn_testnet_deploy/ansible/
provisioning.rs

1// Copyright (c) 2023, MaidSafe.
2// All rights reserved.
3//
4// This SAFE Network Software is licensed under the BSD-3-Clause license.
5// Please see the LICENSE file for more details.
6
7use super::{
8    extra_vars::ExtraVarsDocBuilder,
9    inventory::{
10        generate_full_cone_private_node_static_environment_inventory,
11        generate_symmetric_private_node_static_environment_inventory,
12    },
13    AnsibleInventoryType, AnsiblePlaybook, AnsibleRunner,
14};
15use crate::{
16    ansible::inventory::{
17        generate_custom_environment_inventory,
18        generate_full_cone_nat_gateway_static_environment_inventory,
19    },
20    bootstrap::BootstrapOptions,
21    clients::ClientsDeployOptions,
22    deploy::DeployOptions,
23    error::{Error, Result},
24    funding::FundingOptions,
25    inventory::{DeploymentNodeRegistries, VirtualMachine},
26    print_duration, run_external_command, BinaryOption, CloudProvider, EvmNetwork, LogFormat,
27    NodeType, SshClient, UpgradeOptions,
28};
29use ant_service_management::NodeRegistry;
30use evmlib::common::U256;
31use log::{debug, error, trace};
32use semver::Version;
33use serde::{Deserialize, Serialize};
34use std::{
35    collections::HashMap,
36    net::IpAddr,
37    path::PathBuf,
38    time::{Duration, Instant},
39};
40use walkdir::WalkDir;
41
42use crate::ansible::extra_vars;
43
44pub const DEFAULT_BETA_ENCRYPTION_KEY: &str =
45    "49113d2083f57a976076adbe85decb75115820de1e6e74b47e0429338cef124a";
46
47#[derive(Clone, Serialize, Deserialize)]
48pub struct ProvisionOptions {
49    /// The safe version is also in the binary option, but only for an initial deployment.
50    /// For the upscale, it needs to be provided explicitly, because currently it is not
51    /// recorded in the inventory.
52    pub ant_version: Option<String>,
53    pub binary_option: BinaryOption,
54    pub chunk_size: Option<u64>,
55    pub client_env_variables: Option<Vec<(String, String)>>,
56    pub enable_downloaders: bool,
57    pub enable_telegraf: bool,
58    pub evm_data_payments_address: Option<String>,
59    pub evm_network: EvmNetwork,
60    pub evm_payment_token_address: Option<String>,
61    pub evm_rpc_url: Option<String>,
62    pub full_cone_private_node_count: u16,
63    pub funding_wallet_secret_key: Option<String>,
64    pub gas_amount: Option<U256>,
65    pub interval: Option<Duration>,
66    pub log_format: Option<LogFormat>,
67    pub max_archived_log_files: u16,
68    pub max_log_files: u16,
69    pub max_uploads: Option<u32>,
70    pub name: String,
71    pub network_id: Option<u8>,
72    pub node_count: u16,
73    pub node_env_variables: Option<Vec<(String, String)>>,
74    pub output_inventory_dir_path: PathBuf,
75    pub peer_cache_node_count: u16,
76    pub public_rpc: bool,
77    pub rewards_address: Option<String>,
78    pub symmetric_private_node_count: u16,
79    pub token_amount: Option<U256>,
80    pub uploaders_count: Option<u16>,
81    pub wallet_secret_keys: Option<Vec<String>>,
82}
83
84/// These are obtained by running the inventory playbook
85#[derive(Clone, Debug)]
86pub struct PrivateNodeProvisionInventory {
87    pub full_cone_nat_gateway_vms: Vec<VirtualMachine>,
88    pub full_cone_private_node_vms: Vec<VirtualMachine>,
89    pub symmetric_nat_gateway_vms: Vec<VirtualMachine>,
90    pub symmetric_private_node_vms: Vec<VirtualMachine>,
91}
92
93impl PrivateNodeProvisionInventory {
94    pub fn new(
95        provisioner: &AnsibleProvisioner,
96        full_cone_private_node_vm_count: Option<u16>,
97        symmetric_private_node_vm_count: Option<u16>,
98    ) -> Result<Self> {
99        // All the environment types set private_node_vm count to >0 if not specified.
100        let should_provision_full_cone_private_nodes = full_cone_private_node_vm_count
101            .map(|count| count > 0)
102            .unwrap_or(true);
103        let should_provision_symmetric_private_nodes = symmetric_private_node_vm_count
104            .map(|count| count > 0)
105            .unwrap_or(true);
106
107        let mut inventory = Self {
108            full_cone_nat_gateway_vms: Default::default(),
109            full_cone_private_node_vms: Default::default(),
110            symmetric_nat_gateway_vms: Default::default(),
111            symmetric_private_node_vms: Default::default(),
112        };
113
114        if should_provision_full_cone_private_nodes {
115            let full_cone_private_node_vms = provisioner
116                .ansible_runner
117                .get_inventory(AnsibleInventoryType::FullConePrivateNodes, true)
118                .inspect_err(|err| {
119                    println!("Failed to obtain the inventory of Full Cone private node: {err:?}");
120                })?;
121
122            let full_cone_nat_gateway_inventory = provisioner
123                .ansible_runner
124                .get_inventory(AnsibleInventoryType::FullConeNatGateway, true)
125                .inspect_err(|err| {
126                    println!("Failed to get Full Cone NAT Gateway inventory {err:?}");
127                })?;
128
129            if full_cone_nat_gateway_inventory.len() != full_cone_private_node_vms.len() {
130                println!("The number of Full Cone private nodes does not match the number of Full Cone NAT Gateway VMs");
131                return Err(Error::VmCountMismatch(
132                    Some(AnsibleInventoryType::FullConePrivateNodes),
133                    Some(AnsibleInventoryType::FullConeNatGateway),
134                ));
135            }
136
137            inventory.full_cone_private_node_vms = full_cone_private_node_vms;
138            inventory.full_cone_nat_gateway_vms = full_cone_nat_gateway_inventory;
139        }
140
141        if should_provision_symmetric_private_nodes {
142            let symmetric_private_node_vms = provisioner
143                .ansible_runner
144                .get_inventory(AnsibleInventoryType::SymmetricPrivateNodes, true)
145                .inspect_err(|err| {
146                    println!("Failed to obtain the inventory of Symmetric private node: {err:?}");
147                })?;
148
149            let symmetric_nat_gateway_inventory = provisioner
150                .ansible_runner
151                .get_inventory(AnsibleInventoryType::SymmetricNatGateway, true)
152                .inspect_err(|err| {
153                    println!("Failed to get Symmetric NAT Gateway inventory {err:?}");
154                })?;
155
156            if symmetric_nat_gateway_inventory.len() != symmetric_private_node_vms.len() {
157                println!("The number of Symmetric private nodes does not match the number of Symmetric NAT Gateway VMs");
158                return Err(Error::VmCountMismatch(
159                    Some(AnsibleInventoryType::SymmetricPrivateNodes),
160                    Some(AnsibleInventoryType::SymmetricNatGateway),
161                ));
162            }
163
164            inventory.symmetric_private_node_vms = symmetric_private_node_vms;
165            inventory.symmetric_nat_gateway_vms = symmetric_nat_gateway_inventory;
166        }
167
168        Ok(inventory)
169    }
170
171    pub fn should_provision_full_cone_private_nodes(&self) -> bool {
172        !self.full_cone_private_node_vms.is_empty()
173    }
174
175    pub fn should_provision_symmetric_private_nodes(&self) -> bool {
176        !self.symmetric_private_node_vms.is_empty()
177    }
178
179    pub fn symmetric_private_node_and_gateway_map(
180        &self,
181    ) -> Result<HashMap<VirtualMachine, VirtualMachine>> {
182        Self::match_private_node_vm_and_gateway_vm(
183            &self.symmetric_private_node_vms,
184            &self.symmetric_nat_gateway_vms,
185        )
186    }
187
188    pub fn full_cone_private_node_and_gateway_map(
189        &self,
190    ) -> Result<HashMap<VirtualMachine, VirtualMachine>> {
191        Self::match_private_node_vm_and_gateway_vm(
192            &self.full_cone_private_node_vms,
193            &self.full_cone_nat_gateway_vms,
194        )
195    }
196
197    pub fn match_private_node_vm_and_gateway_vm(
198        private_node_vms: &[VirtualMachine],
199        nat_gateway_vms: &[VirtualMachine],
200    ) -> Result<HashMap<VirtualMachine, VirtualMachine>> {
201        if private_node_vms.len() != nat_gateway_vms.len() {
202            println!(
203            "The number of private node VMs ({}) does not match the number of NAT Gateway VMs ({})",
204            private_node_vms.len(),
205            nat_gateway_vms.len()
206        );
207            error!("The number of private node VMs does not match the number of NAT Gateway VMs: Private VMs: {private_node_vms:?} Nat gateway VMs: {nat_gateway_vms:?}");
208            return Err(Error::VmCountMismatch(None, None));
209        }
210
211        let mut map = HashMap::new();
212        for private_vm in private_node_vms {
213            let nat_gateway = nat_gateway_vms
214                .iter()
215                .find(|vm| {
216                    let private_node_name = private_vm.name.split('-').next_back().unwrap();
217                    let nat_gateway_name = vm.name.split('-').next_back().unwrap();
218                    private_node_name == nat_gateway_name
219                })
220                .ok_or_else(|| {
221                    println!(
222                        "Failed to find a matching NAT Gateway for private node: {}",
223                        private_vm.name
224                    );
225                    error!("Failed to find a matching NAT Gateway for private node: {}. Private VMs: {private_node_vms:?} Nat gateway VMs: {nat_gateway_vms:?}", private_vm.name);
226                    Error::VmCountMismatch(None, None)
227                })?;
228
229            let _ = map.insert(private_vm.clone(), nat_gateway.clone());
230        }
231
232        Ok(map)
233    }
234}
235
236impl From<BootstrapOptions> for ProvisionOptions {
237    fn from(bootstrap_options: BootstrapOptions) -> Self {
238        ProvisionOptions {
239            ant_version: None,
240            binary_option: bootstrap_options.binary_option,
241            chunk_size: bootstrap_options.chunk_size,
242            client_env_variables: None,
243            enable_downloaders: false,
244            enable_telegraf: true,
245            evm_data_payments_address: bootstrap_options.evm_data_payments_address,
246            evm_network: bootstrap_options.evm_network,
247            evm_payment_token_address: bootstrap_options.evm_payment_token_address,
248            evm_rpc_url: bootstrap_options.evm_rpc_url,
249            full_cone_private_node_count: bootstrap_options.full_cone_private_node_count,
250            funding_wallet_secret_key: None,
251            gas_amount: None,
252            interval: Some(bootstrap_options.interval),
253            log_format: bootstrap_options.log_format,
254            max_archived_log_files: bootstrap_options.max_archived_log_files,
255            max_log_files: bootstrap_options.max_log_files,
256            max_uploads: None,
257            name: bootstrap_options.name,
258            network_id: Some(bootstrap_options.network_id),
259            node_count: bootstrap_options.node_count,
260            node_env_variables: bootstrap_options.node_env_variables,
261            output_inventory_dir_path: bootstrap_options.output_inventory_dir_path,
262            peer_cache_node_count: 0,
263            public_rpc: false,
264            rewards_address: Some(bootstrap_options.rewards_address),
265            symmetric_private_node_count: bootstrap_options.symmetric_private_node_count,
266            token_amount: None,
267            uploaders_count: None,
268            wallet_secret_keys: None,
269        }
270    }
271}
272
273impl From<DeployOptions> for ProvisionOptions {
274    fn from(deploy_options: DeployOptions) -> Self {
275        ProvisionOptions {
276            ant_version: None,
277            binary_option: deploy_options.binary_option,
278            chunk_size: deploy_options.chunk_size,
279            client_env_variables: deploy_options.client_env_variables,
280            enable_downloaders: deploy_options.enable_downloaders,
281            enable_telegraf: deploy_options.enable_telegraf,
282            node_env_variables: deploy_options.node_env_variables,
283            evm_data_payments_address: deploy_options.evm_data_payments_address,
284            evm_network: deploy_options.evm_network,
285            evm_payment_token_address: deploy_options.evm_payment_token_address,
286            evm_rpc_url: deploy_options.evm_rpc_url,
287            full_cone_private_node_count: deploy_options.full_cone_private_node_count,
288            funding_wallet_secret_key: deploy_options.funding_wallet_secret_key,
289            gas_amount: deploy_options.initial_gas,
290            interval: Some(deploy_options.interval),
291            log_format: deploy_options.log_format,
292            max_archived_log_files: deploy_options.max_archived_log_files,
293            max_log_files: deploy_options.max_log_files,
294            max_uploads: None,
295            name: deploy_options.name,
296            network_id: deploy_options.network_id,
297            node_count: deploy_options.node_count,
298            output_inventory_dir_path: deploy_options.output_inventory_dir_path,
299            peer_cache_node_count: deploy_options.peer_cache_node_count,
300            public_rpc: deploy_options.public_rpc,
301            rewards_address: Some(deploy_options.rewards_address),
302            symmetric_private_node_count: deploy_options.symmetric_private_node_count,
303            token_amount: deploy_options.initial_tokens,
304            uploaders_count: Some(deploy_options.uploaders_count),
305            wallet_secret_keys: None,
306        }
307    }
308}
309
310impl From<ClientsDeployOptions> for ProvisionOptions {
311    fn from(client_options: ClientsDeployOptions) -> Self {
312        ProvisionOptions {
313            ant_version: None,
314            binary_option: client_options.binary_option,
315            chunk_size: client_options.chunk_size,
316            client_env_variables: client_options.client_env_variables,
317            enable_downloaders: client_options.enable_downloaders,
318            enable_telegraf: client_options.enable_telegraf,
319            evm_data_payments_address: client_options.evm_details.data_payments_address,
320            evm_network: client_options.evm_details.network,
321            evm_payment_token_address: client_options.evm_details.payment_token_address,
322            evm_rpc_url: client_options.evm_details.rpc_url,
323            full_cone_private_node_count: 0,
324            funding_wallet_secret_key: client_options.funding_wallet_secret_key,
325            gas_amount: client_options.initial_gas,
326            interval: None,
327            log_format: None,
328            max_archived_log_files: client_options.max_archived_log_files,
329            max_log_files: client_options.max_log_files,
330            max_uploads: client_options.max_uploads,
331            name: client_options.name,
332            network_id: client_options.network_id,
333            node_count: 0,
334            node_env_variables: None,
335            output_inventory_dir_path: client_options.output_inventory_dir_path,
336            peer_cache_node_count: 0,
337            public_rpc: false,
338            rewards_address: None,
339            symmetric_private_node_count: 0,
340            token_amount: client_options.initial_tokens,
341            uploaders_count: Some(client_options.uploaders_count),
342            wallet_secret_keys: client_options.wallet_secret_keys,
343        }
344    }
345}
346
347#[derive(Clone)]
348pub struct AnsibleProvisioner {
349    pub ansible_runner: AnsibleRunner,
350    pub cloud_provider: CloudProvider,
351    pub ssh_client: SshClient,
352}
353
354impl AnsibleProvisioner {
355    pub fn new(
356        ansible_runner: AnsibleRunner,
357        cloud_provider: CloudProvider,
358        ssh_client: SshClient,
359    ) -> Self {
360        Self {
361            ansible_runner,
362            cloud_provider,
363            ssh_client,
364        }
365    }
366
367    pub fn build_safe_network_binaries(
368        &self,
369        options: &ProvisionOptions,
370        binaries_to_build: Option<Vec<String>>,
371    ) -> Result<()> {
372        let start = Instant::now();
373        println!("Obtaining IP address for build VM...");
374        let build_inventory = self
375            .ansible_runner
376            .get_inventory(AnsibleInventoryType::Build, true)?;
377        let build_ip = build_inventory[0].public_ip_addr;
378        self.ssh_client
379            .wait_for_ssh_availability(&build_ip, &self.cloud_provider.get_ssh_user())?;
380
381        println!("Running ansible against build VM...");
382        let base_extra_vars = extra_vars::build_binaries_extra_vars_doc(options)?;
383
384        let extra_vars = if let Some(binaries) = binaries_to_build {
385            let mut build_ant = false;
386            let mut build_antnode = false;
387            let mut build_antctl = false;
388            let mut build_antctld = false;
389
390            for binary in &binaries {
391                match binary.as_str() {
392                    "ant" => build_ant = true,
393                    "antnode" => build_antnode = true,
394                    "antctl" => build_antctl = true,
395                    "antctld" => build_antctld = true,
396                    _ => return Err(Error::InvalidBinaryName(binary.clone())),
397                }
398            }
399
400            let mut json_value: serde_json::Value = serde_json::from_str(&base_extra_vars)?;
401            if let serde_json::Value::Object(ref mut map) = json_value {
402                map.insert("build_ant".to_string(), serde_json::Value::Bool(build_ant));
403                map.insert(
404                    "build_antnode".to_string(),
405                    serde_json::Value::Bool(build_antnode),
406                );
407                map.insert(
408                    "build_antctl".to_string(),
409                    serde_json::Value::Bool(build_antctl),
410                );
411                map.insert(
412                    "build_antctld".to_string(),
413                    serde_json::Value::Bool(build_antctld),
414                );
415            }
416            json_value.to_string()
417        } else {
418            base_extra_vars
419        };
420
421        self.ansible_runner.run_playbook(
422            AnsiblePlaybook::Build,
423            AnsibleInventoryType::Build,
424            Some(extra_vars),
425        )?;
426        print_duration(start.elapsed());
427        Ok(())
428    }
429
430    pub fn cleanup_node_logs(&self, setup_cron: bool) -> Result<()> {
431        for node_inv_type in AnsibleInventoryType::iter_node_type() {
432            self.ansible_runner.run_playbook(
433                AnsiblePlaybook::CleanupLogs,
434                node_inv_type,
435                Some(format!("{{ \"setup_cron\": \"{setup_cron}\" }}")),
436            )?;
437        }
438
439        Ok(())
440    }
441
442    pub fn copy_logs(&self, name: &str, resources_only: bool) -> Result<()> {
443        for node_inv_type in AnsibleInventoryType::iter_node_type() {
444            self.ansible_runner.run_playbook(
445                AnsiblePlaybook::CopyLogs,
446                node_inv_type,
447                Some(format!(
448                    "{{ \"env_name\": \"{name}\", \"resources_only\" : \"{resources_only}\" }}"
449                )),
450            )?;
451        }
452        Ok(())
453    }
454
455    pub fn get_all_node_inventory(&self) -> Result<Vec<VirtualMachine>> {
456        let mut all_node_inventory = Vec::new();
457        for node_inv_type in AnsibleInventoryType::iter_node_type() {
458            all_node_inventory.extend(self.ansible_runner.get_inventory(node_inv_type, false)?);
459        }
460
461        Ok(all_node_inventory)
462    }
463
464    pub fn get_symmetric_nat_gateway_inventory(&self) -> Result<Vec<VirtualMachine>> {
465        self.ansible_runner
466            .get_inventory(AnsibleInventoryType::SymmetricNatGateway, false)
467    }
468
469    pub fn get_full_cone_nat_gateway_inventory(&self) -> Result<Vec<VirtualMachine>> {
470        self.ansible_runner
471            .get_inventory(AnsibleInventoryType::FullConeNatGateway, false)
472    }
473
474    pub fn get_client_inventory(&self) -> Result<Vec<VirtualMachine>> {
475        self.ansible_runner
476            .get_inventory(AnsibleInventoryType::Clients, false)
477    }
478
479    pub fn get_node_registries(
480        &self,
481        inventory_type: &AnsibleInventoryType,
482    ) -> Result<DeploymentNodeRegistries> {
483        debug!("Fetching node manager inventory for {inventory_type:?}");
484        let temp_dir_path = tempfile::tempdir()?.into_path();
485        let temp_dir_json = serde_json::to_string(&temp_dir_path)?;
486
487        self.ansible_runner.run_playbook(
488            AnsiblePlaybook::AntCtlInventory,
489            *inventory_type,
490            Some(format!("{{ \"dest\": {temp_dir_json} }}")),
491        )?;
492
493        let node_registry_paths = WalkDir::new(temp_dir_path)
494            .into_iter()
495            .flatten()
496            .filter_map(|entry| {
497                if entry.file_type().is_file()
498                    && entry.path().extension().is_some_and(|ext| ext == "json")
499                {
500                    // tempdir/<testnet_name>-node/var/safenode-manager/node_registry.json
501                    let mut vm_name = entry.path().to_path_buf();
502                    trace!("Found file with json extension: {vm_name:?}");
503                    vm_name.pop();
504                    vm_name.pop();
505                    vm_name.pop();
506                    // Extract the <testnet_name>-node string
507                    trace!("Extracting the vm name from the path");
508                    let vm_name = vm_name.file_name()?.to_str()?;
509                    trace!("Extracted vm name from path: {vm_name}");
510                    Some((vm_name.to_string(), entry.path().to_path_buf()))
511                } else {
512                    None
513                }
514            })
515            .collect::<Vec<(String, PathBuf)>>();
516
517        let mut node_registries = Vec::new();
518        let mut failed_vms = Vec::new();
519        for (vm_name, file_path) in node_registry_paths {
520            match NodeRegistry::load(&file_path) {
521                Ok(node_registry) => node_registries.push((vm_name.clone(), node_registry)),
522                Err(_) => failed_vms.push(vm_name.clone()),
523            }
524        }
525
526        let deployment_registries = DeploymentNodeRegistries {
527            inventory_type: *inventory_type,
528            retrieved_registries: node_registries,
529            failed_vms,
530        };
531        Ok(deployment_registries)
532    }
533
534    pub fn provision_evm_nodes(&self, options: &ProvisionOptions) -> Result<()> {
535        let start = Instant::now();
536        println!("Obtaining IP address for EVM nodes...");
537        let evm_node_inventory = self
538            .ansible_runner
539            .get_inventory(AnsibleInventoryType::EvmNodes, true)?;
540        let evm_node_ip = evm_node_inventory[0].public_ip_addr;
541        self.ssh_client
542            .wait_for_ssh_availability(&evm_node_ip, &self.cloud_provider.get_ssh_user())?;
543
544        println!("Running ansible against EVM nodes...");
545        self.ansible_runner.run_playbook(
546            AnsiblePlaybook::EvmNodes,
547            AnsibleInventoryType::EvmNodes,
548            Some(extra_vars::build_evm_nodes_extra_vars_doc(
549                &options.name,
550                &self.cloud_provider,
551                &options.binary_option,
552            )),
553        )?;
554        print_duration(start.elapsed());
555        Ok(())
556    }
557
558    pub fn provision_genesis_node(&self, options: &ProvisionOptions) -> Result<()> {
559        let start = Instant::now();
560        let genesis_inventory = self
561            .ansible_runner
562            .get_inventory(AnsibleInventoryType::Genesis, true)?;
563        let genesis_ip = genesis_inventory[0].public_ip_addr;
564        self.ssh_client
565            .wait_for_ssh_availability(&genesis_ip, &self.cloud_provider.get_ssh_user())?;
566        self.ansible_runner.run_playbook(
567            AnsiblePlaybook::Genesis,
568            AnsibleInventoryType::Genesis,
569            Some(extra_vars::build_node_extra_vars_doc(
570                &self.cloud_provider.to_string(),
571                options,
572                NodeType::Genesis,
573                None,
574                None,
575                1,
576                options.evm_network.clone(),
577                false,
578            )?),
579        )?;
580
581        print_duration(start.elapsed());
582
583        Ok(())
584    }
585
586    pub fn provision_full_cone(
587        &self,
588        options: &ProvisionOptions,
589        initial_contact_peer: Option<String>,
590        initial_network_contacts_url: Option<String>,
591        private_node_inventory: PrivateNodeProvisionInventory,
592        new_full_cone_nat_gateway_new_vms_for_upscale: Option<Vec<VirtualMachine>>,
593    ) -> Result<()> {
594        // Step 1 of Full Cone NAT Gateway
595        let start = Instant::now();
596        self.print_ansible_run_banner("Provision Full Cone NAT Gateway - Step 1");
597
598        for vm in new_full_cone_nat_gateway_new_vms_for_upscale
599            .as_ref()
600            .unwrap_or(&private_node_inventory.full_cone_nat_gateway_vms)
601            .iter()
602        {
603            println!(
604                "Checking SSH availability for Full Cone NAT Gateway: {}",
605                vm.public_ip_addr
606            );
607            self.ssh_client
608                .wait_for_ssh_availability(&vm.public_ip_addr, &self.cloud_provider.get_ssh_user())
609                .map_err(|e| {
610                    println!(
611                        "Failed to establish SSH connection to Full Cone NAT Gateway: {}",
612                        e
613                    );
614                    e
615                })?;
616        }
617
618        let mut modified_private_node_inventory = private_node_inventory.clone();
619
620        // If we are upscaling, then we cannot access the gateway VMs which are already deployed.
621        if let Some(new_full_cone_nat_gateway_new_vms_for_upscale) =
622            &new_full_cone_nat_gateway_new_vms_for_upscale
623        {
624            debug!("Removing existing full cone NAT Gateway and private node VMs from the inventory. Old inventory: {modified_private_node_inventory:?}");
625            let mut names_to_keep = Vec::new();
626
627            for vm in new_full_cone_nat_gateway_new_vms_for_upscale.iter() {
628                let nat_gateway_name = vm.name.split('-').next_back().unwrap();
629                names_to_keep.push(nat_gateway_name);
630            }
631
632            modified_private_node_inventory
633                .full_cone_nat_gateway_vms
634                .retain(|vm| {
635                    let nat_gateway_name = vm.name.split('-').next_back().unwrap();
636                    names_to_keep.contains(&nat_gateway_name)
637                });
638            modified_private_node_inventory
639                .full_cone_private_node_vms
640                .retain(|vm| {
641                    let nat_gateway_name = vm.name.split('-').next_back().unwrap();
642                    names_to_keep.contains(&nat_gateway_name)
643                });
644            debug!("New inventory after removing existing full cone NAT Gateway and private node VMs: {modified_private_node_inventory:?}");
645        }
646
647        if modified_private_node_inventory
648            .full_cone_nat_gateway_vms
649            .is_empty()
650        {
651            error!("There are no full cone NAT Gateway VMs available to upscale");
652            return Ok(());
653        }
654
655        let private_node_ip_map = modified_private_node_inventory
656            .full_cone_private_node_and_gateway_map()?
657            .into_iter()
658            .map(|(k, v)| {
659                let gateway_name = if new_full_cone_nat_gateway_new_vms_for_upscale.is_some() {
660                    debug!("Upscaling, using public IP address for gateway name");
661                    v.public_ip_addr.to_string()
662                } else {
663                    v.name.clone()
664                };
665                (gateway_name, k.private_ip_addr)
666            })
667            .collect::<HashMap<String, IpAddr>>();
668
669        if private_node_ip_map.is_empty() {
670            println!("There are no full cone private node VM available to be routed through the full cone NAT Gateway");
671            return Err(Error::EmptyInventory(
672                AnsibleInventoryType::FullConePrivateNodes,
673            ));
674        }
675
676        let vars = extra_vars::build_nat_gateway_extra_vars_doc(
677            &options.name,
678            private_node_ip_map.clone(),
679            "step1",
680        );
681        debug!("Provisioning Full Cone NAT Gateway - Step 1 with vars: {vars}");
682        let gateway_inventory = if new_full_cone_nat_gateway_new_vms_for_upscale.is_some() {
683            debug!("Upscaling, using static inventory for full cone nat gateway.");
684            generate_full_cone_nat_gateway_static_environment_inventory(
685                &modified_private_node_inventory.full_cone_nat_gateway_vms,
686                &options.name,
687                &options.output_inventory_dir_path,
688            )?;
689
690            AnsibleInventoryType::FullConeNatGatewayStatic
691        } else {
692            AnsibleInventoryType::FullConeNatGateway
693        };
694        self.ansible_runner.run_playbook(
695            AnsiblePlaybook::FullConeNatGateway,
696            gateway_inventory,
697            Some(vars),
698        )?;
699
700        // setup private node config
701        self.print_ansible_run_banner("Provisioning Full Cone Private Node Config");
702
703        generate_full_cone_private_node_static_environment_inventory(
704            &options.name,
705            &options.output_inventory_dir_path,
706            &private_node_inventory.full_cone_private_node_vms,
707            &private_node_inventory.full_cone_nat_gateway_vms,
708            &self.ssh_client.private_key_path,
709        )
710        .inspect_err(|err| {
711            error!("Failed to generate full cone private node static inv with err: {err:?}")
712        })?;
713
714        // For a new deployment, it's quite probable that SSH is available, because this part occurs
715        // after the genesis node has been provisioned. However, for a bootstrap deploy, we need to
716        // check that SSH is available before proceeding.
717        println!("Obtaining IP addresses for nodes...");
718        let inventory = self
719            .ansible_runner
720            .get_inventory(AnsibleInventoryType::FullConePrivateNodes, true)?;
721
722        println!("Waiting for SSH availability on Symmetric Private nodes...");
723        for vm in inventory.iter() {
724            println!(
725                "Checking SSH availability for {}: {}",
726                vm.name, vm.public_ip_addr
727            );
728            self.ssh_client
729                .wait_for_ssh_availability(&vm.public_ip_addr, &self.cloud_provider.get_ssh_user())
730                .map_err(|e| {
731                    println!("Failed to establish SSH connection to {}: {}", vm.name, e);
732                    e
733                })?;
734        }
735
736        println!("SSH is available on all nodes. Proceeding with provisioning...");
737
738        self.ansible_runner.run_playbook(
739            AnsiblePlaybook::PrivateNodeConfig,
740            AnsibleInventoryType::FullConePrivateNodes,
741            Some(
742                extra_vars::build_full_cone_private_node_config_extra_vars_docs(
743                    &private_node_inventory,
744                )?,
745            ),
746        )?;
747
748        // Step 2 of Full Cone NAT Gateway
749
750        let vars = extra_vars::build_nat_gateway_extra_vars_doc(
751            &options.name,
752            private_node_ip_map,
753            "step2",
754        );
755
756        self.print_ansible_run_banner("Provisioning Full Cone NAT Gateway - Step 2");
757        debug!("Provisioning Full Cone NAT Gateway - Step 2 with vars: {vars}");
758        self.ansible_runner.run_playbook(
759            AnsiblePlaybook::FullConeNatGateway,
760            gateway_inventory,
761            Some(vars),
762        )?;
763
764        // provision the nodes
765
766        let home_dir = std::env::var("HOME").inspect_err(|err| {
767            println!("Failed to get home directory with error: {err:?}",);
768        })?;
769        let known_hosts_path = format!("{}/.ssh/known_hosts", home_dir);
770        debug!("Cleaning up known hosts file at {known_hosts_path} ");
771        run_external_command(
772            PathBuf::from("rm"),
773            std::env::current_dir()?,
774            vec![known_hosts_path],
775            false,
776            false,
777        )?;
778
779        self.print_ansible_run_banner("Provision Full Cone Private Nodes");
780
781        self.ssh_client.set_full_cone_nat_routed_vms(
782            &private_node_inventory.full_cone_private_node_vms,
783            &private_node_inventory.full_cone_nat_gateway_vms,
784        )?;
785
786        self.provision_nodes(
787            options,
788            initial_contact_peer,
789            initial_network_contacts_url,
790            NodeType::FullConePrivateNode,
791        )?;
792
793        print_duration(start.elapsed());
794        Ok(())
795    }
796    pub fn provision_symmetric_nat_gateway(
797        &self,
798        options: &ProvisionOptions,
799        private_node_inventory: &PrivateNodeProvisionInventory,
800    ) -> Result<()> {
801        let start = Instant::now();
802        for vm in &private_node_inventory.symmetric_nat_gateway_vms {
803            println!(
804                "Checking SSH availability for Symmetric NAT Gateway: {}",
805                vm.public_ip_addr
806            );
807            self.ssh_client
808                .wait_for_ssh_availability(&vm.public_ip_addr, &self.cloud_provider.get_ssh_user())
809                .map_err(|e| {
810                    println!(
811                        "Failed to establish SSH connection to Symmetric NAT Gateway: {}",
812                        e
813                    );
814                    e
815                })?;
816        }
817
818        let private_node_ip_map = private_node_inventory
819            .symmetric_private_node_and_gateway_map()?
820            .into_iter()
821            .map(|(k, v)| (v.name.clone(), k.private_ip_addr))
822            .collect::<HashMap<String, IpAddr>>();
823
824        if private_node_ip_map.is_empty() {
825            println!("There are no Symmetric private node VM available to be routed through the Symmetric NAT Gateway");
826            return Err(Error::EmptyInventory(
827                AnsibleInventoryType::SymmetricPrivateNodes,
828            ));
829        }
830
831        let vars = extra_vars::build_nat_gateway_extra_vars_doc(
832            &options.name,
833            private_node_ip_map,
834            "symmetric",
835        );
836        debug!("Provisioning Symmetric NAT Gateway with vars: {vars}");
837        self.ansible_runner.run_playbook(
838            AnsiblePlaybook::SymmetricNatGateway,
839            AnsibleInventoryType::SymmetricNatGateway,
840            Some(vars),
841        )?;
842
843        print_duration(start.elapsed());
844        Ok(())
845    }
846
847    pub fn provision_nodes(
848        &self,
849        options: &ProvisionOptions,
850        initial_contact_peer: Option<String>,
851        initial_network_contacts_url: Option<String>,
852        node_type: NodeType,
853    ) -> Result<()> {
854        let start = Instant::now();
855        let mut relay = false;
856        let (inventory_type, node_count) = match &node_type {
857            NodeType::FullConePrivateNode => {
858                relay = true;
859                (
860                    node_type.to_ansible_inventory_type(),
861                    options.full_cone_private_node_count,
862                )
863            }
864            // use provision_genesis_node fn
865            NodeType::Generic => (node_type.to_ansible_inventory_type(), options.node_count),
866            NodeType::Genesis => return Err(Error::InvalidNodeType(node_type)),
867            NodeType::PeerCache => (
868                node_type.to_ansible_inventory_type(),
869                options.peer_cache_node_count,
870            ),
871            NodeType::SymmetricPrivateNode => {
872                relay = true;
873                (
874                    node_type.to_ansible_inventory_type(),
875                    options.symmetric_private_node_count,
876                )
877            }
878        };
879
880        // For a new deployment, it's quite probable that SSH is available, because this part occurs
881        // after the genesis node has been provisioned. However, for a bootstrap deploy, we need to
882        // check that SSH is available before proceeding.
883        println!("Obtaining IP addresses for nodes...");
884        let inventory = self.ansible_runner.get_inventory(inventory_type, true)?;
885
886        println!("Waiting for SSH availability on {node_type:?} nodes...");
887        for vm in inventory.iter() {
888            println!(
889                "Checking SSH availability for {}: {}",
890                vm.name, vm.public_ip_addr
891            );
892            self.ssh_client
893                .wait_for_ssh_availability(&vm.public_ip_addr, &self.cloud_provider.get_ssh_user())
894                .map_err(|e| {
895                    println!("Failed to establish SSH connection to {}: {}", vm.name, e);
896                    e
897                })?;
898        }
899
900        println!("SSH is available on all nodes. Proceeding with provisioning...");
901
902        let playbook = match node_type {
903            NodeType::Generic => AnsiblePlaybook::Nodes,
904            NodeType::PeerCache => AnsiblePlaybook::PeerCacheNodes,
905            NodeType::FullConePrivateNode => AnsiblePlaybook::Nodes,
906            NodeType::SymmetricPrivateNode => AnsiblePlaybook::Nodes,
907            _ => return Err(Error::InvalidNodeType(node_type.clone())),
908        };
909        self.ansible_runner.run_playbook(
910            playbook,
911            inventory_type,
912            Some(extra_vars::build_node_extra_vars_doc(
913                &self.cloud_provider.to_string(),
914                options,
915                node_type.clone(),
916                initial_contact_peer,
917                initial_network_contacts_url,
918                node_count,
919                options.evm_network.clone(),
920                relay,
921            )?),
922        )?;
923
924        print_duration(start.elapsed());
925        Ok(())
926    }
927
928    pub fn provision_symmetric_private_nodes(
929        &self,
930        options: &mut ProvisionOptions,
931        initial_contact_peer: Option<String>,
932        initial_network_contacts_url: Option<String>,
933        private_node_inventory: &PrivateNodeProvisionInventory,
934    ) -> Result<()> {
935        let start = Instant::now();
936        self.print_ansible_run_banner("Provision Symmetric Private Node Config");
937
938        generate_symmetric_private_node_static_environment_inventory(
939            &options.name,
940            &options.output_inventory_dir_path,
941            &private_node_inventory.symmetric_private_node_vms,
942            &private_node_inventory.symmetric_nat_gateway_vms,
943            &self.ssh_client.private_key_path,
944        )
945        .inspect_err(|err| {
946            error!("Failed to generate symmetric private node static inv with err: {err:?}")
947        })?;
948
949        self.ssh_client.set_symmetric_nat_routed_vms(
950            &private_node_inventory.symmetric_private_node_vms,
951            &private_node_inventory.symmetric_nat_gateway_vms,
952        )?;
953
954        let inventory_type = AnsibleInventoryType::SymmetricPrivateNodes;
955
956        // For a new deployment, it's quite probable that SSH is available, because this part occurs
957        // after the genesis node has been provisioned. However, for a bootstrap deploy, we need to
958        // check that SSH is available before proceeding.
959        println!("Obtaining IP addresses for nodes...");
960        let inventory = self.ansible_runner.get_inventory(inventory_type, true)?;
961
962        println!("Waiting for SSH availability on Symmetric Private nodes...");
963        for vm in inventory.iter() {
964            println!(
965                "Checking SSH availability for {}: {}",
966                vm.name, vm.public_ip_addr
967            );
968            self.ssh_client
969                .wait_for_ssh_availability(&vm.public_ip_addr, &self.cloud_provider.get_ssh_user())
970                .map_err(|e| {
971                    println!("Failed to establish SSH connection to {}: {}", vm.name, e);
972                    e
973                })?;
974        }
975
976        println!("SSH is available on all nodes. Proceeding with provisioning...");
977
978        self.ansible_runner.run_playbook(
979            AnsiblePlaybook::PrivateNodeConfig,
980            inventory_type,
981            Some(
982                extra_vars::build_symmetric_private_node_config_extra_vars_doc(
983                    private_node_inventory,
984                )?,
985            ),
986        )?;
987
988        println!("Provisioned Symmetric Private Node Config");
989        print_duration(start.elapsed());
990
991        self.provision_nodes(
992            options,
993            initial_contact_peer,
994            initial_network_contacts_url,
995            NodeType::SymmetricPrivateNode,
996        )?;
997
998        Ok(())
999    }
1000
1001    pub async fn provision_downloaders(
1002        &self,
1003        options: &ProvisionOptions,
1004        genesis_multiaddr: Option<String>,
1005        genesis_network_contacts_url: Option<String>,
1006    ) -> Result<()> {
1007        let start = Instant::now();
1008
1009        println!("Running ansible against Client machine to start the downloader script.");
1010        debug!("Running ansible against Client machine to start the downloader script.");
1011
1012        self.ansible_runner.run_playbook(
1013            AnsiblePlaybook::Downloaders,
1014            AnsibleInventoryType::Clients,
1015            Some(extra_vars::build_downloaders_extra_vars_doc(
1016                &self.cloud_provider.to_string(),
1017                options,
1018                genesis_multiaddr,
1019                genesis_network_contacts_url,
1020            )?),
1021        )?;
1022        print_duration(start.elapsed());
1023        Ok(())
1024    }
1025
1026    pub async fn provision_clients(
1027        &self,
1028        options: &ProvisionOptions,
1029        genesis_multiaddr: Option<String>,
1030        genesis_network_contacts_url: Option<String>,
1031    ) -> Result<()> {
1032        let start = Instant::now();
1033
1034        let sk_map = if let Some(wallet_keys) = &options.wallet_secret_keys {
1035            self.prepare_pre_funded_wallets(wallet_keys).await?
1036        } else {
1037            self.deposit_funds_to_clients(&FundingOptions {
1038                evm_data_payments_address: options.evm_data_payments_address.clone(),
1039                evm_network: options.evm_network.clone(),
1040                evm_payment_token_address: options.evm_payment_token_address.clone(),
1041                evm_rpc_url: options.evm_rpc_url.clone(),
1042                funding_wallet_secret_key: options.funding_wallet_secret_key.clone(),
1043                gas_amount: options.gas_amount,
1044                token_amount: options.token_amount,
1045                uploaders_count: options.uploaders_count,
1046            })
1047            .await?
1048        };
1049
1050        self.ansible_runner.run_playbook(
1051            AnsiblePlaybook::Uploaders,
1052            AnsibleInventoryType::Clients,
1053            Some(extra_vars::build_clients_extra_vars_doc(
1054                &self.cloud_provider.to_string(),
1055                options,
1056                genesis_multiaddr,
1057                genesis_network_contacts_url,
1058                &sk_map,
1059            )?),
1060        )?;
1061        print_duration(start.elapsed());
1062        Ok(())
1063    }
1064
1065    pub fn start_nodes(
1066        &self,
1067        environment_name: &str,
1068        interval: Duration,
1069        node_type: Option<NodeType>,
1070        custom_inventory: Option<Vec<VirtualMachine>>,
1071    ) -> Result<()> {
1072        let mut extra_vars = ExtraVarsDocBuilder::default();
1073        extra_vars.add_variable("interval", &interval.as_millis().to_string());
1074
1075        if let Some(node_type) = node_type {
1076            println!("Running the start nodes playbook for {node_type:?} nodes");
1077            self.ansible_runner.run_playbook(
1078                AnsiblePlaybook::StartNodes,
1079                node_type.to_ansible_inventory_type(),
1080                Some(extra_vars.build()),
1081            )?;
1082            return Ok(());
1083        }
1084
1085        if let Some(custom_inventory) = custom_inventory {
1086            println!("Running the start nodes playbook with a custom inventory");
1087            generate_custom_environment_inventory(
1088                &custom_inventory,
1089                environment_name,
1090                &self.ansible_runner.working_directory_path.join("inventory"),
1091            )?;
1092            self.ansible_runner.run_playbook(
1093                AnsiblePlaybook::StartNodes,
1094                AnsibleInventoryType::Custom,
1095                Some(extra_vars.build()),
1096            )?;
1097            return Ok(());
1098        }
1099
1100        println!("Running the start nodes playbook for all node types");
1101        for node_inv_type in AnsibleInventoryType::iter_node_type() {
1102            self.ansible_runner.run_playbook(
1103                AnsiblePlaybook::StartNodes,
1104                node_inv_type,
1105                Some(extra_vars.build()),
1106            )?;
1107        }
1108        Ok(())
1109    }
1110
1111    pub fn status(&self) -> Result<()> {
1112        for node_inv_type in AnsibleInventoryType::iter_node_type() {
1113            self.ansible_runner
1114                .run_playbook(AnsiblePlaybook::Status, node_inv_type, None)?;
1115        }
1116        Ok(())
1117    }
1118
1119    pub fn start_telegraf(
1120        &self,
1121        environment_name: &str,
1122        node_type: Option<NodeType>,
1123        custom_inventory: Option<Vec<VirtualMachine>>,
1124    ) -> Result<()> {
1125        if let Some(node_type) = node_type {
1126            println!("Running the start telegraf playbook for {node_type:?} nodes");
1127            self.ansible_runner.run_playbook(
1128                AnsiblePlaybook::StartTelegraf,
1129                node_type.to_ansible_inventory_type(),
1130                None,
1131            )?;
1132            return Ok(());
1133        }
1134
1135        if let Some(custom_inventory) = custom_inventory {
1136            println!("Running the start telegraf playbook with a custom inventory");
1137            generate_custom_environment_inventory(
1138                &custom_inventory,
1139                environment_name,
1140                &self.ansible_runner.working_directory_path.join("inventory"),
1141            )?;
1142            self.ansible_runner.run_playbook(
1143                AnsiblePlaybook::StartTelegraf,
1144                AnsibleInventoryType::Custom,
1145                None,
1146            )?;
1147            return Ok(());
1148        }
1149
1150        println!("Running the start telegraf playbook for all node types");
1151        for node_inv_type in AnsibleInventoryType::iter_node_type() {
1152            self.ansible_runner.run_playbook(
1153                AnsiblePlaybook::StartTelegraf,
1154                node_inv_type,
1155                None,
1156            )?;
1157        }
1158
1159        Ok(())
1160    }
1161
1162    pub fn stop_nodes(
1163        &self,
1164        environment_name: &str,
1165        interval: Duration,
1166        node_type: Option<NodeType>,
1167        custom_inventory: Option<Vec<VirtualMachine>>,
1168        delay: Option<u64>,
1169        service_names: Option<Vec<String>>,
1170    ) -> Result<()> {
1171        let mut extra_vars = ExtraVarsDocBuilder::default();
1172        extra_vars.add_variable("interval", &interval.as_millis().to_string());
1173        if let Some(delay) = delay {
1174            extra_vars.add_variable("delay", &delay.to_string());
1175        }
1176        if let Some(service_names) = service_names {
1177            extra_vars.add_list_variable("service_names", service_names);
1178        }
1179        let extra_vars = extra_vars.build();
1180
1181        if let Some(node_type) = node_type {
1182            println!("Running the stop nodes playbook for {node_type:?} nodes");
1183            self.ansible_runner.run_playbook(
1184                AnsiblePlaybook::StopNodes,
1185                node_type.to_ansible_inventory_type(),
1186                Some(extra_vars),
1187            )?;
1188            return Ok(());
1189        }
1190
1191        if let Some(custom_inventory) = custom_inventory {
1192            println!("Running the stop nodes playbook with a custom inventory");
1193            generate_custom_environment_inventory(
1194                &custom_inventory,
1195                environment_name,
1196                &self.ansible_runner.working_directory_path.join("inventory"),
1197            )?;
1198            self.ansible_runner.run_playbook(
1199                AnsiblePlaybook::StopNodes,
1200                AnsibleInventoryType::Custom,
1201                Some(extra_vars),
1202            )?;
1203            return Ok(());
1204        }
1205
1206        println!("Running the stop nodes playbook for all node types");
1207        for node_inv_type in AnsibleInventoryType::iter_node_type() {
1208            self.ansible_runner.run_playbook(
1209                AnsiblePlaybook::StopNodes,
1210                node_inv_type,
1211                Some(extra_vars.clone()),
1212            )?;
1213        }
1214
1215        Ok(())
1216    }
1217
1218    pub fn stop_telegraf(
1219        &self,
1220        environment_name: &str,
1221        node_type: Option<NodeType>,
1222        custom_inventory: Option<Vec<VirtualMachine>>,
1223    ) -> Result<()> {
1224        if let Some(node_type) = node_type {
1225            println!("Running the stop telegraf playbook for {node_type:?} nodes");
1226            self.ansible_runner.run_playbook(
1227                AnsiblePlaybook::StopTelegraf,
1228                node_type.to_ansible_inventory_type(),
1229                None,
1230            )?;
1231            return Ok(());
1232        }
1233
1234        if let Some(custom_inventory) = custom_inventory {
1235            println!("Running the stop telegraf playbook with a custom inventory");
1236            generate_custom_environment_inventory(
1237                &custom_inventory,
1238                environment_name,
1239                &self.ansible_runner.working_directory_path.join("inventory"),
1240            )?;
1241            self.ansible_runner.run_playbook(
1242                AnsiblePlaybook::StopTelegraf,
1243                AnsibleInventoryType::Custom,
1244                None,
1245            )?;
1246            return Ok(());
1247        }
1248
1249        println!("Running the stop telegraf playbook for all node types");
1250        for node_inv_type in AnsibleInventoryType::iter_node_type() {
1251            self.ansible_runner
1252                .run_playbook(AnsiblePlaybook::StopTelegraf, node_inv_type, None)?;
1253        }
1254
1255        Ok(())
1256    }
1257
1258    pub fn upgrade_node_telegraf(&self, name: &str) -> Result<()> {
1259        self.ansible_runner.run_playbook(
1260            AnsiblePlaybook::UpgradeNodeTelegrafConfig,
1261            AnsibleInventoryType::PeerCacheNodes,
1262            Some(extra_vars::build_node_telegraf_upgrade(
1263                name,
1264                &NodeType::PeerCache,
1265            )?),
1266        )?;
1267        self.ansible_runner.run_playbook(
1268            AnsiblePlaybook::UpgradeNodeTelegrafConfig,
1269            AnsibleInventoryType::Nodes,
1270            Some(extra_vars::build_node_telegraf_upgrade(
1271                name,
1272                &NodeType::Generic,
1273            )?),
1274        )?;
1275
1276        self.ansible_runner.run_playbook(
1277            AnsiblePlaybook::UpgradeNodeTelegrafConfig,
1278            AnsibleInventoryType::SymmetricPrivateNodes,
1279            Some(extra_vars::build_node_telegraf_upgrade(
1280                name,
1281                &NodeType::SymmetricPrivateNode,
1282            )?),
1283        )?;
1284
1285        self.ansible_runner.run_playbook(
1286            AnsiblePlaybook::UpgradeNodeTelegrafConfig,
1287            AnsibleInventoryType::FullConePrivateNodes,
1288            Some(extra_vars::build_node_telegraf_upgrade(
1289                name,
1290                &NodeType::FullConePrivateNode,
1291            )?),
1292        )?;
1293        Ok(())
1294    }
1295
1296    pub fn upgrade_client_telegraf(&self, name: &str) -> Result<()> {
1297        self.ansible_runner.run_playbook(
1298            AnsiblePlaybook::UpgradeClientTelegrafConfig,
1299            AnsibleInventoryType::Clients,
1300            Some(extra_vars::build_client_telegraf_upgrade(name)?),
1301        )?;
1302        Ok(())
1303    }
1304
1305    pub fn upgrade_nodes(&self, options: &UpgradeOptions) -> Result<()> {
1306        if let Some(custom_inventory) = &options.custom_inventory {
1307            println!("Running the UpgradeNodes with a custom inventory");
1308            generate_custom_environment_inventory(
1309                custom_inventory,
1310                &options.name,
1311                &self.ansible_runner.working_directory_path.join("inventory"),
1312            )?;
1313            match self.ansible_runner.run_playbook(
1314                AnsiblePlaybook::UpgradeNodes,
1315                AnsibleInventoryType::Custom,
1316                Some(options.get_ansible_vars()),
1317            ) {
1318                Ok(()) => println!("All nodes were successfully upgraded"),
1319                Err(_) => {
1320                    println!("WARNING: some nodes may not have been upgraded or restarted");
1321                }
1322            }
1323            return Ok(());
1324        }
1325
1326        if let Some(node_type) = &options.node_type {
1327            println!("Running the UpgradeNodes playbook for {node_type:?} nodes");
1328            match self.ansible_runner.run_playbook(
1329                AnsiblePlaybook::UpgradeNodes,
1330                node_type.to_ansible_inventory_type(),
1331                Some(options.get_ansible_vars()),
1332            ) {
1333                Ok(()) => println!("All {node_type:?} nodes were successfully upgraded"),
1334                Err(_) => {
1335                    println!(
1336                        "WARNING: some {node_type:?} nodes may not have been upgraded or restarted"
1337                    );
1338                }
1339            }
1340            return Ok(());
1341        }
1342
1343        println!("Running the UpgradeNodes playbook for all node types");
1344
1345        match self.ansible_runner.run_playbook(
1346            AnsiblePlaybook::UpgradeNodes,
1347            AnsibleInventoryType::PeerCacheNodes,
1348            Some(options.get_ansible_vars()),
1349        ) {
1350            Ok(()) => println!("All Peer Cache nodes were successfully upgraded"),
1351            Err(_) => {
1352                println!("WARNING: some Peer Cacche nodes may not have been upgraded or restarted");
1353            }
1354        }
1355        match self.ansible_runner.run_playbook(
1356            AnsiblePlaybook::UpgradeNodes,
1357            AnsibleInventoryType::Nodes,
1358            Some(options.get_ansible_vars()),
1359        ) {
1360            Ok(()) => println!("All generic nodes were successfully upgraded"),
1361            Err(_) => {
1362                println!("WARNING: some nodes may not have been upgraded or restarted");
1363            }
1364        }
1365        match self.ansible_runner.run_playbook(
1366            AnsiblePlaybook::UpgradeNodes,
1367            AnsibleInventoryType::SymmetricPrivateNodes,
1368            Some(options.get_ansible_vars()),
1369        ) {
1370            Ok(()) => println!("All private nodes were successfully upgraded"),
1371            Err(_) => {
1372                println!("WARNING: some nodes may not have been upgraded or restarted");
1373            }
1374        }
1375        // Don't use AnsibleInventoryType::iter_node_type() here, because the genesis node should be upgraded last
1376        match self.ansible_runner.run_playbook(
1377            AnsiblePlaybook::UpgradeNodes,
1378            AnsibleInventoryType::Genesis,
1379            Some(options.get_ansible_vars()),
1380        ) {
1381            Ok(()) => println!("The genesis nodes was successfully upgraded"),
1382            Err(_) => {
1383                println!("WARNING: the genesis node may not have been upgraded or restarted");
1384            }
1385        }
1386        Ok(())
1387    }
1388
1389    pub fn upgrade_antctl(
1390        &self,
1391        environment_name: &str,
1392        version: &Version,
1393        node_type: Option<NodeType>,
1394        custom_inventory: Option<Vec<VirtualMachine>>,
1395    ) -> Result<()> {
1396        let mut extra_vars = ExtraVarsDocBuilder::default();
1397        extra_vars.add_variable("version", &version.to_string());
1398
1399        if let Some(node_type) = node_type {
1400            println!("Running the upgrade safenode-manager playbook for {node_type:?} nodes");
1401            self.ansible_runner.run_playbook(
1402                AnsiblePlaybook::UpgradeAntctl,
1403                node_type.to_ansible_inventory_type(),
1404                Some(extra_vars.build()),
1405            )?;
1406            return Ok(());
1407        }
1408
1409        if let Some(custom_inventory) = custom_inventory {
1410            println!("Running the upgrade safenode-manager playbook with a custom inventory");
1411            generate_custom_environment_inventory(
1412                &custom_inventory,
1413                environment_name,
1414                &self.ansible_runner.working_directory_path.join("inventory"),
1415            )?;
1416            self.ansible_runner.run_playbook(
1417                AnsiblePlaybook::UpgradeAntctl,
1418                AnsibleInventoryType::Custom,
1419                Some(extra_vars.build()),
1420            )?;
1421            return Ok(());
1422        }
1423
1424        println!("Running the upgrade safenode-manager playbook for all node types");
1425        for node_inv_type in AnsibleInventoryType::iter_node_type() {
1426            self.ansible_runner.run_playbook(
1427                AnsiblePlaybook::UpgradeAntctl,
1428                node_inv_type,
1429                Some(extra_vars.build()),
1430            )?;
1431        }
1432
1433        Ok(())
1434    }
1435
1436    pub fn print_ansible_run_banner(&self, s: &str) {
1437        let ansible_run_msg = "Ansible Run: ";
1438        let line = "=".repeat(s.len() + ansible_run_msg.len());
1439        println!("{}\n{}{}\n{}", line, ansible_run_msg, s, line);
1440    }
1441}