korasi_cli/
lib.rs

1pub mod create;
2pub mod ec2;
3pub mod opt;
4pub mod ssh;
5pub mod util;
6
7use anyhow::Context;
8use aws_config::{
9    self, meta::region::RegionProviderChain, timeout::TimeoutConfig, BehaviorVersion,
10};
11use aws_sdk_ec2::types::{InstanceStateName, InstanceType};
12use aws_types::{region::Region, SdkConfig as AwsSdkConfig};
13use inquire::{Select, Text};
14use termion::raw::IntoRawMode;
15use tokio::time::Duration;
16
17use create::CreateCommand;
18use ec2::{EC2Impl as EC2, SSH_KEY_NAME, SSH_SECURITY_GROUP};
19use opt::{Commands, Opt};
20use ssh::Session;
21use util::{ids_to_str, multi_select_instances, select_instance};
22
23/// Loads an AWS config from default environments.
24pub async fn load_config(
25    region: Option<String>,
26    profile_name: Option<String>,
27    operation_timeout: Option<Duration>,
28) -> AwsSdkConfig {
29    tracing::info!("loading config for the region {:?}", region);
30
31    // if region is None, it automatically detects iff it's running inside the EC2 instance
32    let reg_provider = RegionProviderChain::first_try(region.map(Region::new))
33        .or_default_provider()
34        .or_else(Region::new("ap-southeast-1"));
35
36    let mut builder = TimeoutConfig::builder().connect_timeout(Duration::from_secs(5));
37    if let Some(to) = &operation_timeout {
38        if !to.is_zero() {
39            builder = builder.operation_timeout(*to);
40        }
41    }
42    let timeout_cfg = builder.build();
43
44    let mut cfg = aws_config::defaults(BehaviorVersion::v2024_03_28())
45        .region(reg_provider)
46        .profile_name(profile_name.as_ref().unwrap_or(&"default".to_string()))
47        .timeout_config(timeout_cfg);
48    if let Some(p) = profile_name {
49        tracing::info!("loading the aws profile '{p}'");
50        cfg = cfg.profile_name(p);
51    }
52
53    cfg.load().await
54}
55
56pub async fn run(opts: Opt) -> anyhow::Result<()> {
57    let Opt {
58        profile,
59        region,
60        ssh_key,
61        tag,
62        ..
63    } = opts;
64
65    let ssh_path = std::env::var("HOME")
66        .map(|h| {
67            if let Some(ssh_key) = ssh_key {
68                ssh_key
69            } else {
70                format!("{}/.ssh/{SSH_KEY_NAME}.pem", h)
71            }
72        })
73        .unwrap();
74    tracing::info!("Using SSH key at = {}", ssh_path);
75
76    let shared_config = load_config(Some(region), Some(profile), None).await;
77    let client = aws_sdk_ec2::Client::new(&shared_config);
78    let ec2 = EC2::new(client, tag);
79
80    match opts.commands {
81        Commands::Create { ami_id } => {
82            let machine: InstanceType =
83                Select::new("Select the machine type:", InstanceType::values().to_vec())
84                    .prompt()
85                    .unwrap()
86                    .into();
87            tracing::info!("Launching {machine} instance...");
88            CreateCommand
89                .launch(&ec2, machine, ami_id, ssh_path, "start_up.sh".into())
90                .await?;
91        }
92        Commands::List => {
93            let res = ec2.describe_instance(vec![]).await.unwrap();
94            if res.is_empty() {
95                tracing::warn!("There are no active instances.");
96                return Ok(());
97            }
98            for (i, instance) in res.iter().enumerate() {
99                let tags = instance.tags();
100                let mut name = "";
101                for t in tags {
102                    if t.key() == Some("Name") {
103                        name = t.value().unwrap();
104                    }
105                }
106
107                let mut host = "".to_string();
108                if let Some(dns) = instance.public_dns_name() {
109                    if !dns.is_empty() {
110                        host = dns.into();
111                    }
112                }
113
114                tracing::info!(
115                    "{}. {:?}, type = {}, state = {:?}, {:?}",
116                    i + 1,
117                    name,
118                    instance.instance_type.as_ref().unwrap(),
119                    instance.state().unwrap().name().unwrap(),
120                    host,
121                );
122            }
123        }
124        Commands::Delete { wait } => {
125            if let Ok(chosen) =
126                multi_select_instances(&ec2, "Choose the instance(s):", vec![]).await
127            {
128                let instance_ids = ids_to_str(chosen);
129                if instance_ids.is_empty() {
130                    tracing::warn!("Nothing is selected. Use [space] to select option.");
131                } else {
132                    ec2.delete_instances(&instance_ids, wait).await?;
133                }
134            }
135        }
136        Commands::Start => {
137            if let Ok(chosen) = multi_select_instances(
138                &ec2,
139                "Choose the instance(s):",
140                vec![InstanceStateName::Stopped],
141            )
142            .await
143            {
144                let instance_ids = ids_to_str(chosen);
145                if instance_ids.is_empty() {
146                    tracing::warn!("Nothing is selected. Use [space] to select option.");
147                } else {
148                    ec2.start_instances(&instance_ids).await?;
149                }
150            }
151        }
152        Commands::Stop { wait } => {
153            if let Ok(chosen) = multi_select_instances(
154                &ec2,
155                "Choose the instance(s):",
156                vec![InstanceStateName::Running],
157            )
158            .await
159            {
160                let instance_ids = ids_to_str(chosen);
161                if instance_ids.is_empty() {
162                    tracing::warn!("Nothing is selected. Use [space] to select option.");
163                } else {
164                    ec2.stop_instances(&instance_ids, wait).await?;
165                }
166            }
167        }
168        Commands::Upload { src, dst, user } => {
169            if let Ok(chosen) = select_instance(
170                &ec2,
171                "Choose running instance to upload files to:",
172                vec![InstanceStateName::Running],
173            )
174            .await
175            {
176                tracing::info!("Chosen instance: {} = {}", chosen.name, chosen.instance_id);
177                // Refresh inbound IP.
178                ec2.get_ssh_security_group().await?;
179                let session =
180                    Session::connect(&user, chosen.public_dns_name.unwrap(), ssh_path).await?;
181                session.upload(src, dst).await?;
182            } else {
183                tracing::warn!("No active running instances to upload to.");
184            }
185        }
186        Commands::Run { command, user } => {
187            if command.is_empty() {
188                tracing::warn!("Please enter a command to run.");
189                return Ok(());
190            }
191
192            let chosen = select_instance(
193                &ec2,
194                "Choose running instance to execute remote command:",
195                vec![InstanceStateName::Running],
196            )
197            .await
198            .unwrap();
199            tracing::info!(
200                "Chosen instance: name = {}, instance_id = {}",
201                chosen.name,
202                chosen.instance_id
203            );
204
205            // Refresh inbound IP.
206            ec2.get_ssh_security_group().await?;
207
208            let mut session =
209                Session::connect(&user, chosen.public_dns_name.unwrap(), ssh_path).await?;
210            let _raw_term = std::io::stdout().into_raw_mode()?;
211            // TODO: On centos, nothing is printed to stdout (message is received on SDK client).
212            session
213                .exec(
214                    &command
215                        .into_iter()
216                        // arguments are escaped manually since the SSH protocol doesn't support quoting
217                        .map(|cmd_part| shell_escape::escape(cmd_part.into()))
218                        .collect::<Vec<_>>()
219                        .join(" "),
220                )
221                .await?;
222            session.close().await?;
223        }
224        Commands::Shell { user } => {
225            let chosen = select_instance(
226                &ec2,
227                "Choose running instance to ssh:",
228                vec![InstanceStateName::Running],
229            )
230            .await;
231
232            if let Ok(chosen) = chosen {
233                tracing::info!(
234                    "Chosen instance: name = {}, instance_id = {}",
235                    chosen.name,
236                    chosen.instance_id
237                );
238
239                // Refresh inbound IP.
240                ec2.get_ssh_security_group().await?;
241
242                let mut session =
243                    Session::connect(&user, chosen.public_dns_name.unwrap(), ssh_path).await?;
244                let _raw_term = std::io::stdout().into_raw_mode()?;
245                session
246                    .exec(
247                        &vec!["bash"]
248                            .into_iter()
249                            .map(|cmd_part| shell_escape::escape(cmd_part.into()))
250                            .collect::<Vec<_>>()
251                            .join(" "),
252                    )
253                    .await?;
254                session.close().await?;
255            } else {
256                tracing::warn!("There are no active instances to SSH into.");
257            }
258        }
259        Commands::Obliterate => {
260            let yes = Text::new("Do you want to obliterate all resources [y/n]?:").prompt()?;
261            if !(yes == "y" || yes == "Y") {
262                tracing::warn!("Aborting obliterate.");
263                return Ok(());
264            }
265
266            // Passing empty vec means all non-terminated instances are returned.
267            let instances = ec2.describe_instance(vec![]).await?;
268            let select_all = instances.into_iter().map(|i| i.into()).collect();
269            let instance_ids = ids_to_str(select_all);
270
271            let grp = ec2.describe_security_group(SSH_SECURITY_GROUP).await?;
272            let grp_id = grp.as_ref().unwrap().group_id().unwrap();
273
274            let key_pairs = ec2.list_key_pair(SSH_KEY_NAME).await?;
275            let key_pair_ids: Vec<_> = key_pairs.iter().map(|k| k.key_pair_id().unwrap()).collect();
276
277            tracing::info!("instance_ids = {:?}", instance_ids);
278            tracing::info!("grp_id = {:?}", grp_id);
279            tracing::info!("key pairs = {:?}", key_pair_ids);
280
281            ec2.delete_instances(&instance_ids, true).await?;
282            ec2.delete_security_group(grp_id).await?;
283            for id in key_pair_ids {
284                ec2.delete_key_pair(id).await?;
285            }
286
287            // Remove SSH key. PK is useless when key pair is deleted.
288            std::fs::remove_file(&ssh_path)
289                .with_context(|| format!("Failed to remove pk file at {ssh_path}."))?;
290        }
291    };
292
293    Ok(())
294}
korasi_cli/lib.rs

korasi_cli/
lib.rs
