unki/lib/u_lib/src/models/agent.rs

use diesel::{AsChangeset, Identifiable, Insertable, Queryable};
use diesel_derive_enum::DbEnum;
use serde::{Deserialize, Serialize};
use std::{fmt, time::SystemTime};
use strum::Display;

#[cfg(not(target_arch = "wasm32"))]
use crate::builder::NamedJobBuilder;
use crate::{
    config::get_self_uid,
    messaging::Reportable,
    models::schema::*,
    unwrap_enum,
    utils::{systime_to_string, Platform},
};

use uuid::Uuid;

#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, DbEnum, Display)]
#[PgType = "AgentState"]
#[DieselType = "Agentstate"]
pub enum AgentState {
    New,
    Active,
    Banned,
}

//belongs_to
#[derive(
    Clone,
    Debug,
    Serialize,
    Deserialize,
    Identifiable,
    Queryable,
    Insertable,
    AsChangeset,
    PartialEq,
)]
#[table_name = "agents"]
pub struct Agent {
    pub alias: Option<String>,
    pub hostname: String,
    pub id: Uuid,
    pub is_root: bool,
    pub is_root_allowed: bool,
    pub last_active: SystemTime,
    pub platform: String,
    pub regtime: SystemTime,
    pub state: AgentState,
    pub token: Option<String>,
    pub username: String,
}

impl fmt::Display for Agent {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut out = format!("Agent: {}", self.id);
        if let Some(ref alias) = self.alias {
            out += &format!(" ({})", alias)
        }
        out += &format!("\nUsername: {}", self.username);
        out += &format!("\nHostname: {}", self.hostname);
        out += &format!("\nIs root: {}", self.is_root);
        out += &format!("\nRoot allowed: {}", self.is_root_allowed);
        out += &format!("\nLast active: {}", systime_to_string(&self.last_active));
        out += &format!("\nPlatform: {}", self.platform);
        out += &format!("\nState: {}", self.state);
        write!(f, "{}", out)
    }
}

#[cfg(not(target_arch = "wasm32"))]
impl Agent {
    pub fn with_id(uid: Uuid) -> Self {
        Self {
            id: uid,
            ..Default::default()
        }
    }

    #[cfg(unix)]
    pub async fn gather() -> Self {
        let mut builder = NamedJobBuilder::from_shell(vec![
            ("hostname", "hostname"),
            ("is_root", "id -u"),
            ("username", "id -un"),
        ])
        .unwrap_one()
        .wait()
        .await;
        let decoder = |job_result: Reportable| {
            let assoc_job = unwrap_enum!(job_result, Reportable::Assigned);
            assoc_job.to_string_result().trim().to_string()
        };

        Self {
            hostname: decoder(builder.pop("hostname")),
            is_root: &decoder(builder.pop("is_root")) == "0",
            username: decoder(builder.pop("username")),
            platform: Platform::current().into_string(),
            ..Default::default()
        }
    }

    #[cfg(not(unix))]
    pub async fn gather() -> Self {
        Self::default()
    }

    pub async fn run() -> Reportable {
        Reportable::Agent(Agent::gather().await)
    }
}

impl Default for Agent {
    fn default() -> Self {
        Self {
            alias: None,
            id: get_self_uid(),
            hostname: String::new(),
            is_root: false,
            is_root_allowed: false,
            last_active: SystemTime::now(),
            platform: String::new(),
            regtime: SystemTime::now(),
            state: AgentState::New,
            token: None,
            username: String::new(),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_gather() {
        let cli_info = Agent::gather().await;
        assert_eq!(cli_info.alias, None)
    }
}