internal/vulnerability/runner.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"
	"time"

	"github.com/aquasecurity/trivy/pkg/commands/artifact"
	"github.com/aquasecurity/trivy/pkg/commands/clean"
	fanal_types "github.com/aquasecurity/trivy/pkg/fanal/types"
	"github.com/aquasecurity/trivy/pkg/flag"
	trivy_types "github.com/aquasecurity/trivy/pkg/types"
	"github.com/google/go-containerregistry/pkg/name"

	"github.com/elastic/cloudbeat/internal/infra/clog"
)

type VulnerabilityRunner struct {
	log *clog.Logger
	artifact.Runner
}

func NewVulnerabilityRunner(ctx context.Context, log *clog.Logger) (VulnerabilityRunner, error) {
	log.Debug("NewVulnerabilityRunner: New")

	if err := clearTrivyCache(ctx, log); err != nil {
		log.Errorf("error during runner cache clearing %s", err.Error())
	}

	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{
			Scanners: []trivy_types.Scanner{trivy_types.VulnerabilityScanner},
			RekorURL: "https://rekor.sigstore.dev",
		},
		DBOptions: flag.DBOptions{
			NoProgress:         true,
			DBRepositories:     []name.Reference{name.MustParseReference("public.ecr.aws/aquasecurity/trivy-db:2")},
			JavaDBRepositories: []name.Reference{name.MustParseReference("public.ecr.aws/aquasecurity/trivy-java-db:1")},
		},
	}

	runner, err := artifact.NewRunner(ctx, opts)
	if err != nil {
		log.Error("NewVulnerabilityRunner: NewRunner error: ", err)
		return VulnerabilityRunner{}, err
	}

	return VulnerabilityRunner{
		log:    log,
		Runner: runner,
	}, nil
}

func (r VulnerabilityRunner) Close(ctx context.Context) error {
	if r.Runner == nil {
		return nil
	}
	return r.Runner.Close(ctx)
}

func clearTrivyCache(ctx context.Context, log *clog.Logger) error {
	log.Info("Starting VulnerabilityRunner.ClearCache")
	defer log.Info("Ending VulnerabilityRunner.ClearCache")

	return clean.Run(ctx, flag.Options{
		CleanOptions: flag.CleanOptions{
			CleanScanCache: true,
		},
		GlobalOptions: flag.GlobalOptions{
			Timeout: 5 * time.Second,
		},
	})
}