callgraph/main.go

package main

import (
	"flag"
	"fmt"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"syscall"
	"time"

	"./rancher"
	"./sharelatex"
	"./sysdig"
)

var (
	skipTest         = flag.Bool("skip-test", false, "skip selenium integration test")
	rancherHost      = flag.String("rancher-host", "localhost:8080", "host of rancher server")
	rancherAccessKey = flag.String("rancher-access-key", os.Getenv("RANCHER_ACCESS_KEY"), "api access key")
	rancherSecretKey = flag.String("rancher-secret-key", os.Getenv("RANCHER_SECRET_KEY"), "api secret key")
)

func parseArgs() (sharelatexUrl string, hosts []string) {
	flag.Parse()
	args := flag.Args()
	if len(args) < 1 {
		fmt.Fprintf(os.Stderr, "USAGE: %s SHARELATEXURL HOSTS...", os.Args[0])
		os.Exit(1)
	}
	if *rancherAccessKey == "" {
		fmt.Fprintf(os.Stderr, "Set RANCHER_ACCESS_KEY environment variable")
		os.Exit(1)
	}
	if *rancherSecretKey == "" {
		fmt.Fprintf(os.Stderr, "Set RANCHER_SECRET_KEY environment variable")
		os.Exit(1)
	}
	return args[0], args[1:]
}

func rancherServices() ([]sysdig.Service, error) {
	c := rancher.New(*rancherHost, *rancherAccessKey, *rancherSecretKey)
	var api rancher.Api
	err := c.Get(fmt.Sprintf("http://%s/v1", c.Host), &api)
	if err != nil {
		return nil, fmt.Errorf("failed to request rancher api: %v", err)
	}
	if api.Links == nil {
		return nil, fmt.Errorf("no links object found in api response")
	}
	val, ok := api.Links["containers"]
	if !ok {
		return nil, fmt.Errorf("no container link found in api response")
	}
	var containers rancher.Containers
	err = c.Get(val, &containers)
	if err != nil {
		return nil, fmt.Errorf("could not get container list from api: %s", err)
	}
	serviceMap := make(map[string][]net.IP)
	for _, container := range containers.Data {
		var name string
		if container.Labels.ProjectServiceName == "" {
			name = container.Name
		} else {
			name = container.Labels.ProjectServiceName
		}
		if name == "Network Agent" {
			continue
		}
		ip := net.ParseIP(container.PrimaryIpAddress)
		if ip == nil {
			return nil, fmt.Errorf("Got invalid IP from api for '%s' container: %s", container.Name, container.PrimaryIpAddress)
		}
		serviceMap[name] = append(serviceMap[name], ip)
	}
	var services []sysdig.Service
	for name, ips := range serviceMap {
		services = append(services, sysdig.Service{name, ips})
	}
	return services, nil
}

func die(message string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, message, args...)
	fmt.Fprint(os.Stderr, "\n")
	os.Exit(1)
}

func main() {
	url, hosts := parseArgs()

	cmd, err := sysdig.Start(hosts)
	if err != nil {
		die("failed to start sysdig: %v", err)
	}

	if *skipTest {
		signalChan := make(chan os.Signal, 1)
		signal.Notify(signalChan, os.Interrupt)
		signal.Notify(signalChan, syscall.SIGTERM)
		<-signalChan
	} else {
		time.Sleep(4 * time.Second)
		err = sharelatex.IntegrationTest(url)
		if err != nil {
			die("selenium not running?: %v", err)
		}
		time.Sleep(1 * time.Second)
	}

	cmd.SendStopSignal()

	var records []sysdig.Record
	for range hosts {
		select {
		case result := <-cmd.Results:
			if exitError, ok := result.Error.(*exec.ExitError); ok {
				waitStatus := exitError.Sys().(syscall.WaitStatus)
				if waitStatus != 130 { // Interrupt
					fmt.Printf("[%s] sysdig failed with %v\n", result.Host.Name, result.Error, result.Output)
				}
			}
			r, err := sysdig.ParseOutput(result.Output)
			if err != nil {
				die("Error parsing results from %s", result.Host.Name)
			}
			records = append(records, r...)
		}
	}
	services, err := rancherServices()
	if err != nil {
		die("getting rancher services failed: %s", err)
	}
	(sysdig.Records(records)).PrintGraph(services)
}