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