dcoshc/main.go

package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"log"
	"os/exec"
	"strings"
	"sync"

	"github.com/fatih/color"
)

type command struct {
	node string
	hc   string
}

type healthCheck struct {
	Name     string `json:"name"`
	PodSpecs []struct {
		Type      string `json:"type"`
		TaskSpecs []struct {
			HealthCheckSpec struct {
				Command string `json:"command"`
			} `json:"health-check-spec"`
		} `json:"task-specs"`
	} `json:"pod-specs"`
}

type serviceIPs struct {
	Name     string `json:"name"`
	Statuses []struct {
		ContainerStatus struct {
			NetworkInfos []struct {
				IPAddresses []struct {
					IPAddress string `json:"ip_address"`
				} `json:"ip_addresses"`
			} `json:"network_infos"`
		} `json:"container_status"`
	} `json:"statuses"`
}

func getHealthCheck() healthCheck {
	cmd := exec.Command("dcos", "kubernetes", "config", "target")
	var out bytes.Buffer
	cmd.Stdout = &out
	err := cmd.Run()
	if err != nil {
		log.Fatal(err)
	}

	h := healthCheck{}
	if err := json.Unmarshal([]byte(out.String()), &h); err != nil {
		panic(err)
	}
	return h
}

func getServiceIPs() []serviceIPs {
	cmd := exec.Command("dcos", "task", "--json")
	var out bytes.Buffer
	cmd.Stdout = &out
	err := cmd.Run()
	if err != nil {
		log.Fatal(err)
	}

	s := []serviceIPs{}
	if err := json.Unmarshal([]byte(out.String()), &s); err != nil {
		panic(err)
	}

	// get rid of kubernetes pod
	for i, n := range s {
		if n.Name == "kubernetes" {
			s = append(s[:i], s[i+1:]...)
		}
	}
	return s
}

func addAPIServer(n string, ips []serviceIPs) string {
	for _, i := range ips {
		if strings.Contains(i.Name, n) {
			return i.Statuses[0].ContainerStatus.NetworkInfos[0].IPAddresses[0].IPAddress
		}
	}
	return ""
}

func stripHealthCheck(n string, hcs healthCheck) string {
	for _, j := range hcs.PodSpecs {
		if strings.Contains(n, j.Type) {
			if len(j.TaskSpecs) != 0 {
				return j.TaskSpecs[0].HealthCheckSpec.Command
			}
			return ""
		}
	}
	return ""
}

func buildCommands(hcs healthCheck, ips []serviceIPs) []command {
	commands := []command{}
	for _, s := range ips {
		if strings.Contains(s.Name, "apiserver") || strings.Contains(s.Name, "etcd") {
			ip := s.Statuses[0].ContainerStatus.NetworkInfos[0].IPAddresses[0].IPAddress
			c := strings.Replace(stripHealthCheck(s.Name, hcs), "$LIBPROCESS_IP", ip, -1)
			commands = append(commands, command{node: s.Name, hc: c})
		} else {
			commands = append(commands, command{node: s.Name, hc: stripHealthCheck(s.Name, hcs)})
		}
	}
	fmt.Println(commands)
	return commands
}

func main() {
	var wg sync.WaitGroup
	hcs := getHealthCheck()
	ips := getServiceIPs()
	hosts := buildCommands(hcs, ips)

	for _, n := range hosts {
		wg.Add(1)
		go func(c command) {
			defer wg.Done()
			s := strings.Split(c.hc, " ")
			cmd := exec.Command(
				"dcos",
				"task",
				"exec",
				c.node,
				s[0],
				s[1],
				s[2],
			)
			var out bytes.Buffer
			cmd.Stdout = &out
			err := cmd.Run()
			if err != nil {
				log.Fatal(err)
			}
			if strings.Contains(out.String(), "ok") || strings.Contains(out.String(), "true") {
				color.Green(fmt.Sprintf("%v: %q\n", c, out.String()))
			} else {
				color.Red(fmt.Sprintf("%v: %q\n", c, out.String()))
			}
		}(n)
	}
	wg.Wait()
}