use std::io::prelude::*;
use std::net::TcpStream;
use std::path::Path;
use async_trait::async_trait;
use aws_config::meta::region::RegionProviderChain;
use aws_sdk_ec2::client::Waiters;
use aws_sdk_ec2::config::Region;
use aws_sdk_ec2::types::{InstanceType, ResourceType, Tag, TagSpecification};
use aws_sdk_ec2::Client;
use ssh2::Session;
use super::Provider;
use crate::config::Config;
use crate::jobs::Job;
#[derive(Clone)]
pub struct AWSProvider {}
impl Default for AWSProvider {
fn default() -> Self {
Self::new()
}
}
impl AWSProvider {
#[must_use]
pub const fn new() -> Self {
Self {}
}
}
#[async_trait]
impl Provider for AWSProvider {
async fn start_job(&self, name: &str) -> Result<Job, Box<dyn std::error::Error + Send + Sync>> {
let cfg = config();
let region_provider =
RegionProviderChain::first_try(Some(Region::new(cfg.location.clone())))
.or_default_provider();
let shared_config = aws_config::from_env().region(region_provider).load().await;
let client = Client::new(&shared_config);
let tag_spec = TagSpecification::builder()
.resource_type(ResourceType::Instance)
.tags(Tag::builder().key("Name").value(name).build())
.build();
let run_out = client
.run_instances()
.image_id(cfg.image.clone())
.instance_type(InstanceType::from(cfg.server_type.as_str()))
.min_count(1)
.max_count(1)
.key_name(cfg.ssh_key_name.clone())
.tag_specifications(tag_spec)
.send()
.await?;
let instance = run_out.instances().first().ok_or("no instance created")?;
let instance_id = instance
.instance_id()
.ok_or("missing instance id")?
.to_string();
let _ = client
.wait_until_instance_running()
.instance_ids(instance_id.clone())
.wait(std::time::Duration::from_secs(120))
.await;
let desc = client
.describe_instances()
.instance_ids(instance_id.clone())
.send()
.await?;
let ipv4 = desc
.reservations()
.first()
.and_then(|res| res.instances().first())
.and_then(|inst| inst.public_ip_address())
.unwrap_or("")
.to_string();
let job = Job {
id: instance_id.clone(),
ipv4: ipv4.clone(),
name: Some(name.to_string()),
};
let key_path = cfg.ssh_key_path.clone();
tokio::spawn(async move {
let _ = tokio::task::spawn_blocking(move || super::install_over_ssh(&ipv4, &key_path))
.await;
});
Ok(job)
}
async fn get_job(
&self,
job_id: &str,
) -> Result<Option<Job>, Box<dyn std::error::Error + Send + Sync>> {
let cfg = config();
let region_provider =
RegionProviderChain::first_try(Some(Region::new(cfg.location.clone())))
.or_default_provider();
let shared_config = aws_config::from_env().region(region_provider).load().await;
let client = Client::new(&shared_config);
let desc = client
.describe_instances()
.instance_ids(job_id)
.send()
.await?;
if let Some(reservation) = desc.reservations().first() {
if let Some(instance) = reservation.instances().first() {
let name_tag = instance
.tags()
.iter()
.find(|t| t.key() == Some("Name"))
.and_then(|t| t.value())
.map(ToString::to_string);
return Ok(Some(Job {
id: instance.instance_id().unwrap_or_default().to_string(),
ipv4: instance.public_ip_address().unwrap_or_default().to_string(),
name: name_tag,
}));
}
}
Ok(None)
}
async fn stop_job(
&self,
job_id: &str,
) -> Result<Job, Box<dyn std::error::Error + Send + Sync>> {
let cfg = config();
let region_provider =
RegionProviderChain::first_try(Some(Region::new(cfg.location.clone())))
.or_default_provider();
let shared_config = aws_config::from_env().region(region_provider).load().await;
let client = Client::new(&shared_config);
client
.terminate_instances()
.instance_ids(job_id)
.send()
.await?;
Ok(Job {
id: job_id.to_string(),
ipv4: String::new(),
name: None,
})
}
async fn list_jobs(&self) -> Result<Vec<Job>, Box<dyn std::error::Error + Send + Sync>> {
let cfg = config();
let region_provider =
RegionProviderChain::first_try(Some(Region::new(cfg.location.clone())))
.or_default_provider();
let shared_config = aws_config::from_env().region(region_provider).load().await;
let client = Client::new(&shared_config);
let desc = client.describe_instances().send().await?;
let mut jobs = Vec::new();
for reservation in desc.reservations() {
for instance in reservation.instances() {
let name_tag = instance
.tags()
.iter()
.find(|t| t.key() == Some("Name"))
.and_then(|t| t.value())
.map(ToString::to_string);
jobs.push(Job {
id: instance.instance_id().unwrap_or_default().to_string(),
ipv4: instance.public_ip_address().unwrap_or_default().to_string(),
name: name_tag,
});
}
}
Ok(jobs)
}
async fn tail(
&self,
job_id: &str,
filename: &str,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
let cfg = config();
if let Some(job) = self.get_job(job_id).await? {
let tcp = TcpStream::connect((job.ipv4.as_str(), 22))?;
let mut sess = Session::new()?;
sess.set_tcp_stream(tcp);
sess.handshake()?;
sess.userauth_pubkey_file("root", None, Path::new(&cfg.ssh_key_path), None)?;
let mut channel = sess.channel_session()?;
channel.exec(&format!("cat {filename}"))?;
let mut s = String::new();
channel.read_to_string(&mut s)?;
println!("{s}");
channel.wait_close()?;
println!("{}", channel.exit_status()?);
}
Ok(())
}
async fn scp(
&self,
_job_id: &str,
_filename: &str,
_destination: &str,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
Ok(())
}
}
fn config() -> Config {
Config::new()
}