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;

use crate::{
    builder::NamedJobBuilder,
    models::{schema::*, ExecResult},
    unwrap_enum,
    utils::systime_to_string,
    UID,
};

use guess_host_triple::guess_host_triple;
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,
)]
#[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 self.alias.is_some() {
            out += &format!(" ({})", self.alias.as_ref().unwrap())
        }
        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)
    }
}

impl Agent {
    pub fn with_id(uid: Uuid) -> Self {
        Self {
            id: uid,
            ..Default::default()
        }
    }

    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: ExecResult| {
            let assoc_job = unwrap_enum!(job_result, ExecResult::Assigned);
            assoc_job.to_string_result().unwrap().trim().to_string()
        };

        Self {
            hostname: decoder(builder.pop("hostname")),
            is_root: &decoder(builder.pop("is_root")) == "0",
            username: decoder(builder.pop("username")),
            platform: guess_host_triple().unwrap_or("unknown").to_string(),
            ..Default::default()
        }
    }

    pub async fn run() -> ExecResult {
        #[cfg(unix)]
        ExecResult::Agent(Agent::gather().await)
    }
}

impl Default for Agent {
    fn default() -> Self {
        Self {
            alias: None,
            id: UID.clone(),
            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)
    }
}