internal/vulnerability/scanner.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 vulnerability

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"time"

	db_types "github.com/aquasecurity/trivy-db/pkg/types"
	fanal_types "github.com/aquasecurity/trivy/pkg/fanal/types"
	"github.com/aquasecurity/trivy/pkg/flag"
	trivy_types "github.com/aquasecurity/trivy/pkg/types"

	cb_config "github.com/elastic/cloudbeat/internal/config"
	"github.com/elastic/cloudbeat/internal/infra/clog"
	"github.com/elastic/cloudbeat/internal/resources/providers/awslib/ec2"
)

type VulnerabilityScanner struct {
	log     *clog.Logger
	ch      chan []Result
	runner  runnerInterface
	cfg     *cb_config.Config
	seq     time.Time
	manager *SnapshotManager
}

const shutdownGracePeriod = 30 * time.Second

type Result struct {
	reportResult  trivy_types.Result
	vulnerability trivy_types.DetectedVulnerability
	snapshot      ec2.EBSSnapshot
	seq           time.Time
}

type runnerInterface interface {
	ScanVM(ctx context.Context, opts flag.Options) (trivy_types.Report, error)
	Filter(ctx context.Context, opts flag.Options, report trivy_types.Report) (trivy_types.Report, error)
	Report(ctx context.Context, opts flag.Options, report trivy_types.Report) error
	Close(ctx context.Context) error
}

// TODO: Replace sequence with more generic approach
func NewVulnerabilityScanner(log *clog.Logger, runner runnerInterface, manager *SnapshotManager, c *cb_config.Config, seq time.Time) (VulnerabilityScanner, error) {
	log.Debug("VulnerabilityScanner: New")
	ch := make(chan []Result)

	return VulnerabilityScanner{
		log:    log,
		ch:     ch,
		runner: runner,
		cfg:    c,
		// TODO: Replace sequence with more generic approach
		seq:     seq,
		manager: manager,
	}, nil
}

func (f VulnerabilityScanner) ScanSnapshot(ctx context.Context, snapCh chan ec2.EBSSnapshot) {
	f.log.Info("Starting NewVulnerabilityScanner.ScanSnapshot")

	defer close(f.ch)
	for {
		select {
		case <-ctx.Done():
			f.log.Info("NewVulnerabilityScanner.ScanSnapshot context canceled")
			return
		case snap, ok := <-snapCh:
			if !ok {
				f.log.Info("NewVulnerabilityScanner.ScanSnapshot channel is closed")
				return
			}

			f.scan(ctx, snap)
			f.manager.DeleteSnapshot(ctx, snap)
		}
	}
}

func (f VulnerabilityScanner) scan(ctx context.Context, snap ec2.EBSSnapshot) {
	f.log.Infof("Starting VulnerabilityScanner.scan, %s", snap.SnapshotId)
	defer func() {
		if r := recover(); r != nil {
			f.log.Errorf("vulnerability scanner recovered from panic: %v", r)
		}
	}()

	o, err := os.CreateTemp("", "")
	if err != nil {
		f.log.Error("VulnerabilityScanner.scan.TempFile error: ", err)
		return
	}
	defer func(name string) {
		err = os.Remove(name)
		if err != nil {
			f.log.Warnf("Failed to remove temporary file %s: %v", name, err)
		}
	}(o.Name())

	opts := flag.Options{
		GlobalOptions: flag.GlobalOptions{
			// TODO: Make configurable
			Timeout: 1 * time.Hour,
			Quiet:   false,
			Debug:   true,
		},
		PackageOptions: flag.PackageOptions{
			PkgTypes:         []string{trivy_types.PkgTypeOS, trivy_types.PkgTypeLibrary},
			PkgRelationships: fanal_types.Relationships,
		},
		ScanOptions: flag.ScanOptions{
			Target:   fmt.Sprint("ebs:", snap.SnapshotId),
			Scanners: []trivy_types.Scanner{trivy_types.VulnerabilityScanner},
			RekorURL: "https://rekor.sigstore.dev",
		},
		AWSOptions: flag.AWSOptions{
			Region: snap.Region,
		},
		DBOptions: flag.DBOptions{
			SkipDBUpdate:     true,
			SkipJavaDBUpdate: true,
		},
		ReportOptions: flag.ReportOptions{
			Output:     o.Name(),
			Format:     "json",
			Severities: []db_types.Severity{0, 1, 2, 3, 4},
		},
	}
	now := time.Now()
	report, err := f.runner.ScanVM(ctx, opts)
	f.log.Infof(
		"VulnerabilityScanner.scan.ScanVM took %s to scan %s, volume size: %d, isEncrypted: %t, instanceId: %s",
		time.Since(now),
		snap.SnapshotId,
		snap.VolumeSize,
		snap.IsEncrypted,
		*snap.Instance.InstanceId,
	)

	if err != nil {
		f.log.Errorf("VulnerabilityScanner.scan.ScanVM, snapshotId: %s, instanceId: %s, error: %v", snap.SnapshotId, *snap.Instance.InstanceId, err)
		return
	}

	f.log.Info("VulnerabilityScanner.scan.Filter")
	report, err = f.runner.Filter(ctx, opts, report)
	if err != nil {
		f.log.Error("VulnerabilityScanner.scan.Filter error: ", err)
		return
	}

	f.log.Info("VulnerabilityScanner.scan.Report")
	err = f.runner.Report(ctx, opts, report)
	if err != nil {
		f.log.Error("VulnerabilityScanner.scan.Report error: ", err)
		return
	}

	f.log.Info("VulnerabilityScanner.scan.jsonFile")
	jsonFile, err := os.Open(o.Name())
	if err != nil {
		f.log.Error("VulnerabilityScanner.scan.jsonFile error: ", err)
		return
	}

	defer jsonFile.Close()
	f.log.Info("VulnerabilityScanner.scan.ReadAll")
	byteValue, _ := io.ReadAll(jsonFile)
	var unmarshalledReport trivy_types.Report
	err = json.Unmarshal(byteValue, &unmarshalledReport)
	if err != nil {
		f.log.Error("VulnerabilityScanner.scan.Unmarshal error: ", err)
		return
	}

	results := []Result{}
	for _, result := range unmarshalledReport.Results {
		for _, vul := range result.Vulnerabilities {
			// TODO: Replace sequence with more generic approach
			result := Result{
				reportResult:  result,
				vulnerability: vul,
				snapshot:      snap,
				seq:           f.seq,
			}
			results = append(results, result)
		}
	}

	select {
	case <-ctx.Done():
		f.log.Info("VulnerabilityScanner.scan context canceled while sending vulnerabilities")
		return
	case f.ch <- results:
	}

	f.log.Info("VulnerabilityScanner.scan.DONE")
}

func (f VulnerabilityScanner) GetChan() chan []Result {
	return f.ch
}