use std::collections::HashSet;
use std::str::FromStr;
use std::sync::atomic::{AtomicBool, Ordering};
use super::aws::{AwsCredentials, SigV4Request, format_utc, sign_request};
use super::{ProviderError, map_ureq_error};
const SSM_SERVICE: &str = "ssm";
const SSM_TARGET_DESCRIBE: &str = "AmazonSSM.DescribeInstanceInformation";
const SSM_CONTENT_TYPE: &str = "application/x-amz-json-1.1";
const SSM_MAX_RESULTS: u32 = 50;
const SSM_MAX_PAGES: usize = 500;
const SSM_SSH_DOCUMENT: &str = "AWS-StartSSHSession";
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum SsmMode {
#[default]
Off,
Auto,
Always,
}
impl SsmMode {
pub fn as_str(self) -> &'static str {
match self {
SsmMode::Off => "off",
SsmMode::Auto => "auto",
SsmMode::Always => "always",
}
}
pub const ALL: &'static [SsmMode] = &[SsmMode::Off, SsmMode::Auto, SsmMode::Always];
pub fn next(self) -> SsmMode {
match self {
SsmMode::Off => SsmMode::Auto,
SsmMode::Auto => SsmMode::Always,
SsmMode::Always => SsmMode::Off,
}
}
pub fn is_enabled(self) -> bool {
!matches!(self, SsmMode::Off)
}
}
impl FromStr for SsmMode {
type Err = ();
fn from_str(s: &str) -> Result<Self, Self::Err> {
match s.trim().to_ascii_lowercase().as_str() {
"auto" => Ok(SsmMode::Auto),
"always" | "true" | "yes" | "on" => Ok(SsmMode::Always),
_ => Ok(SsmMode::Off),
}
}
}
impl std::fmt::Display for SsmMode {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(self.as_str())
}
}
pub fn is_safe_profile_name(profile: &str) -> bool {
!profile.is_empty()
&& profile
.chars()
.all(|c| c.is_ascii_alphanumeric() || c == '.' || c == '-' || c == '_')
}
pub fn proxy_command(profile: &str, region: &str) -> String {
let profile_part = if profile.is_empty() {
String::new()
} else {
format!(" --profile {}", profile)
};
format!("{}{}{}", command_head(), profile_part, command_tail(region))
}
fn command_head() -> String {
format!(
"sh -c \"aws ssm start-session --target %h --document-name {} --parameters 'portNumber=%p'",
SSM_SSH_DOCUMENT
)
}
fn command_tail(region: &str) -> String {
if region.is_empty() {
"\"".to_string()
} else {
format!(" --region {}\"", region)
}
}
pub(super) fn is_generated_proxy_command(value: &str, region: &str) -> bool {
let Some(rest) = value.strip_prefix(&command_head()) else {
return false;
};
let Some(middle) = rest.strip_suffix(&command_tail(region)) else {
return false;
};
middle.is_empty()
|| middle
.strip_prefix(" --profile ")
.is_some_and(is_safe_profile_name)
}
pub(super) fn region_endpoint(region: &str) -> String {
format!("https://ssm.{}.amazonaws.com", region)
}
pub(super) fn online_nodes_with_endpoint(
agent: &ureq::Agent,
creds: &AwsCredentials,
region: &str,
cancel: &AtomicBool,
endpoint: &str,
) -> Result<HashSet<String>, ProviderError> {
let mut ids = HashSet::new();
let mut next_token: Option<String> = None;
for _ in 0..SSM_MAX_PAGES {
if cancel.load(Ordering::Relaxed) {
return Err(ProviderError::Cancelled);
}
let body = request_body(next_token.as_deref());
let text = post(agent, creds, region, endpoint, &body)?;
let parsed: DescribeInstanceInformationResponse = serde_json::from_str(&text)
.map_err(|e| ProviderError::Parse(format!("{}: {}", region, e)))?;
for node in parsed.instance_information_list {
let (Some(id), Some(status)) = (node.instance_id, node.ping_status) else {
continue;
};
if status == "Online" {
ids.insert(id);
}
}
match parsed.next_token {
Some(token) if !token.is_empty() => next_token = Some(token),
_ => return Ok(ids),
}
}
Ok(ids)
}
fn request_body(next_token: Option<&str>) -> String {
let filters = format!(
r#""Filters":[{{"Key":"PingStatus","Values":["Online"]}}],"MaxResults":{}"#,
SSM_MAX_RESULTS
);
match next_token {
Some(token) => format!(r#"{{{},"NextToken":{}}}"#, filters, json_string(token)),
None => format!("{{{}}}", filters),
}
}
fn json_string(value: &str) -> String {
let mut out = String::with_capacity(value.len() + 2);
out.push('"');
for c in value.chars() {
match c {
'"' => out.push_str("\\\""),
'\\' => out.push_str("\\\\"),
'\n' => out.push_str("\\n"),
'\r' => out.push_str("\\r"),
'\t' => out.push_str("\\t"),
c if (c as u32) < 0x20 => out.push_str(&format!("\\u{:04x}", c as u32)),
c => out.push(c),
}
}
out.push('"');
out
}
fn post(
agent: &ureq::Agent,
creds: &AwsCredentials,
region: &str,
endpoint: &str,
body: &str,
) -> Result<String, ProviderError> {
let host = endpoint
.split_once("://")
.map(|(_, authority)| authority)
.unwrap_or(endpoint)
.to_string();
let epoch = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs();
let (timestamp, datestamp) = format_utc(epoch);
let auth = sign_request(
creds,
region,
&SigV4Request {
method: "POST",
service: SSM_SERVICE,
host: &host,
query_string: "",
payload: body.as_bytes(),
extra_headers: &[
("content-type", SSM_CONTENT_TYPE),
("x-amz-target", SSM_TARGET_DESCRIBE),
],
},
×tamp,
&datestamp,
);
let url = format!("{}/", endpoint);
let mut req = agent
.post(&url)
.header("Authorization", &auth)
.header("x-amz-date", ×tamp)
.header("content-type", SSM_CONTENT_TYPE)
.header("x-amz-target", SSM_TARGET_DESCRIBE);
if let Some(token) = &creds.session_token {
req = req.header("x-amz-security-token", token);
}
let mut resp = req
.config()
.http_status_as_error(false)
.build()
.send(body)
.map_err(map_ureq_error)?;
let status = resp.status();
let text = resp
.body_mut()
.read_to_string()
.map_err(|e| ProviderError::Parse(e.to_string()))?;
if !status.is_success() {
let code = status.as_u16();
if code == super::HTTP_TOO_MANY_REQUESTS {
log::warn!("[external] aws ssm: HTTP {} rate limited", code);
return Err(ProviderError::RateLimited);
}
let detail = parse_error(&text).unwrap_or_else(|| format!("HTTP {}", code));
log::warn!("[external] aws ssm: {}", detail);
return Err(ProviderError::Execute(detail));
}
Ok(text)
}
fn capped(text: String) -> String {
if text.chars().count() > super::aws_sts::MAX_ERROR_MESSAGE {
return text
.chars()
.take(super::aws_sts::MAX_ERROR_MESSAGE)
.collect();
}
text
}
fn parse_error(body: &str) -> Option<String> {
let parsed: serde_json::Value = serde_json::from_str(body).ok()?;
let code = parsed.get("__type")?.as_str()?;
let code = code.rsplit('#').next().unwrap_or(code).trim();
if code.is_empty() {
return None;
}
let message = parsed
.get("message")
.or_else(|| parsed.get("Message"))
.and_then(serde_json::Value::as_str)
.unwrap_or_default()
.trim();
if message.is_empty() {
return Some(capped(code.to_string()));
}
Some(capped(format!("{}: {}", code, message)))
}
#[derive(serde::Deserialize, Debug, Default)]
struct DescribeInstanceInformationResponse {
#[serde(rename = "InstanceInformationList", default)]
instance_information_list: Vec<InstanceInformation>,
#[serde(rename = "NextToken", default)]
next_token: Option<String>,
}
#[derive(serde::Deserialize, Debug)]
struct InstanceInformation {
#[serde(rename = "InstanceId", default)]
instance_id: Option<String>,
#[serde(rename = "PingStatus", default)]
ping_status: Option<String>,
}
#[cfg(test)]
#[path = "aws_ssm_tests.rs"]
mod tests;