Add organization/site/tags support for machine grouping
- Added organization, site, tags columns to connect_machines table - Agent now sends org/site/tags from embedded config in AgentStatus - Server stores org/site/tags metadata in database - Enables grouping machines by client/site/tag in dashboard
This commit is contained in:
@@ -175,6 +175,9 @@ impl SessionManager {
|
||||
display_count: self.get_display_count(),
|
||||
is_streaming: self.state == SessionState::Streaming,
|
||||
agent_version: crate::build_info::short_version(),
|
||||
organization: self.config.company.clone().unwrap_or_default(),
|
||||
site: self.config.site.clone().unwrap_or_default(),
|
||||
tags: self.config.tags.clone(),
|
||||
};
|
||||
|
||||
let msg = Message {
|
||||
|
||||
@@ -277,6 +277,9 @@ message AgentStatus {
|
||||
int32 display_count = 5;
|
||||
bool is_streaming = 6;
|
||||
string agent_version = 7; // Agent version (e.g., "0.1.0-abc123")
|
||||
string organization = 8; // Company/organization name
|
||||
string site = 9; // Site/location name
|
||||
repeated string tags = 10; // Tags for categorization
|
||||
}
|
||||
|
||||
// Server commands agent to uninstall itself
|
||||
|
||||
@@ -116,3 +116,34 @@ pub async fn delete_machine(pool: &PgPool, agent_id: &str) -> Result<(), sqlx::E
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Update machine organization, site, and tags
|
||||
pub async fn update_machine_metadata(
|
||||
pool: &PgPool,
|
||||
agent_id: &str,
|
||||
organization: Option<&str>,
|
||||
site: Option<&str>,
|
||||
tags: &[String],
|
||||
) -> Result<(), sqlx::Error> {
|
||||
// Only update if at least one value is provided
|
||||
if organization.is_none() && site.is_none() && tags.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
UPDATE connect_machines SET
|
||||
organization = COALESCE($1, organization),
|
||||
site = COALESCE($2, site),
|
||||
tags = CASE WHEN $3::text[] = '{}' THEN tags ELSE $3 END
|
||||
WHERE agent_id = $4
|
||||
"#,
|
||||
)
|
||||
.bind(organization)
|
||||
.bind(site)
|
||||
.bind(tags)
|
||||
.bind(agent_id)
|
||||
.execute(pool)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -313,6 +313,16 @@ async fn handle_agent_connection(
|
||||
} else {
|
||||
Some(status.agent_version.clone())
|
||||
};
|
||||
let organization = if status.organization.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(status.organization.clone())
|
||||
};
|
||||
let site = if status.site.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(status.site.clone())
|
||||
};
|
||||
sessions_status.update_agent_status(
|
||||
session_id,
|
||||
Some(status.os_version.clone()),
|
||||
@@ -321,6 +331,9 @@ async fn handle_agent_connection(
|
||||
status.display_count,
|
||||
status.is_streaming,
|
||||
agent_version.clone(),
|
||||
organization.clone(),
|
||||
site.clone(),
|
||||
status.tags.clone(),
|
||||
).await;
|
||||
|
||||
// Update version in database if present
|
||||
@@ -328,8 +341,19 @@ async fn handle_agent_connection(
|
||||
let _ = crate::db::releases::update_machine_version(db.pool(), &agent_id, version).await;
|
||||
}
|
||||
|
||||
info!("Agent status update: {} - streaming={}, uptime={}s, version={:?}",
|
||||
status.hostname, status.is_streaming, status.uptime_secs, agent_version);
|
||||
// Update organization/site/tags in database if present
|
||||
if let Some(ref db) = db {
|
||||
let _ = crate::db::machines::update_machine_metadata(
|
||||
db.pool(),
|
||||
&agent_id,
|
||||
organization.as_deref(),
|
||||
site.as_deref(),
|
||||
&status.tags,
|
||||
).await;
|
||||
}
|
||||
|
||||
info!("Agent status update: {} - streaming={}, uptime={}s, version={:?}, org={:?}, site={:?}",
|
||||
status.hostname, status.is_streaming, status.uptime_secs, agent_version, organization, site);
|
||||
}
|
||||
Some(proto::message::Payload::Heartbeat(_)) => {
|
||||
// Update heartbeat timestamp
|
||||
|
||||
@@ -48,6 +48,9 @@ pub struct Session {
|
||||
pub uptime_secs: i64,
|
||||
pub display_count: i32,
|
||||
pub agent_version: Option<String>, // Agent software version
|
||||
pub organization: Option<String>, // Company/organization name
|
||||
pub site: Option<String>, // Site/location name
|
||||
pub tags: Vec<String>, // Tags for categorization
|
||||
}
|
||||
|
||||
/// Channel for sending frames from agent to viewers
|
||||
@@ -140,6 +143,9 @@ impl SessionManager {
|
||||
uptime_secs: 0,
|
||||
display_count: 1,
|
||||
agent_version: None,
|
||||
organization: None,
|
||||
site: None,
|
||||
tags: Vec::new(),
|
||||
};
|
||||
|
||||
let session_data = SessionData {
|
||||
@@ -170,6 +176,9 @@ impl SessionManager {
|
||||
display_count: i32,
|
||||
is_streaming: bool,
|
||||
agent_version: Option<String>,
|
||||
organization: Option<String>,
|
||||
site: Option<String>,
|
||||
tags: Vec<String>,
|
||||
) {
|
||||
let mut sessions = self.sessions.write().await;
|
||||
if let Some(session_data) = sessions.get_mut(&session_id) {
|
||||
@@ -185,6 +194,15 @@ impl SessionManager {
|
||||
if let Some(version) = agent_version {
|
||||
session_data.info.agent_version = Some(version);
|
||||
}
|
||||
if let Some(org) = organization {
|
||||
session_data.info.organization = Some(org);
|
||||
}
|
||||
if let Some(s) = site {
|
||||
session_data.info.site = Some(s);
|
||||
}
|
||||
if !tags.is_empty() {
|
||||
session_data.info.tags = tags;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -461,6 +479,9 @@ impl SessionManager {
|
||||
uptime_secs: 0,
|
||||
display_count: 1,
|
||||
agent_version: None,
|
||||
organization: None,
|
||||
site: None,
|
||||
tags: Vec::new(),
|
||||
};
|
||||
|
||||
// Create placeholder channels (will be replaced on reconnect)
|
||||
|
||||
Reference in New Issue
Block a user