aws_ssm_bridge/
interactive.rs1use 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#[derive(Debug, Clone)]
46pub struct InteractiveConfig {
47 pub terminal: TerminalConfig,
49 pub send_initial_size: bool,
51 pub show_banner: bool,
53 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
68pub struct InteractiveShell {
70 config: InteractiveConfig,
71 terminal: Terminal,
72 session: Option<Session>,
73 session_manager: Option<SessionManager>,
74}
75
76impl InteractiveShell {
77 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 #[instrument(skip(self))]
90 pub async fn connect(&mut self, target: &str) -> Result<()> {
91 info!(target = %target, "Connecting to instance");
92
93 let manager = SessionManager::new().await?;
95
96 let session_config = SessionConfig {
98 target: target.to_string(),
99 ..Default::default()
100 };
101
102 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 if self.config.send_initial_size {
117 self.send_terminal_size().await?;
118 }
119
120 Ok(())
121 }
122
123 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 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 #[instrument(skip(self))]
139 pub async fn run(&mut self) -> Result<()> {
140 let mut session = self
142 .session
143 .take()
144 .ok_or_else(|| Error::Config("Not connected".to_string()))?;
145
146 let _raw_guard = self.terminal.enable_raw_mode()?;
148
149 let mut input_rx = self.terminal.start_input_reader();
151
152 let mut output = session.output();
154
155 let result = loop {
157 tokio::select! {
158 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 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 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 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 pub fn session(&self) -> Option<&Session> {
226 self.session.as_ref()
227 }
228
229 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
245pub 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}