1  use anyhow::Result;
  2  use aws_sdk_ec2::types::Instance;
  3  use serde::{Deserialize, Serialize};
  4  
  5  use crate::opts::Opts;
  6  
  7  pub const CACHE_FILE: &str = "/tmp/blaze_ssh_cache.json";
  8  
  9  #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
 10  pub struct InstanceSet {
 11      pub instances: Vec<InstanceDetails>,
 12  }
 13  
 14  impl InstanceSet {
 15      pub fn new(instances: Vec<InstanceDetails>) -> Result<Self> {
 16          Ok(Self { instances })
 17      }
 18  
 19      async fn fetch_remote() -> Result<Self> {
 20          let aws_config = aws_config::load_from_env().await;
 21          let aws_client = aws_sdk_ec2::Client::new(&aws_config);
 22          let response = aws_client.describe_instances().send().await?;
 23          let instances = response
 24              .reservations()
 25              .iter()
 26              .flat_map(|reservation| reservation.instances())
 27              .map(InstanceDetails::from_instance)
 28              .flat_map(Result::ok)
 29              .collect::<Vec<_>>();
 30  
 31          let instance_set = InstanceSet::new(instances)?;
 32          instance_set.write()?;
 33  
 34          Ok(instance_set)
 35      }
 36  
 37      pub fn is_non_selectable(&self) -> bool {
 38          self.instances.len() == 1
 39      }
 40  
 41      pub async fn fetch(opts: &Opts) -> Result<Self> {
 42          match opts.no_cache {
 43              false => {
 44                  let cache_result = std::fs::read_to_string(CACHE_FILE);
 45                  match cache_result {
 46                      Ok(cache) => serde_json::from_str(&cache).map_err(|e| e.into()),
 47                      Err(_e) => Self::fetch_remote().await,
 48                  }
 49              }
 50              true => Self::fetch_remote().await,
 51          }
 52      }
 53  
 54      pub fn write(&self) -> Result<()> {
 55          std::fs::write(CACHE_FILE, serde_json::to_string(self)?)?;
 56  
 57          Ok(())
 58      }
 59  
 60      pub fn filter(&self, search: &Option<String>) -> Result<Self> {
 61          if search.is_none() {
 62              Self::new(self.instances.clone())
 63          } else {
 64              Self::new(
 65                  self.instances
 66                      .clone()
 67                      .into_iter()
 68                      .filter(|instance| {
 69                          instance
 70                              .instance_name
 71                              .clone()
 72                              .unwrap_or("".to_string())
 73                              .contains(search.clone().unwrap().as_str())
 74                      })
 75                      .collect::<Vec<_>>(),
 76              )
 77          }
 78      }
 79  }
 80  
 81  #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash, Default)]
 82  pub struct InstanceDetails {
 83      pub public_ip: Option<String>,
 84      pub private_ip: Option<String>,
 85      pub instance_id: Option<String>,
 86      pub instance_name: Option<String>,
 87  }
 88  
 89  impl InstanceDetails {
 90      pub fn from_instance(instance: &Instance) -> Result<Self> {
 91          let instance_clone = instance.clone();
 92          Ok(InstanceDetails {
 93              public_ip: instance_clone.public_ip_address,
 94              private_ip: instance_clone.private_ip_address,
 95              instance_id: instance_clone.instance_id,
 96              instance_name: Self::extract_instance_name(&instance.clone()),
 97          })
 98      }
 99  
100      pub fn is_empty(&self) -> bool {
101          self.public_ip.is_none()
102              && self.private_ip.is_none()
103              && self.instance_id.is_none()
104              && self.instance_name.is_none()
105      }
106  
107      pub fn display_name(&self) -> Result<String> {
108          let cloned_instance = self.clone();
109          Ok(format!(
110              "{:<32} | priv_ip: {:>16} | pub_ip: {:>16} | {:<32}",
111              Self::truncate_string(cloned_instance.instance_name.unwrap_or("None".to_string())),
112              cloned_instance.private_ip.unwrap_or("None".to_string()),
113              cloned_instance.public_ip.unwrap_or("None".to_string()),
114              cloned_instance.instance_id.unwrap_or("None".to_string())
115          ))
116      }
117  
118      fn extract_instance_name(instance: &Instance) -> Option<String> {
119          instance
120              .tags()
121              .iter()
122              .find(|tag| tag.key == Some("Name".to_string()))
123              .cloned()
124              .and_then(|tag| tag.value)
125      }
126  
127      fn truncate_string(input: String) -> String {
128          if input.len() <= 32 {
129              return input;
130          }
131  
132          format!("{}...", &input[..29])
133      }
134  }