fix: added SSM support also

Author: Sandesh Grangdan

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "e2s"
-version = "0.1.2"
+version = "0.1.3"
 edition = "2021"
 
 # Github Repo

Cargo.toml

@@ -3,7 +3,7 @@ use clap::Parser;
 use rand::Rng;
 
 use crate::app::{
-    aws::ec2::Ec2Client,
+    aws::ec2::{ConnectMode, Ec2Client},
     input::{ssh_keys::SshKeys, ssh_user::SshUsers},
 };
 use input::user_input::{
@@ -49,7 +49,8 @@ pub struct App {
     pub ec2_client: Ec2Client,
     pub ssh_keys: SshKeys,
     pub ssh_user: SshUsers,
-    pub private: bool,
+    pub connect_mode: ConnectMode,
+
 }
 
 // ANCHOR: application_impl
@@ -70,7 +71,7 @@ impl App {
             ec2_client: Ec2Client::None,
             ssh_keys: SshKeys::load(),
             ssh_user: SshUsers::load(),
-            private: true,
+            connect_mode: ConnectMode::Public,
         }
     }
 

src/app.rs

@@ -45,6 +45,31 @@ impl Data {
     }
 }
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum ConnectMode {
+    Public,
+    Private,
+    Ssm,
+}
+
+impl ConnectMode {
+    pub fn next(self) -> Self {
+        match self {
+            ConnectMode::Public => ConnectMode::Private,
+            ConnectMode::Private => ConnectMode::Ssm,
+            ConnectMode::Ssm => ConnectMode::Public,
+        }
+    }
+
+    pub fn toggle(&mut self) {
+        *self = match *self {
+            ConnectMode::Public => ConnectMode::Private,
+            ConnectMode::Private => ConnectMode::Ssm,
+            ConnectMode::Ssm => ConnectMode::Public,
+        };
+    }
+}
+
 fn generate_instance_id() -> String {
     let hex: String = (0..17)
         .map(|_| format!("{:x}", rand::random::<u8>() % 16))
@@ -253,30 +278,52 @@ impl App {
     pub async fn ssh(&mut self) -> io::Result<()> {
         if let Some(selected) = self.state.selected() {
             if let Some(item) = self.display_items.get(selected) {
-                let key_path = match &self.ssh_keys.selected_key {
-                    Some(key) => key,
-                    None => {
-                        println!("No SSH key selected.");
-                        return Ok(());
+                let mut cmd = match self.connect_mode {
+                    ConnectMode::Public | ConnectMode::Private => {
+                        let key_path = match &self.ssh_keys.selected_key {
+                            Some(key) => key,
+                            None => {
+                                eprintln!("No SSH key selected.");
+                                return Ok(());
+                            }
+                        };
+                        let user = self
+                            .ssh_user
+                            .selected_user
+                            .as_deref()
+                            .unwrap_or("ec2-user");
+                        let ip = match self.connect_mode {
+                            ConnectMode::Public => &item.public_ipv4,
+                            ConnectMode::Private => &item.private_ipv4,
+                            _ => unreachable!(),  
+                        };
+                        let mut ssh_cmd = Command::new("ssh");
+                        ssh_cmd.args(["-i", key_path]);
+                        ssh_cmd.arg(format!("{}@{}", user, ip));
+                        ssh_cmd
                     }
-                };
+                    ConnectMode::Ssm => {
+                        let mut ssm_cmd = Command::new("aws");
+                        ssm_cmd.args([
+                            "ssm",
+                            "start-session",
+                            "--target",
+                            &item.instance_id,
+                        ]);
 
-                let user = self.ssh_user.selected_user.as_deref().unwrap_or("ec2-user");
-                let ip = if self.private {
-                    &item.private_ipv4
-                } else {
-                    &item.public_ipv4
+                        if self.args.region != *"None" {
+                            ssm_cmd.args(["--region", &self.args.region]);
+                        } 
+
+                        if self.args.profile != *"None" {
+                            ssm_cmd.args(["--profile", &self.args.profile]);
+                        } 
+                        ssm_cmd
+                    }
                 };
 
-                let status = Command::new("ssh")
-                .arg("-i")
-                .arg(key_path)
-                .arg(format!("{}@{}", user, ip))
-                .stdin(std::process::Stdio::inherit())
-                .stdout(std::process::Stdio::inherit())
-                .stderr(std::process::Stdio::inherit())
-                .status()?;
-               
+                let status = cmd.status()?;
+
                 if !status.success() {
                     eprintln!("Failed to launch SSH session.");
                 }

src/app/aws/ec2.rs

@@ -26,15 +26,15 @@ use update::update;
 // ANCHOR: main
 #[tokio::main]
 async fn main() -> Result<()> {
-    use std::time::Instant;
+    // use std::time::Instant;
     // Create an application.
     let mut app = App::new(app::Args::parse());
     app.set_ec2_client().await;
 
-    let start = Instant::now();
+    // let start = Instant::now();
     app.fetch_ec2_data().await;
-    let duration = start.elapsed();
-    println!("fetch_ec2_data took: {:?}", duration);
+    // let duration = start.elapsed();
+    // println!("fetch_ec2_data took: {:?}", duration);
 
     // Initialize the terminal user interface.
     let terminal = ratatui::init();

src/main.rs

@@ -4,7 +4,7 @@ use ratatui::{
     widgets::{Block, BorderType, Borders, Padding, Paragraph, Wrap},
 };
 
-use crate::app::App;
+use crate::app::{App, aws::ec2::ConnectMode};
 
 pub fn render(app: &mut App, f: &mut Frame, layout: Rect) {
     let title_block = Block::default()
@@ -13,7 +13,11 @@ pub fn render(app: &mut App, f: &mut Frame, layout: Rect) {
         .padding(Padding::new(2, 1, 0, 0))
         .border_type(BorderType::Plain);
 
-    let ssh_from_private = if app.private { "Private" } else { "Public" };
+    let ssh_from_private = match app.connect_mode {
+        ConnectMode::Private => "Private",
+        ConnectMode::Public => "Public",
+        ConnectMode::Ssm => "SSM",
+    };
 
     let selected_ssh_key = if Some(app.ssh_keys.selected_key.clone().unwrap_or_default()) != None {
         app.ssh_keys.selected_key.clone().unwrap_or_default()

src/ui/ui_block/header.rs

@@ -47,7 +47,7 @@ pub async fn update(app: &mut App, key_event: KeyEvent, tui: &mut Tui) {
                 app.ssh_user.previous();
             }
             KeyCode::Char('p') => {
-                app.private = !app.private;
+                app.connect_mode.toggle();
             }
             KeyCode::Char('?') => {
                 app.show_help = !app.show_help;