lda/wakey
1
0
Fork 0
Code Issues 5 Pull Requests Actions Packages Projects Releases 18 Wiki Activity

uhhh

Browse Source
This commit is contained in:
lda
2026-04-11 19:08:22 +07:00 Unverified
parent 3dfc2b4832
commit 306cce0e6f
6 changed files with 213 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
wakey-control-plane/src/api.rs
+36 -2
View File
@@ -37,19 +37,28 @@ pub struct ListEnrollTokenQuery {
pub include_expired: Option<bool>,
}
#[derive(Debug, Serialize)]
#[derive(Debug, Serialize, Deserialize)]
pub struct EnrollTokenStatus {
pub enroll_token: String,
pub expires_at_unix: u64,
pub expired: bool,
}
#[derive(Debug, Serialize)]
#[derive(Debug, Serialize, Deserialize)]
pub struct RevokeEnrollTokenResponse {
pub token: String,
pub revoked: bool,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct StateStatsResponse {
pub db_path: String,
pub schema_version: u32,
pub agent_count: usize,
pub enroll_token_count: usize,
pub expired_enroll_token_count: usize,
}
#[derive(Debug, Serialize)]
pub struct AgentStatus {
pub agent_id: String,
@@ -201,6 +210,31 @@ pub async fn revoke_enroll_token(
}
}
pub async fn state_stats(
State(state): State<AppState>,
) -> Result<impl IntoResponse, (StatusCode, Json<serde_json::Value>)> {
match state.store.stats().await {
Ok(stats) => Ok((
StatusCode::OK,
Json(StateStatsResponse {
db_path: stats.db_path.display().to_string(),
schema_version: stats.schema_version,
agent_count: stats.agent_count,
enroll_token_count: stats.enroll_token_count,
expired_enroll_token_count: stats.expired_enroll_token_count,
}),
)),
Err(err) => {
warn!(error = %err, "failed to read state stats");
Err(json_error(
StatusCode::INTERNAL_SERVER_ERROR,
"state_stats_failed",
&err.to_string(),
))
}
}
}
pub async fn run_command(
State(state): State<AppState>,
AxumPath(agent_id): AxumPath<String>,
wakey-control-plane/src/cli.rs
+12
View File
@@ -64,6 +64,9 @@ pub struct ServeArgs {
#[arg(long, default_value = DEFAULT_CONFIG_FILE)]
pub config_file: PathBuf,
#[arg(long)]
pub bootstrap_config: bool,
}
#[derive(Args, Clone)]
@@ -140,6 +143,9 @@ pub struct ListEnrollTokensArgs {
#[arg(long)]
pub data_dir: Option<PathBuf>,
#[arg(long)]
pub public_url: Option<String>,
#[arg(long)]
pub include_expired: bool,
@@ -158,6 +164,9 @@ pub struct RevokeEnrollTokenArgs {
#[arg(long)]
pub data_dir: Option<PathBuf>,
#[arg(long)]
pub public_url: Option<String>,
#[arg(long)]
pub token: String,
}
@@ -173,6 +182,9 @@ pub struct StateStatsArgs {
#[arg(long)]
pub data_dir: Option<PathBuf>,
#[arg(long)]
pub public_url: Option<String>,
#[arg(long)]
pub json: bool,
}
wakey-control-plane/src/config.rs
+56 -4
View File
@@ -268,6 +268,31 @@ pub fn write_init_config(args: &InitConfigArgs) -> Result<()> {
Ok(())
}
pub fn bootstrap_config_if_missing(args: &ServeArgs) -> Result<bool> {
if args.config_file.exists() {
return Ok(false);
}
let init = InitConfigArgs {
config_file: args.config_file.clone(),
data_dir: args.data_dir.clone(),
bind: args.bind,
public_url: args.public_url.clone(),
state_file: args.state_file.clone(),
pid_file: args.pid_file.clone(),
command_timeout_ms: args.command_timeout_ms,
enroll_token_ttl_seconds: args.enroll_token_ttl_seconds,
enroll_tokens: args.enroll_tokens.clone(),
telemetry_otlp_endpoint: None,
telemetry_service_name: None,
telemetry_json_logs: false,
force: false,
};
write_init_config(&init)?;
Ok(true)
}
pub fn resolve_path(data_dir: &Path, candidate: PathBuf) -> PathBuf {
if candidate.is_absolute() {
candidate
@@ -285,11 +310,17 @@ pub struct IssueTokenSettings {
pub struct StateAccessSettings {
pub data_dir: PathBuf,
pub state_file: PathBuf,
pub public_url: Option<String>,
}
pub fn resolve_issue_token_settings(args: &IssueEnrollTokenArgs) -> Result<IssueTokenSettings> {
let file = load_file_config(&args.config_file)?;
let state = resolve_state_access(&args.config_file, args.data_dir.clone(), args.state_file.clone())?;
let state = resolve_state_access(
&args.config_file,
args.data_dir.clone(),
args.state_file.clone(),
args.public_url.clone(),
)?;
let ttl = Duration::from_secs(
args.ttl_seconds
@@ -306,21 +337,37 @@ pub fn resolve_issue_token_settings(args: &IssueEnrollTokenArgs) -> Result<Issue
}
pub fn resolve_list_enroll_token_settings(args: &ListEnrollTokensArgs) -> Result<StateAccessSettings> {
resolve_state_access(&args.config_file, args.data_dir.clone(), args.state_file.clone())
resolve_state_access(
&args.config_file,
args.data_dir.clone(),
args.state_file.clone(),
args.public_url.clone(),
)
}
pub fn resolve_revoke_enroll_token_settings(args: &RevokeEnrollTokenArgs) -> Result<StateAccessSettings> {
resolve_state_access(&args.config_file, args.data_dir.clone(), args.state_file.clone())
resolve_state_access(
&args.config_file,
args.data_dir.clone(),
args.state_file.clone(),
args.public_url.clone(),
)
}
pub fn resolve_state_stats_settings(args: &StateStatsArgs) -> Result<StateAccessSettings> {
resolve_state_access(&args.config_file, args.data_dir.clone(), args.state_file.clone())
resolve_state_access(
&args.config_file,
args.data_dir.clone(),
args.state_file.clone(),
args.public_url.clone(),
)
}
fn resolve_state_access(
config_file: &Path,
cli_data_dir: Option<PathBuf>,
cli_state_file: Option<PathBuf>,
cli_public_url: Option<String>,
) -> Result<StateAccessSettings> {
let file = load_file_config(config_file)?;
@@ -335,8 +382,13 @@ fn resolve_state_access(
.unwrap_or_else(|| PathBuf::from("state.db")),
);
let public_url = cli_public_url
.or(file.public_url)
.map(|url| normalize_public_url(&url));
Ok(StateAccessSettings {
data_dir,
state_file,
public_url,
})
}
wakey-control-plane/src/main.rs
+9
View File
@@ -16,6 +16,15 @@ async fn main() -> Result<()> {
match cli.command {
Command::Serve(args) => {
if args.bootstrap_config {
let created = config::bootstrap_config_if_missing(&args)?;
if created {
eprintln!(
"bootstrapped missing config at {}",
args.config_file.display()
);
}
}
let daemon = config::DaemonConfig::from_serve_args(&args)?;
tracing::init(cli.verbose, &daemon.telemetry)?;
runtime::serve(daemon).await
wakey-control-plane/src/runtime.rs
+94
View File
@@ -63,6 +63,7 @@ pub async fn serve(daemon: config::DaemonConfig) -> Result<()> {
"/api/v1/control/enroll-tokens/{token}",
axum::routing::delete(api::revoke_enroll_token),
)
.route("/api/v1/control/state-stats", get(api::state_stats))
.route("/api/v1/agent/ws", get(ws::agent_ws))
.route("/api/v1/control/agents", get(api::list_agents))
.route(
@@ -186,6 +187,43 @@ pub async fn issue_enroll_token(args: IssueEnrollTokenArgs) -> Result<()> {
pub async fn list_enroll_tokens(args: ListEnrollTokensArgs) -> Result<()> {
let settings = config::resolve_list_enroll_token_settings(&args)?;
if let Some(base) = settings.public_url.as_deref() {
let url = format!(
"{}/api/v1/control/enroll-tokens?include_expired={}",
base,
args.include_expired
);
let response = reqwest::get(&url)
.await
.with_context(|| format!("failed to call {url}"))?;
if !response.status().is_success() {
let status = response.status();
let body = response
.text()
.await
.unwrap_or_else(|_| "<unreadable error body>".to_string());
anyhow::bail!("live list-enroll-tokens failed with {status}: {body}");
}
let body: Vec<api::EnrollTokenStatus> = response
.json()
.await
.context("failed to decode list-enroll-tokens response")?;
if args.json {
println!(
"{}",
serde_json::to_string_pretty(&body).context("failed to render json")?
);
return Ok(());
}
for token in body {
println!(
"token={} expires_at_unix={} expired={}",
token.enroll_token, token.expires_at_unix, token.expired
);
}
return Ok(());
}
let store = state::Store::load_or_init(&settings.state_file, Vec::new(), Duration::from_secs(1))
.await
.with_context(|| format!("failed to initialize store {}", settings.state_file.display()))?;
@@ -208,6 +246,30 @@ pub async fn list_enroll_tokens(args: ListEnrollTokensArgs) -> Result<()> {
pub async fn revoke_enroll_token(args: RevokeEnrollTokenArgs) -> Result<()> {
let settings = config::resolve_revoke_enroll_token_settings(&args)?;
if let Some(base) = settings.public_url.as_deref() {
let url = format!("{}/api/v1/control/enroll-tokens/{}", base, args.token);
let client = reqwest::Client::new();
let response = client
.delete(&url)
.send()
.await
.with_context(|| format!("failed to call {url}"))?;
if !response.status().is_success() {
let status = response.status();
let body = response
.text()
.await
.unwrap_or_else(|_| "<unreadable error body>".to_string());
anyhow::bail!("live revoke-enroll-token failed with {status}: {body}");
}
let body: api::RevokeEnrollTokenResponse = response
.json()
.await
.context("failed to decode revoke-enroll-token response")?;
println!("token={} revoked={}", body.token, body.revoked);
return Ok(());
}
let store = state::Store::load_or_init(&settings.state_file, Vec::new(), Duration::from_secs(1))
.await
.with_context(|| format!("failed to initialize store {}", settings.state_file.display()))?;
@@ -218,6 +280,38 @@ pub async fn revoke_enroll_token(args: RevokeEnrollTokenArgs) -> Result<()> {
pub async fn state_stats(args: StateStatsArgs) -> Result<()> {
let settings = config::resolve_state_stats_settings(&args)?;
if let Some(base) = settings.public_url.as_deref() {
let url = format!("{}/api/v1/control/state-stats", base);
let response = reqwest::get(&url)
.await
.with_context(|| format!("failed to call {url}"))?;
if !response.status().is_success() {
let status = response.status();
let body = response
.text()
.await
.unwrap_or_else(|_| "<unreadable error body>".to_string());
anyhow::bail!("live state-stats failed with {status}: {body}");
}
let body: api::StateStatsResponse = response
.json()
.await
.context("failed to decode state-stats response")?;
if args.json {
println!(
"{}",
serde_json::to_string_pretty(&body).context("failed to render json")?
);
return Ok(());
}
println!("db_path={}", body.db_path);
println!("schema_version={}", body.schema_version);
println!("agent_count={}", body.agent_count);
println!("enroll_token_count={}", body.enroll_token_count);
println!("expired_enroll_token_count={}", body.expired_enroll_token_count);
return Ok(());
}
let store = state::Store::load_or_init(&settings.state_file, Vec::new(), Duration::from_secs(1))
.await
.with_context(|| format!("failed to initialize store {}", settings.state_file.display()))?;