internal/flavors/vulnerability.go

// Licensed to Elasticsearch B.V. under one or more contributor
// license agreements. See the NOTICE file distributed with
// this work for additional information regarding copyright
// ownership. Elasticsearch B.V. licenses this file to you under
// the Apache License, Version 2.0 (the "License"); you may
// not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

package flavors

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

	dlog "github.com/aquasecurity/go-dep-parser/pkg/log"
	tlog "github.com/aquasecurity/trivy/pkg/log"
	"github.com/elastic/beats/v7/libbeat/beat"
	agentconfig "github.com/elastic/elastic-agent-libs/config"
	xlog "github.com/masahiro331/go-xfs-filesystem/log"
	"go.uber.org/zap"
	"go.uber.org/zap/exp/zapslog"

	"github.com/elastic/cloudbeat/internal/config"
	"github.com/elastic/cloudbeat/internal/dataprovider"
	"github.com/elastic/cloudbeat/internal/dataprovider/providers/cloud"
	"github.com/elastic/cloudbeat/internal/dataprovider/providers/common"
	"github.com/elastic/cloudbeat/internal/infra/clog"
	"github.com/elastic/cloudbeat/internal/resources/providers/awslib"
	vuln "github.com/elastic/cloudbeat/internal/vulnerability"
	"github.com/elastic/cloudbeat/version"
)

// vulnerability configuration.
type vulnerability struct {
	flavorBase
	bdp dataprovider.CommonDataProvider
	cdp *common.DataProvider
}

// NewVulnerability creates an instance of vulnerability.
func NewVulnerability(b *beat.Beat, cfg *agentconfig.C) (beat.Beater, error) {
	log := clog.NewLogger("vulnerability")

	// Override trivy's logger
	scanLog := zap.New(log.Core()).Sugar()
	dlog.SetLogger(scanLog)
	xlog.SetLogger(scanLog)
	tlog.SetDefault(slog.New(zapslog.NewHandler(log.Core())))

	ctx, cancel := context.WithCancel(context.Background())

	c, err := config.New(cfg)
	if err != nil {
		cancel()
		return nil, fmt.Errorf("error reading config file: %w", err)
	}
	log.Info("Config initiated with cycle period of ", c.Period)

	bdp, err := getAWSDataProvider(ctx, *c, log)
	if err != nil {
		cancel()
		return nil, fmt.Errorf("failed to init common data provider: %w", err)
	}

	cdp, err := common.New(version.CloudbeatVersionInfo{Version: version.CloudbeatVersion()}, c)
	if err != nil {
		cancel()
		return nil, err
	}

	client, err := NewClient(b.Publisher, c.Processors)
	if err != nil {
		cancel()
		return nil, fmt.Errorf("failed to init client: %w", err)
	}
	log.Infof("vulnerability configured %d processors", len(c.Processors))

	publisher := NewPublisher(log, flushInterval, eventsThreshold, client)

	base := flavorBase{
		ctx:       ctx,
		cancel:    cancel,
		publisher: publisher,
		config:    c,
		log:       log,
		client:    client,
	}

	bt := &vulnerability{
		flavorBase: base,
		bdp:        bdp,
		cdp:        cdp,
	}

	return bt, nil
}

func getAWSDataProvider(ctx context.Context, cfg config.Config, logger *clog.Logger) (dataprovider.CommonDataProvider, error) {
	awsConfig, err := awslib.InitializeAWSConfig(cfg.CloudConfig.Aws.Cred)
	if err != nil {
		return nil, fmt.Errorf("failed to initialize AWS credentials: %w", err)
	}

	identity, err := awslib.IdentityProvider{Logger: logger}.GetIdentity(ctx, *awsConfig)
	if err != nil {
		return nil, fmt.Errorf("failed to get AWS identity: %w", err)
	}

	return cloud.NewDataProvider(cloud.WithAccount(*identity)), nil
}

// Run starts vulnerability.
func (bt *vulnerability) Run(*beat.Beat) error {
	bt.log.Info("vulnerability is running! Hit CTRL-C to stop it.")
	repeater := NewRepeater(bt.log, bt.config.Period)
	return repeater.Run(bt.ctx, bt.runIteration)
}

func (bt *vulnerability) runIteration() error {
	worker, err := vuln.NewVulnerabilityWorker(bt.ctx, bt.log, bt.config, bt.bdp, bt.cdp)
	if err != nil {
		bt.log.Error("vulnerability.runIteration worker creation failed")
		bt.cancel()
		return err
	}

	go bt.publisher.HandleEvents(bt.ctx, worker.GetChan())

	worker.Run(bt.ctx)

	bt.log.Info("vulnerability.runIteration cycle finished")
	return nil
}

// Stop stops vulnerability.
func (bt *vulnerability) Stop() {
	bt.cancel()
}