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