use super::{
AgentPolicySpec, AgentPolicySummary, BTreeMap, BTreeSet, Error, ErrorKind, LiveAgentPolicy,
MutationPlan, Path, Result, Serialize, Transport, Value, agent_policies, collect, http, json,
live_from_policy, read_live, validate,
};
#[derive(Debug, Clone, PartialEq)]
pub struct AgentPolicyImportPlan {
pub preview: MutationPlan,
pub skipped: Vec<Value>,
pub package_installs: Vec<String>,
pub total: usize,
source: std::path::PathBuf,
specs: Vec<AgentPolicySpec>,
before: BTreeMap<String, Option<LiveAgentPolicy>>,
bodies: BTreeMap<String, Value>,
monitoring_package: Option<agent_policies::PackageStatus>,
overwrite: bool,
}
#[derive(Debug, Clone, PartialEq, Serialize)]
pub struct AgentPolicyImportReport {
pub applied: bool,
pub succeeded: Vec<Value>,
pub unchanged: Vec<Value>,
pub skipped: Vec<Value>,
pub failed: Vec<Value>,
pub total: usize,
pub affected_agents: u64,
pub package_installs: Vec<String>,
}
const MONITORING_PACKAGE: &str = "elastic_agent";
const SERVER_SELECTED_INSTALL: &str = "elastic_agent@server-selected";
pub fn build_replace_body(current: &AgentPolicySpec, desired: &AgentPolicySpec) -> Result<Value> {
current.validate()?;
desired.validate()?;
if current.id != desired.id {
return unsupported(
"changing agent policy id is not supported by the agent-policy update API",
);
}
let removed = [
(
"description",
current.description.is_some() && desired.description.is_none(),
),
(
"unenroll_timeout",
current.unenroll_timeout.is_some() && desired.unenroll_timeout.is_none(),
),
(
"monitoring_pprof_enabled",
current.monitoring_pprof_enabled.is_some()
&& desired.monitoring_pprof_enabled.is_none(),
),
(
"advanced_settings",
current.advanced_settings.is_some() && desired.advanced_settings.is_none(),
),
(
"monitoring_http",
current.monitoring_http.is_some() && desired.monitoring_http.is_none(),
),
(
"monitoring_diagnostics",
current.monitoring_diagnostics.is_some() && desired.monitoring_diagnostics.is_none(),
),
];
if let Some((field, _)) = removed.iter().find(|(_, gone)| *gone) {
return unsupported(format!(
"removing {field} is not supported by the agent-policy update API"
));
}
let mut body = serde_json::to_value(desired)
.map_err(|error| Error::new(ErrorKind::Error, format!("encoding agent policy: {error}")))?
.as_object()
.cloned()
.expect("specs serialize to objects");
body.remove("id");
if current.overrides.is_some() && desired.overrides.is_none() {
body.insert("overrides".into(), Value::Null);
}
if current.keep_monitoring_alive.is_some() && desired.keep_monitoring_alive.is_none() {
body.insert("keep_monitoring_alive".into(), Value::Null);
}
Ok(Value::Object(body))
}
fn unsupported<T>(message: impl Into<String>) -> Result<T> {
Err(Error::new(ErrorKind::Unsupported, message))
}
pub async fn plan_import(
transport: &Transport,
path: &Path,
overwrite: bool,
skip_existing: bool,
) -> Result<AgentPolicyImportPlan> {
let mut specs = validate(path)?;
if specs.is_empty() {
return Err(Error::new(
ErrorKind::Error,
"agent-policy import needs at least one agent policy",
));
}
if overwrite && skip_existing {
return Err(Error::new(
ErrorKind::Error,
"--overwrite and --skip-existing cannot be used together",
));
}
let total = specs.len();
let mut before_raw = BTreeMap::new();
let mut conflicts = Vec::new();
for spec in &specs {
match agent_policies::get(transport, &spec.id).await {
Ok(policy) => {
if !overwrite && !skip_existing {
conflicts.push(spec.id.clone());
}
before_raw.insert(spec.id.clone(), Some(policy));
}
Err(error) if error.kind == ErrorKind::NotFound => {
before_raw.insert(spec.id.clone(), None);
}
Err(error) => return Err(error),
}
}
let mut live_names = BTreeMap::new();
for item in collect(transport).await? {
let row = AgentPolicySummary::from_item(&item)?;
if live_names
.insert(row.name.clone(), row.id.clone())
.is_some()
{
return Err(http(format!(
"decoding agent policies list: duplicate name '{}'",
row.name
)));
}
}
let taken: Vec<String> = specs
.iter()
.filter_map(|spec| {
live_names
.get(&spec.name)
.filter(|owner| **owner != spec.id)
.map(|owner| format!("{} ({owner})", spec.name))
})
.collect();
if !taken.is_empty() {
return Err(Error::new(
ErrorKind::Conflict,
format!("agent policy names already exist: {}", taken.join(", ")),
));
}
if !conflicts.is_empty() {
return Err(Error::new(
ErrorKind::Conflict,
format!("agent policies already exist: {}", conflicts.join(", ")),
));
}
let mut skipped = Vec::new();
if skip_existing {
specs.retain(|spec| match before_raw.get(&spec.id) {
Some(Some(_)) => {
skipped.push(json!({"id": spec.id, "reason": "exists"}));
false
}
_ => true,
});
before_raw.retain(|id, _| specs.iter().any(|spec| spec.id == *id));
}
let mut before = BTreeMap::new();
for (id, raw) in before_raw {
let live = match raw {
Some(policy) => Some(live_from_policy(&policy, &id, transport.space())?),
None => None,
};
before.insert(id, live);
}
let mut bodies = BTreeMap::new();
for spec in &specs {
if let Some(Some(current)) = before.get(&spec.id)
&& current.spec != *spec
{
bodies.insert(spec.id.clone(), build_replace_body(¤t.spec, spec)?);
}
}
let mut package_installs = Vec::new();
let needs_monitoring = specs
.iter()
.any(|spec| monitoring_can_install(before.get(&spec.id).and_then(Option::as_ref), spec));
let monitoring_package = if needs_monitoring {
let status = agent_policies::package_status(transport, MONITORING_PACKAGE).await?;
if status.status != "installed" {
package_installs.push(SERVER_SELECTED_INSTALL.to_string());
}
Some(status)
} else {
None
};
let preview = MutationPlan {
preview_action: format!(
"Import {} agent policy(ies) from {}",
specs.len(),
path.display()
),
preview_details: import_details(&specs, &before, &package_installs),
targets: specs.iter().map(|spec| spec.id.clone()).collect(),
};
Ok(AgentPolicyImportPlan {
preview,
skipped,
package_installs,
total,
source: path.to_path_buf(),
specs,
before,
bodies,
monitoring_package,
overwrite,
})
}
fn monitoring_can_install(current: Option<&LiveAgentPolicy>, desired: &AgentPolicySpec) -> bool {
!desired.monitoring_enabled.is_empty()
&& current.is_none_or(|current| current.spec.monitoring_enabled.is_empty())
}
fn import_details(
specs: &[AgentPolicySpec],
before: &BTreeMap<String, Option<LiveAgentPolicy>>,
package_installs: &[String],
) -> Vec<String> {
let mut details: Vec<String> = specs
.iter()
.filter_map(|spec| match before.get(&spec.id) {
Some(None) => Some(format!("{} create {}", spec.id, spec.name)),
Some(Some(current)) if current.spec == *spec => {
Some(format!("{} unchanged {}", spec.id, spec.name))
}
Some(Some(current)) => {
let name = if current.spec.name == spec.name {
spec.name.clone()
} else {
format!("{} -> {}", current.spec.name, spec.name)
};
Some(format!(
"{} replace {name} agents {}",
spec.id, current.agents
))
}
None => None,
})
.collect();
details.extend(
package_installs
.iter()
.map(|install| format!("package install {install}")),
);
details
}
pub async fn apply_import(
transport: &Transport,
plan: &AgentPolicyImportPlan,
) -> Result<AgentPolicyImportReport> {
validate_import_plan(plan)?;
let mut succeeded = Vec::new();
let mut unchanged = Vec::new();
let mut failed = Vec::new();
let mut affected_agents = 0;
let mut expected_package = plan.monitoring_package.clone();
let mut package_installs = Vec::new();
for desired in &plan.specs {
let Some(before) = plan.before.get(&desired.id) else {
failed.push(failed_row(&desired.id, false, "missing preflight snapshot"));
continue;
};
let current = match read_live(transport, &desired.id).await {
Ok(live) => Some(live),
Err(error) if error.kind == ErrorKind::NotFound => None,
Err(error) => {
failed.push(failed_row(&desired.id, false, error.message));
continue;
}
};
match (before, current) {
(None, Some(_)) => failed.push(failed_row(
&desired.id,
false,
"agent policy appeared since preview",
)),
(Some(_), None) => failed.push(failed_row(
&desired.id,
false,
"agent policy disappeared since preview",
)),
(Some(before), Some(live)) if before != &live => failed.push(failed_row(
&desired.id,
false,
"agent policy changed since preview",
)),
(before, current) => {
let package_can_change = monitoring_can_install(before.as_ref(), desired);
if package_can_change {
let Some(expected) = expected_package.as_ref() else {
failed.push(failed_row(
&desired.id,
false,
"missing monitoring package snapshot",
));
continue;
};
match agent_policies::package_status(transport, MONITORING_PACKAGE).await {
Ok(actual) if actual == *expected => {}
Ok(_) => {
failed.push(failed_row(
&desired.id,
false,
"elastic_agent package changed since preview",
));
continue;
}
Err(error) => {
failed.push(failed_row(&desired.id, false, error.message));
continue;
}
}
}
let (action, applied, route_error) = match (before, current) {
(None, None) => {
match other_owner_of_name(transport, &desired.name).await {
Ok(Some(owner)) => {
failed.push(failed_row(
&desired.id,
false,
format!(
"agent policy name appeared since preview: {} ({owner})",
desired.name
),
));
continue;
}
Ok(None) => {}
Err(error) => {
failed.push(failed_row(&desired.id, false, error));
continue;
}
}
match agent_policies::create(transport, desired).await {
Ok(_) => ("created", true, None),
Err(error) => ("created", false, Some(error.message)),
}
}
(Some(before), Some(_)) if before.spec == *desired => {
unchanged.push(json!({"id": desired.id}));
continue;
}
(Some(before), Some(_)) => {
let body = plan
.bodies
.get(&desired.id)
.expect("validated replacement body");
match agent_policies::update(transport, &desired.id, body).await {
Ok(_) => {
affected_agents += before.agents;
("replaced", true, None)
}
Err(error) => ("replaced", false, Some(error.message)),
}
}
_ => unreachable!("appearance and disappearance handled above"),
};
let stored_error = if applied {
verify_stored(transport, desired).await.err()
} else {
None
};
let package_error = if package_can_change {
let expected = expected_package
.as_ref()
.expect("validated package snapshot");
match observe_package_after_write(transport, expected).await {
Ok((after, installed)) => {
expected_package = Some(after);
if let Some(installed) = installed
&& !package_installs.contains(&installed)
{
package_installs.push(installed);
}
None
}
Err(error) => Some(error),
}
} else {
None
};
let errors = [route_error, stored_error, package_error]
.into_iter()
.flatten()
.collect::<Vec<_>>();
if errors.is_empty() {
succeeded.push(json!({"id": desired.id, "action": action}));
} else {
failed.push(failed_row(&desired.id, applied, errors.join("; ")));
}
}
}
}
Ok(AgentPolicyImportReport {
applied: true,
succeeded,
unchanged,
skipped: plan.skipped.clone(),
failed,
total: plan.total,
affected_agents,
package_installs,
})
}
async fn verify_stored(
transport: &Transport,
desired: &AgentPolicySpec,
) -> std::result::Result<(), String> {
match read_live(transport, &desired.id).await {
Ok(live) if live.spec == *desired => Ok(()),
Ok(_) => Err("server stored a different agent-policy spec".into()),
Err(error) => Err(error.message),
}
}
async fn other_owner_of_name(
transport: &Transport,
name: &str,
) -> std::result::Result<Option<String>, String> {
let items = collect(transport).await.map_err(|error| error.message)?;
for item in &items {
let row = AgentPolicySummary::from_item(item).map_err(|error| error.message)?;
if row.name == name {
return Ok(Some(row.id));
}
}
Ok(None)
}
async fn observe_package_after_write(
transport: &Transport,
before: &agent_policies::PackageStatus,
) -> std::result::Result<(agent_policies::PackageStatus, Option<String>), String> {
let after = agent_policies::package_status(transport, MONITORING_PACKAGE)
.await
.map_err(|error| error.message)?;
if before.status != "installed" && after.status == "installed" {
let version = after
.installed_version
.clone()
.expect("decoder requires installed version");
return Ok((after, Some(format!("{MONITORING_PACKAGE}@{version}"))));
}
Ok((after, None))
}
fn failed_row(id: &str, applied: bool, error: impl Into<String>) -> Value {
json!({"id": id, "applied": applied, "error": error.into()})
}
fn validate_import_plan(plan: &AgentPolicyImportPlan) -> Result<()> {
let invalid = |message: &str| {
Err(Error::new(
ErrorKind::Error,
format!("invalid agent-policy import plan: {message}"),
))
};
if plan.total == 0 || plan.total != plan.specs.len() + plan.skipped.len() {
return invalid("total does not equal pending and skipped agent policies");
}
let mut previous_id: Option<&str> = None;
let mut names = BTreeSet::new();
let mut expected_body_ids = BTreeSet::new();
for spec in &plan.specs {
spec.validate()?;
if previous_id.is_some_and(|previous| previous >= spec.id.as_str()) {
return invalid("pending agent policies must be unique and sorted by id");
}
previous_id = Some(&spec.id);
if !names.insert(spec.name.as_str()) {
return invalid("pending agent-policy names must be unique");
}
let Some(before) = plan.before.get(&spec.id) else {
return invalid("preflight snapshots do not match pending agent policies");
};
if let Some(current) = before {
current.spec.validate()?;
if current.spec.id != spec.id {
return invalid("live snapshot id does not match its pending policy");
}
if current.attached.windows(2).any(|ids| ids[0] >= ids[1]) {
return invalid("attached integration ids must be unique and sorted");
}
}
match before {
None if plan.bodies.contains_key(&spec.id) => {
return invalid("planned creates must not carry a replacement body");
}
None => {}
Some(current) if current.spec == *spec => {
if plan.bodies.contains_key(&spec.id) {
return invalid("unchanged agent policies must not carry a replacement body");
}
}
Some(current) => {
expected_body_ids.insert(spec.id.as_str());
if !plan.overwrite {
return invalid("replacement plan requires overwrite");
}
if plan.bodies.get(&spec.id) != Some(&build_replace_body(¤t.spec, spec)?) {
return invalid("replacement body does not match its snapshots");
}
}
}
}
if plan.before.len() != plan.specs.len() {
return invalid("preflight snapshots do not match pending agent policies");
}
if plan
.bodies
.keys()
.map(String::as_str)
.collect::<BTreeSet<_>>()
!= expected_body_ids
{
return invalid("replacement bodies do not match changed agent policies");
}
let mut previous_skipped: Option<&str> = None;
for skipped in &plan.skipped {
let object = skipped.as_object().ok_or_else(|| {
Error::new(
ErrorKind::Error,
"invalid agent-policy import plan: skipped row must be an object",
)
})?;
if object.len() != 2 || object.get("reason").and_then(Value::as_str) != Some("exists") {
return invalid("skipped rows must contain only id and reason exists");
}
let id = object
.get("id")
.and_then(Value::as_str)
.filter(|id| !id.is_empty())
.ok_or_else(|| {
Error::new(
ErrorKind::Error,
"invalid agent-policy import plan: skipped id must be non-empty",
)
})?;
if previous_skipped.is_some_and(|previous| previous >= id) {
return invalid("skipped agent policies must be unique and sorted by id");
}
if plan.before.contains_key(id) {
return invalid("an agent policy cannot be both pending and skipped");
}
previous_skipped = Some(id);
}
let needs_monitoring = plan.specs.iter().any(|spec| {
monitoring_can_install(plan.before.get(&spec.id).and_then(Option::as_ref), spec)
});
match (&plan.monitoring_package, needs_monitoring) {
(Some(status), true) if status.name == MONITORING_PACKAGE => {
let expected = if status.status == "installed" {
Vec::new()
} else {
vec![SERVER_SELECTED_INSTALL.to_string()]
};
if status.status == "installed" && status.installed_version.is_none() {
return invalid("installed monitoring package needs an exact version");
}
if plan.package_installs != expected {
return invalid("monitoring package preview does not match its snapshot");
}
}
(None, false) if plan.package_installs.is_empty() => {}
_ => return invalid("monitoring package snapshot does not match pending transitions"),
}
let expected_preview = MutationPlan {
preview_action: format!(
"Import {} agent policy(ies) from {}",
plan.specs.len(),
plan.source.display()
),
preview_details: import_details(&plan.specs, &plan.before, &plan.package_installs),
targets: plan.specs.iter().map(|spec| spec.id.clone()).collect(),
};
if plan.preview != expected_preview {
return invalid("preview does not match the canonical plan");
}
Ok(())
}