Oculus/main.go

183 lines
5.4 KiB
Go

package main
import (
"context"
"encoding/json"
"fmt"
"io/ioutil"
"log"
"os"
"os/exec"
"strings"
"time"
"github.com/docker/docker/api/types"
"github.com/docker/docker/api/types/events"
"github.com/docker/docker/api/types/filters"
"github.com/docker/docker/client"
"github.com/robfig/cron/v3"
)
const (
notifyScript = "/notify.sh"
)
// ContainerDiff represents the state of a container
type ContainerDiff struct {
ID string
Image string
Labels map[string]string
}
// getRunningContainers fetches the list of running containers with relevant labels
func getRunningContainers(cli *client.Client) ([]ContainerDiff, error) {
containers, err := cli.ContainerList(context.Background(), types.ContainerListOptions{})
if err != nil {
return nil, err
}
var diffs []ContainerDiff
for _, container := range containers {
if _, ok := container.Labels["oculus.containerid"]; ok {
diffs = append(diffs, ContainerDiff{
ID: container.ID,
Image: container.Image,
Labels: container.Labels,
})
}
}
return diffs, nil
}
// saveDiffs writes the container diffs to a file
func saveDiffs(diffs []ContainerDiff, filePath string) error {
data, err := json.Marshal(diffs)
if err != nil {
return err
}
return ioutil.WriteFile(filePath, data, 0644)
}
// loadDiffs reads the container diffs from a file
func loadDiffs(filePath string) ([]ContainerDiff, error) {
data, err := ioutil.ReadFile(filePath)
if err != nil {
return nil, err
}
var diffs []ContainerDiff
err = json.Unmarshal(data, &diffs)
return diffs, err
}
// profileContainers logs the current state of containers to a file
func profileContainers(cli *client.Client) {
diffs, err := getRunningContainers(cli)
if err != nil {
log.Fatalf("Error profiling containers: %v", err)
}
for _, diff := range diffs {
if diff.Labels["oculus.mode"] == "profile" {
filePath := fmt.Sprintf("%s.json", diff.Labels["oculus.containerid"])
err := saveDiffs([]ContainerDiff{diff}, filePath)
if err != nil {
log.Fatalf("Error saving diffs: %v", err)
}
}
}
log.Println("Profiled current containers and saved to file")
}
// compareContainers compares the current state of containers with the saved state
func compareContainers(cli *client.Client) {
currentDiffs, err := getRunningContainers(cli)
if err != nil {
log.Fatalf("Error getting current containers: %v", err)
}
for _, current := range currentDiffs {
if current.Labels["oculus.mode"] == "monitor" {
filePath := fmt.Sprintf("%s.json", current.Labels["oculus.containerid"])
savedDiffs, err := loadDiffs(filePath)
if err != nil {
log.Printf("Error loading saved diffs for %s: %v", current.ID, err)
continue
}
if len(savedDiffs) > 0 {
saved := savedDiffs[0]
if current.Image != saved.Image {
log.Printf("Container %s changed: Image %s -> %s", current.ID, saved.Image, current.Image)
runNotifyScript(current.ID, current.Image, saved.Image)
}
if current.Labels["oculus.ignorelist"] != saved.Labels["oculus.ignorelist"] {
log.Printf("Container %s ignore list changed: %s -> %s", current.ID, saved.Labels["oculus.ignorelist"], current.Labels["oculus.ignorelist"])
runNotifyScript(current.ID, current.Image, saved.Image)
}
} else {
log.Printf("New container detected: ID %s, Image %s", current.ID, current.Image)
runNotifyScript(current.ID, current.Image, "")
}
}
}
}
// runNotifyScript executes the notification script with the container details
func runNotifyScript(containerID, newImage, oldImage string) {
cmd := exec.Command(notifyScript, containerID, newImage, oldImage)
err := cmd.Run()
if err != nil {
log.Printf("Error running notify script: %v", err)
}
}
// watchDockerEvents listens for Docker events and triggers comparisons
func watchDockerEvents(cli *client.Client) {
eventFilter := filters.NewArgs()
eventFilter.Add("type", "container")
eventFilter.Add("event", "start")
eventFilter.Add("event", "stop")
eventFilter.Add("event", "destroy")
messages, errs := cli.Events(context.Background(), types.EventsOptions{
Filters: eventFilter,
})
for {
select {
case event := <-messages:
log.Printf("Docker event received: %s for container %s", event.Action, event.ID)
compareContainers(cli)
case err := <-errs:
log.Printf("Error watching Docker events: %v", err)
return
}
}
}
// scheduler sets up the cron jobs for profiling
func scheduler(cli *client.Client) {
c := cron.New()
c.AddFunc("@every 1h", func() { profileContainers(cli) })
c.Start()
}
func main() {
cli, err := client.NewClientWithOpts(client.FromEnv, client.WithAPIVersionNegotiation())
if (err != nil) {
log.Fatalf("Error creating Docker client: %v", err)
}
go scheduler(cli)
go watchDockerEvents(cli)
// Keep the program running
select {}
}