1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107

use std::{ffi::OsStr, process::Stdio};

use tokio::process::Command;

use super::Result;
use crate::util::run_command_output;

fn prepare_base_command<Cmd, NS, Pod>(
    kubectl_command: Cmd,
    namespace: NS,
    pod: Option<Pod>,
) -> Command
where
    Cmd: AsRef<OsStr>,
    NS: AsRef<OsStr>,
    Pod: AsRef<OsStr>,
{
    let mut command = Command::new(&kubectl_command);

    command
        .stdin(Stdio::null())
        .stdout(Stdio::piped())
        .stderr(Stdio::piped());

    command.arg("get").arg("pod");

    if let Some(pod) = pod {
        command.arg(pod);
    }

    command.arg("-n").arg(namespace).arg("-o").arg("json");

    command
}

/// Returns the node the given pod is on
pub async fn get_node<Cmd, NS, Pod>(kubectl_command: Cmd, namespace: NS, pod: Pod) -> Result<String>
where
    Pod: AsRef<OsStr>,
    Cmd: AsRef<OsStr>,
    NS: AsRef<OsStr>,
{
    let command = prepare_base_command(kubectl_command, namespace, Some(pod));
    let pod = run_command_output(command).await?;
    let pod: serde_json::Value = serde_json::from_str(&pod)?;

    let node = pod["spec"]["nodeName"]
        .as_str()
        .ok_or("nodename must be a string")?;

    Ok(node.to_string())
}

/// Set label on all nodes to discover this label from the pod.
pub async fn label_nodes<Cmd, Label>(kubectl_command: Cmd, label: Label) -> Result<String>
where
    Cmd: AsRef<OsStr>,
    Label: AsRef<OsStr>,
{
    let mut command = Command::new(&kubectl_command);
    command
        .arg("label")
        .arg("node")
        .arg(label)
        .arg("--all")
        .arg("--overwrite");

    let res = run_command_output(command).await?;
    Ok(res)
}

pub async fn get_pod_on_node<Cmd, NS>(
    kubectl_command: Cmd,
    namespace: NS,
    node: &str,
    service: &str,
) -> Result<String>
where
    Cmd: AsRef<OsStr>,
    NS: AsRef<OsStr>,
{
    let nopod: Option<&str> = None;
    let command = prepare_base_command(kubectl_command, namespace, nopod);
    let pods = run_command_output(command).await?;
    let pods: serde_json::Value = serde_json::from_str(&pods)?;

    let pods = pods["items"].as_array().ok_or("items should be an array")?;

    for pod in pods {
        if pod["spec"]["nodeName"]
            .as_str()
            .ok_or("nodeName must be a string")?
            == node
            && pod["spec"]["serviceAccount"]
                .as_str()
                .ok_or("serviceAccount must be a string")?
                == service
        {
            return Ok(pod["metadata"]["name"]
                .as_str()
                .ok_or("name must be a string")?
                .to_string());
        }
    }

    Err("No pod on this node".into())
}