From c1235449734a4f74761ab1d6de51083e31831915 Mon Sep 17 00:00:00 2001 From: Kyrylo Kulyhin Date: Tue, 1 Oct 2024 22:31:24 +0300 Subject: [PATCH] new interactive terminal to access history, use nano, vim (#22) --- Cargo.lock | 241 +++++++++++++++++++++++++++++++++++++++++++++++++++- Cargo.toml | 4 + src/main.rs | 182 +++++++++++++++++++++------------------ src/pty.rs | 101 ++++++++++++++++++++++ 4 files changed, 445 insertions(+), 83 deletions(-) create mode 100644 src/pty.rs diff --git a/Cargo.lock b/Cargo.lock index d9d7d28..9307f39 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -66,6 +66,12 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "anyhow" +version = "1.0.89" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "86fdf8605db99b54d3cd748a44c6d04df638eb5dafb219b135d0149bd0db01f6" + [[package]] name = "autocfg" version = "1.4.0" @@ -426,6 +432,12 @@ dependencies = [ "vsimd", ] +[[package]] +name = "bitflags" +version = "1.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" + [[package]] name = "bitflags" version = "2.6.0" @@ -472,6 +484,12 @@ version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" +[[package]] +name = "cfg_aliases" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" + [[package]] name = "clap" version = "4.5.18" @@ -573,13 +591,23 @@ dependencies = [ "subtle", ] +[[package]] +name = "downcast-rs" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "75b325c5dbd37f80359721ad39aca5a29fb04c89279657cffdda8736d0c0b9d2" + [[package]] name = "ecs-exec" version = "0.1.0" dependencies = [ "aws-config", "aws-sdk-ecs", + "aws-sdk-sts", "clap", + "nix 0.29.0", + "portable-pty", + "termios 0.3.3", "tokio", ] @@ -601,6 +629,17 @@ version = "2.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e8c02a5121d4ea3eb16a80748c74f5549a5665e4c21333c6098f283870fbdea6" +[[package]] +name = "filedescriptor" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7199d965852c3bac31f779ef99cbb4537f80e952e2d6aa0ffeb30cce00f4f46e" +dependencies = [ + "libc", + "thiserror", + "winapi", +] + [[package]] name = "fnv" version = "1.0.7" @@ -862,6 +901,15 @@ dependencies = [ "hashbrown", ] +[[package]] +name = "ioctl-rs" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f7970510895cee30b3e9128319f2cefd4bde883a39f38baa279567ba3a7eb97d" +dependencies = [ + "libc", +] + [[package]] name = "is_terminal_polyfill" version = "1.70.1" @@ -874,6 +922,12 @@ version = "1.0.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "49f1f14873335454500d59611f1cf4a4b0f786f9ac11f4312a78e4cf2566695b" +[[package]] +name = "lazy_static" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" + [[package]] name = "libc" version = "0.2.159" @@ -902,6 +956,15 @@ version = "2.7.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3" +[[package]] +name = "memoffset" +version = "0.6.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5aa361d4faea93603064a027415f07bd8e1d5c88c9fbf68bf56a285428fd79ce" +dependencies = [ + "autocfg", +] + [[package]] name = "miniz_oxide" version = "0.8.0" @@ -923,6 +986,32 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "nix" +version = "0.25.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f346ff70e7dbfd675fe90590b92d59ef2de15a8779ae305ebcbfd3f0caf59be4" +dependencies = [ + "autocfg", + "bitflags 1.3.2", + "cfg-if", + "libc", + "memoffset", + "pin-utils", +] + +[[package]] +name = "nix" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "71e2746dc3a24dd78b3cfcb7be93368c6de9963d30f43a6a73998a9cf4b17b46" +dependencies = [ + "bitflags 2.6.0", + "cfg-if", + "cfg_aliases", + "libc", +] + [[package]] name = "num-conv" version = "0.1.0" @@ -1024,6 +1113,27 @@ version = "1.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cc9c68a3f6da06753e9335d63e27f6b9754dd1920d941135b7ea8224f141adb2" +[[package]] +name = "portable-pty" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "806ee80c2a03dbe1a9fb9534f8d19e4c0546b790cde8fd1fea9d6390644cb0be" +dependencies = [ + "anyhow", + "bitflags 1.3.2", + "downcast-rs", + "filedescriptor", + "lazy_static", + "libc", + "log", + "nix 0.25.1", + "serial", + "shared_library", + "shell-words", + "winapi", + "winreg", +] + [[package]] name = "powerfmt" version = "0.2.0" @@ -1054,7 +1164,7 @@ version = "0.5.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b6dfecf2c74bce2466cabf93f6664d6998a69eb21e39f4207930065b27b771f" dependencies = [ - "bitflags", + "bitflags 2.6.0", ] [[package]] @@ -1173,7 +1283,7 @@ version = "2.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "897b2245f0b511c87893af39b033e5ca9cce68824c4d7e7630b5a1d339658d02" dependencies = [ - "bitflags", + "bitflags 2.6.0", "core-foundation", "core-foundation-sys", "libc", @@ -1216,6 +1326,48 @@ dependencies = [ "syn", ] +[[package]] +name = "serial" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a1237a96570fc377c13baa1b88c7589ab66edced652e43ffb17088f003db3e86" +dependencies = [ + "serial-core", + "serial-unix", + "serial-windows", +] + +[[package]] +name = "serial-core" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f46209b345401737ae2125fe5b19a77acce90cd53e1658cda928e4fe9a64581" +dependencies = [ + "libc", +] + +[[package]] +name = "serial-unix" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f03fbca4c9d866e24a459cbca71283f545a37f8e3e002ad8c70593871453cab7" +dependencies = [ + "ioctl-rs", + "libc", + "serial-core", + "termios 0.2.2", +] + +[[package]] +name = "serial-windows" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "15c6d3b776267a75d31bbdfd5d36c0ca051251caafc285827052bc53bcdc8162" +dependencies = [ + "libc", + "serial-core", +] + [[package]] name = "sha2" version = "0.10.8" @@ -1227,6 +1379,22 @@ dependencies = [ "digest", ] +[[package]] +name = "shared_library" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a9e7e0f2bfae24d8a5b5a66c5b257a83c7412304311512a0c054cd5e619da11" +dependencies = [ + "lazy_static", + "libc", +] + +[[package]] +name = "shell-words" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24188a676b6ae68c3b2cb3a01be17fbf7240ce009799bb56d5b1409051e78fde" + [[package]] name = "shlex" version = "1.3.0" @@ -1296,6 +1464,44 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "termios" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d5d9cf598a6d7ce700a4e6a9199da127e6819a61e64b68609683cc9a01b5683a" +dependencies = [ + "libc", +] + +[[package]] +name = "termios" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "411c5bf740737c7918b8b1fe232dca4dc9f8e754b8ad5e20966814001ed0ac6b" +dependencies = [ + "libc", +] + +[[package]] +name = "thiserror" +version = "1.0.64" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d50af8abc119fb8bb6dbabcfa89656f46f84aa0ac7688088608076ad2b459a84" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.64" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08904e7672f5eb876eaaf87e0ce17857500934f4981c4a0ab2b4aa98baac7fc3" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "time" version = "0.3.36" @@ -1525,6 +1731,28 @@ version = "0.11.0+wasi-snapshot-preview1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423" +[[package]] +name = "winapi" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" +dependencies = [ + "winapi-i686-pc-windows-gnu", + "winapi-x86_64-pc-windows-gnu", +] + +[[package]] +name = "winapi-i686-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" + +[[package]] +name = "winapi-x86_64-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" + [[package]] name = "windows-sys" version = "0.52.0" @@ -1607,6 +1835,15 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" +[[package]] +name = "winreg" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "80d0f4e272c85def139476380b12f9ac60926689dd2e01d4923222f40580869d" +dependencies = [ + "winapi", +] + [[package]] name = "xmlparser" version = "0.13.6" diff --git a/Cargo.toml b/Cargo.toml index a697881..f5187da 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -8,3 +8,7 @@ aws-config = "1.5.6" aws-sdk-ecs = "1.46.0" tokio = { version = "1.40.0", features = ["full"] } clap = { version = "4.0", features = ["derive"] } +aws-sdk-sts = "1.44.0" +nix = { version = "0.29.0", features = ["term", "process", "fs"] } +portable-pty = "0.8.1" +termios = "0.3" diff --git a/src/main.rs b/src/main.rs index 4ebe02f..8caab58 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,16 +1,20 @@ +mod pty; + use clap::Parser; use aws_config::meta::region::RegionProviderChain; use aws_sdk_ecs::Client as EcsClient; use aws_sdk_ecs::Error as EcsError; use aws_sdk_ecs::config::Region as Region; +use aws_sdk_sts::Client as StsClient; +use aws_sdk_sts::operation::get_caller_identity::GetCallerIdentityError; use std::error::Error; use std::io::{self, Write}; -use std::process::{Command, Stdio}; -use std::io::Read; +use std::process::{Command}; /// A CLI tool to interactively run ECS `execute-command` #[derive(Parser, Debug)] -#[command(author = "Kyrylo Kulyhin", version = "0.1.0", about = "ECS execute-command CLI tool", long_about = None)] +#[command(author = "Kyrylo Kulyhin", version = "0.1.3", about = "ECS execute-command CLI tool", long_about = None +)] struct Cli { /// The AWS profile to use #[arg(long, short = 'p', default_value = "dt-infra")] @@ -41,6 +45,41 @@ async fn main() -> Result<(), Box> { let aws_profile = args.profile.clone(); let container_name = args.container.clone(); + + match check_sso_session(&aws_profile).await { + Ok(_) => { + // Proceed with ECS client operations if the session exists + // let region_provider = RegionProviderChain::first_try(Region::new(aws_region.clone())) + // .or_default_provider(); + + // let shared_config = aws_config::defaults(aws_config::BehaviorVersion::latest()) + // .region(region_provider) + // .profile_name(aws_profile.clone()) + // .load() + // .await; + + println!("SSO session is active. Proceeding with ECS operations for container: {}", container_name); + } + Err(_) => { + if prompt_user_for_login() { + let status = Command::new("aws") + .arg("--profile") + .arg(aws_profile.clone()) + .arg("sso") + .arg("login") + .status()?; // Wait for the command to complete + + if !status.success() { + println!("AWS SSO login failed. Exiting the program."); + return Ok(()); // Exit if login fails + } + println!("Please run `aws sso login` in another terminal, then re-run this program."); + } else { + println!("Exiting the program."); + } + } + } + let region_provider = RegionProviderChain::first_try(Region::new(aws_region.clone())).or_default_provider(); let shared_config = aws_config::defaults(aws_config::BehaviorVersion::latest()) @@ -48,94 +87,75 @@ async fn main() -> Result<(), Box> { .profile_name(aws_profile.clone()) .load() .await; - let ecs_client = EcsClient::new(&shared_config); + let ecs_client = EcsClient::new(&shared_config); let cluster_arn = show_clusters(&ecs_client).await?; let task_arn = show_tasks(&ecs_client, &cluster_arn, &args.service).await?; - let mut child = Command::new("aws") - .arg("--profile") - .arg(aws_profile) - .arg("--region") - .arg(aws_region) - .arg("ecs") - .arg("execute-command") - .arg("--cluster") - .arg(cluster_arn) - .arg("--task") - .arg(task_arn) - .arg("--container") - .arg(container_name) - .arg("--interactive") - .arg("--command") - .arg(args.command) - .stdin(Stdio::piped()) // Attach stdin for interactive input - .stdout(Stdio::piped()) // Capture stdout for interactive output - .stderr(Stdio::piped()) // Capture stderr for error handling - .spawn()?; - - // Handles for child's stdin, stdout, and stderr - let stdin = child.stdin.as_mut().expect("Failed to open stdin"); - let mut stdout = child.stdout.take().expect("Failed to open stdout"); - let mut stderr = child.stderr.take().expect("Failed to open stderr"); - - // Spawn a thread to read from stdout and print to the user's console - let stdout_thread = std::thread::spawn(move || { - let mut stdout_buf = [0; 1024]; - let stdout = &mut stdout; - - loop { - match stdout.read(&mut stdout_buf) { - Ok(0) => break, // End of output - Ok(n) => { - print!("{}", String::from_utf8_lossy(&stdout_buf[..n])); - io::stdout().flush().unwrap(); // Ensure output is printed immediately - } - Err(err) => { - eprintln!("Error reading stdout: {:?}", err); - break; - } - } - } - }); - - // Spawn a thread to read from stderr and print errors to the console - let stderr_thread = std::thread::spawn(move || { - let mut stderr_buf = [0; 1024]; - let stderr = &mut stderr; - - loop { - match stderr.read(&mut stderr_buf) { - Ok(0) => break, // End of error output - Ok(n) => { - eprint!("{}", String::from_utf8_lossy(&stderr_buf[..n])); - io::stderr().flush().unwrap(); // Ensure error output is printed immediately - } - Err(err) => { - eprintln!("Error reading stderr: {:?}", err); - break; - } + let cmd = "aws"; + let args = [ + "--profile", &*aws_profile, + "--region", &*aws_region, + "ecs", "execute-command", + "--cluster", &*cluster_arn, + "--task", &*task_arn, + "--container", &*container_name, + "--interactive", + "--command", &*args.command, + ]; + + pty::spawn_pty_shell(cmd, &args)?; + + Ok(()) +} + +async fn check_sso_session(profile: &str) -> Result<(), aws_sdk_sts::error::SdkError> { + let region_provider = RegionProviderChain::default_provider(); + let shared_config = aws_config::defaults(aws_config::BehaviorVersion::latest()) + .region(region_provider) + .profile_name(profile) + .load() + .await; + + let sts_client = StsClient::new(&shared_config); + + // Call `get_caller_identity` to verify the session is active + match sts_client.get_caller_identity().send().await { + Ok(_) => Ok(()), + Err(err) => match err { + aws_sdk_sts::error::SdkError::ServiceError { .. } => { + println!("Service error occurred: {:?}", err); + Err(err) } - } - }); - // Reading input from the user and writing it to the command's stdin - let mut input = String::new(); - while io::stdin().read_line(&mut input).unwrap() > 0 { - stdin.write_all(input.as_bytes())?; - stdin.flush()?; - input.clear(); + aws_sdk_sts::error::SdkError::TimeoutError(_) => { + println!("The request timed out. Please check your connection."); + Err(err) + } + aws_sdk_sts::error::SdkError::DispatchFailure(_) => { + println!("Network error. Please check your internet connection."); + Err(err) + } + _ => { + println!("An unknown error occurred: {:?}", err); + Err(err) + } + }, } +} - // Wait for stdout and stderr threads to finish - stdout_thread.join().expect("Failed to join stdout thread"); - stderr_thread.join().expect("Failed to join stderr thread"); +fn prompt_user_for_login() -> bool { + println!("No active AWS SSO session found."); + println!("Please run the following command to login via AWS SSO:"); + println!(" aws sso login"); + print!("Would you like to retry after logging in? (Y/n): "); + io::stdout().flush().unwrap(); // Make sure prompt gets printed immediately - // Wait for the command to finish - let status = child.wait()?; - println!("Command exited with status: {}", status); + let mut input = String::new(); + io::stdin().read_line(&mut input).expect("Failed to read input"); + let input = input.trim().to_lowercase(); - Ok(()) + input == "y" || input == "yes" } // List your clusters. diff --git a/src/pty.rs b/src/pty.rs new file mode 100644 index 0000000..1899387 --- /dev/null +++ b/src/pty.rs @@ -0,0 +1,101 @@ +use portable_pty::{CommandBuilder, NativePtySystem, PtySize, PtySystem}; +use std::io::{Read, Write}; +use std::thread; +use termios::{Termios, TCSANOW, ECHO, ICANON, ISIG}; +use std::os::unix::io::AsRawFd; + +pub(crate) fn spawn_pty_shell(cmd: &str, args: &[&str]) -> Result<(), Box> { + // Create a new PTY system + let pty_system = NativePtySystem::default(); + let pair = pty_system.openpty(PtySize { + rows: 24, + cols: 80, + pixel_width: 0, + pixel_height: 0, + })?; + + // Save the original terminal attributes of stdin + let stdin_fd = std::io::stdin().as_raw_fd(); + let orig_termios = Termios::from_fd(stdin_fd)?; + + // Set stdin to raw mode + let mut raw_termios = orig_termios.clone(); + raw_termios.c_lflag &= !(ECHO | ICANON | ISIG); + termios::tcsetattr(stdin_fd, TCSANOW, &raw_termios)?; + + // Ensure that stdin is set back to original mode when the program exits + let _termios_guard = TermiosGuard::new(stdin_fd, orig_termios.clone()); + + // Build your command that connects to the ECS task + // Replace this with your actual command and arguments + + let mut cmd_builder = CommandBuilder::new(cmd); + cmd_builder.args(args); + cmd_builder.env( + "TERM", + std::env::var("TERM").unwrap_or_else(|_| "xterm-256color".to_string()), + ); + + // Spawn the child process within the PTY + let mut child = pair.slave.spawn_command(cmd_builder)?; + + // Close the slave side of the PTY in the parent process + drop(pair.slave); + + // Set up reader and writer for the PTY master + let mut reader = pair.master.try_clone_reader()?; + let mut writer = pair.master.take_writer()?; + + // Spawn a thread to handle input from stdin to the PTY + thread::spawn(move || { + let stdin = std::io::stdin(); + let mut stdin = stdin.lock(); + let mut buffer = [0u8; 1024]; + loop { + let n = stdin.read(&mut buffer).expect("Failed to read from stdin"); + if n == 0 { + break; + } + writer + .write_all(&buffer[..n]) + .expect("Failed to write to PTY"); + writer.flush().expect("Failed to flush PTY writer"); + } + }); + + // Read output from the PTY and write to stdout + let stdout = std::io::stdout(); + let mut stdout = stdout.lock(); + let mut buffer = [0u8; 1024]; + loop { + let n = reader.read(&mut buffer)?; + if n == 0 { + break; + } + stdout.write_all(&buffer[..n])?; + stdout.flush()?; + } + + // Wait for the child process to exit + child.wait()?; + + Ok(()) +} + +// Helper struct to restore terminal attributes on exit +struct TermiosGuard { + fd: i32, + termios: Termios, +} + +impl TermiosGuard { + fn new(fd: i32, termios: Termios) -> Self { + TermiosGuard { fd, termios } + } +} + +impl Drop for TermiosGuard { + fn drop(&mut self) { + let _ = termios::tcsetattr(self.fd, TCSANOW, &self.termios); + } +} \ No newline at end of file