1use std::{2 env::current_dir,3 ffi::{OsStr, OsString},4 io::Write,5 ops::Deref,6 path::PathBuf,7 sync::{Arc, Mutex, MutexGuard, OnceLock},8};910use anyhow::{anyhow, bail, Context, Result};11use clap::{ArgGroup, Parser};12use openssh::SessionBuilder;13use tempfile::NamedTempFile;1415use crate::{16 better_nix_eval::{Field, NixSessionPool},17 command::MyCommand,18 fleetdata::{FleetData, FleetSecret, FleetSharedSecret},19 nix_go, nix_go_json,20};2122pub struct FleetConfigInternals {23 pub local_system: String,24 pub directory: PathBuf,25 pub opts: FleetOpts,26 pub data: Mutex<FleetData>,27 pub nix_args: Vec<OsString>,28 29 pub fleet_field: Field,30 31 pub config_field: Field,32 33 pub config_unchecked_field: Field,34}3536#[derive(Clone)]37pub struct Config(Arc<FleetConfigInternals>);3839impl Deref for Config {40 type Target = FleetConfigInternals;4142 fn deref(&self) -> &Self::Target {43 &self.044 }45}4647pub struct ConfigHost {48 pub name: String,49 pub session: OnceLock<Arc<openssh::Session>>,50}51impl ConfigHost {52 pub async fn open_session(&self) -> Result<Arc<openssh::Session>> {53 54 if let Some(session) = &self.session.get() {55 return Ok((*session).clone());56 };57 let session = SessionBuilder::default();5859 let session = session60 .connect(&self.name)61 .await62 .map_err(|e| anyhow!("ssh error: {e}"))?;63 let session = Arc::new(session);64 self.session.set(session.clone()).expect("TOCTOU happened");65 Ok(session)66 }67 pub async fn mktemp_dir(&self) -> Result<String> {68 let mut cmd = self.cmd("mktemp").await?;69 cmd.arg("-d");70 let path = cmd.run_string().await?;71 Ok(path.trim_end().to_owned())72 }73 pub async fn read_file_bin(&self, path: impl AsRef<OsStr>) -> Result<Vec<u8>> {74 let mut cmd = self.cmd("cat").await?;75 cmd.arg(path);76 cmd.run_bytes().await77 }78 pub async fn read_file_text(&self, path: impl AsRef<OsStr>) -> Result<String> {79 let mut cmd = self.cmd("cat").await?;80 cmd.arg(path);81 cmd.run_string().await82 }83 pub async fn cmd(&self, cmd: impl AsRef<OsStr>) -> Result<MyCommand> {84 let session = self.open_session().await?;85 Ok(MyCommand::new_on(cmd, session))86 }8788 pub async fn decrypt(&self, data: Vec<u8>) -> Result<Vec<u8>> {89 let mut cmd = self.cmd("fleet-install-secrets").await?;90 cmd.arg("decrypt").eqarg("--secret", z85::encode(&data));91 let encoded = cmd92 .sudo()93 .run_string()94 .await95 .context("failed to call remote host for decrypt")?;96 z85::decode(encoded.trim_end()).context("bad encoded data? outdated host?")97 }98}99100impl Config {101 pub fn should_skip(&self, host: &str) -> bool {102 if !self.opts.skip.is_empty() {103 self.opts.skip.iter().any(|h| h as &str == host)104 } else if !self.opts.only.is_empty() {105 !self.opts.only.iter().any(|h| h as &str == host)106 } else {107 false108 }109 }110 pub fn is_local(&self, host: &str) -> bool {111 self.opts.localhost.as_ref().map(|s| s as &str) == Some(host)112 }113114 pub async fn run_on(&self, host: &str, mut command: MyCommand, sudo: bool) -> Result<()> {115 if sudo {116 command = command.sudo();117 }118 if !self.is_local(host) {119 command = command.ssh(host);120 }121 command.run().await122 }123 pub async fn run_string_on(124 &self,125 host: &str,126 mut command: MyCommand,127 sudo: bool,128 ) -> Result<String> {129 if sudo {130 command = command.sudo();131 }132 if !self.is_local(host) {133 command = command.ssh(host);134 }135 command.run_string().await136 }137138 pub async fn host(&self, name: &str) -> Result<ConfigHost> {139 Ok(ConfigHost {140 name: name.to_owned(),141 session: OnceLock::new(),142 })143 }144 pub async fn list_hosts(&self) -> Result<Vec<ConfigHost>> {145 let fleet_field = &self.fleet_field;146 let names = nix_go!(fleet_field.configuredHosts).list_fields().await?;147 let mut out = vec![];148 for name in names {149 out.push(ConfigHost {150 name,151 session: OnceLock::new(),152 })153 }154 Ok(out)155 }156 pub async fn system_config(&self, host: &str) -> Result<Field> {157 let fleet_field = &self.fleet_field;158 Ok(nix_go!(fleet_field.configuredSystems[{ host }].config))159 }160161 pub(super) fn data(&self) -> MutexGuard<FleetData> {162 self.data.lock().unwrap()163 }164 pub(super) fn data_mut(&self) -> MutexGuard<FleetData> {165 self.data.lock().unwrap()166 }167 168 pub async fn list_configured_shared(&self) -> Result<Vec<String>> {169 let config_field = &self.config_field;170 nix_go!(config_field.sharedSecrets).list_fields().await171 }172 173 pub fn list_shared(&self) -> Vec<String> {174 let data = self.data();175 data.shared_secrets.keys().cloned().collect()176 }177 pub fn has_shared(&self, name: &str) -> bool {178 let data = self.data();179 data.shared_secrets.contains_key(name)180 }181 pub fn replace_shared(&self, name: String, shared: FleetSharedSecret) {182 let mut data = self.data_mut();183 data.shared_secrets.insert(name.to_owned(), shared);184 }185 pub fn remove_shared(&self, secret: &str) {186 let mut data = self.data_mut();187 data.shared_secrets.remove(secret);188 }189190 pub fn has_secret(&self, host: &str, secret: &str) -> bool {191 let data = self.data();192 let Some(host_secrets) = data.host_secrets.get(host) else {193 return false;194 };195 host_secrets.contains_key(secret)196 }197 pub fn insert_secret(&self, host: &str, secret: String, value: FleetSecret) {198 let mut data = self.data_mut();199 let host_secrets = data.host_secrets.entry(host.to_owned()).or_default();200 host_secrets.insert(secret, value);201 }202203 pub async fn reencrypt_on_host(204 &self,205 host: &str,206 data: Vec<u8>,207 targets: Vec<String>,208 ) -> Result<Vec<u8>> {209 let data = z85::encode(&data);210 let mut recmd = MyCommand::new("fleet-install-secrets");211 recmd.arg("reencrypt").eqarg("--secret", data);212 for target in targets {213 recmd.eqarg("--targets", target);214 }215 recmd = recmd.sudo().ssh(host);216 let encoded = recmd217 .run_string()218 .await219 .context("failed to call remote host for decrypt")?220 .trim()221 .to_owned();222 z85::decode(encoded).context("bad encoded data? outdated host?")223 }224225 pub fn host_secret(&self, host: &str, secret: &str) -> Result<FleetSecret> {226 let data = self.data();227 let Some(host_secrets) = data.host_secrets.get(host) else {228 bail!("no secrets for machine {host}");229 };230 let Some(secret) = host_secrets.get(secret) else {231 bail!("machine {host} has no secret {secret}");232 };233 Ok(secret.clone())234 }235 pub fn shared_secret(&self, secret: &str) -> Result<FleetSharedSecret> {236 let data = self.data();237 let Some(secret) = data.shared_secrets.get(secret) else {238 bail!("no shared secret {secret}");239 };240 Ok(secret.clone())241 }242 pub async fn shared_secret_expected_owners(&self, secret: &str) -> Result<Vec<String>> {243 let config_field = &self.config_field;244 Ok(nix_go_json!(245 config_field.sharedSecrets[{ secret }].expectedOwners246 ))247 }248249 pub fn save(&self) -> Result<()> {250 let mut tempfile = NamedTempFile::new_in(self.directory.clone())?;251 let data = nixlike::serialize(&self.data() as &FleetData)?;252 tempfile.write_all(253 format!(254 "# This file contains fleet state and shouldn't be edited by hand\n\n{}\n\n# vim: ts=2 et nowrap\n",255 data256 )257 .as_bytes(),258 )?;259 let mut fleet_data_path = self.directory.clone();260 fleet_data_path.push("fleet.nix");261 tempfile.persist(fleet_data_path)?;262 Ok(())263 }264}265266#[derive(Parser, Clone)]267#[clap(group = ArgGroup::new("target_hosts"))]268pub struct FleetOpts {269 270 #[clap(long, number_of_values = 1, group = "target_hosts")]271 only: Vec<String>,272273 274 #[clap(long, number_of_values = 1, group = "target_hosts")]275 skip: Vec<String>,276277 278 #[clap(long)]279 pub localhost: Option<String>,280281 282 283 #[clap(long, default_value = "detect")]284 pub local_system: String,285}286287impl FleetOpts {288 pub async fn build(mut self, nix_args: Vec<OsString>) -> Result<Config> {289 if self.localhost.is_none() {290 self.localhost291 .replace(hostname::get().unwrap().to_str().unwrap().to_owned());292 }293 let directory = current_dir()?;294295 let pool = NixSessionPool::new(directory.as_os_str().to_owned(), nix_args.clone()).await?;296 let root_field = pool.get().await?;297298 if self.local_system == "detect" {299 let builtins_field = Field::field(root_field.clone(), "builtins").await?;300 self.local_system = nix_go_json!(builtins_field.currentSystem);301 }302 let local_system = self.local_system.clone();303304 let fleet_root = Field::field(root_field, "fleetConfigurations").await?;305306 let fleet_field = nix_go!(fleet_root.default);307 let config_field = nix_go!(fleet_field.configUnchecked);308 let config_unchecked_field = nix_go!(fleet_field.unchecked);309310 let mut fleet_data_path = directory.clone();311 fleet_data_path.push("fleet.nix");312 let bytes = std::fs::read_to_string(fleet_data_path)?;313 let data = nixlike::parse_str(&bytes)?;314315 Ok(Config(Arc::new(FleetConfigInternals {316 opts: self,317 directory,318 data,319 local_system,320 nix_args,321 fleet_field,322 config_field,323 config_unchecked_field,324 })))325 }326}