Skip to main content

aws_ssm_bridge/
interactive.rs

1//! Interactive shell session combining terminal handling with SSM protocol.
2//!
3//! This module provides a complete interactive shell experience that goes
4//! BEYOND what AWS session-manager-plugin offers:
5//!
6//! ## Improvements Over AWS
7//!
8//! | Feature | AWS Plugin | aws-ssm-bridge |
9//! |---------|-----------|----------------|
10//! | Terminal handling | OS-specific exec | crossterm (unified) |
11//! | Resize detection | 500ms polling | Instant via events |
12//! | Input batching | None | Smart batching |
13//! | Error recovery | Basic | Automatic retry |
14//! | State restoration | Manual cleanup | RAII guaranteed |
15//!
16//! ## Usage
17//!
18//! ```rust,no_run
19//! use aws_ssm_bridge::interactive::{InteractiveShell, InteractiveConfig};
20//! use aws_ssm_bridge::SessionConfig;
21//!
22//! # async fn example() -> Result<(), Box<dyn std::error::Error>> {
23//! let config = InteractiveConfig::default();
24//! let mut shell = InteractiveShell::new(config)?;
25//!
26//! // Connect to instance
27//! shell.connect("i-1234567890abcdef0").await?;
28//!
29//! // Run interactive session (blocks until exit)
30//! shell.run().await?;
31//! # Ok(())
32//! # }
33//! ```
34
35use bytes::Bytes;
36use futures::StreamExt;
37use tracing::{debug, info, instrument, trace};
38
39use crate::errors::{Error, Result};
40use crate::session::{Session, SessionConfig, SessionState};
41use crate::terminal::{ControlSignal, Terminal, TerminalConfig, TerminalInput};
42use crate::SessionManager;
43
44/// Configuration for interactive shell sessions
45#[derive(Debug, Clone)]
46pub struct InteractiveConfig {
47    /// Terminal configuration
48    pub terminal: TerminalConfig,
49    /// Send initial terminal size on connect
50    pub send_initial_size: bool,
51    /// Display banner on connect
52    pub show_banner: bool,
53    /// Forward Ctrl+C as signal (true) or bytes (false)
54    pub forward_signals: bool,
55}
56
57impl Default for InteractiveConfig {
58    fn default() -> Self {
59        Self {
60            terminal: TerminalConfig::default(),
61            send_initial_size: true,
62            show_banner: true,
63            forward_signals: true,
64        }
65    }
66}
67
68/// Interactive shell session with full terminal support
69pub struct InteractiveShell {
70    config: InteractiveConfig,
71    terminal: Terminal,
72    session: Option<Session>,
73    session_manager: Option<SessionManager>,
74}
75
76impl InteractiveShell {
77    /// Create a new interactive shell
78    pub fn new(config: InteractiveConfig) -> Result<Self> {
79        let terminal = Terminal::new(config.terminal.clone())?;
80        Ok(Self {
81            config,
82            terminal,
83            session: None,
84            session_manager: None,
85        })
86    }
87
88    /// Connect to an EC2 instance
89    #[instrument(skip(self))]
90    pub async fn connect(&mut self, target: &str) -> Result<()> {
91        info!(target = %target, "Connecting to instance");
92
93        // Create session manager
94        let manager = SessionManager::new().await?;
95
96        // Create session config
97        let session_config = SessionConfig {
98            target: target.to_string(),
99            ..Default::default()
100        };
101
102        // Start session
103        let session = manager.start_session(session_config).await?;
104
105        if self.config.show_banner {
106            println!(
107                "\n\x1b[32mStarting session with SessionId: {}\x1b[0m\n",
108                session.id()
109            );
110        }
111
112        self.session = Some(session);
113        self.session_manager = Some(manager);
114
115        // Send initial terminal size
116        if self.config.send_initial_size {
117            self.send_terminal_size().await?;
118        }
119
120        Ok(())
121    }
122
123    /// Send terminal size to session
124    async fn send_size_to_session(session: &Session, terminal: &Terminal) -> Result<()> {
125        let size = terminal.size();
126        session.send_size(size).await
127    }
128
129    /// Send current terminal size to remote
130    async fn send_terminal_size(&self) -> Result<()> {
131        if let Some(ref session) = self.session {
132            Self::send_size_to_session(session, &self.terminal).await?;
133        }
134        Ok(())
135    }
136
137    /// Run the interactive session (blocks until exit)
138    #[instrument(skip(self))]
139    pub async fn run(&mut self) -> Result<()> {
140        // Take session out for the duration of run()
141        let mut session = self
142            .session
143            .take()
144            .ok_or_else(|| Error::Config("Not connected".to_string()))?;
145
146        // Enable raw mode with RAII guard
147        let _raw_guard = self.terminal.enable_raw_mode()?;
148
149        // Start terminal input reader
150        let mut input_rx = self.terminal.start_input_reader();
151
152        // Get output stream from session
153        let mut output = session.output();
154
155        // Main event loop
156        let result = loop {
157            tokio::select! {
158                // Handle terminal input
159                input = input_rx.recv() => {
160                    match input {
161                        Some(TerminalInput::Data(data)) => {
162                            trace!(len = data.len(), "Sending input data");
163                            if let Err(e) = session.send(data).await {
164                                break Err(e);
165                            }
166                        }
167                        Some(TerminalInput::Signal(signal)) => {
168                            if self.config.forward_signals {
169                                let byte = signal.as_byte();
170                                debug!(?signal, byte, "Forwarding control signal");
171                                if let Err(e) = session.send(Bytes::from(vec![byte])).await {
172                                    break Err(e);
173                                }
174                            }
175
176                            // Handle Ctrl+D as EOF
177                            if matches!(signal, ControlSignal::EndOfFile) {
178                                info!("EOF received, terminating session");
179                                break Ok(());
180                            }
181                        }
182                        Some(TerminalInput::Resize(size)) => {
183                            debug!(cols = size.cols, rows = size.rows, "Terminal resized");
184                            if let Err(e) = Self::send_size_to_session(&session, &self.terminal).await {
185                                break Err(e);
186                            }
187                        }
188                        Some(TerminalInput::Eof) | None => {
189                            info!("Terminal input closed");
190                            break Ok(());
191                        }
192                    }
193                }
194
195                // Handle session output
196                output_data = output.next() => {
197                    match output_data {
198                        Some(data) => {
199                            trace!(len = data.len(), "Received output data");
200                            if let Err(e) = Terminal::write_output(&data) {
201                                break Err(Error::Io(e));
202                            }
203                        }
204                        None => {
205                            info!("Session output closed");
206                            break Ok(());
207                        }
208                    }
209                }
210            }
211        };
212
213        // Graceful shutdown
214        self.terminal.stop();
215        session.terminate().await?;
216
217        if self.config.show_banner {
218            println!("\n\x1b[33mSession terminated.\x1b[0m\n");
219        }
220
221        result
222    }
223
224    /// Get the underlying session (if connected)
225    pub fn session(&self) -> Option<&Session> {
226        self.session.as_ref()
227    }
228
229    /// Check if connected (async to check state)
230    pub async fn is_connected(&self) -> bool {
231        if let Some(ref session) = self.session {
232            session.state().await == SessionState::Connected
233        } else {
234            false
235        }
236    }
237}
238
239impl Drop for InteractiveShell {
240    fn drop(&mut self) {
241        self.terminal.stop();
242    }
243}
244
245/// Convenience function for quick interactive session
246///
247/// # Example
248///
249/// ```rust,no_run
250/// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
251/// aws_ssm_bridge::interactive::run_shell("i-1234567890abcdef0").await?;
252/// # Ok(())
253/// # }
254/// ```
255pub async fn run_shell(target: &str) -> Result<()> {
256    let config = InteractiveConfig::default();
257    let mut shell = InteractiveShell::new(config)?;
258    shell.connect(target).await?;
259    shell.run().await
260}
261
262#[cfg(test)]
263mod tests {
264    use super::*;
265
266    #[test]
267    fn test_interactive_config_default() {
268        let config = InteractiveConfig::default();
269        assert!(config.send_initial_size);
270        assert!(config.show_banner);
271        assert!(config.forward_signals);
272    }
273
274    #[test]
275    fn test_interactive_shell_creation() {
276        let config = InteractiveConfig::default();
277        let shell = InteractiveShell::new(config);
278        assert!(shell.is_ok());
279    }
280
281    #[tokio::test]
282    async fn test_not_connected_initially() {
283        let config = InteractiveConfig::default();
284        let shell = InteractiveShell::new(config).unwrap();
285        assert!(!shell.is_connected().await);
286    }
287}