johnandersen777.bsky.social/asyncai

Switch repository

Loading…

use std::collections::BTreeSet;

use serde::{Deserialize, Serialize};
use serde_json::Value;

use super::runner::{shell_quote, Runner};

pub const DEFAULT_ROOTS: &[&str] = &[
    "/var/lib/socialweb-computer",
    "/var/tmp/swcit",
    "~/.local/state/socialweb-computer",
];

#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
pub struct KcpInfo {
    pub id: String,
    pub host: String,
    pub kubeconfig: String,
    pub server: String,
    pub workspace: String,
    pub context: String,
    pub running: bool,
    pub error: Option<String>,
}

#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
pub struct WorkspaceInfo {
    pub name: String,
    pub path: String,
    pub url: String,
    pub phase: String,
    pub kind: String,
}

#[derive(Clone, Debug)]
pub struct AdminConfig {
    pub ca_data: Option<String>,
    pub tls_server_name: Option<String>,
}

#[derive(Clone, Debug)]
pub struct Discovery {
    runner: Runner,
    roots: Vec<String>,
    kubeconfigs: Vec<String>,
}

impl Discovery {
    pub fn new(runner: Runner, roots: Vec<String>, kubeconfigs: Vec<String>) -> Self {
        Self {
            runner,
            roots,
            kubeconfigs,
        }
    }

    pub async fn discover(&self) -> (Vec<KcpInfo>, Vec<String>) {
        let mut errors = Vec::new();
        let mut paths = BTreeSet::new();

        for path in &self.kubeconfigs {
            paths.insert(path.clone());
        }

        if self.kubeconfigs.is_empty() {
            let roots = if self.roots.is_empty() {
                DEFAULT_ROOTS
                    .iter()
                    .map(|root| expand_home(root))
                    .collect::<Vec<_>>()
            } else {
                self.roots.iter().map(|root| expand_home(root)).collect()
            };
            let script = find_script(&roots);
            match self.runner.shell(&script).await {
                Ok(output) => {
                    for path in prefer_kcp_admin_kubeconfigs(
                        output
                            .lines()
                            .map(str::trim)
                            .filter(|line| !line.is_empty())
                            .map(str::to_string),
                    ) {
                        paths.insert(path);
                    }
                }
                Err(error) => errors.push(format!("kubeconfig discovery failed: {error}")),
            }
        }

        let mut kcps = Vec::new();
        for path in paths {
            match self.inspect_kubeconfig(&path).await {
                Ok(info) => kcps.push(info),
                Err(error) => errors.push(format!("{path}: {error}")),
            }
        }

        (kcps, errors)
    }

    async fn inspect_kubeconfig(&self, path: &str) -> Result<KcpInfo, String> {
        let view = self
            .runner
            .kubectl(
                &[
                    "--kubeconfig".to_string(),
                    path.to_string(),
                    "config".to_string(),
                    "view".to_string(),
                    "--raw".to_string(),
                    "-o".to_string(),
                    "json".to_string(),
                ],
                None,
            )
            .await?;
        let current = parse_current_context(&view)?;

        Ok(KcpInfo {
            id: format!("{}:{}", self.host_label(), path),
            host: self.host_label(),
            kubeconfig: path.to_string(),
            server: current.server_base,
            workspace: current.workspace,
            context: current.context,
            running: false,
            error: None,
        })
    }

    pub async fn admin_config(&self, kcp: &KcpInfo) -> Result<AdminConfig, String> {
        let view = self
            .runner
            .kubectl(
                &[
                    "--kubeconfig".to_string(),
                    kcp.kubeconfig.clone(),
                    "config".to_string(),
                    "view".to_string(),
                    "--raw".to_string(),
                    "-o".to_string(),
                    "json".to_string(),
                ],
                None,
            )
            .await?;
        let value: Value = serde_json::from_str(&view)
            .map_err(|error| format!("kubeconfig is not JSON: {error}"))?;
        let cluster = current_cluster(&value)?;
        Ok(AdminConfig {
            ca_data: cluster
                .get("certificate-authority-data")
                .and_then(Value::as_str)
                .map(str::to_string),
            tls_server_name: cluster
                .get("tls-server-name")
                .and_then(Value::as_str)
                .map(str::to_string),
        })
    }

    fn host_label(&self) -> String {
        self.runner_host_label()
    }

    fn runner_host_label(&self) -> String {
        self.runner
            .ssh_host()
            .map(str::to_string)
            .unwrap_or_else(|| "local".to_string())
    }
}

#[derive(Clone, Debug)]
struct CurrentContext {
    context: String,
    workspace: String,
    server_base: String,
}

fn parse_current_context(view: &str) -> Result<CurrentContext, String> {
    let value: Value =
        serde_json::from_str(view).map_err(|error| format!("kubeconfig is not JSON: {error}"))?;
    let context_name = value
        .get("current-context")
        .and_then(Value::as_str)
        .filter(|name| !name.is_empty())
        .ok_or_else(|| "kubeconfig has no current-context".to_string())?
        .to_string();
    let server = current_cluster(&value)?
        .get("server")
        .and_then(Value::as_str)
        .filter(|server| !server.is_empty())
        .ok_or_else(|| "kubeconfig current cluster has no server".to_string())?
        .to_string();
    let workspace = workspace_from_server(&server);
    Ok(CurrentContext {
        context: context_name,
        workspace,
        server_base: server_base(&server),
    })
}

fn current_cluster(value: &Value) -> Result<&Value, String> {
    let context_name = value
        .get("current-context")
        .and_then(Value::as_str)
        .ok_or_else(|| "kubeconfig has no current-context".to_string())?;
    let context = value
        .get("contexts")
        .and_then(Value::as_array)
        .and_then(|contexts| {
            contexts
                .iter()
                .find(|context| context.get("name").and_then(Value::as_str) == Some(context_name))
        })
        .and_then(|context| context.get("context"))
        .ok_or_else(|| format!("kubeconfig has no context {context_name:?}"))?;
    let cluster_name = context
        .get("cluster")
        .and_then(Value::as_str)
        .ok_or_else(|| format!("kubeconfig context {context_name:?} has no cluster"))?;
    value
        .get("clusters")
        .and_then(Value::as_array)
        .and_then(|clusters| {
            clusters
                .iter()
                .find(|cluster| cluster.get("name").and_then(Value::as_str) == Some(cluster_name))
        })
        .and_then(|cluster| cluster.get("cluster"))
        .ok_or_else(|| format!("kubeconfig has no cluster {cluster_name:?}"))
}

pub fn server_base(server: &str) -> String {
    match server.find("/clusters/") {
        Some(index) => server[..index].trim_end_matches('/').to_string(),
        None => server.trim_end_matches('/').to_string(),
    }
}

pub fn workspace_from_server(server: &str) -> String {
    match server.find("/clusters/") {
        Some(index) => server[index + "/clusters/".len()..]
            .trim_end_matches('/')
            .to_string(),
        None => "root".to_string(),
    }
}

pub fn workspace_url(server_base: &str, workspace: &str) -> String {
    if workspace.is_empty() || workspace == "root" {
        return format!("{}/clusters/root", server_base.trim_end_matches('/'));
    }
    format!(
        "{}/clusters/{}",
        server_base.trim_end_matches('/'),
        workspace
    )
}

pub(crate) fn parse_workspace_list(
    output: &str,
    parent: &str,
    server_base: &str,
) -> Result<Vec<WorkspaceInfo>, String> {
    if output.trim().is_empty() {
        return Ok(Vec::new());
    }
    let value: Value = serde_json::from_str(output)
        .map_err(|error| format!("workspace list is not JSON: {error}"))?;
    let items = value
        .get("items")
        .and_then(Value::as_array)
        .cloned()
        .unwrap_or_default();

    let mut workspaces = Vec::new();
    for item in items {
        let name = item
            .pointer("/metadata/name")
            .and_then(Value::as_str)
            .unwrap_or_default()
            .to_string();
        if name.is_empty() {
            continue;
        }
        let path = if parent.is_empty() {
            name.clone()
        } else {
            format!("{parent}:{name}")
        };
        let url = item
            .pointer("/status/url")
            .and_then(Value::as_str)
            .filter(|url| !url.is_empty())
            .map(str::to_string)
            .unwrap_or_else(|| workspace_url(server_base, &path));
        let phase = item
            .pointer("/status/phase")
            .and_then(Value::as_str)
            .unwrap_or("Unknown")
            .to_string();
        let kind = item
            .pointer("/spec/type/name")
            .and_then(Value::as_str)
            .unwrap_or("")
            .to_string();
        workspaces.push(WorkspaceInfo {
            name,
            path,
            url,
            phase,
            kind,
        });
    }
    Ok(workspaces)
}

fn find_script(roots: &[String]) -> String {
    let mut script = String::from("for root in");
    for root in roots {
        script.push(' ');
        script.push_str(&shell_quote(root));
    }
    script.push_str("; do [ -d \"$root\" ] && { find \"$root\" -type f -name admin.kubeconfig -path '*/kcp/*' ! -path '*/shared/*' 2>/dev/null || true; }; done; true");
    script
}

fn prefer_kcp_admin_kubeconfigs(paths: impl Iterator<Item = String>) -> Vec<String> {
    let mut grouped: std::collections::BTreeMap<String, (usize, String)> =
        std::collections::BTreeMap::new();
    for path in paths {
        let Some(index) = path.find("/kcp/") else {
            continue;
        };
        let kcp_root = path[..index + "/kcp/".len()].to_string();
        let suffix_len = path.len() - kcp_root.len();
        match grouped.get(&kcp_root) {
            Some((best, _)) if *best <= suffix_len => continue,
            _ => {
                grouped.insert(kcp_root, (suffix_len, path));
            }
        }
    }
    grouped.into_values().map(|(_, path)| path).collect()
}

fn expand_home(path: &str) -> String {
    if let Some(rest) = path.strip_prefix("~/") {
        if let Ok(home) = std::env::var("HOME") {
            return format!("{home}/{rest}");
        }
    }
    path.to_string()
}

#[cfg(test)]
mod tests {
    use super::{parse_current_context, server_base, workspace_from_server, workspace_url};

    const VIEW: &str = r#"{
        "current-context": "root",
        "clusters": [{"name": "root", "cluster": {"server": "https://127.0.0.1:33059/clusters/root", "certificate-authority-data": "Y2E=", "tls-server-name": "127.0.0.1"}}],
        "contexts": [{"name": "root", "context": {"cluster": "root", "user": "shard-admin"}}],
        "users": [{"name": "shard-admin", "user": {"token": "secret"}}]
    }"#;

    #[test]
    fn reads_the_current_context_without_exposing_credentials() {
        let current = parse_current_context(VIEW).unwrap();
        assert_eq!(current.context, "root");
        assert_eq!(current.workspace, "root");
        assert_eq!(current.server_base, "https://127.0.0.1:33059");
    }

    #[test]
    fn workspace_urls_are_derived_from_the_base() {
        assert_eq!(
            server_base("https://host:1/clusters/root:app"),
            "https://host:1"
        );
        assert_eq!(
            workspace_from_server("https://host:1/clusters/root:app"),
            "root:app"
        );
        assert_eq!(
            workspace_url("https://host:1", "root:app"),
            "https://host:1/clusters/root:app"
        );
    }
}