sshpoke/internal/docker/api.go

package docker

import (
	"context"
	"errors"
	"time"

	"github.com/Neur0toxine/sshpoke/internal/config"
	"github.com/Neur0toxine/sshpoke/internal/logger"
	"github.com/Neur0toxine/sshpoke/pkg/dto"
	"github.com/Neur0toxine/sshpoke/pkg/smarttypes"
	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/filters"
	"github.com/docker/docker/client"
)

var Default *Docker

type Docker struct {
	cli           *client.Client
	ctx           context.Context
	wait          chan struct{}
	services      map[smarttypes.MatchableString]config.ServiceLabels
	defaultServer string
}

func New(ctx context.Context, services []config.Service, defaultServer string) (*Docker, error) {
	cli, err := client.NewClientWithOpts(config.Default.Docker.Opts)
	if err != nil {
		return nil, err
	}
	servicesMap := make(map[smarttypes.MatchableString]config.ServiceLabels)
	for _, svc := range services {
		servicesMap[svc.Name] = svc.Params
	}
	return &Docker{
		cli:           cli,
		ctx:           ctx,
		wait:          make(chan struct{}),
		services:      servicesMap,
		defaultServer: defaultServer,
	}, nil
}

func (d *Docker) findServiceLabels(id string, names []string) *config.ServiceLabels {
	if labels, ok := d.services[smarttypes.MatchableString(id)]; ok {
		return &labels
	}
	for matcher, labels := range d.services {
		for _, name := range names {
			if matcher.Match(name) {
				return &labels
			}
		}
	}
	return nil
}

func (d *Docker) Containers() map[string]dto.Container {
	items, err := d.cli.ContainerList(d.ctx, types.ContainerListOptions{
		Filters: filters.NewArgs(filters.Arg("status", "running")),
	})
	if err != nil {
		logger.Sugar.Errorf("cannot get containers list: %s", err)
		return nil
	}
	containers := map[string]dto.Container{}
	for _, item := range items {
		container, ok := dockerContainerToInternal(item, d.findServiceLabels(item.ID, item.Names), d.defaultServer)
		if !ok {
			continue
		}
		containers[item.ID] = container
	}
	return containers
}

func (d *Docker) GetContainer(id string, all bool) (dto.Container, bool) {
	container, err := d.cli.ContainerList(d.ctx, types.ContainerListOptions{
		Filters: filters.NewArgs(filters.Arg("id", id)),
		All:     all,
	})
	if err != nil || len(container) != 1 {
		return dto.Container{}, false
	}
	converted, ok := dockerContainerToInternal(
		container[0], d.findServiceLabels(container[0].ID, container[0].Names), d.defaultServer)
	if !ok {
		return dto.Container{}, false
	}
	return converted, true
}

func (d *Docker) Listen() (chan dto.Event, error) {
	cli, err := client.NewClientWithOpts(config.Default.Docker.Opts)
	if err != nil {
		return nil, err
	}

	output := make(chan dto.Event)
	go func() {
		for {
			eventSource, errSource := cli.Events(d.ctx, types.EventsOptions{
				Filters: filters.NewArgs(filters.Arg("type", "container")),
			})
			select {
			case event := <-eventSource:
				eventType := dto.TypeFromAction(event.Action)
				if (eventType != dto.EventStart && eventType != dto.EventStop) || !actorEnabled(event.Actor).Bool() {
					continue
				}
				container, err := d.cli.ContainerList(d.ctx, types.ContainerListOptions{
					Filters: filters.NewArgs(filters.Arg("id", event.Actor.ID)),
					All:     true,
				})
				if err != nil || len(container) != 1 {
					logger.Sugar.Errorw("cannot get container info",
						"id", event.Actor.ID, "err", err)
					continue
				}
				converted, ok := dockerContainerToInternal(container[0],
					d.findServiceLabels(container[0].ID, container[0].Names), d.defaultServer)
				if !ok {
					continue
				}
				newEvent := dto.Event{
					Type:      eventType,
					Container: converted,
				}
				msg := "exposing container"
				if eventType == dto.EventStop {
					msg = "stopping container"
				}
				logger.Sugar.Debugw(msg,
					"type", event.Action,
					"container.id", event.Actor.ID,
					"container.ip", converted.IP.String(),
					"container.port", converted.Port,
					"container.server", converted.Server,
					"container.remote_host", converted.RemoteHost)
				output <- newEvent
			case err := <-errSource:
				if errors.Is(err, context.Canceled) {
					d.wait <- struct{}{}
					logger.Sugar.Debug("stopping docker event listener...")
					return
				}
				logger.Sugar.Errorf("docker error, reconnect in 1s: %s", err)
				time.Sleep(time.Second)
				continue
			}
		}
	}()

	return output, nil
}

func (d *Docker) Wait() {
	<-d.wait
}