sn_testnet_deploy/ansible/
provisioning.rs

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