sn_testnet_deploy/ansible/
provisioning.rs

1// Copyright (c) 2023, MaidSafe.
2// All rights reserved.
3//
4// This SAFE Network Software is licensed under the BSD-3-Clause license.
5// Please see the LICENSE file for more details.
6
7use super::{
8    extra_vars::ExtraVarsDocBuilder,
9    inventory::{
10        generate_full_cone_private_node_static_environment_inventory,
11        generate_symmetric_private_node_static_environment_inventory,
12    },
13    AnsibleInventoryType, AnsiblePlaybook, AnsibleRunner,
14};
15use crate::{
16    ansible::inventory::{
17        generate_custom_environment_inventory,
18        generate_full_cone_nat_gateway_static_environment_inventory,
19    },
20    bootstrap::BootstrapOptions,
21    clients::ClientsDeployOptions,
22    deploy::DeployOptions,
23    error::{Error, Result},
24    funding::FundingOptions,
25    inventory::{DeploymentNodeRegistries, VirtualMachine},
26    print_duration, run_external_command, BinaryOption, CloudProvider, EvmNetwork, LogFormat,
27    NodeType, SshClient, UpgradeOptions,
28};
29use ant_service_management::NodeRegistry;
30use evmlib::common::U256;
31use log::{debug, error, trace};
32use semver::Version;
33use serde::{Deserialize, Serialize};
34use std::{
35    collections::HashMap,
36    net::IpAddr,
37    path::PathBuf,
38    time::{Duration, Instant},
39};
40use walkdir::WalkDir;
41
42use crate::ansible::extra_vars;
43
44pub const DEFAULT_BETA_ENCRYPTION_KEY: &str =
45    "49113d2083f57a976076adbe85decb75115820de1e6e74b47e0429338cef124a";
46
47#[derive(Clone, Serialize, Deserialize)]
48pub struct ProvisionOptions {
49    /// The safe version is also in the binary option, but only for an initial deployment.
50    /// For the upscale, it needs to be provided explicitly, because currently it is not
51    /// recorded in the inventory.
52    pub ant_version: Option<String>,
53    pub binary_option: BinaryOption,
54    pub chunk_size: Option<u64>,
55    pub client_env_variables: Option<Vec<(String, String)>>,
56    pub enable_download_verifier: bool,
57    pub enable_random_verifier: bool,
58    pub enable_performance_verifier: bool,
59    pub enable_telegraf: bool,
60    pub enable_uploaders: bool,
61    pub evm_data_payments_address: Option<String>,
62    pub evm_network: EvmNetwork,
63    pub evm_payment_token_address: Option<String>,
64    pub evm_rpc_url: Option<String>,
65    pub expected_hash: Option<String>,
66    pub expected_size: Option<u64>,
67    pub file_address: Option<String>,
68    pub full_cone_private_node_count: u16,
69    pub funding_wallet_secret_key: Option<String>,
70    pub gas_amount: Option<U256>,
71    pub interval: Option<Duration>,
72    pub log_format: Option<LogFormat>,
73    pub max_archived_log_files: u16,
74    pub max_log_files: u16,
75    pub max_uploads: Option<u32>,
76    pub name: String,
77    pub network_id: Option<u8>,
78    pub network_dashboard_branch: Option<String>,
79    pub node_count: u16,
80    pub node_env_variables: Option<Vec<(String, String)>>,
81    pub output_inventory_dir_path: PathBuf,
82    pub peer_cache_node_count: u16,
83    pub public_rpc: bool,
84    pub rewards_address: Option<String>,
85    pub symmetric_private_node_count: u16,
86    pub token_amount: Option<U256>,
87    pub upload_size: Option<u16>,
88    pub upload_interval: Option<u16>,
89    pub uploaders_count: Option<u16>,
90    pub wallet_secret_keys: Option<Vec<String>>,
91}
92
93/// These are obtained by running the inventory playbook
94#[derive(Clone, Debug)]
95pub struct PrivateNodeProvisionInventory {
96    pub full_cone_nat_gateway_vms: Vec<VirtualMachine>,
97    pub full_cone_private_node_vms: Vec<VirtualMachine>,
98    pub symmetric_nat_gateway_vms: Vec<VirtualMachine>,
99    pub symmetric_private_node_vms: Vec<VirtualMachine>,
100}
101
102impl PrivateNodeProvisionInventory {
103    pub fn new(
104        provisioner: &AnsibleProvisioner,
105        full_cone_private_node_vm_count: Option<u16>,
106        symmetric_private_node_vm_count: Option<u16>,
107    ) -> Result<Self> {
108        // All the environment types set private_node_vm count to >0 if not specified.
109        let should_provision_full_cone_private_nodes = full_cone_private_node_vm_count
110            .map(|count| count > 0)
111            .unwrap_or(true);
112        let should_provision_symmetric_private_nodes = symmetric_private_node_vm_count
113            .map(|count| count > 0)
114            .unwrap_or(true);
115
116        let mut inventory = Self {
117            full_cone_nat_gateway_vms: Default::default(),
118            full_cone_private_node_vms: Default::default(),
119            symmetric_nat_gateway_vms: Default::default(),
120            symmetric_private_node_vms: Default::default(),
121        };
122
123        if should_provision_full_cone_private_nodes {
124            let full_cone_private_node_vms = provisioner
125                .ansible_runner
126                .get_inventory(AnsibleInventoryType::FullConePrivateNodes, true)
127                .inspect_err(|err| {
128                    println!("Failed to obtain the inventory of Full Cone private node: {err:?}");
129                })?;
130
131            let full_cone_nat_gateway_inventory = provisioner
132                .ansible_runner
133                .get_inventory(AnsibleInventoryType::FullConeNatGateway, true)
134                .inspect_err(|err| {
135                    println!("Failed to get Full Cone NAT Gateway inventory {err:?}");
136                })?;
137
138            if full_cone_nat_gateway_inventory.len() != full_cone_private_node_vms.len() {
139                println!("The number of Full Cone private nodes does not match the number of Full Cone NAT Gateway VMs");
140                return Err(Error::VmCountMismatch(
141                    Some(AnsibleInventoryType::FullConePrivateNodes),
142                    Some(AnsibleInventoryType::FullConeNatGateway),
143                ));
144            }
145
146            inventory.full_cone_private_node_vms = full_cone_private_node_vms;
147            inventory.full_cone_nat_gateway_vms = full_cone_nat_gateway_inventory;
148        }
149
150        if should_provision_symmetric_private_nodes {
151            let symmetric_private_node_vms = provisioner
152                .ansible_runner
153                .get_inventory(AnsibleInventoryType::SymmetricPrivateNodes, true)
154                .inspect_err(|err| {
155                    println!("Failed to obtain the inventory of Symmetric private node: {err:?}");
156                })?;
157
158            let symmetric_nat_gateway_inventory = provisioner
159                .ansible_runner
160                .get_inventory(AnsibleInventoryType::SymmetricNatGateway, true)
161                .inspect_err(|err| {
162                    println!("Failed to get Symmetric NAT Gateway inventory {err:?}");
163                })?;
164
165            if symmetric_nat_gateway_inventory.len() != symmetric_private_node_vms.len() {
166                println!("The number of Symmetric private nodes does not match the number of Symmetric NAT Gateway VMs");
167                return Err(Error::VmCountMismatch(
168                    Some(AnsibleInventoryType::SymmetricPrivateNodes),
169                    Some(AnsibleInventoryType::SymmetricNatGateway),
170                ));
171            }
172
173            inventory.symmetric_private_node_vms = symmetric_private_node_vms;
174            inventory.symmetric_nat_gateway_vms = symmetric_nat_gateway_inventory;
175        }
176
177        Ok(inventory)
178    }
179
180    pub fn should_provision_full_cone_private_nodes(&self) -> bool {
181        !self.full_cone_private_node_vms.is_empty()
182    }
183
184    pub fn should_provision_symmetric_private_nodes(&self) -> bool {
185        !self.symmetric_private_node_vms.is_empty()
186    }
187
188    pub fn symmetric_private_node_and_gateway_map(
189        &self,
190    ) -> Result<HashMap<VirtualMachine, VirtualMachine>> {
191        Self::match_private_node_vm_and_gateway_vm(
192            &self.symmetric_private_node_vms,
193            &self.symmetric_nat_gateway_vms,
194        )
195    }
196
197    pub fn full_cone_private_node_and_gateway_map(
198        &self,
199    ) -> Result<HashMap<VirtualMachine, VirtualMachine>> {
200        Self::match_private_node_vm_and_gateway_vm(
201            &self.full_cone_private_node_vms,
202            &self.full_cone_nat_gateway_vms,
203        )
204    }
205
206    pub fn match_private_node_vm_and_gateway_vm(
207        private_node_vms: &[VirtualMachine],
208        nat_gateway_vms: &[VirtualMachine],
209    ) -> Result<HashMap<VirtualMachine, VirtualMachine>> {
210        if private_node_vms.len() != nat_gateway_vms.len() {
211            println!(
212            "The number of private node VMs ({}) does not match the number of NAT Gateway VMs ({})",
213            private_node_vms.len(),
214            nat_gateway_vms.len()
215        );
216            error!("The number of private node VMs does not match the number of NAT Gateway VMs: Private VMs: {private_node_vms:?} Nat gateway VMs: {nat_gateway_vms:?}");
217            return Err(Error::VmCountMismatch(None, None));
218        }
219
220        let mut map = HashMap::new();
221        for private_vm in private_node_vms {
222            let nat_gateway = nat_gateway_vms
223                .iter()
224                .find(|vm| {
225                    let private_node_name = private_vm.name.split('-').next_back().unwrap();
226                    let nat_gateway_name = vm.name.split('-').next_back().unwrap();
227                    private_node_name == nat_gateway_name
228                })
229                .ok_or_else(|| {
230                    println!(
231                        "Failed to find a matching NAT Gateway for private node: {}",
232                        private_vm.name
233                    );
234                    error!("Failed to find a matching NAT Gateway for private node: {}. Private VMs: {private_node_vms:?} Nat gateway VMs: {nat_gateway_vms:?}", private_vm.name);
235                    Error::VmCountMismatch(None, None)
236                })?;
237
238            let _ = map.insert(private_vm.clone(), nat_gateway.clone());
239        }
240
241        Ok(map)
242    }
243}
244
245impl From<BootstrapOptions> for ProvisionOptions {
246    fn from(bootstrap_options: BootstrapOptions) -> Self {
247        ProvisionOptions {
248            ant_version: None,
249            binary_option: bootstrap_options.binary_option,
250            chunk_size: bootstrap_options.chunk_size,
251            client_env_variables: None,
252            enable_download_verifier: false,
253            enable_random_verifier: false,
254            enable_performance_verifier: false,
255            enable_telegraf: true,
256            enable_uploaders: false,
257            evm_data_payments_address: bootstrap_options.evm_data_payments_address,
258            evm_network: bootstrap_options.evm_network,
259            evm_payment_token_address: bootstrap_options.evm_payment_token_address,
260            evm_rpc_url: bootstrap_options.evm_rpc_url,
261            expected_hash: None,
262            expected_size: None,
263            file_address: None,
264            full_cone_private_node_count: bootstrap_options.full_cone_private_node_count,
265            funding_wallet_secret_key: None,
266            gas_amount: None,
267            interval: Some(bootstrap_options.interval),
268            log_format: bootstrap_options.log_format,
269            max_archived_log_files: bootstrap_options.max_archived_log_files,
270            max_log_files: bootstrap_options.max_log_files,
271            max_uploads: None,
272            name: bootstrap_options.name,
273            network_id: Some(bootstrap_options.network_id),
274            network_dashboard_branch: None,
275            node_count: bootstrap_options.node_count,
276            node_env_variables: bootstrap_options.node_env_variables,
277            output_inventory_dir_path: bootstrap_options.output_inventory_dir_path,
278            peer_cache_node_count: 0,
279            public_rpc: false,
280            rewards_address: Some(bootstrap_options.rewards_address),
281            symmetric_private_node_count: bootstrap_options.symmetric_private_node_count,
282            token_amount: None,
283            upload_size: None,
284            upload_interval: None,
285            uploaders_count: None,
286            wallet_secret_keys: None,
287        }
288    }
289}
290
291impl From<DeployOptions> for ProvisionOptions {
292    fn from(deploy_options: DeployOptions) -> Self {
293        ProvisionOptions {
294            ant_version: None,
295            binary_option: deploy_options.binary_option,
296            chunk_size: deploy_options.chunk_size,
297            client_env_variables: deploy_options.client_env_variables,
298            enable_download_verifier: deploy_options.enable_download_verifier,
299            enable_performance_verifier: deploy_options.enable_performance_verifier,
300            enable_random_verifier: deploy_options.enable_random_verifier,
301            enable_telegraf: deploy_options.enable_telegraf,
302            enable_uploaders: true,
303            node_env_variables: deploy_options.node_env_variables,
304            evm_data_payments_address: deploy_options.evm_data_payments_address,
305            evm_network: deploy_options.evm_network,
306            evm_payment_token_address: deploy_options.evm_payment_token_address,
307            evm_rpc_url: deploy_options.evm_rpc_url,
308            expected_hash: None,
309            expected_size: None,
310            file_address: None,
311            full_cone_private_node_count: deploy_options.full_cone_private_node_count,
312            funding_wallet_secret_key: deploy_options.funding_wallet_secret_key,
313            gas_amount: deploy_options.initial_gas,
314            interval: Some(deploy_options.interval),
315            log_format: deploy_options.log_format,
316            max_archived_log_files: deploy_options.max_archived_log_files,
317            max_log_files: deploy_options.max_log_files,
318            max_uploads: None,
319            name: deploy_options.name,
320            network_id: Some(deploy_options.network_id),
321            network_dashboard_branch: deploy_options.network_dashboard_branch,
322            node_count: deploy_options.node_count,
323            output_inventory_dir_path: deploy_options.output_inventory_dir_path,
324            peer_cache_node_count: deploy_options.peer_cache_node_count,
325            public_rpc: deploy_options.public_rpc,
326            rewards_address: Some(deploy_options.rewards_address),
327            symmetric_private_node_count: deploy_options.symmetric_private_node_count,
328            token_amount: deploy_options.initial_tokens,
329            upload_size: None,
330            upload_interval: Some(deploy_options.upload_interval),
331            uploaders_count: Some(deploy_options.uploaders_count),
332            wallet_secret_keys: None,
333        }
334    }
335}
336
337impl From<ClientsDeployOptions> for ProvisionOptions {
338    fn from(client_options: ClientsDeployOptions) -> Self {
339        Self {
340            ant_version: None,
341            binary_option: client_options.binary_option,
342            chunk_size: client_options.chunk_size,
343            client_env_variables: client_options.client_env_variables,
344            enable_download_verifier: client_options.enable_download_verifier,
345            enable_random_verifier: client_options.enable_random_verifier,
346            enable_performance_verifier: client_options.enable_performance_verifier,
347            enable_telegraf: client_options.enable_telegraf,
348            enable_uploaders: client_options.enable_uploaders,
349            evm_data_payments_address: client_options.evm_details.data_payments_address,
350            evm_network: client_options.evm_details.network,
351            evm_payment_token_address: client_options.evm_details.payment_token_address,
352            evm_rpc_url: client_options.evm_details.rpc_url,
353            expected_hash: client_options.expected_hash,
354            expected_size: client_options.expected_size,
355            file_address: client_options.file_address,
356            full_cone_private_node_count: 0,
357            funding_wallet_secret_key: client_options.funding_wallet_secret_key,
358            gas_amount: client_options.initial_gas,
359            interval: None,
360            log_format: None,
361            max_archived_log_files: client_options.max_archived_log_files,
362            max_log_files: client_options.max_log_files,
363            max_uploads: client_options.max_uploads,
364            name: client_options.name,
365            network_id: client_options.network_id,
366            network_dashboard_branch: None,
367            node_count: 0,
368            node_env_variables: None,
369            output_inventory_dir_path: client_options.output_inventory_dir_path,
370            peer_cache_node_count: 0,
371            public_rpc: false,
372            rewards_address: None,
373            symmetric_private_node_count: 0,
374            token_amount: client_options.initial_tokens,
375            upload_size: client_options.upload_size,
376            upload_interval: None,
377            uploaders_count: Some(client_options.uploaders_count),
378            wallet_secret_keys: client_options.wallet_secret_keys,
379        }
380    }
381}
382
383#[derive(Clone)]
384pub struct AnsibleProvisioner {
385    pub ansible_runner: AnsibleRunner,
386    pub cloud_provider: CloudProvider,
387    pub ssh_client: SshClient,
388}
389
390impl AnsibleProvisioner {
391    pub fn new(
392        ansible_runner: AnsibleRunner,
393        cloud_provider: CloudProvider,
394        ssh_client: SshClient,
395    ) -> Self {
396        Self {
397            ansible_runner,
398            cloud_provider,
399            ssh_client,
400        }
401    }
402
403    pub fn build_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::FullConeNatGateway,
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::FullConeNatGateway,
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_clients(
1054        &self,
1055        options: &ProvisionOptions,
1056        genesis_multiaddr: Option<String>,
1057        genesis_network_contacts_url: Option<String>,
1058    ) -> Result<()> {
1059        let start = Instant::now();
1060
1061        let sk_map = if let Some(wallet_keys) = &options.wallet_secret_keys {
1062            self.prepare_pre_funded_wallets(wallet_keys).await?
1063        } else {
1064            self.deposit_funds_to_clients(&FundingOptions {
1065                evm_data_payments_address: options.evm_data_payments_address.clone(),
1066                evm_network: options.evm_network.clone(),
1067                evm_payment_token_address: options.evm_payment_token_address.clone(),
1068                evm_rpc_url: options.evm_rpc_url.clone(),
1069                funding_wallet_secret_key: options.funding_wallet_secret_key.clone(),
1070                gas_amount: options.gas_amount,
1071                token_amount: options.token_amount,
1072                uploaders_count: options.uploaders_count,
1073            })
1074            .await?
1075        };
1076
1077        self.ansible_runner.run_playbook(
1078            AnsiblePlaybook::Uploaders,
1079            AnsibleInventoryType::Clients,
1080            Some(extra_vars::build_clients_extra_vars_doc(
1081                &self.cloud_provider.to_string(),
1082                options,
1083                genesis_multiaddr,
1084                genesis_network_contacts_url,
1085                &sk_map,
1086            )?),
1087        )?;
1088        print_duration(start.elapsed());
1089        Ok(())
1090    }
1091
1092    pub fn start_nodes(
1093        &self,
1094        environment_name: &str,
1095        interval: Duration,
1096        node_type: Option<NodeType>,
1097        custom_inventory: Option<Vec<VirtualMachine>>,
1098    ) -> Result<()> {
1099        let mut extra_vars = ExtraVarsDocBuilder::default();
1100        extra_vars.add_variable("interval", &interval.as_millis().to_string());
1101
1102        if let Some(node_type) = node_type {
1103            println!("Running the start nodes playbook for {node_type:?} nodes");
1104            self.ansible_runner.run_playbook(
1105                AnsiblePlaybook::StartNodes,
1106                node_type.to_ansible_inventory_type(),
1107                Some(extra_vars.build()),
1108            )?;
1109            return Ok(());
1110        }
1111
1112        if let Some(custom_inventory) = custom_inventory {
1113            println!("Running the start nodes playbook with a custom inventory");
1114            generate_custom_environment_inventory(
1115                &custom_inventory,
1116                environment_name,
1117                &self.ansible_runner.working_directory_path.join("inventory"),
1118            )?;
1119            self.ansible_runner.run_playbook(
1120                AnsiblePlaybook::StartNodes,
1121                AnsibleInventoryType::Custom,
1122                Some(extra_vars.build()),
1123            )?;
1124            return Ok(());
1125        }
1126
1127        println!("Running the start nodes playbook for all node types");
1128        for node_inv_type in AnsibleInventoryType::iter_node_type() {
1129            self.ansible_runner.run_playbook(
1130                AnsiblePlaybook::StartNodes,
1131                node_inv_type,
1132                Some(extra_vars.build()),
1133            )?;
1134        }
1135        Ok(())
1136    }
1137
1138    pub fn status(&self) -> Result<()> {
1139        for node_inv_type in AnsibleInventoryType::iter_node_type() {
1140            self.ansible_runner
1141                .run_playbook(AnsiblePlaybook::Status, node_inv_type, None)?;
1142        }
1143        Ok(())
1144    }
1145
1146    pub fn start_telegraf(
1147        &self,
1148        environment_name: &str,
1149        node_type: Option<NodeType>,
1150        custom_inventory: Option<Vec<VirtualMachine>>,
1151    ) -> Result<()> {
1152        if let Some(node_type) = node_type {
1153            println!("Running the start telegraf playbook for {node_type:?} nodes");
1154            self.ansible_runner.run_playbook(
1155                AnsiblePlaybook::StartTelegraf,
1156                node_type.to_ansible_inventory_type(),
1157                None,
1158            )?;
1159            return Ok(());
1160        }
1161
1162        if let Some(custom_inventory) = custom_inventory {
1163            println!("Running the start telegraf playbook with a custom inventory");
1164            generate_custom_environment_inventory(
1165                &custom_inventory,
1166                environment_name,
1167                &self.ansible_runner.working_directory_path.join("inventory"),
1168            )?;
1169            self.ansible_runner.run_playbook(
1170                AnsiblePlaybook::StartTelegraf,
1171                AnsibleInventoryType::Custom,
1172                None,
1173            )?;
1174            return Ok(());
1175        }
1176
1177        println!("Running the start telegraf playbook for all node types");
1178        for node_inv_type in AnsibleInventoryType::iter_node_type() {
1179            self.ansible_runner.run_playbook(
1180                AnsiblePlaybook::StartTelegraf,
1181                node_inv_type,
1182                None,
1183            )?;
1184        }
1185
1186        Ok(())
1187    }
1188
1189    pub fn stop_nodes(
1190        &self,
1191        environment_name: &str,
1192        interval: Duration,
1193        node_type: Option<NodeType>,
1194        custom_inventory: Option<Vec<VirtualMachine>>,
1195        delay: Option<u64>,
1196        service_names: Option<Vec<String>>,
1197    ) -> Result<()> {
1198        let mut extra_vars = ExtraVarsDocBuilder::default();
1199        extra_vars.add_variable("interval", &interval.as_millis().to_string());
1200        if let Some(delay) = delay {
1201            extra_vars.add_variable("delay", &delay.to_string());
1202        }
1203        if let Some(service_names) = service_names {
1204            extra_vars.add_list_variable("service_names", service_names);
1205        }
1206        let extra_vars = extra_vars.build();
1207
1208        if let Some(node_type) = node_type {
1209            println!("Running the stop nodes playbook for {node_type:?} nodes");
1210            self.ansible_runner.run_playbook(
1211                AnsiblePlaybook::StopNodes,
1212                node_type.to_ansible_inventory_type(),
1213                Some(extra_vars),
1214            )?;
1215            return Ok(());
1216        }
1217
1218        if let Some(custom_inventory) = custom_inventory {
1219            println!("Running the stop nodes playbook with a custom inventory");
1220            generate_custom_environment_inventory(
1221                &custom_inventory,
1222                environment_name,
1223                &self.ansible_runner.working_directory_path.join("inventory"),
1224            )?;
1225            self.ansible_runner.run_playbook(
1226                AnsiblePlaybook::StopNodes,
1227                AnsibleInventoryType::Custom,
1228                Some(extra_vars),
1229            )?;
1230            return Ok(());
1231        }
1232
1233        println!("Running the stop nodes playbook for all node types");
1234        for node_inv_type in AnsibleInventoryType::iter_node_type() {
1235            self.ansible_runner.run_playbook(
1236                AnsiblePlaybook::StopNodes,
1237                node_inv_type,
1238                Some(extra_vars.clone()),
1239            )?;
1240        }
1241
1242        Ok(())
1243    }
1244
1245    pub fn stop_telegraf(
1246        &self,
1247        environment_name: &str,
1248        node_type: Option<NodeType>,
1249        custom_inventory: Option<Vec<VirtualMachine>>,
1250    ) -> Result<()> {
1251        if let Some(node_type) = node_type {
1252            println!("Running the stop telegraf playbook for {node_type:?} nodes");
1253            self.ansible_runner.run_playbook(
1254                AnsiblePlaybook::StopTelegraf,
1255                node_type.to_ansible_inventory_type(),
1256                None,
1257            )?;
1258            return Ok(());
1259        }
1260
1261        if let Some(custom_inventory) = custom_inventory {
1262            println!("Running the stop telegraf playbook with a custom inventory");
1263            generate_custom_environment_inventory(
1264                &custom_inventory,
1265                environment_name,
1266                &self.ansible_runner.working_directory_path.join("inventory"),
1267            )?;
1268            self.ansible_runner.run_playbook(
1269                AnsiblePlaybook::StopTelegraf,
1270                AnsibleInventoryType::Custom,
1271                None,
1272            )?;
1273            return Ok(());
1274        }
1275
1276        println!("Running the stop telegraf playbook for all node types");
1277        for node_inv_type in AnsibleInventoryType::iter_node_type() {
1278            self.ansible_runner
1279                .run_playbook(AnsiblePlaybook::StopTelegraf, node_inv_type, None)?;
1280        }
1281
1282        Ok(())
1283    }
1284
1285    pub fn upgrade_node_telegraf(&self, name: &str) -> Result<()> {
1286        self.ansible_runner.run_playbook(
1287            AnsiblePlaybook::UpgradeNodeTelegrafConfig,
1288            AnsibleInventoryType::PeerCacheNodes,
1289            Some(extra_vars::build_node_telegraf_upgrade(
1290                name,
1291                &NodeType::PeerCache,
1292            )?),
1293        )?;
1294        self.ansible_runner.run_playbook(
1295            AnsiblePlaybook::UpgradeNodeTelegrafConfig,
1296            AnsibleInventoryType::Nodes,
1297            Some(extra_vars::build_node_telegraf_upgrade(
1298                name,
1299                &NodeType::Generic,
1300            )?),
1301        )?;
1302
1303        self.ansible_runner.run_playbook(
1304            AnsiblePlaybook::UpgradeNodeTelegrafConfig,
1305            AnsibleInventoryType::SymmetricPrivateNodes,
1306            Some(extra_vars::build_node_telegraf_upgrade(
1307                name,
1308                &NodeType::SymmetricPrivateNode,
1309            )?),
1310        )?;
1311
1312        self.ansible_runner.run_playbook(
1313            AnsiblePlaybook::UpgradeNodeTelegrafConfig,
1314            AnsibleInventoryType::FullConePrivateNodes,
1315            Some(extra_vars::build_node_telegraf_upgrade(
1316                name,
1317                &NodeType::FullConePrivateNode,
1318            )?),
1319        )?;
1320        Ok(())
1321    }
1322
1323    pub fn upgrade_client_telegraf(&self, name: &str) -> Result<()> {
1324        self.ansible_runner.run_playbook(
1325            AnsiblePlaybook::UpgradeClientTelegrafConfig,
1326            AnsibleInventoryType::Clients,
1327            Some(extra_vars::build_client_telegraf_upgrade(name)?),
1328        )?;
1329        Ok(())
1330    }
1331
1332    pub fn upgrade_nodes(&self, options: &UpgradeOptions) -> Result<()> {
1333        if let Some(custom_inventory) = &options.custom_inventory {
1334            println!("Running the UpgradeNodes with a custom inventory");
1335            generate_custom_environment_inventory(
1336                custom_inventory,
1337                &options.name,
1338                &self.ansible_runner.working_directory_path.join("inventory"),
1339            )?;
1340            match self.ansible_runner.run_playbook(
1341                AnsiblePlaybook::UpgradeNodes,
1342                AnsibleInventoryType::Custom,
1343                Some(options.get_ansible_vars()),
1344            ) {
1345                Ok(()) => println!("All nodes were successfully upgraded"),
1346                Err(_) => {
1347                    println!("WARNING: some nodes may not have been upgraded or restarted");
1348                }
1349            }
1350            return Ok(());
1351        }
1352
1353        if let Some(node_type) = &options.node_type {
1354            println!("Running the UpgradeNodes playbook for {node_type:?} nodes");
1355            match self.ansible_runner.run_playbook(
1356                AnsiblePlaybook::UpgradeNodes,
1357                node_type.to_ansible_inventory_type(),
1358                Some(options.get_ansible_vars()),
1359            ) {
1360                Ok(()) => println!("All {node_type:?} nodes were successfully upgraded"),
1361                Err(_) => {
1362                    println!(
1363                        "WARNING: some {node_type:?} nodes may not have been upgraded or restarted"
1364                    );
1365                }
1366            }
1367            return Ok(());
1368        }
1369
1370        println!("Running the UpgradeNodes playbook for all node types");
1371
1372        match self.ansible_runner.run_playbook(
1373            AnsiblePlaybook::UpgradeNodes,
1374            AnsibleInventoryType::PeerCacheNodes,
1375            Some(options.get_ansible_vars()),
1376        ) {
1377            Ok(()) => println!("All Peer Cache nodes were successfully upgraded"),
1378            Err(_) => {
1379                println!("WARNING: some Peer Cacche nodes may not have been upgraded or restarted");
1380            }
1381        }
1382        match self.ansible_runner.run_playbook(
1383            AnsiblePlaybook::UpgradeNodes,
1384            AnsibleInventoryType::Nodes,
1385            Some(options.get_ansible_vars()),
1386        ) {
1387            Ok(()) => println!("All generic nodes were successfully upgraded"),
1388            Err(_) => {
1389                println!("WARNING: some nodes may not have been upgraded or restarted");
1390            }
1391        }
1392        match self.ansible_runner.run_playbook(
1393            AnsiblePlaybook::UpgradeNodes,
1394            AnsibleInventoryType::SymmetricPrivateNodes,
1395            Some(options.get_ansible_vars()),
1396        ) {
1397            Ok(()) => println!("All private nodes were successfully upgraded"),
1398            Err(_) => {
1399                println!("WARNING: some nodes may not have been upgraded or restarted");
1400            }
1401        }
1402        // Don't use AnsibleInventoryType::iter_node_type() here, because the genesis node should be upgraded last
1403        match self.ansible_runner.run_playbook(
1404            AnsiblePlaybook::UpgradeNodes,
1405            AnsibleInventoryType::Genesis,
1406            Some(options.get_ansible_vars()),
1407        ) {
1408            Ok(()) => println!("The genesis nodes was successfully upgraded"),
1409            Err(_) => {
1410                println!("WARNING: the genesis node may not have been upgraded or restarted");
1411            }
1412        }
1413        Ok(())
1414    }
1415
1416    pub fn upgrade_antctl(
1417        &self,
1418        environment_name: &str,
1419        version: &Version,
1420        node_type: Option<NodeType>,
1421        custom_inventory: Option<Vec<VirtualMachine>>,
1422    ) -> Result<()> {
1423        let mut extra_vars = ExtraVarsDocBuilder::default();
1424        extra_vars.add_variable("version", &version.to_string());
1425
1426        if let Some(node_type) = node_type {
1427            println!("Running the upgrade safenode-manager playbook for {node_type:?} nodes");
1428            self.ansible_runner.run_playbook(
1429                AnsiblePlaybook::UpgradeAntctl,
1430                node_type.to_ansible_inventory_type(),
1431                Some(extra_vars.build()),
1432            )?;
1433            return Ok(());
1434        }
1435
1436        if let Some(custom_inventory) = custom_inventory {
1437            println!("Running the upgrade safenode-manager playbook with a custom inventory");
1438            generate_custom_environment_inventory(
1439                &custom_inventory,
1440                environment_name,
1441                &self.ansible_runner.working_directory_path.join("inventory"),
1442            )?;
1443            self.ansible_runner.run_playbook(
1444                AnsiblePlaybook::UpgradeAntctl,
1445                AnsibleInventoryType::Custom,
1446                Some(extra_vars.build()),
1447            )?;
1448            return Ok(());
1449        }
1450
1451        println!("Running the upgrade safenode-manager playbook for all node types");
1452        for node_inv_type in AnsibleInventoryType::iter_node_type() {
1453            self.ansible_runner.run_playbook(
1454                AnsiblePlaybook::UpgradeAntctl,
1455                node_inv_type,
1456                Some(extra_vars.build()),
1457            )?;
1458        }
1459
1460        Ok(())
1461    }
1462
1463    pub fn upgrade_nginx_config(
1464        &self,
1465        environment_name: &str,
1466        custom_inventory: Option<Vec<VirtualMachine>>,
1467    ) -> Result<()> {
1468        if let Some(custom_inventory) = custom_inventory {
1469            println!("Running the upgrade nginx config playbook with a custom inventory");
1470            generate_custom_environment_inventory(
1471                &custom_inventory,
1472                environment_name,
1473                &self.ansible_runner.working_directory_path.join("inventory"),
1474            )?;
1475            self.ansible_runner.run_playbook(
1476                AnsiblePlaybook::UpgradeNginx,
1477                AnsibleInventoryType::Custom,
1478                None,
1479            )?;
1480            return Ok(());
1481        }
1482
1483        println!("Running the upgrade nginx config playbook for peer cache nodes");
1484        self.ansible_runner.run_playbook(
1485            AnsiblePlaybook::UpgradeNginx,
1486            AnsibleInventoryType::PeerCacheNodes,
1487            None,
1488        )?;
1489        Ok(())
1490    }
1491
1492    pub fn upgrade_geoip_telegraf(&self, name: &str) -> Result<()> {
1493        self.ansible_runner.run_playbook(
1494            AnsiblePlaybook::UpgradeGeoIpTelegrafConfig,
1495            AnsibleInventoryType::PeerCacheNodes,
1496            Some(extra_vars::build_node_telegraf_upgrade(
1497                name,
1498                &NodeType::PeerCache,
1499            )?),
1500        )?;
1501        Ok(())
1502    }
1503
1504    pub fn print_ansible_run_banner(&self, s: &str) {
1505        let ansible_run_msg = "Ansible Run: ";
1506        let line = "=".repeat(s.len() + ansible_run_msg.len());
1507        println!("{line}\n{ansible_run_msg}{s}\n{line}");
1508    }
1509}