1use std::{1use std::{
2 cell::OnceCell,
3 collections::BTreeMap,
2 env::current_dir,4 env::current_dir,
3 ffi::{OsStr, OsString},5 ffi::{OsStr, OsString},
4 fmt::Display,6 fmt::Display,
10};12};
1113
12use anyhow::{anyhow, bail, ensure, Context, Result};14use anyhow::{anyhow, bail, ensure, Context, Result};
13use clap::{ArgGroup, Parser};15use clap::Parser;
14use fleet_shared::SecretData;16use fleet_shared::SecretData;
15use nix_eval::{nix_go, nix_go_json, NixSessionPool, Value};17use nix_eval::{nix_go, nix_go_json, NixSessionPool, Value};
18use nom::{
19 bytes::complete::take_while1,
20 character::complete::char,
21 combinator::{map, opt},
22 multi::separated_list1,
23 sequence::{preceded, separated_pair},
24};
16use openssh::SessionBuilder;25use openssh::SessionBuilder;
17use serde::de::DeserializeOwned;26use serde::de::DeserializeOwned;
18use tempfile::NamedTempFile;27use tempfile::NamedTempFile;
53 pub name: String,62 pub name: String,
54 pub local: bool,63 pub local: bool,
55 pub session: OnceLock<Arc<openssh::Session>>,64 pub session: OnceLock<Arc<openssh::Session>>,
65 groups: OnceCell<Vec<String>>,
5666
57 pub nixos_config: Option<Value>,67 pub nixos_config: Option<Value>,
58}68}
59impl ConfigHost {69impl ConfigHost {
70 pub async fn tags(&self) -> Result<Vec<String>> {
71 if let Some(v) = self.groups.get() {
72 return Ok(v.clone());
73 }
74 // TOCTOU is possible here in case if config is changed, but this case is not handled anywhere anyway,
75 // assuming getting tags always returns the same value.
76 let Some(nixos_config) = &self.nixos_config else {
77 return Ok(vec![]);
78 };
79 let tags: Vec<String> = nix_go_json!(nixos_config.tags);
80
81 let _ = self.groups.set(tags.clone());
82
83 Ok(tags)
84 }
60 async fn open_session(&self) -> Result<Arc<openssh::Session>> {85 async fn open_session(&self) -> Result<Arc<openssh::Session>> {
61 assert!(!self.local, "do not open ssh connection to local session");86 assert!(!self.local, "do not open ssh connection to local session");
62 // FIXME: TOCTOU87 // FIXME: TOCTOU
217}242}
218243
219impl Config {244impl Config {
220 pub fn should_skip(&self, host: &str) -> bool {245 pub async fn should_skip(&self, host: &ConfigHost) -> Result<bool> {
221 if !self.opts.skip.is_empty() {246 if !self.opts.skip.is_empty() && self.opts.skip.iter().any(|h| h as &str == host.name) {
222 self.opts.skip.iter().any(|h| h as &str == host)
223 } else if !self.opts.only.is_empty() {
224 !self.opts.only.iter().any(|h| h as &str == host)247 return Ok(true);
225 } else {
226 false
227 }248 }
249 if self.opts.only.is_empty() {
250 return Ok(false);
251 }
252 let mut have_group_matches = false;
253 for item in self.opts.only.iter() {
254 match item {
255 HostItem::Host { name, .. } if *name == host.name => {
256 return Ok(false);
257 }
258 HostItem::Tag { .. } => {
259 have_group_matches = true;
260 }
261 _ => {}
262 }
263 }
264 if have_group_matches {
265 let host_tags = host.tags().await?;
266 for item in self.opts.only.iter() {
267 match item {
268 HostItem::Tag { name, .. } if host_tags.contains(name) => {
269 return Ok(false);
270 }
271 _ => {}
272 }
273 }
274 }
275 Ok(true)
228 }276 }
277 pub async fn action_attr(&self, host: &ConfigHost, attr: &str) -> Result<Option<String>> {
278 if self.opts.only.is_empty() {
279 return Ok(None);
280 }
281 let mut have_group_matches = false;
282 for item in self.opts.only.iter() {
283 match item {
284 HostItem::Host { name, attrs }
285 if *name == host.name && attrs.contains_key(attr) =>
286 {
287 return Ok(attrs.get(attr).cloned());
288 }
289 HostItem::Tag { attrs, .. } if attrs.contains_key(attr) => {
290 have_group_matches = true;
291 }
292 _ => {}
293 }
294 }
295 if have_group_matches {
296 let host_tags = host.tags().await?;
297 for item in self.opts.only.iter() {
298 match item {
299 HostItem::Tag { name, attrs }
300 if host_tags.contains(name) && attrs.contains_key(attr) =>
301 {
302 return Ok(attrs.get(attr).cloned());
303 }
304 _ => {}
305 }
306 }
307 }
308 Ok(None)
309 }
229 pub fn is_local(&self, host: &str) -> bool {310 pub fn is_local(&self, host: &str) -> bool {
230 self.opts.localhost.as_ref().map(|s| s as &str) == Some(host)311 self.opts.localhost.as_ref().map(|s| s as &str) == Some(host)
231 }312 }
237 local: true,318 local: true,
238 session: OnceLock::new(),319 session: OnceLock::new(),
239 nixos_config: None,320 nixos_config: None,
321 groups: {
322 let cell = OnceCell::new();
323 let _ = cell.set(vec![]);
324 cell
325 },
240 }326 }
241 }327 }
242328
249 local: self.is_local(name),335 local: self.is_local(name),
250 session: OnceLock::new(),336 session: OnceLock::new(),
251 nixos_config: Some(nixos_config),337 nixos_config: Some(nixos_config),
338 groups: OnceCell::new(),
252 })339 })
253 }340 }
254 pub async fn list_hosts(&self) -> Result<Vec<ConfigHost>> {341 pub async fn list_hosts(&self) -> Result<Vec<ConfigHost>> {
353 fleet_data_path.push("fleet.nix");440 fleet_data_path.push("fleet.nix");
354 tempfile.persist(fleet_data_path)?;441 tempfile.persist(fleet_data_path)?;
355 Ok(())442 Ok(())
443 }
444}
445
446#[derive(Clone)]
447enum HostItem {
448 Host {
449 name: String,
450 attrs: BTreeMap<String, String>,
451 },
452 Tag {
453 name: String,
454 attrs: BTreeMap<String, String>,
455 },
456}
457fn host_item_parser(input: &str) -> Result<HostItem, String> {
458 fn err_to_string(err: nom::Err<nom::error::Error<&str>>) -> String {
459 err.to_string()
460 }
461
462 let (input, is_tag) = map(opt(char('@')), |c| c.is_some())(input).map_err(err_to_string)?;
463 let (input, name) = map(
464 take_while1(|v| v != ',' && v != '?' && v != '@'),
465 str::to_owned,
466 )(input)
467 .map_err(err_to_string)?;
468
469 let kw_item = separated_pair(
470 map(take_while1(|v| v != '&' && v != '='), str::to_owned),
471 char('='),
472 map(take_while1(|v| v != '&'), str::to_owned),
473 );
474 let kw = map(separated_list1(char('&'), kw_item), |vec| {
475 vec.into_iter().collect::<BTreeMap<_, _>>()
476 });
477 let mut opt_kw = map(opt(preceded(char('?'), kw)), Option::unwrap_or_default);
478
479 let (input, attrs) = opt_kw(input).map_err(err_to_string)?;
480
481 if !input.is_empty() {
482 return Err(format!("unexpected trailing input: {input:?}"));
356 }483 }
484 Ok(if is_tag {
485 HostItem::Tag { name, attrs }
486 } else {
487 HostItem::Host { name, attrs }
488 })
357}489}
358490
359#[derive(Parser, Clone)]491#[derive(Parser, Clone)]
360#[clap(group = ArgGroup::new("target_hosts"))]
361pub struct FleetOpts {492pub struct FleetOpts {
362 /// All hosts except those would be skipped493 /// All hosts except those would be skipped
363 #[clap(long, number_of_values = 1, group = "target_hosts")]494 #[clap(long, number_of_values = 1, value_parser = host_item_parser)]
364 only: Vec<String>,495 only: Vec<HostItem>,
365496
366 /// Hosts to skip497 /// Hosts to skip
367 #[clap(long, number_of_values = 1, group = "target_hosts")]498 #[clap(long, number_of_values = 1)]
368 skip: Vec<String>,499 skip: Vec<String>,
369500
370 /// Host, which should be threaten as current machine501 /// Host, which should be threaten as current machine