kratactl/cli/zone/
exec.rs

1use std::collections::HashMap;
2
3use anyhow::Result;
4
5use clap::Parser;
6use crossterm::tty::IsTty;
7use krata::v1::{
8    common::{TerminalSize, ZoneTaskSpec, ZoneTaskSpecEnvVar},
9    control::{control_service_client::ControlServiceClient, ExecInsideZoneRequest},
10};
11
12use tokio::io::stdin;
13use tonic::{transport::Channel, Request};
14
15use crate::console::StdioConsoleStream;
16
17use crate::cli::resolve_zone;
18
19#[derive(Parser)]
20#[command(about = "Execute a command inside the zone")]
21pub struct ZoneExecCommand {
22    #[arg[short, long, help = "Environment variables"]]
23    env: Option<Vec<String>>,
24    #[arg(short = 'w', long, help = "Working directory")]
25    working_directory: Option<String>,
26    #[arg(short = 't', long, help = "Allocate tty")]
27    tty: bool,
28    #[arg(help = "Zone to exec inside, either the name or the uuid")]
29    zone: String,
30    #[arg(
31        allow_hyphen_values = true,
32        trailing_var_arg = true,
33        help = "Command to run inside the zone"
34    )]
35    command: Vec<String>,
36}
37
38impl ZoneExecCommand {
39    pub async fn run(self, mut client: ControlServiceClient<Channel>) -> Result<()> {
40        let zone_id: String = resolve_zone(&mut client, &self.zone).await?;
41        let should_map_tty = self.tty && stdin().is_tty();
42        let initial = ExecInsideZoneRequest {
43            zone_id,
44            task: Some(ZoneTaskSpec {
45                environment: env_map(&self.env.unwrap_or_default())
46                    .iter()
47                    .map(|(key, value)| ZoneTaskSpecEnvVar {
48                        key: key.clone(),
49                        value: value.clone(),
50                    })
51                    .collect(),
52                command: self.command,
53                working_directory: self.working_directory.unwrap_or_default(),
54                tty: self.tty,
55            }),
56            stdin: vec![],
57            stdin_closed: false,
58            terminal_size: if should_map_tty {
59                let size = crossterm::terminal::size().ok();
60                size.map(|(columns, rows)| TerminalSize {
61                    rows: rows as u32,
62                    columns: columns as u32,
63                })
64            } else {
65                None
66            },
67        };
68
69        let stream = StdioConsoleStream::input_stream_exec(initial, should_map_tty).await;
70
71        let response = client
72            .exec_inside_zone(Request::new(stream))
73            .await?
74            .into_inner();
75
76        let code = StdioConsoleStream::exec_output(response, should_map_tty).await?;
77        std::process::exit(code);
78    }
79}
80
81fn env_map(env: &[String]) -> HashMap<String, String> {
82    let mut map = HashMap::<String, String>::new();
83    for item in env {
84        if let Some((key, value)) = item.split_once('=') {
85            map.insert(key.to_string(), value.to_string());
86        }
87    }
88    map
89}