Skip to main content

regent_sdk/hosts/
managed_host.rs

1//! Managed host and builder
2//!
3//! This module provides the [`ManagedHost`] and [`ManagedHostBuilder`] types for
4//! managing connections to target hosts and executing compliance operations.
5
6use serde::{Deserialize, Serialize};
7use std::collections::HashMap;
8use std::time::Duration;
9use tracing::Instrument;
10use tracing::Level;
11use tracing::span;
12#[allow(unused)]
13use tracing::{debug, error, info, trace, warn};
14
15use crate::LocalHostHandler;
16use crate::Ssh2AuthMethod;
17use crate::Ssh2HostHandler;
18use crate::SshoqHostHandler;
19use crate::WhichUser;
20use crate::error::RegentError;
21use crate::hosts::handlers::ConnectionMethod;
22use crate::hosts::handlers::Handler;
23use crate::hosts::handlers::HostHandler;
24use crate::hosts::handlers::TargetUserKind;
25use crate::hosts::handlers::ssh2::Ssh2AuthReference;
26use crate::hosts::handlers::sshoq::{SshoqAuthReference, SshoqConnectionParams, SshoqTlsConfig};
27use crate::hosts::privilege::Credentials;
28use crate::hosts::privilege::LoginKey;
29use crate::hosts::privilege::Privilege;
30use crate::hosts::properties::HostProperties;
31use crate::secrets::SecretProvidersPool;
32use crate::state::ExpectedState;
33use crate::state::attribute::Remediation;
34use crate::state::compliance::Action;
35use crate::state::compliance::AttributeComplianceAssessment;
36use crate::state::compliance::HostStatus;
37use crate::state::compliance::ManagedHostStatus;
38
39/// Represents the connection state of a managed host.
40///
41/// This enum tracks whether a host is currently connected, disconnected, or in an error state,
42/// enabling proper state management and preventing issues like double-connection or
43/// reconnection failures due to stale state.
44#[derive(Clone, Serialize, Deserialize, Debug, PartialEq)]
45pub enum ConnectionState {
46    /// The host is currently disconnected.
47    Disconnected,
48    /// The host is currently connected and ready for operations.
49    Connected,
50    /// The host connection state is unknown (e.g., after a failed connection attempt).
51    Unknown,
52}
53
54impl Default for ConnectionState {
55    fn default() -> Self {
56        ConnectionState::Disconnected
57    }
58}
59
60/// Builder for creating [`ManagedHost`] instances.
61///
62/// This struct is used to configure a managed host before it's built and connected.
63/// It holds all the configuration needed to establish a connection and manage the host.
64///
65/// # Serialization
66///
67/// `ManagedHostBuilder` can be serialized and deserialized as JSON or YAML:
68///
69/// ```no_run
70/// use regent_sdk::hosts::managed_host::ManagedHostBuilder;
71/// use regent_sdk::hosts::handlers::{ConnectionMethod, TargetUser};
72///
73/// let builder = ManagedHostBuilder::new(
74///     "web-server-01",
75///     "192.168.1.100:22",
76///     Some(ConnectionMethod::Localhost(TargetUser::current_user())),
77/// );
78///
79/// let yaml = serde_yaml::to_string(&builder).unwrap();
80/// let json = serde_json::to_string(&builder).unwrap();
81/// ```
82///
83/// # Example
84///
85/// ```no_run
86/// use regent_sdk::hosts::managed_host::ManagedHostBuilder;
87/// use regent_sdk::hosts::handlers::{ConnectionMethod, TargetUser};
88/// use regent_sdk::secrets::{SecretProvider, SecretProvidersPoolBuilder};
89///
90/// let builder = ManagedHostBuilder::new(
91///     "web-server-01",
92///     "192.168.1.100:22",
93///     Some(ConnectionMethod::Localhost(TargetUser::current_user())),
94/// );
95///
96/// let secrets_pool = SecretProvidersPoolBuilder::new()
97///     .with_default_provider(&SecretProvider::files("files"))
98///     .build()
99///     .unwrap();
100///
101/// let managed_host = builder.build(Some(secrets_pool)).await.unwrap();
102/// ```
103#[derive(Debug, Clone, Serialize, Deserialize)]
104#[serde(rename_all = "PascalCase")]
105#[serde(deny_unknown_fields)]
106pub struct ManagedHostBuilder {
107    /// Unique identifier for this managed host.
108    ///
109    /// Used for tracking and logging purposes.
110    pub id: String,
111    /// The network endpoint for the host (e.g., "192.168.1.100:22" or "localhost").
112    endpoint: String,
113    /// The connection method to use for connecting to the host.
114    pub host_connection_method: Option<ConnectionMethod>,
115    /// Optional host properties that have been collected.
116    host_properties: Option<HostProperties>,
117    /// Optional host-specific variables for templating.
118    ///
119    /// These variables are available during template rendering and can be used
120    /// to customize attribute behavior per host.
121    pub host_vars: Option<HashMap<String, String>>,
122}
123
124impl ManagedHostBuilder {
125    /// Create a new managed host builder.
126    ///
127    /// # Example
128    ///
129    /// ```no_run
130    /// use regent_sdk::hosts::managed_host::ManagedHostBuilder;
131    /// use regent_sdk::hosts::handlers::{ConnectionMethod, TargetUser};
132    ///
133    /// let builder = ManagedHostBuilder::new(
134    ///     "web-server-01",
135    ///     "192.168.1.100:22",
136    ///     Some(ConnectionMethod::Localhost(TargetUser::current_user())),
137    /// );
138    /// ```
139    pub fn new(id: &str, endpoint: &str, connection_method: Option<ConnectionMethod>) -> Self {
140        Self {
141            id: id.to_string(),
142            endpoint: endpoint.to_string(),
143            host_connection_method: connection_method,
144            host_properties: None,
145            host_vars: None,
146        }
147    }
148
149    pub fn id(&self) -> &str {
150        &self.id
151    }
152
153    /// Set the connection method for this host.
154    ///
155    /// # Example
156    ///
157    /// ```no_run
158    /// use regent_sdk::hosts::managed_host::ManagedHostBuilder;
159    /// use regent_sdk::hosts::handlers::{ConnectionMethod, TargetUser};
160    ///
161    /// let mut builder = ManagedHostBuilder::new("host-01", "localhost", None);
162    /// builder.set_connection_method(ConnectionMethod::Localhost(TargetUser::current_user()));
163    /// ```
164    pub fn set_connection_method(&mut self, connection_method: ConnectionMethod) {
165        self.host_connection_method = Some(connection_method);
166    }
167
168    /// Set the host variables for templating.
169    ///
170    /// # Example
171    ///
172    /// ```no_run
173    /// use regent_sdk::hosts::managed_host::ManagedHostBuilder;
174    /// use std::collections::HashMap;
175    ///
176    /// let mut builder = ManagedHostBuilder::new("host-01", "localhost", None);
177    ///
178    /// let mut vars = HashMap::new();
179    /// vars.insert("env".to_string(), "production".to_string());
180    /// vars.insert("region".to_string(), "us-east-1".to_string());
181    ///
182    /// builder.set_host_vars(Some(vars));
183    /// ```
184    pub fn set_host_vars(&mut self, host_vars: Option<HashMap<String, String>>) {
185        self.host_vars = host_vars;
186    }
187
188    /// Parse a managed host builder from raw YAML content.
189    ///
190    /// # Example
191    ///
192    /// ```no_run
193    /// use regent_sdk::hosts::managed_host::ManagedHostBuilder;
194    ///
195    /// let yaml = r#"
196    /// Id: web-server-01
197    /// Endpoint: 192.168.1.100:22
198    /// HostConnectionMethod:
199    ///   Localhost:
200    ///     UserKind: CurrentUser
201    /// "#;
202    ///
203    /// let builder = ManagedHostBuilder::from_raw_yaml(yaml).unwrap();
204    /// ```
205    pub fn from_raw_yaml(raw_yaml: &str) -> Result<Self, RegentError> {
206        match yaml_serde::from_str::<Self>(raw_yaml) {
207            Ok(managed_host_builder) => Ok(managed_host_builder),
208            Err(details) => Err(RegentError::FailureToParseContent(format!("{:?}", details))),
209        }
210    }
211
212    /// Parse a managed host builder from raw JSON content.
213    ///
214    /// # Example
215    ///
216    /// ```no_run
217    /// use regent_sdk::hosts::managed_host::ManagedHostBuilder;
218    ///
219    /// let json = r#"{"Id":"web-server-01","Endpoint":"192.168.1.100:22"}"#;
220    /// let builder = ManagedHostBuilder::from_raw_json(json).unwrap();
221    /// ```
222    pub fn from_raw_json(raw_json: &str) -> Result<Self, RegentError> {
223        match serde_json::from_str::<Self>(raw_json) {
224            Ok(managed_host_builder) => Ok(managed_host_builder),
225            Err(details) => Err(RegentError::FailureToParseContent(format!("{:?}", details))),
226        }
227    }
228
229    pub async fn build(
230        self,
231        optional_secret_provider: Option<SecretProvidersPool>,
232    ) -> Result<ManagedHost, RegentError> {
233        // Check that each required field is set
234        if let None = self.host_connection_method {
235            return Err(RegentError::WrongInitialization(format!(
236                "connection method unset"
237            )));
238        }
239
240        // Retrieve connection secrets when needed
241        match self.host_connection_method {
242            Some(connection) => {
243                match connection {
244                    ConnectionMethod::Localhost(target_user) => {
245                        match target_user.user_kind {
246                            TargetUserKind::CurrentUser => {
247                                // No secret required
248                                Ok(ManagedHost::new(
249                                    self.id,
250                                    &self.endpoint,
251                                    Handler::localhost(LocalHostHandler::from(
252                                        WhichUser::CurrentUser,
253                                    )),
254                                    self.host_vars,
255                                    self.host_properties,
256                                    optional_secret_provider,
257                                ))
258                            }
259                            TargetUserKind::User(secret_reference) => {
260                                match &optional_secret_provider {
261                                    Some(secret_provider) => {
262                                        match secret_provider
263                                            .get_secret_typed::<Credentials>(&secret_reference)
264                                            .await
265                                        {
266                                            Ok(secret) => Ok(ManagedHost::new(
267                                                self.id,
268                                                &self.endpoint,
269                                                Handler::localhost(LocalHostHandler::from(
270                                                    WhichUser::UsernamePassword(secret.inner()),
271                                                )),
272                                                self.host_vars,
273                                                self.host_properties,
274                                                optional_secret_provider,
275                                            )),
276                                            Err(details) => Err(details),
277                                        }
278                                    }
279                                    None => Err(RegentError::WrongInitialization(format!(
280                                        "secret required to connect to host but secret_provider unset"
281                                    ))),
282                                }
283                            }
284                            TargetUserKind::RawUser(credentials) => {
285                                // No secret required
286                                Ok(ManagedHost::new(
287                                    self.id,
288                                    &self.endpoint,
289                                    Handler::localhost(LocalHostHandler::from(
290                                        WhichUser::UsernamePassword(credentials),
291                                    )),
292                                    self.host_vars,
293                                    self.host_properties,
294                                    optional_secret_provider,
295                                ))
296                            }
297                        }
298                    }
299                    ConnectionMethod::Ssh2(ssh2_auth_reference) => {
300                        match ssh2_auth_reference.auth_method {
301                            Ssh2AuthReference::UsernamePassword(secret_reference) => {
302                                match &optional_secret_provider {
303                                    Some(secret_provider) => {
304                                        match secret_provider
305                                            .get_secret_typed::<Credentials>(&secret_reference)
306                                            .await
307                                        {
308                                            Ok(secret) => Ok(ManagedHost::new(
309                                                self.id,
310                                                &self.endpoint,
311                                                Handler::ss2(Ssh2HostHandler::NotConnected(
312                                                    Ssh2AuthMethod::UsernamePassword(
313                                                        secret.inner(),
314                                                    ),
315                                                )),
316                                                self.host_vars,
317                                                self.host_properties,
318                                                optional_secret_provider,
319                                            )),
320                                            Err(details) => Err(details),
321                                        }
322                                    }
323                                    None => Err(RegentError::WrongInitialization(format!(
324                                        "secret required to connect to host but secret_provider unset"
325                                    ))),
326                                }
327                            }
328                            Ssh2AuthReference::Key(login_key_ref) => {
329                                match &optional_secret_provider {
330                                    Some(secret_provider) => {
331                                        match secret_provider
332                                            .get_secret_raw(login_key_ref.key_ref())
333                                            .await
334                                        {
335                                            Ok(secret) => Ok(ManagedHost::new(
336                                                self.id,
337                                                &self.endpoint,
338                                                Handler::ss2(Ssh2HostHandler::NotConnected(
339                                                    Ssh2AuthMethod::Key(LoginKey::from(
340                                                        login_key_ref.username().to_string(),
341                                                        secret.inner(),
342                                                    )),
343                                                )),
344                                                self.host_vars,
345                                                self.host_properties,
346                                                optional_secret_provider,
347                                            )),
348                                            Err(details) => Err(details),
349                                        }
350                                    }
351                                    None => Err(RegentError::WrongInitialization(format!(
352                                        "secret required to connect to host but secret_provider unset"
353                                    ))),
354                                }
355                            }
356                        }
357                    }
358                    ConnectionMethod::Sshoq(sshoq_auth_reference) => {
359                        let sshoq_tls = sshoq_auth_reference
360                            .tls_config
361                            .clone()
362                            .unwrap_or_default()
363                            .to_client_tls_config()?;
364                        match sshoq_auth_reference.auth_method {
365                            SshoqAuthReference::UsernamePassword(secret_reference) => {
366                                match &optional_secret_provider {
367                                    Some(secret_provider) => {
368                                        match secret_provider
369                                            .get_secret_typed::<Credentials>(&secret_reference)
370                                            .await
371                                        {
372                                            Ok(secret) => {
373                                                let mut params = SshoqConnectionParams::username_password(
374                                                    secret.inner(),
375                                                    sshoq_tls,
376                                                );
377                                                if let Some(p) = sshoq_auth_reference.path.clone() {
378                                                    params = params.with_path(&p);
379                                                }
380                                                Ok(ManagedHost::new(
381                                                    self.id,
382                                                    &self.endpoint,
383                                                    Handler::sshoq(SshoqHostHandler::NotConnected(
384                                                        params,
385                                                    )),
386                                                    self.host_vars,
387                                                    self.host_properties,
388                                                    optional_secret_provider,
389                                                ))
390                                            }
391                                            Err(details) => Err(details),
392                                        }
393                                    }
394                                    None => Err(RegentError::WrongInitialization(format!(
395                                        "secret required to connect to host but secret_provider unset"
396                                    ))),
397                                }
398                            }
399                            SshoqAuthReference::Key(login_key_ref) => {
400                                match &optional_secret_provider {
401                                    Some(secret_provider) => {
402                                        match secret_provider
403                                            .get_secret_raw(login_key_ref.key_ref())
404                                            .await
405                                        {
406                                            Ok(secret) => {
407                                                let mut params = SshoqConnectionParams::key(
408                                                    LoginKey::from(
409                                                        login_key_ref.username().to_string(),
410                                                        secret.inner(),
411                                                    ),
412                                                    sshoq_tls,
413                                                );
414                                                if let Some(p) = sshoq_auth_reference.path.clone() {
415                                                    params = params.with_path(&p);
416                                                }
417                                                Ok(ManagedHost::new(
418                                                    self.id,
419                                                    &self.endpoint,
420                                                    Handler::sshoq(SshoqHostHandler::NotConnected(
421                                                        params,
422                                                    )),
423                                                    self.host_vars,
424                                                    self.host_properties,
425                                                    optional_secret_provider,
426                                                ))
427                                            }
428                                            Err(details) => Err(details),
429                                        }
430                                    }
431                                    None => Err(RegentError::WrongInitialization(format!(
432                                        "secret required to connect to host but secret_provider unset"
433                                    ))),
434                                }
435                            }
436                            SshoqAuthReference::BearerToken(secret_reference) => {
437                                match &optional_secret_provider {
438                                    Some(secret_provider) => {
439                                        match secret_provider
440                                            .get_secret_raw(&secret_reference)
441                                            .await
442                                        {
443                                            Ok(secret) => {
444                                                let mut params = SshoqConnectionParams::bearer_token(
445                                                    secret.inner(),
446                                                    sshoq_tls,
447                                                );
448                                                if let Some(p) = sshoq_auth_reference.path.clone() {
449                                                    params = params.with_path(&p);
450                                                }
451                                                Ok(ManagedHost::new(
452                                                    self.id,
453                                                    &self.endpoint,
454                                                    Handler::sshoq(SshoqHostHandler::NotConnected(
455                                                        params,
456                                                    )),
457                                                    self.host_vars,
458                                                    self.host_properties,
459                                                    optional_secret_provider,
460                                                ))
461                                            }
462                                            Err(details) => Err(details),
463                                        }
464                                    }
465                                    None => Err(RegentError::WrongInitialization(format!(
466                                        "secret required to connect to host but secret_provider unset"
467                                    ))),
468                                }
469                            }
470                        }
471                    }
472                }
473            }
474            None => Err(RegentError::WrongInitialization(format!(
475                "connection_method unset"
476            ))),
477        }
478    }
479}
480
481/// A managed host that can execute compliance operations.
482///
483/// `ManagedHost` represents a target system that can be connected to and managed.
484/// It provides methods for assessing and reaching compliance with expected states.
485///
486/// # Lifecycle
487///
488/// 1. Create using [`ManagedHostBuilder::build()`]
489/// 2. Connect using [`ManagedHost::connect()`]
490/// 3. Execute operations (assess/reach compliance)
491/// 4. Disconnect using [`ManagedHost::disconnect()`]
492///
493/// # Example
494///
495/// ```no_run
496/// use regent_sdk::hosts::managed_host::{ManagedHost, ManagedHostBuilder};
497/// use regent_sdk::hosts::handlers::{ConnectionMethod, TargetUser};
498/// use regent_sdk::state::ExpectedState;
499///
500/// let builder = ManagedHostBuilder::new(
501///     "web-server-01",
502///     "192.168.1.100:22",
503///     Some(ConnectionMethod::Localhost(TargetUser::current_user())),
504/// );
505///
506/// let mut host = builder.build(None).await.unwrap();
507/// host.connect().unwrap();
508///
509/// let expected_state = ExpectedState::new();
510/// let status = host.assess_compliance(&expected_state).await.unwrap();
511///
512/// host.disconnect().unwrap();
513/// ```
514// #[derive(Clone)]
515pub struct ManagedHost {
516    /// Unique identifier for this managed host.
517    id: String,
518    /// The network endpoint for the host.
519    endpoint: String,
520    /// The handler used for connecting and executing commands on the host.
521    pub handler: Handler,
522    /// Tera context for template rendering with host variables.
523    context: tera::Context,
524    /// Cached host properties (OS, architecture, etc.).
525    host_properties: Option<HostProperties>,
526    /// Secret providers pool for retrieving secrets.
527    // TODO : how to avoid cloning the whole Pool for each Host ? Is it worth introducting Arc<Mutex<T>> ?
528    secret_providers: Option<SecretProvidersPool>,
529    /// The current connection state of this host.
530    ///
531    /// This tracks whether the host is connected, disconnected, or in an unknown state,
532    /// enabling proper idempotency for connection/disconnection operations.
533    connection_state: ConnectionState,
534}
535
536impl Clone for ManagedHost {
537    fn clone(&self) -> Self {
538        Self {
539            id: self.id.clone(),
540            endpoint: self.endpoint.clone(),
541            handler: self.handler.clone(),
542            context: self.context.clone(),
543            host_properties: self.host_properties.clone(),
544            secret_providers: self.secret_providers.clone(),
545            connection_state: self.connection_state.clone(),
546        }
547    }
548}
549
550impl ManagedHost {
551    /// Create a new managed host.
552    pub fn new(
553        id: String,
554        endpoint: &str,
555        handler: Handler,
556        host_vars: Option<HashMap<String, String>>,
557        host_properties: Option<HostProperties>,
558        secret_providers: Option<SecretProvidersPool>,
559    ) -> ManagedHost {
560        let context = match host_vars {
561            Some(content) => match tera::Context::from_serialize(&content) {
562                Ok(context) => context,
563                Err(details) => {
564                    error!("Failed to create Tera context : {:?}", details);
565                    // TODO : turn this into an Err -> return type of this function -> Result
566                    tera::Context::new()
567                }
568            },
569            None => tera::Context::new(),
570        };
571        ManagedHost {
572            id,
573            endpoint: endpoint.to_string(),
574            handler,
575            context,
576            host_properties,
577            secret_providers: secret_providers.clone(),
578            connection_state: ConnectionState::Disconnected,
579        }
580    }
581
582    /// Create a new managed host from an iterator of variables.
583    ///
584    /// This is a convenience constructor that accepts variables as an iterator.
585    ///
586    pub fn from(
587        id: String,
588        endpoint: &str,
589        handler: Handler,
590        vars: Option<impl IntoIterator<Item = (String, String)>>,
591        host_properties: Option<HostProperties>,
592        secret_providers: Option<SecretProvidersPool>,
593    ) -> ManagedHost {
594        let final_vars = match vars {
595            Some(vars_list) => {
596                let mut final_vars: HashMap<String, String> = HashMap::new();
597
598                for (key, value) in vars_list.into_iter() {
599                    final_vars.insert(key, value);
600                }
601
602                Some(final_vars)
603            }
604            None => None,
605        };
606
607        ManagedHost {
608            id,
609            endpoint: endpoint.to_string(),
610            handler,
611            context: tera::Context::from_serialize(&final_vars).unwrap(),
612            host_properties,
613            secret_providers: secret_providers.clone(),
614            connection_state: ConnectionState::Disconnected,
615        }
616    }
617
618    /// Get the unique identifier for this managed host.
619    pub fn id(&self) -> &str {
620        &self.id
621    }
622
623    /// Add a variable to the host's template context.
624    ///
625    /// This variable will be available during template rendering for attributes.
626    pub fn add_var(&mut self, key: String, value: String) {
627        if self.context.contains_key(&key) {
628            warn!(
629                key,
630                "Writing an already-existing variable to context. Is this expected ?"
631            );
632        }
633        self.context.insert(key, &value);
634    }
635
636    /// Set the host properties.
637    ///
638    /// This replaces any existing host properties with the provided value.
639    pub fn set_host_properties(&mut self, host_properties: Option<HostProperties>) {
640        self.host_properties = host_properties;
641    }
642
643    /// Enable secret caching for this managed host.
644    ///
645    /// This ensures that secrets retrieved during assessment and enforcement
646    /// phases are the same, even if the underlying secrets are rotated.
647    /// This is crucial for idempotency in compliance operations.
648    pub fn enable_secret_caching(&mut self) {
649        if let Some(secret_providers) = &mut self.secret_providers {
650            secret_providers.enable_caching();
651        }
652    }
653
654    /// Disable secret caching for this managed host.
655    pub fn disable_secret_caching(&mut self) {
656        if let Some(secret_providers) = &mut self.secret_providers {
657            secret_providers.disable_caching();
658        }
659    }
660
661    /// Collect host properties dynamically from the host.
662    ///
663    /// This method connects to the host and collects information about its
664    /// operating system, architecture, and other properties.
665    ///
666    pub async fn collect_properties(&mut self) -> Result<(), RegentError> {
667        if matches!(self.host_properties, None) {
668            match HostProperties::collect_dynamically(&mut self.handler).await {
669                Ok(host_properties) => {
670                    self.host_properties = Some(host_properties);
671                }
672                Err(details) => {
673                    return Err(details);
674                }
675            }
676        } else {
677            warn!("HostProperties already collected. Logic error ?");
678        }
679        Ok(())
680    }
681
682    /// Get the host properties.
683    pub fn get_host_properties(&self) -> &Option<HostProperties> {
684        &self.host_properties
685    }
686
687    /// Connect to the host.
688    ///
689    /// This method establishes a connection to the host using the configured handler.
690    ///
691    /// # Example
692    ///
693    /// ```no_run
694    /// use regent_sdk::hosts::managed_host::ManagedHost;
695    ///
696    /// let mut host = /* create host */;
697    /// host.connect().unwrap();
698    /// assert!(host.is_connected());
699    /// ```
700    pub async fn connect(&mut self) -> Result<(), RegentError> {
701        let result = self.handler.connect(&self.endpoint).await;
702        if result.is_ok() {
703            self.connection_state = ConnectionState::Connected;
704        } else {
705            self.connection_state = ConnectionState::Unknown;
706        }
707        result
708    }
709
710    /// Check if the host is currently connected.
711    pub async fn is_connected(&mut self) -> bool {
712        self.handler.is_connected().await
713    }
714
715    /// Check the tracked connection state of this host.
716    ///
717    /// Unlike `is_connected()`, which queries the underlying handler,
718    /// this returns the tracked state that was last set by `connect()` or `disconnect()`.
719    ///
720    pub fn connection_state(&self) -> &ConnectionState {
721        &self.connection_state
722    }
723
724    /// Check if the host is in a tracked connected state.
725    ///
726    /// This returns `true` if the tracked state is `Connected`, regardless of
727    /// the actual underlying handler state. Use `is_connected()` to check
728    /// the actual connection status.
729    ///
730    pub fn is_tracked_connected(&self) -> bool {
731        matches!(self.connection_state, ConnectionState::Connected)
732    }
733
734    /// Check if the host is in a tracked disconnected state.
735    ///
736    /// This returns `true` if the tracked state is `Disconnected`.
737    ///
738    pub fn is_tracked_disconnected(&self) -> bool {
739        matches!(self.connection_state, ConnectionState::Disconnected)
740    }
741
742    /// Disconnect from the host.
743    ///
744    pub async fn disconnect(&mut self) -> Result<(), RegentError> {
745        let result = self.handler.disconnect().await;
746        if result.is_ok() {
747            self.connection_state = ConnectionState::Disconnected;
748        } else {
749            self.connection_state = ConnectionState::Unknown;
750        }
751        result
752    }
753
754    /// Assess compliance of the host with the expected state.
755    ///
756    /// This method checks each attribute in the expected state and determines
757    /// whether the host is compliant. It does not make any changes to the host.
758    ///
759    /// # Example
760    ///
761    /// ```no_run
762    /// use regent_sdk::hosts::managed_host::ManagedHost;
763    /// use regent_sdk::state::{ExpectedState, Attribute};
764    ///
765    /// let mut host = /* create and connect host */;
766    /// let expected_state = ExpectedState::new()
767    ///     .with_attribute(Attribute::service(/* ... */))
768    ///     .build();
769    ///
770    /// let status = host.assess_compliance(&expected_state).await.unwrap();
771    ///
772    /// if status.is_already_compliant() {
773    ///     println!("Host is compliant!");
774    /// } else {
775    ///     println!("Host needs remediation");
776    /// }
777    /// ```
778    // Defaults to sequential assessment
779    pub async fn assess_compliance(
780        &mut self,
781        expected_state: &ExpectedState,
782        logging: bool,
783    ) -> Result<ManagedHostStatus, RegentError> {
784        if !self.is_connected().await {
785            return Err(RegentError::NotConnectedToHost);
786        }
787
788        // Enable secret caching to ensure idempotency
789        self.enable_secret_caching();
790
791        let mut already_compliant = true;
792        let mut final_remediations_list: Vec<Remediation> = Vec::new();
793
794        for attribute in expected_state.attributes.clone().iter_mut() {
795            let attribute_span = span!(Level::INFO, "attribute", attr_name = attribute.name());
796
797            // We execute the entire loop iteration inside an instrumented async block for tracing consistency
798            let iteration_result = async {
799                // Taking context into account before working on the Attribute
800                match attribute.consider_context(&self.context) {
801                    Ok(context_aware_attribute) => {
802                        match context_aware_attribute
803                            .assess(
804                                &mut self.handler,
805                                &self.host_properties,
806                                &self.secret_providers,
807                            )
808                            .await
809                        {
810                            Ok(attribute_compliance_assessment) => {
811                                match attribute_compliance_assessment {
812                                    AttributeComplianceAssessment::Compliant => {
813                                        if logging {
814                                            info!(target: "run", "Compliant");
815                                        }
816                                        Ok(Some(Vec::new())) // Compliant, no remediations
817                                    }
818                                    AttributeComplianceAssessment::NonCompliant(remediations) => {
819                                        if logging {
820                                            warn!(target: "run", remediations = ?remediations.remediations(), "Not compliant");
821                                        }
822                                        Ok(Some(remediations.into_inner()))
823                                    }
824                                    AttributeComplianceAssessment::NonCompliantFatal(details) => {
825                                        if logging {
826                                            error!(target: "run", details, "Not compliant and no remediation possible");
827                                        }
828                                        Ok(None) // Non-compliant fatal, can't remediate
829                                    }
830                                }
831                            }
832                            Err(details) => Err(details),
833                        }
834                    }
835                    Err(details) => {
836                        let content = match &details {
837                            RegentError::FailureToConsiderContext(content) => content,
838                            _ => &format!("{:?}", details),
839                        };
840                        error!("{}", content);
841                        Err(details)
842                    }
843                }
844            }
845            .instrument(attribute_span) // <-- Captures EVERYTHING inside the async block
846            .await;
847
848            // Process the control-flow state safely outside the span tracking lifecycle
849            match iteration_result {
850                Ok(Some(remediations)) => {
851                    if !remediations.is_empty() {
852                        already_compliant = false;
853                        final_remediations_list.extend(remediations);
854                    }
855                }
856                Ok(None) => {
857                    already_compliant = false;
858                }
859                Err(details) => {
860                    return Err(details);
861                }
862            }
863        }
864
865        if already_compliant {
866            Ok(ManagedHostStatus::already_compliant())
867        } else {
868            Ok(ManagedHostStatus::not_compliant(final_remediations_list))
869        }
870    }
871
872    /// Automatically reach compliance with the expected state.
873    ///
874    /// This method assesses compliance and then automatically performs the necessary
875    /// remediations to bring the host into the expected state. It will stop at the
876    /// first failure unless all remediations are marked as allowed to fail.
877    ///
878    /// # Example
879    ///
880    /// ```no_run
881    /// use regent_sdk::hosts::managed_host::ManagedHost;
882    /// use regent_sdk::state::{ExpectedState, Attribute};
883    ///
884    /// let mut host = /* create and connect host */;
885    /// let expected_state = ExpectedState::new()
886    ///     .with_attribute(Attribute::service(/* ... */))
887    ///     .build();
888    ///
889    /// let result = host.reach_compliance(&expected_state).await.unwrap();
890    ///
891    /// if result.is_reach_compliance_success() {
892    ///     println!("Compliance reached successfully!");
893    /// } else if result.is_reach_compliance_failed() {
894    ///     println!("Failed to reach compliance");
895    /// }
896    /// ```
897    pub async fn reach_compliance(
898        &mut self,
899        expected_state: &ExpectedState,
900    ) -> Result<ManagedHostStatus, RegentError> {
901        if !self.is_connected().await {
902            return Err(RegentError::NotConnectedToHost);
903        }
904
905        // Enable secret caching to ensure idempotency
906        self.enable_secret_caching();
907
908        let mut final_host_status = HostStatus::AlreadyCompliant;
909        let mut reaching_compliance_failed = false;
910        let mut actions_taken: Vec<Action> = Vec::new();
911
912        for attribute in &expected_state.attributes {
913            let attribute_span = span!(Level::INFO, "attribute", attr_name = attribute.name());
914
915            // Wrap the logic into an instrumented async block to facilitate tracing
916            let iteration_result = async {
917                match attribute.consider_context(&self.context) {
918                    Ok(context_aware_attribute) => {
919                        let timeout_duration = context_aware_attribute.timeout();
920
921                        match context_aware_attribute
922                            .assess(
923                                &mut self.handler,
924                                &self.host_properties,
925                                &self.secret_providers,
926                            )
927                            .await
928                        {
929                            Ok(attribute_compliance) => {
930                                match attribute_compliance {
931                                    AttributeComplianceAssessment::Compliant => {
932                                        info!(target: "run", "Compliant");
933                                        // Return progress updates back to the parent collector
934                                        Ok(LoopControl::Continue)
935                                    }
936                                    AttributeComplianceAssessment::NonCompliant(remediations) => {
937                                        warn!(target: "run", remediations = ?remediations.remediations(), "Not compliant");
938
939                                        let mut local_failed = false;
940                                        let mut local_actions = Vec::new();
941
942                                        for remediation in remediations.iter() {
943                                            let remediation_span = span!(Level::INFO, "remediation", action = ?remediation);
944
945                                            let remediation_result = remediation
946                                                .reach_compliance(
947                                                    &mut self.handler,
948                                                    &self.host_properties,
949                                                    &self.secret_providers,
950                                                    timeout_duration,
951                                                )
952                                                .instrument(remediation_span.clone())
953                                                .await;
954
955                                            let _remediation_span_guard = remediation_span.enter();
956
957                                            match remediation_result {
958                                                Ok(internal_api_call_outcome) => {
959                                                    local_actions.push(Action::from(
960                                                        remediation.clone(),
961                                                        Some(internal_api_call_outcome.clone()),
962                                                    ));
963
964                                                    match &internal_api_call_outcome {
965                                                        InternalApiCallOutcome::Success(details) => {
966                                                            info!(target: "run", remediation_outcome = "Success", details = details.clone().unwrap_or("none".to_string()));
967                                                        }
968                                                        InternalApiCallOutcome::AllowedFailure(details) => {
969                                                            info!(target: "run", remediation_outcome = "AllowedFailure", details);
970                                                        }
971                                                        InternalApiCallOutcome::Failure(details) => {
972                                                            local_failed = true;
973                                                            warn!(remediation_outcome = "Failure", details);
974                                                            // Break the local remediation sub-loop
975                                                            break;
976                                                        }
977                                                    }
978                                                }
979                                                Err(details) => {
980                                                    warn!("Failed to apply remediation");
981                                                    return Err(details);
982                                                }
983                                            }
984                                        }
985
986                                        Ok(LoopControl::Remediated {
987                                            failed: local_failed,
988                                            actions: local_actions
989                                        })
990                                    }
991                                    AttributeComplianceAssessment::NonCompliantFatal(details) => {
992                                        error!(target: "run", details, "Not compliant and no remediation possible");
993                                        Ok(LoopControl::Fatal)
994                                    }
995                                }
996                            }
997                            Err(details) => {
998                                warn!(reason = ?details, "Failed assessment");
999                                Err(details)
1000                            }
1001                        }
1002                    }
1003                    Err(details) => {
1004                        let content = match &details {
1005                            RegentError::FailureToConsiderContext(content) => content,
1006                            _ => &format!("{:?}", details),
1007                        };
1008                        error!("{}", content);
1009                        Err(details)
1010                    }
1011                }
1012            }
1013            .instrument(attribute_span) // <-- Seamlessly maintains 'attribute' tracing span across everything
1014            .await;
1015
1016            // Process loop control and state changes outside of the tracing block context
1017            match iteration_result {
1018                Ok(LoopControl::Continue) => {}
1019                Ok(LoopControl::Fatal) => {
1020                    final_host_status = HostStatus::ReachComplianceFailed;
1021                }
1022                Ok(LoopControl::Remediated { failed, actions }) => {
1023                    actions_taken.extend(actions);
1024                    if failed {
1025                        // reaching_compliance_failed = true;
1026                        final_host_status = HostStatus::ReachComplianceFailed;
1027                        break; // Stop processing further attributes due to failure
1028                    } else {
1029                        final_host_status = HostStatus::ReachComplianceSuccess;
1030                    }
1031                }
1032                Err(details) => {
1033                    return Err(details);
1034                }
1035            }
1036        }
1037
1038        match final_host_status {
1039            HostStatus::AlreadyCompliant => Ok(ManagedHostStatus::already_compliant()),
1040            HostStatus::ReachComplianceFailed => {
1041                Ok(ManagedHostStatus::reach_compliance_failed(actions_taken))
1042            }
1043            _ => Ok(ManagedHostStatus::reach_compliance_success(actions_taken)),
1044        }
1045    }
1046}
1047
1048/// Trait for types that can assess compliance of a host.
1049///
1050/// Implement this trait for custom attribute types that need to check
1051/// whether a host is compliant with a specific configuration.
1052///
1053/// # Type Parameters
1054///
1055/// * `Handler` - The type of host handler
1056pub trait AssessCompliance<Handler: HostHandler> {
1057    /// Assess whether the host is compliant.
1058    async fn assess_compliance(
1059        &self,
1060        host_handler: &mut Handler,
1061        host_properties: &Option<HostProperties>,
1062        privilege: &Privilege,
1063        optional_secret_provider: &Option<SecretProvidersPool>,
1064    ) -> Result<AttributeComplianceAssessment, RegentError>;
1065}
1066
1067/// Trait for types that can perform remediation to reach compliance.
1068///
1069/// Implement this trait for custom remediation types that need to make
1070/// changes to a host to bring it into compliance.
1071///
1072/// # Type Parameters
1073///
1074/// * `Handler` - The type of host handler
1075pub trait ReachCompliance<Handler: HostHandler> {
1076    /// Perform remediation on the host.
1077    async fn call(
1078        &self,
1079        host_handler: &mut Handler,
1080        host_properties: &Option<HostProperties>,
1081        optional_secret_provider: &Option<SecretProvidersPool>,
1082    ) -> Result<InternalApiCallOutcome, RegentError>;
1083}
1084
1085/// Trait for types that have a default timeout.
1086///
1087/// Implement this trait for operations that should have a configurable timeout.
1088pub trait Timeout {
1089    /// Get the default timeout for this operation.
1090    fn default_timeout(&self) -> Duration;
1091}
1092
1093/// Outcome of an attribute-level operation.
1094///
1095/// This enum represents the possible results of assessing and/or remediating
1096/// a single attribute on a host.
1097///
1098/// # Variants
1099///
1100/// - `AlreadyCompliant`: The attribute was already compliant
1101/// - `NotCompliant`: The attribute was not compliant, with list of required remediations
1102/// - `ReachComplianceFailed`: Attempt to reach compliance failed
1103/// - `ComplianceReachedWithAllowedFailure`: Compliance reached but with allowed failures
1104/// - `ComplianceReached`: Compliance successfully reached with list of actions taken
1105#[derive(Serialize, Deserialize)]
1106pub enum AttributeLevelOperationOutcome {
1107    /// The attribute was already compliant.
1108    AlreadyCompliant,
1109    /// The attribute was not compliant.
1110    ///
1111    /// Contains the list of remediations that would be needed to reach compliance.
1112    NotCompliant(Vec<Remediation>),
1113    /// Attempt to reach compliance failed.
1114    ///
1115    /// Contains the failure outcome.
1116    ReachComplianceFailed(InternalApiCallOutcome),
1117    /// Compliance reached but with allowed failures.
1118    ///
1119    /// Some remediations failed but were marked as allowed to fail.
1120    ComplianceReachedWithAllowedFailure(InternalApiCallOutcome),
1121    /// Compliance successfully reached.
1122    ///
1123    /// Contains the list of remediations performed and their outcomes.
1124    ComplianceReached(Vec<(Remediation, InternalApiCallOutcome)>),
1125}
1126
1127/// Outcome of an internal API call (remediation execution).
1128///
1129/// This enum represents the result of executing a single remediation action.
1130///
1131/// # Variants
1132///
1133/// - `Success`: The remediation succeeded, with optional details
1134/// - `Failure`: The remediation failed, with error details
1135/// - `AllowedFailure`: The remediation failed but was allowed to fail
1136#[derive(Debug, Serialize, Deserialize, Clone)]
1137pub enum InternalApiCallOutcome {
1138    /// The remediation succeeded.
1139    ///
1140    /// Contains optional details about the success.
1141    Success(Option<String>),
1142    /// The remediation failed.
1143    ///
1144    /// Contains error details.
1145    Failure(String),
1146    /// The remediation failed but was marked as allowed to fail.
1147    ///
1148    /// This allows for "best effort" compliance where some failures are acceptable.
1149    AllowedFailure(String),
1150}
1151
1152// Internal state tracking helper for out-of-span control flow
1153enum LoopControl {
1154    Continue,
1155    Remediated { failed: bool, actions: Vec<Action> },
1156    Fatal,
1157}