search-albeast-vuln-albs.go

package main

import (
	"context"
	"fmt"
	"log/slog"
	"os"

	"github.com/aws/aws-sdk-go-v2/config"
	elbv2 "github.com/aws/aws-sdk-go-v2/service/elasticloadbalancingv2"
)

func doRun(ctx context.Context) error {
	cfg, err := config.LoadDefaultConfig(ctx)
	if err != nil {
		return fmt.Errorf("config.LoadDefaultConfig: %w", err)
	}
	client := elbv2.NewFromConfig(cfg)
	listALBsInput := &elbv2.DescribeLoadBalancersInput{}
	affectedALBs := make([]string, 0)
	slog.InfoContext(ctx, "start search")
	for lbOut, err := range describeLoadBalancers(ctx, client, listALBsInput) {
		if err != nil {
			return fmt.Errorf("DescribeLoadBalancers: %w", err)
		}
		for _, lb := range lbOut.LoadBalancers {
			input := &elbv2.DescribeListenersInput{LoadBalancerArn: lb.LoadBalancerArn}
			for listenerOut, err := range describeListeners(ctx, client, input) {
				if err != nil {
					return fmt.Errorf("DescribeListeners: %w", err)
				}
				for _, listener := range listenerOut.Listeners {
					for _, action := range listener.DefaultActions {
						if action.AuthenticateCognitoConfig != nil || action.AuthenticateOidcConfig != nil {
							affectedALBs = append(affectedALBs, *lb.LoadBalancerName)
						}
					}
				}
			}
		}
	}
	slog.InfoContext(ctx, "searched load balancers", slog.Int("affected_albs_num", len(affectedALBs)))
	for _, lb := range affectedALBs {
		slog.InfoContext(ctx, "affected load balancer", slog.String("load_balancer_name", lb))
	}
	return nil
}

func describeListeners(ctx context.Context, client *elbv2.Client, initialInput *elbv2.DescribeListenersInput) func(func(*elbv2.DescribeListenersOutput, error) bool) {
	return newSDKIterator(ctx, func(ctx context.Context, in *elbv2.DescribeListenersInput, opts ...func(*elbv2.Options)) (*elbv2.DescribeListenersOutput, error) {
		var mk string
		if in.Marker != nil {
			mk = *in.Marker
		}
		slog.InfoContext(ctx, "describe listeners", slog.String("load_balancer_arn", *in.LoadBalancerArn), slog.String("marker", mk))
		return client.DescribeListeners(ctx, in, opts...)
	}, describeListenersNext, initialInput)
}

func describeListenersNext(input *elbv2.DescribeListenersInput, output *elbv2.DescribeListenersOutput) bool {
	if output == nil || output.NextMarker == nil || *output.NextMarker == "" {
		return false
	}
	input.Marker = output.NextMarker
	return true
}

func describeLoadBalancers(ctx context.Context, client *elbv2.Client, initialInput *elbv2.DescribeLoadBalancersInput) func(func(*elbv2.DescribeLoadBalancersOutput, error) bool) {
	return newSDKIterator(ctx, func(ctx context.Context, in *elbv2.DescribeLoadBalancersInput, opts ...func(*elbv2.Options)) (*elbv2.DescribeLoadBalancersOutput, error) {
		var mk string
		if in.Marker != nil {
			mk = *in.Marker
		}
		slog.InfoContext(ctx, "describe load balancers", slog.String("marker", mk))
		return client.DescribeLoadBalancers(ctx, in, opts...)
	}, describeLoadBalancersNext, initialInput)
}

func describeLoadBalancersNext(input *elbv2.DescribeLoadBalancersInput, output *elbv2.DescribeLoadBalancersOutput) bool {
	if output == nil || output.NextMarker == nil || *output.NextMarker == "" {
		return false
	}
	input.Marker = output.NextMarker
	return true
}

// refs. https://zenn.dev/fujiwara/scraps/5d68298a83e50f
func newSDKIterator[I any, O any, OP any](ctx context.Context, do func(context.Context, *I, ...OP) (*O, error), next func(*I, *O) bool, input *I, opts ...OP) func(func(*O, error) bool) {
	return func(yield func(*O, error) bool) {
		for {
			out, err := do(ctx, input, opts...)
			if err != nil {
				yield(nil, err)
				return
			}
			if !yield(out, err) {
				return
			}
			if !next(input, out) {
				return
			}
		}
	}
}

func run() int {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	if err := doRun(ctx); err != nil {
		slog.ErrorContext(ctx, "failed", slog.String("error", err.Error()))
		return 1
	}
	return 0
}

func main() {
	os.Exit(run())
}