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    deploy::DeployOptions,
22    error::{Error, Result},
23    funding::FundingOptions,
24    inventory::{DeploymentNodeRegistries, VirtualMachine},
25    print_duration, run_external_command,
26    uploaders::UploaderDeployOptions,
27    BinaryOption, CloudProvider, EvmNetwork, LogFormat, 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 downloaders_count: u16,
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('-').last().unwrap();
217                    let nat_gateway_name = vm.name.split('-').last().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            downloaders_count: 0,
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            downloaders_count: deploy_options.downloaders_count,
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<UploaderDeployOptions> for ProvisionOptions {
311    fn from(uploader_options: UploaderDeployOptions) -> Self {
312        ProvisionOptions {
313            ant_version: None,
314            binary_option: uploader_options.binary_option,
315            chunk_size: uploader_options.chunk_size,
316            client_env_variables: uploader_options.client_env_variables,
317            downloaders_count: 0,
318            enable_telegraf: uploader_options.enable_telegraf,
319            evm_data_payments_address: uploader_options.evm_details.data_payments_address,
320            evm_network: uploader_options.evm_details.network,
321            evm_payment_token_address: uploader_options.evm_details.payment_token_address,
322            evm_rpc_url: uploader_options.evm_details.rpc_url,
323            full_cone_private_node_count: 0,
324            funding_wallet_secret_key: uploader_options.funding_wallet_secret_key,
325            gas_amount: uploader_options.initial_gas,
326            interval: None,
327            log_format: None,
328            max_archived_log_files: uploader_options.max_archived_log_files,
329            max_log_files: uploader_options.max_log_files,
330            max_uploads: uploader_options.max_uploads,
331            name: uploader_options.name,
332            network_id: uploader_options.network_id,
333            node_count: 0,
334            node_env_variables: None,
335            output_inventory_dir_path: uploader_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: uploader_options.initial_tokens,
341            uploaders_count: Some(uploader_options.uploaders_count),
342            wallet_secret_keys: uploader_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_uploader_inventory(&self) -> Result<Vec<VirtualMachine>> {
475        self.ansible_runner
476            .get_inventory(AnsibleInventoryType::Uploaders, 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('-').last().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('-').last().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('-').last().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_uploaders(
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        let sk_map = if let Some(wallet_keys) = &options.wallet_secret_keys {
1010            self.prepare_pre_funded_wallets(wallet_keys).await?
1011        } else {
1012            self.deposit_funds_to_uploaders(&FundingOptions {
1013                evm_data_payments_address: options.evm_data_payments_address.clone(),
1014                evm_network: options.evm_network.clone(),
1015                evm_payment_token_address: options.evm_payment_token_address.clone(),
1016                evm_rpc_url: options.evm_rpc_url.clone(),
1017                funding_wallet_secret_key: options.funding_wallet_secret_key.clone(),
1018                gas_amount: options.gas_amount,
1019                token_amount: options.token_amount,
1020                uploaders_count: options.uploaders_count,
1021            })
1022            .await?
1023        };
1024
1025        self.ansible_runner.run_playbook(
1026            AnsiblePlaybook::Uploaders,
1027            AnsibleInventoryType::Uploaders,
1028            Some(extra_vars::build_uploaders_extra_vars_doc(
1029                &self.cloud_provider.to_string(),
1030                options,
1031                genesis_multiaddr,
1032                genesis_network_contacts_url,
1033                &sk_map,
1034            )?),
1035        )?;
1036        print_duration(start.elapsed());
1037        Ok(())
1038    }
1039
1040    pub fn start_nodes(
1041        &self,
1042        environment_name: &str,
1043        interval: Duration,
1044        node_type: Option<NodeType>,
1045        custom_inventory: Option<Vec<VirtualMachine>>,
1046    ) -> Result<()> {
1047        let mut extra_vars = ExtraVarsDocBuilder::default();
1048        extra_vars.add_variable("interval", &interval.as_millis().to_string());
1049
1050        if let Some(node_type) = node_type {
1051            println!("Running the start nodes playbook for {node_type:?} nodes");
1052            self.ansible_runner.run_playbook(
1053                AnsiblePlaybook::StartNodes,
1054                node_type.to_ansible_inventory_type(),
1055                Some(extra_vars.build()),
1056            )?;
1057            return Ok(());
1058        }
1059
1060        if let Some(custom_inventory) = custom_inventory {
1061            println!("Running the start nodes playbook with a custom inventory");
1062            generate_custom_environment_inventory(
1063                &custom_inventory,
1064                environment_name,
1065                &self.ansible_runner.working_directory_path.join("inventory"),
1066            )?;
1067            self.ansible_runner.run_playbook(
1068                AnsiblePlaybook::StartNodes,
1069                AnsibleInventoryType::Custom,
1070                Some(extra_vars.build()),
1071            )?;
1072            return Ok(());
1073        }
1074
1075        println!("Running the start nodes playbook for all node types");
1076        for node_inv_type in AnsibleInventoryType::iter_node_type() {
1077            self.ansible_runner.run_playbook(
1078                AnsiblePlaybook::StartNodes,
1079                node_inv_type,
1080                Some(extra_vars.build()),
1081            )?;
1082        }
1083        Ok(())
1084    }
1085
1086    pub fn status(&self) -> Result<()> {
1087        for node_inv_type in AnsibleInventoryType::iter_node_type() {
1088            self.ansible_runner
1089                .run_playbook(AnsiblePlaybook::Status, node_inv_type, None)?;
1090        }
1091        Ok(())
1092    }
1093
1094    pub fn start_telegraf(
1095        &self,
1096        environment_name: &str,
1097        node_type: Option<NodeType>,
1098        custom_inventory: Option<Vec<VirtualMachine>>,
1099    ) -> Result<()> {
1100        if let Some(node_type) = node_type {
1101            println!("Running the start telegraf playbook for {node_type:?} nodes");
1102            self.ansible_runner.run_playbook(
1103                AnsiblePlaybook::StartTelegraf,
1104                node_type.to_ansible_inventory_type(),
1105                None,
1106            )?;
1107            return Ok(());
1108        }
1109
1110        if let Some(custom_inventory) = custom_inventory {
1111            println!("Running the start telegraf playbook with a custom inventory");
1112            generate_custom_environment_inventory(
1113                &custom_inventory,
1114                environment_name,
1115                &self.ansible_runner.working_directory_path.join("inventory"),
1116            )?;
1117            self.ansible_runner.run_playbook(
1118                AnsiblePlaybook::StartTelegraf,
1119                AnsibleInventoryType::Custom,
1120                None,
1121            )?;
1122            return Ok(());
1123        }
1124
1125        println!("Running the start telegraf playbook for all node types");
1126        for node_inv_type in AnsibleInventoryType::iter_node_type() {
1127            self.ansible_runner.run_playbook(
1128                AnsiblePlaybook::StartTelegraf,
1129                node_inv_type,
1130                None,
1131            )?;
1132        }
1133
1134        Ok(())
1135    }
1136
1137    pub fn stop_nodes(
1138        &self,
1139        environment_name: &str,
1140        interval: Duration,
1141        node_type: Option<NodeType>,
1142        custom_inventory: Option<Vec<VirtualMachine>>,
1143        delay: Option<u64>,
1144        service_names: Option<Vec<String>>,
1145    ) -> Result<()> {
1146        let mut extra_vars = ExtraVarsDocBuilder::default();
1147        extra_vars.add_variable("interval", &interval.as_millis().to_string());
1148        if let Some(delay) = delay {
1149            extra_vars.add_variable("delay", &delay.to_string());
1150        }
1151        if let Some(service_names) = service_names {
1152            extra_vars.add_list_variable("service_names", service_names);
1153        }
1154        let extra_vars = extra_vars.build();
1155
1156        if let Some(node_type) = node_type {
1157            println!("Running the stop nodes playbook for {node_type:?} nodes");
1158            self.ansible_runner.run_playbook(
1159                AnsiblePlaybook::StopNodes,
1160                node_type.to_ansible_inventory_type(),
1161                Some(extra_vars),
1162            )?;
1163            return Ok(());
1164        }
1165
1166        if let Some(custom_inventory) = custom_inventory {
1167            println!("Running the stop nodes playbook with a custom inventory");
1168            generate_custom_environment_inventory(
1169                &custom_inventory,
1170                environment_name,
1171                &self.ansible_runner.working_directory_path.join("inventory"),
1172            )?;
1173            self.ansible_runner.run_playbook(
1174                AnsiblePlaybook::StopNodes,
1175                AnsibleInventoryType::Custom,
1176                Some(extra_vars),
1177            )?;
1178            return Ok(());
1179        }
1180
1181        println!("Running the stop nodes playbook for all node types");
1182        for node_inv_type in AnsibleInventoryType::iter_node_type() {
1183            self.ansible_runner.run_playbook(
1184                AnsiblePlaybook::StopNodes,
1185                node_inv_type,
1186                Some(extra_vars.clone()),
1187            )?;
1188        }
1189
1190        Ok(())
1191    }
1192
1193    pub fn stop_telegraf(
1194        &self,
1195        environment_name: &str,
1196        node_type: Option<NodeType>,
1197        custom_inventory: Option<Vec<VirtualMachine>>,
1198    ) -> Result<()> {
1199        if let Some(node_type) = node_type {
1200            println!("Running the stop telegraf playbook for {node_type:?} nodes");
1201            self.ansible_runner.run_playbook(
1202                AnsiblePlaybook::StopTelegraf,
1203                node_type.to_ansible_inventory_type(),
1204                None,
1205            )?;
1206            return Ok(());
1207        }
1208
1209        if let Some(custom_inventory) = custom_inventory {
1210            println!("Running the stop telegraf playbook with a custom inventory");
1211            generate_custom_environment_inventory(
1212                &custom_inventory,
1213                environment_name,
1214                &self.ansible_runner.working_directory_path.join("inventory"),
1215            )?;
1216            self.ansible_runner.run_playbook(
1217                AnsiblePlaybook::StopTelegraf,
1218                AnsibleInventoryType::Custom,
1219                None,
1220            )?;
1221            return Ok(());
1222        }
1223
1224        println!("Running the stop telegraf playbook for all node types");
1225        for node_inv_type in AnsibleInventoryType::iter_node_type() {
1226            self.ansible_runner
1227                .run_playbook(AnsiblePlaybook::StopTelegraf, node_inv_type, None)?;
1228        }
1229
1230        Ok(())
1231    }
1232
1233    pub fn upgrade_node_telegraf(&self, name: &str) -> Result<()> {
1234        self.ansible_runner.run_playbook(
1235            AnsiblePlaybook::UpgradeNodeTelegrafConfig,
1236            AnsibleInventoryType::PeerCacheNodes,
1237            Some(extra_vars::build_node_telegraf_upgrade(
1238                name,
1239                &NodeType::PeerCache,
1240            )?),
1241        )?;
1242        self.ansible_runner.run_playbook(
1243            AnsiblePlaybook::UpgradeNodeTelegrafConfig,
1244            AnsibleInventoryType::Nodes,
1245            Some(extra_vars::build_node_telegraf_upgrade(
1246                name,
1247                &NodeType::Generic,
1248            )?),
1249        )?;
1250
1251        self.ansible_runner.run_playbook(
1252            AnsiblePlaybook::UpgradeNodeTelegrafConfig,
1253            AnsibleInventoryType::SymmetricPrivateNodes,
1254            Some(extra_vars::build_node_telegraf_upgrade(
1255                name,
1256                &NodeType::SymmetricPrivateNode,
1257            )?),
1258        )?;
1259
1260        self.ansible_runner.run_playbook(
1261            AnsiblePlaybook::UpgradeNodeTelegrafConfig,
1262            AnsibleInventoryType::FullConePrivateNodes,
1263            Some(extra_vars::build_node_telegraf_upgrade(
1264                name,
1265                &NodeType::FullConePrivateNode,
1266            )?),
1267        )?;
1268        Ok(())
1269    }
1270
1271    pub fn upgrade_uploader_telegraf(&self, name: &str) -> Result<()> {
1272        self.ansible_runner.run_playbook(
1273            AnsiblePlaybook::UpgradeUploaderTelegrafConfig,
1274            AnsibleInventoryType::Uploaders,
1275            Some(extra_vars::build_uploader_telegraf_upgrade(name)?),
1276        )?;
1277        Ok(())
1278    }
1279
1280    pub fn upgrade_nodes(&self, options: &UpgradeOptions) -> Result<()> {
1281        if let Some(custom_inventory) = &options.custom_inventory {
1282            println!("Running the UpgradeNodes with a custom inventory");
1283            generate_custom_environment_inventory(
1284                custom_inventory,
1285                &options.name,
1286                &self.ansible_runner.working_directory_path.join("inventory"),
1287            )?;
1288            match self.ansible_runner.run_playbook(
1289                AnsiblePlaybook::UpgradeNodes,
1290                AnsibleInventoryType::Custom,
1291                Some(options.get_ansible_vars()),
1292            ) {
1293                Ok(()) => println!("All nodes were successfully upgraded"),
1294                Err(_) => {
1295                    println!("WARNING: some nodes may not have been upgraded or restarted");
1296                }
1297            }
1298            return Ok(());
1299        }
1300
1301        if let Some(node_type) = &options.node_type {
1302            println!("Running the UpgradeNodes playbook for {node_type:?} nodes");
1303            match self.ansible_runner.run_playbook(
1304                AnsiblePlaybook::UpgradeNodes,
1305                node_type.to_ansible_inventory_type(),
1306                Some(options.get_ansible_vars()),
1307            ) {
1308                Ok(()) => println!("All {node_type:?} nodes were successfully upgraded"),
1309                Err(_) => {
1310                    println!(
1311                        "WARNING: some {node_type:?} nodes may not have been upgraded or restarted"
1312                    );
1313                }
1314            }
1315            return Ok(());
1316        }
1317
1318        println!("Running the UpgradeNodes playbook for all node types");
1319
1320        match self.ansible_runner.run_playbook(
1321            AnsiblePlaybook::UpgradeNodes,
1322            AnsibleInventoryType::PeerCacheNodes,
1323            Some(options.get_ansible_vars()),
1324        ) {
1325            Ok(()) => println!("All Peer Cache nodes were successfully upgraded"),
1326            Err(_) => {
1327                println!("WARNING: some Peer Cacche nodes may not have been upgraded or restarted");
1328            }
1329        }
1330        match self.ansible_runner.run_playbook(
1331            AnsiblePlaybook::UpgradeNodes,
1332            AnsibleInventoryType::Nodes,
1333            Some(options.get_ansible_vars()),
1334        ) {
1335            Ok(()) => println!("All generic nodes were successfully upgraded"),
1336            Err(_) => {
1337                println!("WARNING: some nodes may not have been upgraded or restarted");
1338            }
1339        }
1340        match self.ansible_runner.run_playbook(
1341            AnsiblePlaybook::UpgradeNodes,
1342            AnsibleInventoryType::SymmetricPrivateNodes,
1343            Some(options.get_ansible_vars()),
1344        ) {
1345            Ok(()) => println!("All private nodes were successfully upgraded"),
1346            Err(_) => {
1347                println!("WARNING: some nodes may not have been upgraded or restarted");
1348            }
1349        }
1350        // Don't use AnsibleInventoryType::iter_node_type() here, because the genesis node should be upgraded last
1351        match self.ansible_runner.run_playbook(
1352            AnsiblePlaybook::UpgradeNodes,
1353            AnsibleInventoryType::Genesis,
1354            Some(options.get_ansible_vars()),
1355        ) {
1356            Ok(()) => println!("The genesis nodes was successfully upgraded"),
1357            Err(_) => {
1358                println!("WARNING: the genesis node may not have been upgraded or restarted");
1359            }
1360        }
1361        Ok(())
1362    }
1363
1364    pub fn upgrade_antctl(
1365        &self,
1366        environment_name: &str,
1367        version: &Version,
1368        node_type: Option<NodeType>,
1369        custom_inventory: Option<Vec<VirtualMachine>>,
1370    ) -> Result<()> {
1371        let mut extra_vars = ExtraVarsDocBuilder::default();
1372        extra_vars.add_variable("version", &version.to_string());
1373
1374        if let Some(node_type) = node_type {
1375            println!("Running the upgrade safenode-manager playbook for {node_type:?} nodes");
1376            self.ansible_runner.run_playbook(
1377                AnsiblePlaybook::UpgradeAntctl,
1378                node_type.to_ansible_inventory_type(),
1379                Some(extra_vars.build()),
1380            )?;
1381            return Ok(());
1382        }
1383
1384        if let Some(custom_inventory) = custom_inventory {
1385            println!("Running the upgrade safenode-manager playbook with a custom inventory");
1386            generate_custom_environment_inventory(
1387                &custom_inventory,
1388                environment_name,
1389                &self.ansible_runner.working_directory_path.join("inventory"),
1390            )?;
1391            self.ansible_runner.run_playbook(
1392                AnsiblePlaybook::UpgradeAntctl,
1393                AnsibleInventoryType::Custom,
1394                Some(extra_vars.build()),
1395            )?;
1396            return Ok(());
1397        }
1398
1399        println!("Running the upgrade safenode-manager playbook for all node types");
1400        for node_inv_type in AnsibleInventoryType::iter_node_type() {
1401            self.ansible_runner.run_playbook(
1402                AnsiblePlaybook::UpgradeAntctl,
1403                node_inv_type,
1404                Some(extra_vars.build()),
1405            )?;
1406        }
1407
1408        Ok(())
1409    }
1410
1411    pub fn print_ansible_run_banner(&self, s: &str) {
1412        let ansible_run_msg = "Ansible Run: ";
1413        let line = "=".repeat(s.len() + ansible_run_msg.len());
1414        println!("{}\n{}{}\n{}", line, ansible_run_msg, s, line);
1415    }
1416}