churn-v2/crates/churn/src/main.rs

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")
    }
}