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