Merge pull request #370 from stakwork/feat/swarm-tag

Feat/swarm tag

src/bin/super/aws_util.rs

@@ -0,0 +1,18 @@
+use anyhow::Error;
+use aws_config::meta::region::RegionProviderChain;
+use aws_config::Region;
+use aws_sdk_ec2::Client;
+use aws_smithy_types::retry::RetryConfig;
+use sphinx_swarm::utils::getenv;
+
+pub async fn make_aws_client() -> Result<Client, Error> {
+    let region = getenv("AWS_S3_REGION_NAME")?;
+    let region_provider = RegionProviderChain::first_try(Some(Region::new(region)));
+    let config = aws_config::from_env()
+        .region(region_provider)
+        .retry_config(RetryConfig::standard().with_max_attempts(10))
+        .load()
+        .await;
+
+    Ok(Client::new(&config))
+}

src/bin/super/ec2.rs

@@ -0,0 +1,45 @@
+use crate::{aws_util::make_aws_client, state::InstanceFromAws};
+use anyhow::{anyhow, Error};
+use aws_sdk_ec2::types::Filter;
+
+pub async fn get_swarms_by_tag(key: &str, value: &str) -> Result<Vec<InstanceFromAws>, Error> {
+    let client = make_aws_client().await?;
+
+    let mut instances: Vec<InstanceFromAws> = vec![];
+
+    let tag_filter = Filter::builder()
+        .name(format!("tag:{}", key))
+        .values(format!("{}", value))
+        .build();
+
+    let response = client
+        .describe_instances()
+        .filters(tag_filter)
+        .send()
+        .await?;
+
+    if response.reservations().is_empty() {
+        log::error!("No instances found with the given tag.");
+        return Err(anyhow!("No instances found with the given tag."));
+    }
+
+    for reservation in response.reservations.unwrap() {
+        if !reservation.instances().is_empty() {
+            for instance in reservation.instances.unwrap() {
+                if instance.public_ip_address.is_some()
+                    && instance.instance_id.is_some()
+                    && instance.instance_type.is_some()
+                {
+                    instances.push(InstanceFromAws {
+                        instacne_id: instance.instance_id.unwrap(),
+                        intance_type: instance.instance_type.unwrap().to_string(),
+                    });
+                }
+            }
+        } else {
+            log::error!("Instances do not exist")
+        }
+    }
+
+    return Ok(instances);
+}

src/bin/super/mod.rs

@@ -1,6 +1,8 @@
 mod auth_token;
+mod aws_util;
 mod checker;
 mod cmd;
+mod ec2;
 mod route53;
 mod routes;
 mod state;
@@ -13,7 +15,7 @@ use state::RemoteStack;
 use state::Super;
 use util::{
     accessing_child_container_controller, add_new_swarm_details, add_new_swarm_from_child_swarm,
-    get_aws_instance_types, get_child_swarm_config, get_child_swarm_containers,
+    get_aws_instance_types, get_child_swarm_config, get_child_swarm_containers, get_config,
     get_swarm_instance_type, update_aws_instance_type,
 };
 
@@ -133,7 +135,8 @@ pub async fn super_handle(
     let ret = match cmd {
         Cmd::Swarm(swarm_cmd) => match swarm_cmd {
             SwarmCmd::GetConfig => {
-                let res = &state.remove_tokens();
+                let res = get_config(&mut state).await?;
+                must_save_stack = true;
                 Some(serde_json::to_string(&res)?)
             }
             SwarmCmd::Login(ld) => match state.users.iter().find(|u| u.username == ld.username) {

src/bin/super/state.rs

@@ -39,6 +39,12 @@ pub struct BotCred {
     pub bot_url: String,
 }
 
+#[derive(Serialize, Deserialize, Debug, Eq, PartialEq, Default, Clone)]
+pub struct InstanceFromAws {
+    pub instacne_id: String,
+    pub intance_type: String,
+}
+
 impl Default for Super {
     fn default() -> Self {
         Self {

src/bin/super/superapp/src/Remotes.svelte

@@ -54,7 +54,7 @@
   let swarm_id = "";
   let delete_host = "";
   let errorMessage = false;
-  let loading = false;
+  let loading = true;
   let errors = [];
   let name = "";
   let vanity_address = "";
@@ -129,6 +129,8 @@
 
     await getConfig();
 
+    loading = false;
+
     await getConfigSortByUnhealthy();
   });
 

src/bin/super/util.rs

@@ -6,7 +6,8 @@ use aws_config::meta::region::RegionProviderChain;
 use aws_config::Region;
 use aws_sdk_ec2::client::Waiters;
 use aws_sdk_ec2::types::{
-    AttributeValue, BlockDeviceMapping, EbsBlockDevice, InstanceType, Tag, TagSpecification,
+    AttributeBooleanValue, AttributeValue, BlockDeviceMapping, EbsBlockDevice, InstanceType, Tag,
+    TagSpecification,
 };
 use aws_sdk_ec2::Client;
 use aws_smithy_types::retry::RetryConfig;
@@ -17,12 +18,14 @@ use sphinx_swarm::cmd::{send_cmd_request, Cmd, LoginInfo, SwarmCmd, UpdateNode};
 use sphinx_swarm::config::Stack;
 use sphinx_swarm::utils::{getenv, make_reqwest_client};
 
+use crate::aws_util::make_aws_client;
 use crate::cmd::{
     AccessNodesInfo, AddSwarmResponse, CreateEc2InstanceInfo, GetInstanceTypeByInstanceId,
     GetInstanceTypeRes, LoginResponse, SuperSwarmResponse, UpdateInstanceDetails,
 };
+use crate::ec2::get_swarms_by_tag;
 use crate::route53::add_domain_name_to_route53;
-use crate::state::{AwsInstanceType, RemoteStack, Super};
+use crate::state::{AwsInstanceType, InstanceFromAws, RemoteStack, Super};
 use rand::Rng;
 use tokio::time::{sleep, Duration};
 
@@ -385,6 +388,10 @@ async fn create_ec2_instance(
 
     let custom_domain = vanity_address.unwrap_or_else(|| String::from(""));
 
+    let key = getenv("SWARM_TAG_KEY")?;
+
+    let value = getenv("SWARM_TAG_VALUE")?;
+
     // Load the AWS configuration
     let config = aws_config::from_env()
         .region(region_provider)
@@ -477,15 +484,15 @@ async fn create_ec2_instance(
       "#
     );
 
-    let tag = Tag::builder()
-        .key("Name")
-        .value(swarm_name) // Replace with the desired instance name
-        .build();
+    let tags = vec![
+        Tag::builder().key("Name").value(swarm_name).build(),
+        Tag::builder().key(key).value(value).build(),
+    ];
 
     // Define the TagSpecification to apply the tags when the instance is created
     let tag_specification = TagSpecification::builder()
-        .resource_type("instance".into()) // Tag the instance
-        .tags(tag)
+        .resource_type("instance".into())
+        .set_tags(Some(tags))
         .build();
 
     let block_device = BlockDeviceMapping::builder()
@@ -510,6 +517,7 @@ async fn create_ec2_instance(
         .block_device_mappings(block_device)
         .tag_specifications(tag_specification)
         .subnet_id(subnet_id)
+        .disable_api_termination(true)
         .send()
         .map_err(|err| {
             log::error!("Error Creating instance instance: {}", err);
@@ -522,7 +530,27 @@ async fn create_ec2_instance(
     }
 
     let instance_id: String = result.instances()[0].instance_id().unwrap().to_string();
-    println!("Created instance with ID: {}", instance_id);
+    log::info!("Created instance with ID: {}", instance_id);
+
+    client
+        .modify_instance_attribute()
+        .instance_id(instance_id.clone())
+        .disable_api_termination(
+            AttributeBooleanValue::builder()
+                .set_value(Some(true))
+                .build(),
+        )
+        .send()
+        .await
+        .map_err(|err| {
+            log::error!("Error enabling termination protection: {}", err);
+            anyhow::anyhow!(err.to_string())
+        })?;
+
+    log::info!(
+        "Instance {} created and termination protection enabled.",
+        instance_id
+    );
 
     Ok((instance_id, swarm_number))
 }
@@ -805,18 +833,6 @@ pub async fn update_ec2_instance_type(
     Ok(())
 }
 
-async fn make_aws_client() -> Result<Client, Error> {
-    let region = getenv("AWS_S3_REGION_NAME")?;
-    let region_provider = RegionProviderChain::first_try(Some(Region::new(region)));
-    let config = aws_config::from_env()
-        .region(region_provider)
-        .retry_config(RetryConfig::standard().with_max_attempts(10))
-        .load()
-        .await;
-
-    Ok(Client::new(&config))
-}
-
 pub fn get_swarm_instance_type(
     info: GetInstanceTypeByInstanceId,
     state: &Super,
@@ -859,3 +875,30 @@ fn get_instance(instance_type: &str) -> Option<AwsInstanceType> {
 
     return Some(instance_types[postion.unwrap()].clone());
 }
+
+pub async fn get_config(state: &mut Super) -> Result<Super, Error> {
+    let key = getenv("SWARM_TAG_KEY")?;
+    let value = getenv("SWARM_TAG_VALUE")?;
+    let aws_instances = get_swarms_by_tag(&key, &value).await?;
+
+    let mut aws_instances_hashmap: HashMap<String, InstanceFromAws> = HashMap::new();
+
+    for aws_instance in aws_instances {
+        aws_instances_hashmap.insert(aws_instance.instacne_id.clone(), aws_instance.clone());
+    }
+
+    for stack in state.stacks.iter_mut() {
+        if aws_instances_hashmap.contains_key(&stack.ec2_instance_id) {
+            let aws_instance_hashmap = aws_instances_hashmap.get(&stack.ec2_instance_id).unwrap();
+            if stack.ec2.is_none() {
+                stack.ec2 = Some(aws_instance_hashmap.intance_type.clone());
+            } else {
+                if aws_instance_hashmap.intance_type != stack.ec2.clone().unwrap() {
+                    stack.ec2 = Some(aws_instance_hashmap.intance_type.clone())
+                }
+            }
+        }
+    }
+    let res = state.remove_tokens();
+    Ok(res)
+}

superadmin.yml

@@ -90,6 +90,9 @@ services:
       - AWS_SECURITY_GROUP_ID=$AWS_SECURITY_GROUP_ID
       - AWS_KEY_NAME=$AWS_KEY_NAME
       - AWS_SUBNET_ID=$AWS_SUBNET_ID
+      - SWARM_TAG_VALUE=$SWARM_TAG_VALUE
+      - SWARM_TAG_KEY=$SWARM_TAG_KEY
+      - SWARM_UPDATER_PASSWORD=$SWARM_UPDATER_PASSWORD
 
 networks:
   sphinx-swarm: