sn_testnet_deploy/ansible/
provisioning.rs

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