churn-v2/crates/churn/src/main.rs
kjuulh 43ed89d0d8
feat: with agent db
Signed-off-by: kjuulh <contact@kjuulh.io>
2023-08-27 19:42:33 +02:00

155 lines
4.9 KiB
Rust

use churn_domain::{AgentEnrollReq, LeaseResp, ServerMonitorResp};
use clap::{Parser, Subcommand};
#[derive(Parser)]
#[command(author, version, about, long_about = None, subcommand_required = true)]
struct Command {
#[command(subcommand)]
command: Option<Commands>,
}
#[derive(Subcommand)]
enum Commands {
Auth {
#[arg(env = "CHURN_SERVER", long)]
server: String,
#[arg(env = "CHURN_SERVER_TOKEN", long)]
server_token: String,
},
Bootstrap {
#[arg(env = "CHURN_AGENT", long)]
agent: String,
#[arg(env = "CHURN_AGENT_NAME", long)]
agent_name: String,
#[arg(env = "CHURN_SERVER", long)]
server: String,
#[arg(env = "CHURN_SERVER_TOKEN", long)]
server_token: String,
},
Health {
#[arg(env = "CHURN_SERVER", long)]
server: String,
#[arg(env = "CHURN_AGENT", long)]
agent: String,
},
Monitor {
#[arg(env = "CHURN_SERVER", long)]
server: String,
#[arg(env = "CHURN_SERVER_TOKEN", long)]
server_token: String,
},
}
#[tokio::main]
async fn main() -> anyhow::Result<()> {
dotenv::dotenv().ok();
tracing_subscriber::fmt::init();
let cli = Command::parse();
handle_command(cli).await?;
Ok(())
}
async fn handle_command(cmd: Command) -> anyhow::Result<()> {
if let Some(cmd) = cmd.command {
match cmd {
Commands::Bootstrap {
agent,
agent_name,
server,
server_token: _,
} => {
tracing::info!("enrolling agent: {} for server: {}", agent, server);
let client = reqwest::Client::new();
let req = client.post(format!("{server}/agent/lease")).build()?;
let lease_resp = client.execute(req).await?;
let lease = lease_resp.json::<LeaseResp>().await?;
let req = client
.post(format!("{agent}/enroll"))
.json(&AgentEnrollReq {
lease: lease.token,
server,
agent_name,
})
.build()?;
let lease_resp = client.execute(req).await?;
if !lease_resp.status().is_success() {
if let Ok(text) = lease_resp.text().await {
tracing::warn!(
"could not enroll because agent server encoutered error: {}",
text
);
anyhow::bail!("encountered error: {}", text);
}
anyhow::bail!("encountered error");
}
Ok(())
}
Commands::Health { server, agent } => {
tracing::info!("connecting to server: {}", server);
reqwest::get(format!("{server}/ping")).await?;
tracing::info!("connected to server successfully");
tracing::info!("connecting to agent: {}", agent);
reqwest::get(format!("{agent}/ping")).await?;
tracing::info!("connected to agent successfully");
Ok(())
}
Commands::Auth {
server: _,
server_token: _,
} => todo!(),
Commands::Monitor {
server,
server_token: _,
} => {
tracing::info!("monitoring server: {}", server);
let mut cursor: Option<uuid::Uuid> = None;
loop {
tracing::debug!("reading logs from server: {}", server);
let resp = reqwest::get(format!(
"{server}/logs{}",
match &cursor {
None => "".to_string(),
Some(cursor) => format!("?cursor={}", cursor),
}
))
.await?;
if !resp.status().is_success() {
if let Ok(text) = resp.text().await {
anyhow::bail!("encountered error: {}", text);
}
anyhow::bail!("encountered error");
}
match resp.json::<ServerMonitorResp>().await {
Ok(resp) => {
for line in resp.logs {
tracing::info!("event: {}", line);
}
cursor = resp.cursor;
}
Err(e) => {
tracing::warn!("failed to call server (error={})", e);
}
}
tokio::time::sleep(std::time::Duration::from_secs(1)).await;
}
}
}
} else {
panic!("no command supplied")
}
}