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

#[cfg(feature = "server")]
use diesel::{AsChangeset, Identifiable, Insertable, Queryable};
#[cfg(feature = "server")]
use diesel_derive_enum::DbEnum;
use serde::{Deserialize, Serialize};
use std::{fmt, time::SystemTime};
use strum::Display;

#[cfg(feature = "server")]
use crate::models::schema::*;

use crate::{
    config::get_self_uid,
    executor::ExecResult,
    runner::NamedJobRunner,
    utils::{systime_to_string, Platform},
};

use uuid::Uuid;

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

//belongs_to
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
#[cfg_attr(
    feature = "server",
    derive(Identifiable, Queryable, Insertable, AsChangeset),
    table_name = "agents"
)]
pub struct Agent {
    pub alias: Option<String>,
    pub hostname: String,
    pub id: Uuid,
    pub ip_gray: Option<String>,
    pub ip_white: Option<String>,
    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 {
        write!(f, "Agent: {}", self.id)?;
        if let Some(ref alias) = self.alias {
            write!(f, " ({})", alias)?
        }
        writeln!(f, "\nUsername: {}", self.username)?;
        writeln!(f, "Hostname: {}", self.hostname)?;
        writeln!(f, "Is root: {}", self.is_root)?;
        writeln!(f, "Root allowed: {}", self.is_root_allowed)?;
        writeln!(f, "Last active: {}", systime_to_string(&self.last_active))?;
        writeln!(f, "Platform: {}", self.platform)?;
        writeln!(f, "State: {}", self.state)
    }
}

#[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 = NamedJobRunner::from_shell(vec![
            ("hostname", "hostname"),
            ("is_root", "id -u"),
            ("username", "id -un"),
        ])
        .unwrap_one()
        .wait()
        .await;
        let decoder =
            |job_result: ExecResult| job_result.unwrap().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 {
        todo!()
    }

    pub async fn run() -> 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(),
            ip_gray: None,
            ip_white: None,
        }
    }
}

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

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