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