1use super::reports::PullReport;
2use crate::codec::Format;
3use crate::exceptions;
4use crate::model::{ExceptionItem, ExceptionList, ListKey, Rule};
5use crate::normalize;
6use crate::rules::RuleSource;
7use crate::state::mirror::{encode_list_file, encode_rule_file};
8use crate::state::transaction::{
9 PullLock, StagedFile, acquire_pull, recover_pull, replace_staged_files,
10};
11use elasticctl_core::{Error, ErrorKind, Result, Transport};
12use serde_json::Value;
13use std::collections::{BTreeMap, BTreeSet};
14use std::path::Path;
15
16pub async fn pull(
17 t: &Transport,
18 dir: &Path,
19 format: Format,
20 selectors: &[String],
21 tag: Option<&str>,
22 search: Option<&str>,
23 source: RuleSource,
24) -> Result<PullReport> {
25 let lock = acquire_pull(dir)?;
26 recover_pull(&lock)?;
27
28 let scope = super::scope_of(t, selectors, tag, search, source, &[], "pull").await?;
29 let mut remote = scope.remote(t).await?;
30 if remote.is_empty()
31 && !scope.is_scoped()
32 && matches!(scope.source, RuleSource::Custom | RuleSource::Prebuilt)
33 {
34 crate::rules::verify_source_partition(t).await?;
35 }
36 normalize::sort_rules(&mut remote);
37
38 let wanted: BTreeSet<ListKey> = super::referenced_keys(&remote);
39 let mut fetched: BTreeMap<ListKey, (ExceptionList, Vec<ExceptionItem>)> = BTreeMap::new();
40 for key in &wanted {
41 let list = match exceptions::get_list(t, key).await {
42 Ok(list) => list,
43 Err(e) if e.kind == ErrorKind::NotFound => {
44 return Err(Error::new(
45 ErrorKind::NotFound,
46 format!(
47 "a rule references exception list \"{}\" ({}), which does not exist on \
48 this stack",
49 key.list_id, key.namespace_type
50 ),
51 ));
52 }
53 Err(e) => return Err(e),
54 };
55 let items = exceptions::find_items(t, key).await?;
56 fetched.insert(key.clone(), (list, items));
57 }
58
59 let all_items: Vec<_> = fetched
60 .values()
61 .flat_map(|(_, items)| items.iter().cloned())
62 .collect();
63 let (value_lists, value_list_warnings) = super::values::pull_lists(t, &all_items).await?;
64
65 let (pulled, item_count) = write_mirror(&lock, &remote, &fetched, &value_lists, format)?;
66
67 let target = super::rules_dir(dir);
68
69 Ok(PullReport {
70 pulled,
71 exception_lists: fetched.len(),
72 exception_items: item_count,
73 value_lists: value_lists.len(),
74 value_list_warnings,
75 dir: target.display().to_string(),
76 selected: scope.is_scoped().then(|| scope.selected()),
77 })
78}
79
80fn write_mirror(
81 lock: &PullLock,
82 remote: &[Rule],
83 fetched: &BTreeMap<ListKey, (ExceptionList, Vec<ExceptionItem>)>,
84 value_lists: &[crate::value_lists::LocalPair],
85 format: Format,
86) -> Result<(usize, usize)> {
87 let ext = match format {
88 Format::Yaml => "yaml",
89 Format::Ndjson => "ndjson",
90 };
91
92 let mut lists_to_write: Vec<(ListKey, ExceptionList)> = Vec::new();
93 let mut inlines: Vec<(ListKey, ExceptionList)> = Vec::new();
94 let mut item_count = 0usize;
95 for (key, (list, items)) in fetched {
96 let mut canonical = normalize::canonical_list(list);
97 let canonical_items: Vec<ExceptionItem> =
98 items.iter().map(normalize::canonical_item).collect();
99 item_count += canonical_items.len();
100 let items_value = Value::Array(
101 canonical_items
102 .iter()
103 .map(|i| i.clone().into_value())
104 .collect(),
105 );
106 canonical.as_map_mut().insert("items".into(), items_value);
107 if list.list_type() == "rule_default" {
108 inlines.push((key.clone(), canonical));
109 } else {
110 lists_to_write.push((key.clone(), canonical));
111 }
112 }
113
114 let mut inline_by_rule: BTreeMap<String, Vec<ExceptionList>> = BTreeMap::new();
115 for (key, list) in &inlines {
116 let owner = remote.iter().find_map(|rule| {
117 let matches = crate::model::exception_refs(rule)
118 .iter()
119 .any(|rf| rf.list_id == key.list_id && rf.namespace_type == key.namespace_type);
120 matches
121 .then(|| rule.rule_id().ok())
122 .flatten()
123 .map(str::to_string)
124 });
125 if let Some(rule_id) = owner {
126 inline_by_rule
127 .entry(rule_id)
128 .or_default()
129 .push(list.clone());
130 }
131 }
132
133 let mut claimed_rules: BTreeMap<String, Vec<String>> = BTreeMap::new();
134 let mut planned_rules: Vec<(String, Rule)> = Vec::with_capacity(remote.len());
135 let mut collisions: Vec<String> = crate::value_lists::stem_collisions(
136 value_lists
137 .iter()
138 .map(|pair| (pair.container.id.as_str(), pair.stem.as_str())),
139 )
140 .into_iter()
141 .map(|(a, b)| {
142 format!(
143 "value lists {a:?} and {b:?} both use stem {:?}",
144 crate::value_lists::file_stem(&a)
145 )
146 })
147 .collect();
148
149 for rule in remote {
150 let canonical = normalize::canonical(rule);
151 let rule_id = canonical.rule_id()?.to_string();
152 let filename = super::safe_filename(&rule_id, ext);
153
154 let previous = claimed_rules.entry(filename.clone()).or_default();
155 for other in previous.iter() {
156 collisions.push(format!(
157 "\"{other}\" and \"{rule_id}\" both sanitise to \"{filename}\""
158 ));
159 }
160 if previous.is_empty() {
161 planned_rules.push((filename, canonical));
162 }
163 previous.push(rule_id);
164 }
165
166 let mut claimed_lists: BTreeMap<String, Vec<String>> = BTreeMap::new();
167 for (key, _) in &lists_to_write {
168 let qualified = format!("{} ({})", key.list_id, key.namespace_type);
169 let filename = super::safe_filename(&key.list_id, ext);
170 let previous = claimed_lists.entry(filename.clone()).or_default();
171 for other in previous.iter() {
172 collisions.push(format!(
173 "\"{other}\" and \"{qualified}\" both sanitise to \"{filename}\""
174 ));
175 }
176 previous.push(qualified);
177 }
178
179 if !collisions.is_empty() {
180 return Err(Error::new(
181 ErrorKind::Conflict,
182 format!(
183 "{} filename collision(s); rename one id in each pair: {}",
184 collisions.len(),
185 collisions.join("; ")
186 ),
187 ));
188 }
189
190 let mut staged = Vec::with_capacity(planned_rules.len() + lists_to_write.len());
191 for (filename, canonical) in &planned_rules {
192 let rule_id = canonical.rule_id()?;
193 let inline = inline_by_rule.get(rule_id).cloned().unwrap_or_default();
194 let body = encode_rule_file(canonical, &inline, format)?;
195 staged.push(StagedFile {
196 relative: Path::new("rules").join(filename),
197 bytes: body.into_bytes(),
198 });
199 }
200
201 for (key, list) in &lists_to_write {
202 let filename = super::safe_filename(&key.list_id, ext);
203 let body = encode_list_file(list, format)?;
204 staged.push(StagedFile {
205 relative: Path::new("exceptions").join(filename),
206 bytes: body.into_bytes(),
207 });
208 }
209 for pair in value_lists {
210 staged.push(StagedFile {
211 relative: Path::new("value-lists").join(format!("{}.yaml", pair.stem)),
212 bytes: crate::value_lists::encode_container(&pair.container)?.into_bytes(),
213 });
214 staged.push(StagedFile {
215 relative: Path::new("value-lists").join(format!("{}.txt", pair.stem)),
216 bytes: crate::value_lists::encode_values(&pair.values).into_bytes(),
217 });
218 }
219 replace_staged_files(lock, &staged)?;
220
221 Ok((planned_rules.len(), item_count))
222}
223
224#[cfg(test)]
225mod tests;