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_delayed_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_delayed_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_delayed_verifier: deploy_options.enable_delayed_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_delayed_verifier: client_options.enable_delayed_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_autonomi_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                false,
614            )?),
615        )?;
616
617        print_duration(start.elapsed());
618
619        Ok(())
620    }
621
622    pub fn provision_full_cone(
623        &self,
624        options: &ProvisionOptions,
625        initial_contact_peer: Option<String>,
626        initial_network_contacts_url: Option<String>,
627        private_node_inventory: PrivateNodeProvisionInventory,
628        new_full_cone_nat_gateway_new_vms_for_upscale: Option<Vec<VirtualMachine>>,
629    ) -> Result<()> {
630        // Step 1 of Full Cone NAT Gateway
631        let start = Instant::now();
632        self.print_ansible_run_banner("Provision Full Cone NAT Gateway - Step 1");
633
634        for vm in new_full_cone_nat_gateway_new_vms_for_upscale
635            .as_ref()
636            .unwrap_or(&private_node_inventory.full_cone_nat_gateway_vms)
637            .iter()
638        {
639            println!(
640                "Checking SSH availability for Full Cone NAT Gateway: {}",
641                vm.public_ip_addr
642            );
643            self.ssh_client
644                .wait_for_ssh_availability(&vm.public_ip_addr, &self.cloud_provider.get_ssh_user())
645                .map_err(|e| {
646                    println!("Failed to establish SSH connection to Full Cone NAT Gateway: {e}");
647                    e
648                })?;
649        }
650
651        let mut modified_private_node_inventory = private_node_inventory.clone();
652
653        // If we are upscaling, then we cannot access the gateway VMs which are already deployed.
654        if let Some(new_full_cone_nat_gateway_new_vms_for_upscale) =
655            &new_full_cone_nat_gateway_new_vms_for_upscale
656        {
657            debug!("Removing existing full cone NAT Gateway and private node VMs from the inventory. Old inventory: {modified_private_node_inventory:?}");
658            let mut names_to_keep = Vec::new();
659
660            for vm in new_full_cone_nat_gateway_new_vms_for_upscale.iter() {
661                let nat_gateway_name = vm.name.split('-').next_back().unwrap();
662                names_to_keep.push(nat_gateway_name);
663            }
664
665            modified_private_node_inventory
666                .full_cone_nat_gateway_vms
667                .retain(|vm| {
668                    let nat_gateway_name = vm.name.split('-').next_back().unwrap();
669                    names_to_keep.contains(&nat_gateway_name)
670                });
671            modified_private_node_inventory
672                .full_cone_private_node_vms
673                .retain(|vm| {
674                    let nat_gateway_name = vm.name.split('-').next_back().unwrap();
675                    names_to_keep.contains(&nat_gateway_name)
676                });
677            debug!("New inventory after removing existing full cone NAT Gateway and private node VMs: {modified_private_node_inventory:?}");
678        }
679
680        if modified_private_node_inventory
681            .full_cone_nat_gateway_vms
682            .is_empty()
683        {
684            error!("There are no full cone NAT Gateway VMs available to upscale");
685            return Ok(());
686        }
687
688        let private_node_ip_map = modified_private_node_inventory
689            .full_cone_private_node_and_gateway_map()?
690            .into_iter()
691            .map(|(k, v)| {
692                let gateway_name = if new_full_cone_nat_gateway_new_vms_for_upscale.is_some() {
693                    debug!("Upscaling, using public IP address for gateway name");
694                    v.public_ip_addr.to_string()
695                } else {
696                    v.name.clone()
697                };
698                (gateway_name, k.private_ip_addr)
699            })
700            .collect::<HashMap<String, IpAddr>>();
701
702        if private_node_ip_map.is_empty() {
703            println!("There are no full cone private node VM available to be routed through the full cone NAT Gateway");
704            return Err(Error::EmptyInventory(
705                AnsibleInventoryType::FullConePrivateNodes,
706            ));
707        }
708
709        let vars = extra_vars::build_nat_gateway_extra_vars_doc(
710            &options.name,
711            private_node_ip_map.clone(),
712            "step1",
713        );
714        debug!("Provisioning Full Cone NAT Gateway - Step 1 with vars: {vars}");
715        let gateway_inventory = if new_full_cone_nat_gateway_new_vms_for_upscale.is_some() {
716            debug!("Upscaling, using static inventory for full cone nat gateway.");
717            generate_full_cone_nat_gateway_static_environment_inventory(
718                &modified_private_node_inventory.full_cone_nat_gateway_vms,
719                &options.name,
720                &options.output_inventory_dir_path,
721            )?;
722
723            AnsibleInventoryType::FullConeNatGatewayStatic
724        } else {
725            AnsibleInventoryType::FullConeNatGateway
726        };
727        self.ansible_runner.run_playbook(
728            AnsiblePlaybook::StaticFullConeNatGateway,
729            gateway_inventory,
730            Some(vars),
731        )?;
732
733        // setup private node config
734        self.print_ansible_run_banner("Provisioning Full Cone Private Node Config");
735
736        generate_full_cone_private_node_static_environment_inventory(
737            &options.name,
738            &options.output_inventory_dir_path,
739            &private_node_inventory.full_cone_private_node_vms,
740            &private_node_inventory.full_cone_nat_gateway_vms,
741            &self.ssh_client.private_key_path,
742        )
743        .inspect_err(|err| {
744            error!("Failed to generate full cone private node static inv with err: {err:?}")
745        })?;
746
747        // For a new deployment, it's quite probable that SSH is available, because this part occurs
748        // after the genesis node has been provisioned. However, for a bootstrap deploy, we need to
749        // check that SSH is available before proceeding.
750        println!("Obtaining IP addresses for nodes...");
751        let inventory = self
752            .ansible_runner
753            .get_inventory(AnsibleInventoryType::FullConePrivateNodes, true)?;
754
755        println!("Waiting for SSH availability on Symmetric Private nodes...");
756        for vm in inventory.iter() {
757            println!(
758                "Checking SSH availability for {}: {}",
759                vm.name, vm.public_ip_addr
760            );
761            self.ssh_client
762                .wait_for_ssh_availability(&vm.public_ip_addr, &self.cloud_provider.get_ssh_user())
763                .map_err(|e| {
764                    println!("Failed to establish SSH connection to {}: {}", vm.name, e);
765                    e
766                })?;
767        }
768
769        println!("SSH is available on all nodes. Proceeding with provisioning...");
770
771        self.ansible_runner.run_playbook(
772            AnsiblePlaybook::PrivateNodeConfig,
773            AnsibleInventoryType::FullConePrivateNodes,
774            Some(
775                extra_vars::build_full_cone_private_node_config_extra_vars_docs(
776                    &private_node_inventory,
777                )?,
778            ),
779        )?;
780
781        // Step 2 of Full Cone NAT Gateway
782
783        let vars = extra_vars::build_nat_gateway_extra_vars_doc(
784            &options.name,
785            private_node_ip_map,
786            "step2",
787        );
788
789        self.print_ansible_run_banner("Provisioning Full Cone NAT Gateway - Step 2");
790        debug!("Provisioning Full Cone NAT Gateway - Step 2 with vars: {vars}");
791        self.ansible_runner.run_playbook(
792            AnsiblePlaybook::StaticFullConeNatGateway,
793            gateway_inventory,
794            Some(vars),
795        )?;
796
797        // provision the nodes
798
799        let home_dir = std::env::var("HOME").inspect_err(|err| {
800            println!("Failed to get home directory with error: {err:?}",);
801        })?;
802        let known_hosts_path = format!("{home_dir}/.ssh/known_hosts");
803        debug!("Cleaning up known hosts file at {known_hosts_path} ");
804        run_external_command(
805            PathBuf::from("rm"),
806            std::env::current_dir()?,
807            vec![known_hosts_path],
808            false,
809            false,
810        )?;
811
812        self.print_ansible_run_banner("Provision Full Cone Private Nodes");
813
814        self.ssh_client.set_full_cone_nat_routed_vms(
815            &private_node_inventory.full_cone_private_node_vms,
816            &private_node_inventory.full_cone_nat_gateway_vms,
817        )?;
818
819        self.provision_nodes(
820            options,
821            initial_contact_peer,
822            initial_network_contacts_url,
823            NodeType::FullConePrivateNode,
824        )?;
825
826        print_duration(start.elapsed());
827        Ok(())
828    }
829    pub fn provision_symmetric_nat_gateway(
830        &self,
831        options: &ProvisionOptions,
832        private_node_inventory: &PrivateNodeProvisionInventory,
833    ) -> Result<()> {
834        let start = Instant::now();
835        for vm in &private_node_inventory.symmetric_nat_gateway_vms {
836            println!(
837                "Checking SSH availability for Symmetric NAT Gateway: {}",
838                vm.public_ip_addr
839            );
840            self.ssh_client
841                .wait_for_ssh_availability(&vm.public_ip_addr, &self.cloud_provider.get_ssh_user())
842                .map_err(|e| {
843                    println!("Failed to establish SSH connection to Symmetric NAT Gateway: {e}");
844                    e
845                })?;
846        }
847
848        let private_node_ip_map = private_node_inventory
849            .symmetric_private_node_and_gateway_map()?
850            .into_iter()
851            .map(|(k, v)| (v.name.clone(), k.private_ip_addr))
852            .collect::<HashMap<String, IpAddr>>();
853
854        if private_node_ip_map.is_empty() {
855            println!("There are no Symmetric private node VM available to be routed through the Symmetric NAT Gateway");
856            return Err(Error::EmptyInventory(
857                AnsibleInventoryType::SymmetricPrivateNodes,
858            ));
859        }
860
861        let vars = extra_vars::build_nat_gateway_extra_vars_doc(
862            &options.name,
863            private_node_ip_map,
864            "symmetric",
865        );
866        debug!("Provisioning Symmetric NAT Gateway with vars: {vars}");
867        self.ansible_runner.run_playbook(
868            AnsiblePlaybook::SymmetricNatGateway,
869            AnsibleInventoryType::SymmetricNatGateway,
870            Some(vars),
871        )?;
872
873        print_duration(start.elapsed());
874        Ok(())
875    }
876
877    pub fn provision_nodes(
878        &self,
879        options: &ProvisionOptions,
880        initial_contact_peer: Option<String>,
881        initial_network_contacts_url: Option<String>,
882        node_type: NodeType,
883    ) -> Result<()> {
884        let start = Instant::now();
885        let mut write_older_cache_files = false;
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                write_older_cache_files = true;
896                (
897                    node_type.to_ansible_inventory_type(),
898                    options.peer_cache_node_count,
899                )
900            }
901            NodeType::SymmetricPrivateNode => (
902                node_type.to_ansible_inventory_type(),
903                options.symmetric_private_node_count,
904            ),
905        };
906
907        // For a new deployment, it's quite probable that SSH is available, because this part occurs
908        // after the genesis node has been provisioned. However, for a bootstrap deploy, we need to
909        // check that SSH is available before proceeding.
910        println!("Obtaining IP addresses for nodes...");
911        let inventory = self.ansible_runner.get_inventory(inventory_type, true)?;
912
913        println!("Waiting for SSH availability on {node_type:?} nodes...");
914        for vm in inventory.iter() {
915            println!(
916                "Checking SSH availability for {}: {}",
917                vm.name, vm.public_ip_addr
918            );
919            self.ssh_client
920                .wait_for_ssh_availability(&vm.public_ip_addr, &self.cloud_provider.get_ssh_user())
921                .map_err(|e| {
922                    println!("Failed to establish SSH connection to {}: {}", vm.name, e);
923                    e
924                })?;
925        }
926
927        println!("SSH is available on all nodes. Proceeding with provisioning...");
928
929        let playbook = match node_type {
930            NodeType::Generic => AnsiblePlaybook::Nodes,
931            NodeType::PeerCache => AnsiblePlaybook::PeerCacheNodes,
932            NodeType::FullConePrivateNode => AnsiblePlaybook::Nodes,
933            NodeType::SymmetricPrivateNode => AnsiblePlaybook::Nodes,
934            _ => return Err(Error::InvalidNodeType(node_type.clone())),
935        };
936        self.ansible_runner.run_playbook(
937            playbook,
938            inventory_type,
939            Some(extra_vars::build_node_extra_vars_doc(
940                &self.cloud_provider.to_string(),
941                options,
942                node_type.clone(),
943                initial_contact_peer,
944                initial_network_contacts_url,
945                node_count,
946                options.evm_network.clone(),
947                write_older_cache_files,
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_static_downloaders(
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        println!("Running ansible against client machine to start the static downloaders.");
1062        debug!("Running ansible against client machine to start the static downloaders.");
1063
1064        self.ansible_runner.run_playbook(
1065            AnsiblePlaybook::StaticDownloaders,
1066            AnsibleInventoryType::Clients,
1067            Some(extra_vars::build_downloaders_extra_vars_doc(
1068                &self.cloud_provider.to_string(),
1069                options,
1070                genesis_multiaddr,
1071                genesis_network_contacts_url,
1072            )?),
1073        )?;
1074        print_duration(start.elapsed());
1075        Ok(())
1076    }
1077
1078    pub async fn provision_clients(
1079        &self,
1080        options: &ProvisionOptions,
1081        genesis_multiaddr: Option<String>,
1082        genesis_network_contacts_url: Option<String>,
1083    ) -> Result<()> {
1084        let start = Instant::now();
1085
1086        let sk_map = if let Some(wallet_keys) = &options.wallet_secret_keys {
1087            self.prepare_pre_funded_wallets(wallet_keys).await?
1088        } else {
1089            self.deposit_funds_to_clients(&FundingOptions {
1090                evm_data_payments_address: options.evm_data_payments_address.clone(),
1091                evm_network: options.evm_network.clone(),
1092                evm_payment_token_address: options.evm_payment_token_address.clone(),
1093                evm_rpc_url: options.evm_rpc_url.clone(),
1094                funding_wallet_secret_key: options.funding_wallet_secret_key.clone(),
1095                gas_amount: options.gas_amount,
1096                token_amount: options.token_amount,
1097                uploaders_count: options.uploaders_count,
1098            })
1099            .await?
1100        };
1101
1102        self.ansible_runner.run_playbook(
1103            AnsiblePlaybook::Uploaders,
1104            AnsibleInventoryType::Clients,
1105            Some(extra_vars::build_clients_extra_vars_doc(
1106                &self.cloud_provider.to_string(),
1107                options,
1108                genesis_multiaddr,
1109                genesis_network_contacts_url,
1110                &sk_map,
1111            )?),
1112        )?;
1113        print_duration(start.elapsed());
1114        Ok(())
1115    }
1116
1117    pub fn start_nodes(
1118        &self,
1119        environment_name: &str,
1120        interval: Duration,
1121        node_type: Option<NodeType>,
1122        custom_inventory: Option<Vec<VirtualMachine>>,
1123    ) -> Result<()> {
1124        let mut extra_vars = ExtraVarsDocBuilder::default();
1125        extra_vars.add_variable("interval", &interval.as_millis().to_string());
1126
1127        if let Some(node_type) = node_type {
1128            println!("Running the start nodes playbook for {node_type:?} nodes");
1129            self.ansible_runner.run_playbook(
1130                AnsiblePlaybook::StartNodes,
1131                node_type.to_ansible_inventory_type(),
1132                Some(extra_vars.build()),
1133            )?;
1134            return Ok(());
1135        }
1136
1137        if let Some(custom_inventory) = custom_inventory {
1138            println!("Running the start nodes playbook with a custom inventory");
1139            generate_custom_environment_inventory(
1140                &custom_inventory,
1141                environment_name,
1142                &self.ansible_runner.working_directory_path.join("inventory"),
1143            )?;
1144            self.ansible_runner.run_playbook(
1145                AnsiblePlaybook::StartNodes,
1146                AnsibleInventoryType::Custom,
1147                Some(extra_vars.build()),
1148            )?;
1149            return Ok(());
1150        }
1151
1152        println!("Running the start nodes playbook for all node types");
1153        for node_inv_type in AnsibleInventoryType::iter_node_type() {
1154            self.ansible_runner.run_playbook(
1155                AnsiblePlaybook::StartNodes,
1156                node_inv_type,
1157                Some(extra_vars.build()),
1158            )?;
1159        }
1160        Ok(())
1161    }
1162
1163    pub fn status(&self) -> Result<()> {
1164        for node_inv_type in AnsibleInventoryType::iter_node_type() {
1165            self.ansible_runner
1166                .run_playbook(AnsiblePlaybook::Status, node_inv_type, None)?;
1167        }
1168        Ok(())
1169    }
1170
1171    pub fn start_telegraf(
1172        &self,
1173        environment_name: &str,
1174        node_type: Option<NodeType>,
1175        custom_inventory: Option<Vec<VirtualMachine>>,
1176    ) -> Result<()> {
1177        if let Some(node_type) = node_type {
1178            println!("Running the start telegraf playbook for {node_type:?} nodes");
1179            self.ansible_runner.run_playbook(
1180                AnsiblePlaybook::StartTelegraf,
1181                node_type.to_ansible_inventory_type(),
1182                None,
1183            )?;
1184            return Ok(());
1185        }
1186
1187        if let Some(custom_inventory) = custom_inventory {
1188            println!("Running the start telegraf playbook with a custom inventory");
1189            generate_custom_environment_inventory(
1190                &custom_inventory,
1191                environment_name,
1192                &self.ansible_runner.working_directory_path.join("inventory"),
1193            )?;
1194            self.ansible_runner.run_playbook(
1195                AnsiblePlaybook::StartTelegraf,
1196                AnsibleInventoryType::Custom,
1197                None,
1198            )?;
1199            return Ok(());
1200        }
1201
1202        println!("Running the start telegraf playbook for all node types");
1203        for node_inv_type in AnsibleInventoryType::iter_node_type() {
1204            self.ansible_runner.run_playbook(
1205                AnsiblePlaybook::StartTelegraf,
1206                node_inv_type,
1207                None,
1208            )?;
1209        }
1210
1211        Ok(())
1212    }
1213
1214    pub fn stop_nodes(
1215        &self,
1216        environment_name: &str,
1217        interval: Duration,
1218        node_type: Option<NodeType>,
1219        custom_inventory: Option<Vec<VirtualMachine>>,
1220        delay: Option<u64>,
1221        service_names: Option<Vec<String>>,
1222    ) -> Result<()> {
1223        let mut extra_vars = ExtraVarsDocBuilder::default();
1224        extra_vars.add_variable("interval", &interval.as_millis().to_string());
1225        if let Some(delay) = delay {
1226            extra_vars.add_variable("delay", &delay.to_string());
1227        }
1228        if let Some(service_names) = service_names {
1229            extra_vars.add_list_variable("service_names", service_names);
1230        }
1231        let extra_vars = extra_vars.build();
1232
1233        if let Some(node_type) = node_type {
1234            println!("Running the stop nodes playbook for {node_type:?} nodes");
1235            self.ansible_runner.run_playbook(
1236                AnsiblePlaybook::StopNodes,
1237                node_type.to_ansible_inventory_type(),
1238                Some(extra_vars),
1239            )?;
1240            return Ok(());
1241        }
1242
1243        if let Some(custom_inventory) = custom_inventory {
1244            println!("Running the stop nodes playbook with a custom inventory");
1245            generate_custom_environment_inventory(
1246                &custom_inventory,
1247                environment_name,
1248                &self.ansible_runner.working_directory_path.join("inventory"),
1249            )?;
1250            self.ansible_runner.run_playbook(
1251                AnsiblePlaybook::StopNodes,
1252                AnsibleInventoryType::Custom,
1253                Some(extra_vars),
1254            )?;
1255            return Ok(());
1256        }
1257
1258        println!("Running the stop nodes playbook for all node types");
1259        for node_inv_type in AnsibleInventoryType::iter_node_type() {
1260            self.ansible_runner.run_playbook(
1261                AnsiblePlaybook::StopNodes,
1262                node_inv_type,
1263                Some(extra_vars.clone()),
1264            )?;
1265        }
1266
1267        Ok(())
1268    }
1269
1270    pub fn stop_telegraf(
1271        &self,
1272        environment_name: &str,
1273        node_type: Option<NodeType>,
1274        custom_inventory: Option<Vec<VirtualMachine>>,
1275    ) -> Result<()> {
1276        if let Some(node_type) = node_type {
1277            println!("Running the stop telegraf playbook for {node_type:?} nodes");
1278            self.ansible_runner.run_playbook(
1279                AnsiblePlaybook::StopTelegraf,
1280                node_type.to_ansible_inventory_type(),
1281                None,
1282            )?;
1283            return Ok(());
1284        }
1285
1286        if let Some(custom_inventory) = custom_inventory {
1287            println!("Running the stop telegraf playbook with a custom inventory");
1288            generate_custom_environment_inventory(
1289                &custom_inventory,
1290                environment_name,
1291                &self.ansible_runner.working_directory_path.join("inventory"),
1292            )?;
1293            self.ansible_runner.run_playbook(
1294                AnsiblePlaybook::StopTelegraf,
1295                AnsibleInventoryType::Custom,
1296                None,
1297            )?;
1298            return Ok(());
1299        }
1300
1301        println!("Running the stop telegraf playbook for all node types");
1302        for node_inv_type in AnsibleInventoryType::iter_node_type() {
1303            self.ansible_runner
1304                .run_playbook(AnsiblePlaybook::StopTelegraf, node_inv_type, None)?;
1305        }
1306
1307        Ok(())
1308    }
1309
1310    pub fn upgrade_node_telegraf(&self, name: &str) -> Result<()> {
1311        self.ansible_runner.run_playbook(
1312            AnsiblePlaybook::UpgradeNodeTelegrafConfig,
1313            AnsibleInventoryType::PeerCacheNodes,
1314            Some(extra_vars::build_node_telegraf_upgrade(
1315                name,
1316                &NodeType::PeerCache,
1317            )?),
1318        )?;
1319        self.ansible_runner.run_playbook(
1320            AnsiblePlaybook::UpgradeNodeTelegrafConfig,
1321            AnsibleInventoryType::Nodes,
1322            Some(extra_vars::build_node_telegraf_upgrade(
1323                name,
1324                &NodeType::Generic,
1325            )?),
1326        )?;
1327
1328        self.ansible_runner.run_playbook(
1329            AnsiblePlaybook::UpgradeNodeTelegrafConfig,
1330            AnsibleInventoryType::SymmetricPrivateNodes,
1331            Some(extra_vars::build_node_telegraf_upgrade(
1332                name,
1333                &NodeType::SymmetricPrivateNode,
1334            )?),
1335        )?;
1336
1337        self.ansible_runner.run_playbook(
1338            AnsiblePlaybook::UpgradeNodeTelegrafConfig,
1339            AnsibleInventoryType::FullConePrivateNodes,
1340            Some(extra_vars::build_node_telegraf_upgrade(
1341                name,
1342                &NodeType::FullConePrivateNode,
1343            )?),
1344        )?;
1345        Ok(())
1346    }
1347
1348    pub fn upgrade_client_telegraf(&self, name: &str) -> Result<()> {
1349        self.ansible_runner.run_playbook(
1350            AnsiblePlaybook::UpgradeClientTelegrafConfig,
1351            AnsibleInventoryType::Clients,
1352            Some(extra_vars::build_client_telegraf_upgrade(name)?),
1353        )?;
1354        Ok(())
1355    }
1356
1357    pub fn upgrade_nodes(&self, options: &UpgradeOptions) -> Result<()> {
1358        if let Some(custom_inventory) = &options.custom_inventory {
1359            println!("Running the UpgradeNodes with a custom inventory");
1360            generate_custom_environment_inventory(
1361                custom_inventory,
1362                &options.name,
1363                &self.ansible_runner.working_directory_path.join("inventory"),
1364            )?;
1365            match self.ansible_runner.run_playbook(
1366                AnsiblePlaybook::UpgradeNodes,
1367                AnsibleInventoryType::Custom,
1368                Some(options.get_ansible_vars()),
1369            ) {
1370                Ok(()) => println!("All nodes were successfully upgraded"),
1371                Err(_) => {
1372                    println!("WARNING: some nodes may not have been upgraded or restarted");
1373                }
1374            }
1375            return Ok(());
1376        }
1377
1378        if let Some(node_type) = &options.node_type {
1379            println!("Running the UpgradeNodes playbook for {node_type:?} nodes");
1380            match self.ansible_runner.run_playbook(
1381                AnsiblePlaybook::UpgradeNodes,
1382                node_type.to_ansible_inventory_type(),
1383                Some(options.get_ansible_vars()),
1384            ) {
1385                Ok(()) => println!("All {node_type:?} nodes were successfully upgraded"),
1386                Err(_) => {
1387                    println!(
1388                        "WARNING: some {node_type:?} nodes may not have been upgraded or restarted"
1389                    );
1390                }
1391            }
1392            return Ok(());
1393        }
1394
1395        println!("Running the UpgradeNodes playbook for all node types");
1396
1397        match self.ansible_runner.run_playbook(
1398            AnsiblePlaybook::UpgradeNodes,
1399            AnsibleInventoryType::PeerCacheNodes,
1400            Some(options.get_ansible_vars()),
1401        ) {
1402            Ok(()) => println!("All Peer Cache nodes were successfully upgraded"),
1403            Err(_) => {
1404                println!("WARNING: some Peer Cacche nodes may not have been upgraded or restarted");
1405            }
1406        }
1407        match self.ansible_runner.run_playbook(
1408            AnsiblePlaybook::UpgradeNodes,
1409            AnsibleInventoryType::Nodes,
1410            Some(options.get_ansible_vars()),
1411        ) {
1412            Ok(()) => println!("All generic nodes were successfully upgraded"),
1413            Err(_) => {
1414                println!("WARNING: some nodes may not have been upgraded or restarted");
1415            }
1416        }
1417        match self.ansible_runner.run_playbook(
1418            AnsiblePlaybook::UpgradeNodes,
1419            AnsibleInventoryType::SymmetricPrivateNodes,
1420            Some(options.get_ansible_vars()),
1421        ) {
1422            Ok(()) => println!("All private nodes were successfully upgraded"),
1423            Err(_) => {
1424                println!("WARNING: some nodes may not have been upgraded or restarted");
1425            }
1426        }
1427        // Don't use AnsibleInventoryType::iter_node_type() here, because the genesis node should be upgraded last
1428        match self.ansible_runner.run_playbook(
1429            AnsiblePlaybook::UpgradeNodes,
1430            AnsibleInventoryType::Genesis,
1431            Some(options.get_ansible_vars()),
1432        ) {
1433            Ok(()) => println!("The genesis nodes was successfully upgraded"),
1434            Err(_) => {
1435                println!("WARNING: the genesis node may not have been upgraded or restarted");
1436            }
1437        }
1438        Ok(())
1439    }
1440
1441    pub fn upgrade_antctl(
1442        &self,
1443        environment_name: &str,
1444        version: &Version,
1445        node_type: Option<NodeType>,
1446        custom_inventory: Option<Vec<VirtualMachine>>,
1447    ) -> Result<()> {
1448        let mut extra_vars = ExtraVarsDocBuilder::default();
1449        extra_vars.add_variable("version", &version.to_string());
1450
1451        if let Some(node_type) = node_type {
1452            println!("Running the upgrade safenode-manager playbook for {node_type:?} nodes");
1453            self.ansible_runner.run_playbook(
1454                AnsiblePlaybook::UpgradeAntctl,
1455                node_type.to_ansible_inventory_type(),
1456                Some(extra_vars.build()),
1457            )?;
1458            return Ok(());
1459        }
1460
1461        if let Some(custom_inventory) = custom_inventory {
1462            println!("Running the upgrade safenode-manager playbook with a custom inventory");
1463            generate_custom_environment_inventory(
1464                &custom_inventory,
1465                environment_name,
1466                &self.ansible_runner.working_directory_path.join("inventory"),
1467            )?;
1468            self.ansible_runner.run_playbook(
1469                AnsiblePlaybook::UpgradeAntctl,
1470                AnsibleInventoryType::Custom,
1471                Some(extra_vars.build()),
1472            )?;
1473            return Ok(());
1474        }
1475
1476        println!("Running the upgrade safenode-manager playbook for all node types");
1477        for node_inv_type in AnsibleInventoryType::iter_node_type() {
1478            self.ansible_runner.run_playbook(
1479                AnsiblePlaybook::UpgradeAntctl,
1480                node_inv_type,
1481                Some(extra_vars.build()),
1482            )?;
1483        }
1484
1485        Ok(())
1486    }
1487
1488    pub fn upgrade_nginx_config(
1489        &self,
1490        environment_name: &str,
1491        custom_inventory: Option<Vec<VirtualMachine>>,
1492    ) -> Result<()> {
1493        if let Some(custom_inventory) = custom_inventory {
1494            println!("Running the upgrade nginx config playbook with a custom inventory");
1495            generate_custom_environment_inventory(
1496                &custom_inventory,
1497                environment_name,
1498                &self.ansible_runner.working_directory_path.join("inventory"),
1499            )?;
1500            self.ansible_runner.run_playbook(
1501                AnsiblePlaybook::UpgradeNginx,
1502                AnsibleInventoryType::Custom,
1503                None,
1504            )?;
1505            return Ok(());
1506        }
1507
1508        println!("Running the upgrade nginx config playbook for peer cache nodes");
1509        self.ansible_runner.run_playbook(
1510            AnsiblePlaybook::UpgradeNginx,
1511            AnsibleInventoryType::PeerCacheNodes,
1512            None,
1513        )?;
1514        Ok(())
1515    }
1516
1517    pub fn upgrade_geoip_telegraf(&self, name: &str) -> Result<()> {
1518        self.ansible_runner.run_playbook(
1519            AnsiblePlaybook::UpgradeGeoIpTelegrafConfig,
1520            AnsibleInventoryType::PeerCacheNodes,
1521            Some(extra_vars::build_node_telegraf_upgrade(
1522                name,
1523                &NodeType::PeerCache,
1524            )?),
1525        )?;
1526        Ok(())
1527    }
1528
1529    pub fn print_ansible_run_banner(&self, s: &str) {
1530        let ansible_run_msg = "Ansible Run: ";
1531        let line = "=".repeat(s.len() + ansible_run_msg.len());
1532        println!("{line}\n{ansible_run_msg}{s}\n{line}");
1533    }
1534}