god-eye/internal/modules/takeover/takeover.go

// Package takeover runs v1 takeover detection on every host with a CNAME.
// Reads from the store; listens for late DNSResolved events for concurrent
// modules.
package takeover

import (
	"context"
	"sync"
	"time"

	"god-eye/internal/eventbus"
	"god-eye/internal/module"
	"god-eye/internal/scanner"
	"god-eye/internal/store"
)

const ModuleName = "takeover.cname"

type takeoverModule struct{}

func Register() { module.Register(&takeoverModule{}) }

func (*takeoverModule) Name() string                   { return ModuleName }
func (*takeoverModule) Phase() module.Phase            { return module.PhaseAnalysis }
func (*takeoverModule) Consumes() []eventbus.EventType { return []eventbus.EventType{eventbus.EventDNSResolved} }
func (*takeoverModule) Produces() []eventbus.EventType {
	return []eventbus.EventType{eventbus.EventTakeoverCandidate}
}
func (*takeoverModule) DefaultEnabled() bool { return true }

func (*takeoverModule) Run(mctx module.Context) error {
	if mctx.Config.Bool("no_takeover", false) {
		return nil
	}
	conc := mctx.Config.Int("concurrency", 100)
	if conc <= 0 {
		conc = 100
	}
	timeout := mctx.Config.Int("timeout", 5)

	processed := make(map[string]struct{})
	var processedMu sync.Mutex
	shouldProcess := func(host string) bool {
		processedMu.Lock()
		defer processedMu.Unlock()
		if _, dup := processed[host]; dup {
			return false
		}
		processed[host] = struct{}{}
		return true
	}

	work := make(chan string, conc*2)
	var wg sync.WaitGroup
	for i := 0; i < conc; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for host := range work {
				if mctx.Ctx.Err() != nil {
					return
				}
				service := scanner.CheckTakeover(host, timeout)
				if service == "" {
					continue
				}
				_ = mctx.Store.Upsert(mctx.Ctx, host, func(h *store.Host) {
					h.Takeover = &store.Takeover{
						Service:   service,
						CNAME:     h.CNAME,
						Confirmed: false,
						FoundAt:   time.Now(),
					}
				})
				mctx.Bus.Publish(mctx.Ctx, eventbus.TakeoverCandidate{
					EventMeta: eventbus.EventMeta{At: time.Now(), Source: ModuleName, Target: host},
					Subdomain: host,
					Service:   service,
				})
			}
		}()
	}

	// Drain: every host with a CNAME is a takeover candidate.
	for _, h := range mctx.Store.All(mctx.Ctx) {
		if h == nil || h.CNAME == "" {
			continue
		}
		if !shouldProcess(h.Subdomain) {
			continue
		}
		select {
		case work <- h.Subdomain:
		case <-mctx.Ctx.Done():
			close(work)
			wg.Wait()
			return nil
		}
	}

	sub := mctx.Bus.Subscribe(eventbus.EventDNSResolved, func(_ context.Context, e eventbus.Event) {
		ev, ok := e.(eventbus.DNSResolved)
		if !ok || ev.CNAME == "" {
			return
		}
		if !shouldProcess(ev.Subdomain) {
			return
		}
		select {
		case work <- ev.Subdomain:
		case <-mctx.Ctx.Done():
		}
	})
	defer sub.Unsubscribe()

	select {
	case <-time.After(500 * time.Millisecond):
	case <-mctx.Ctx.Done():
	}

	close(work)
	wg.Wait()
	return nil
}