1#![forbid(unsafe_code)]
7
8mod bulk;
11mod cli;
12mod clients;
13mod commands;
14mod debug;
15mod env;
16mod error;
17mod files;
18mod input;
19mod json;
20mod output;
21mod pagination;
22mod policy;
23mod redteam;
24mod render;
25mod report;
26mod signals;
27mod store;
28mod tenant;
29mod views;
30mod writes;
31
32use clap::FromArgMatches;
33use commands::{Invocation, Plan};
34use error::{CliError, EXIT_OK, EXIT_OPERATIONAL, EXIT_USAGE};
35use output::Output;
36use std::ffi::OsString;
37use std::io::{ErrorKind, Write};
38
39pub const VERSION: &str = env!("CARGO_PKG_VERSION");
41
42pub fn run<I, T>(args: I) -> i32
44where
45 I: IntoIterator<Item = T>,
46 T: Into<OsString>,
47{
48 let args: Vec<OsString> = args.into_iter().map(Into::into).collect();
49 if version_requested(&args) {
50 return print_version();
51 }
52 let parsed = cli::command()
53 .try_get_matches_from(&args)
54 .and_then(|matches| cli::Cli::from_arg_matches(&matches));
55 let cli = match parsed {
56 Ok(cli) => cli,
57 Err(error) => {
58 let error = cli::redacted_parse_error(&args, &error).unwrap_or(error);
61 let _ = error.print();
62 return error.exit_code();
63 }
64 };
65 let mut output = Output::stdio(cli.global.quiet);
66 output.set_debug(cli.global.debug);
67 if commands::uses_format_alias(&cli) {
69 output.warning(commands::FORMAT_DEPRECATED);
70 }
71 if let Some(cli::TopCommand::Redteam(group)) = &cli.command {
73 for warning in commands::redteam::deprecation_warnings(group) {
74 output.warning(warning);
75 }
76 }
77 if commands::uses_input_alias(&cli) {
79 output.warning(commands::INPUT_DEPRECATED);
80 }
81 let plan = match commands::plan(&cli) {
82 Ok(plan) => plan,
83 Err(error) => {
84 output.report(&error);
85 return error.exit_code(None);
86 }
87 };
88 if let Plan::GroupHelp(path) = plan {
89 output.write_stderr(&cli::group_help(path));
90 return EXIT_USAGE;
91 }
92 let mut invocation = Invocation::new(output, cli.global.tenant.clone());
93 if cli.global.debug
95 && plan.debug_policy() == debug::DebugPolicy::Log
96 && let Err(error) = invocation.open_debug_log()
97 {
98 invocation.output.report(&error);
99 return error.exit_code(None);
100 }
101 execute(plan, invocation)
102}
103
104pub fn version_requested(args: &[OsString]) -> bool {
107 args.iter()
108 .skip(1)
109 .take_while(|arg| *arg != "--")
110 .any(|arg| arg == "-V" || arg == "--version")
111}
112
113fn stdout_write_failed() -> CliError {
116 CliError::operational("AISEC_SDK_ERROR:Output could not be written to stdout")
117}
118
119fn print_version() -> i32 {
120 let mut stdout = std::io::stdout().lock();
121 match writeln!(stdout, "{VERSION}").and_then(|()| stdout.flush()) {
122 Err(error) if error.kind() != ErrorKind::BrokenPipe => {
123 let text = stdout_write_failed().render_colored(render::style::stderr_color());
124 let _ = std::io::stderr().write_all(text.as_bytes());
125 EXIT_OPERATIONAL
126 }
127 _ => EXIT_OK,
128 }
129}
130
131fn execute(plan: Plan, mut invocation: Invocation) -> i32 {
132 let runtime = match tokio::runtime::Builder::new_current_thread()
133 .enable_all()
134 .build()
135 {
136 Ok(runtime) => runtime,
137 Err(_) => {
138 let error = CliError::operational("AISEC_SDK_ERROR:The async runtime could not start");
139 invocation.output.report(&error);
140 return EXIT_OPERATIONAL;
141 }
142 };
143 let signals = signals::install(invocation.session.clone());
144 signals.set_quiet(invocation.output.quiet());
145 invocation.output.attach_signals(signals.clone());
146 let fallback = plan.interrupt_context();
147 let result = runtime.block_on(commands::dispatch(&mut invocation, plan));
148 let fallback = if invocation.committed {
150 error::InterruptContext::ChangeSaved
151 } else {
152 invocation.interrupt.take().unwrap_or(fallback)
153 };
154 runtime.shutdown_background();
157 let signal = signals.code();
158 let result = match (result, signal) {
160 (Err(CliError::Interrupted(context)), _) => Err(CliError::Interrupted(context)),
161 (_, Some(_)) => Err(CliError::Interrupted(fallback)),
162 (result, None) => result,
163 };
164 finish(&mut invocation.output, result, signal)
165}
166
167fn finish(output: &mut Output, result: Result<(), CliError>, signal: Option<i32>) -> i32 {
168 match result {
169 Ok(()) => match output.commit() {
170 Ok(()) => EXIT_OK,
171 Err(error) if error.kind() == ErrorKind::BrokenPipe => EXIT_OK,
173 Err(_) => {
174 let error = match output.pending_signal() {
176 Some(_) => CliError::Interrupted(error::InterruptContext::OutputIncomplete),
177 None => stdout_write_failed(),
178 };
179 output.report(&error);
180 error.exit_code(output.pending_signal())
181 }
182 },
183 Err(error) => {
184 output.discard();
185 output.report(&error);
186 error.exit_code(signal)
187 }
188 }
189}
190
191#[cfg(test)]
192mod tests {
193 use super::*;
194
195 fn args(list: &[&str]) -> Vec<OsString> {
196 list.iter().map(OsString::from).collect()
197 }
198
199 #[test]
200 fn version_pre_scan() {
201 assert!(version_requested(&args(&["airs-cli-rs", "-V"])));
202 assert!(version_requested(&args(&[
203 "airs-cli-rs",
204 "runtime",
205 "scan",
206 "--profile",
207 "-V",
208 "hi"
209 ])));
210 assert!(version_requested(&args(&[
211 "airs-cli-rs",
212 "doctor",
213 "--version"
214 ])));
215 assert!(!version_requested(&args(&[
216 "airs-cli-rs",
217 "runtime",
218 "scan",
219 "--profile",
220 "p",
221 "--",
222 "-V"
223 ])));
224 assert!(
225 !version_requested(&args(&["-V"])),
226 "the program name is skipped"
227 );
228 assert!(!version_requested(&args(&["airs-cli-rs", "-v"])));
229 }
230
231 #[test]
232 fn finish_maps_results_to_exit_codes() {
233 let (mut out, stdout, stderr) = output::testing::capture(false);
234 out.data_line("data");
235 assert_eq!(finish(&mut out, Ok(()), None), 0);
236 assert_eq!(stdout.text(), "data\n");
237
238 let (mut out, stdout, stderr2) = output::testing::capture(false);
239 out.data_line("partial");
240 assert_eq!(finish(&mut out, Err(CliError::usage("bad")), None), 2);
241 assert_eq!(stdout.text(), "");
242 assert_eq!(stderr2.text(), "\n ✗ bad\n\n");
243
244 let (mut out, _, stderr3) = output::testing::capture(false);
245 let code = finish(
246 &mut out,
247 Err(CliError::Interrupted(
248 error::InterruptContext::NoDataWritten,
249 )),
250 Some(143),
251 );
252 assert_eq!(code, 143);
253 assert!(stderr3.text().contains("Interrupted; no data was written"));
254 assert_eq!(stderr.text(), "");
255 }
256
257 #[test]
258 fn broken_pipe_is_success() {
259 struct Closed;
260 impl Write for Closed {
261 fn write(&mut self, _: &[u8]) -> std::io::Result<usize> {
262 Err(std::io::Error::from(ErrorKind::BrokenPipe))
263 }
264 fn flush(&mut self) -> std::io::Result<()> {
265 Ok(())
266 }
267 }
268 let mut out = Output::with_sinks(false, Box::new(Closed), Box::new(std::io::sink()));
269 out.data_line("data");
270 assert_eq!(finish(&mut out, Ok(()), None), 0);
271 }
272
273 #[test]
274 fn other_stdout_write_failures_are_reported() {
275 struct Full;
276 impl Write for Full {
277 fn write(&mut self, _: &[u8]) -> std::io::Result<usize> {
278 Err(std::io::Error::from(ErrorKind::StorageFull))
279 }
280 fn flush(&mut self) -> std::io::Result<()> {
281 Ok(())
282 }
283 }
284 let stderr = output::testing::Capture::default();
285 let mut out = Output::with_sinks(false, Box::new(Full), Box::new(stderr.clone()));
286 out.data_line("data");
287 assert_eq!(finish(&mut out, Ok(()), None), 1);
288 assert_eq!(
289 stderr.text(),
290 "\n ✗ Error: AISEC_SDK_ERROR:Output could not be written to stdout\n\n"
291 );
292 }
293
294 #[test]
295 fn a_signal_before_the_final_write_reports_the_interruption() {
296 let (mut out, stdout, stderr) = output::testing::capture(false);
297 let signals = std::sync::Arc::new(signals::SignalState::default());
298 assert_eq!(signals.on_signal(143), signals::Action::Cancel);
299 out.attach_signals(signals);
300 out.data_line("data");
301 assert_eq!(finish(&mut out, Ok(()), None), 143);
302 assert_eq!(stdout.text(), "", "nothing is written after a signal");
303 assert_eq!(
304 stderr.text(),
305 "\n ✗ Error: Interrupted; the output was not completely written\n\n"
306 );
307 }
308}