trivy/internal/hooktest/hook.go

package hooktest

import (
	"context"
	"errors"
	"testing"

	"github.com/aquasecurity/trivy/pkg/extension"
	"github.com/aquasecurity/trivy/pkg/flag"
	"github.com/aquasecurity/trivy/pkg/types"
)

type testHook struct{}

func (*testHook) Name() string {
	return "test"
}

func (*testHook) Version() int {
	return 1
}

// RunHook implementation
func (*testHook) PreRun(_ context.Context, opts flag.Options) error {
	if opts.GlobalOptions.ConfigFile == "bad-config" {
		return errors.New("bad pre-run")
	}
	return nil
}

func (*testHook) PostRun(_ context.Context, opts flag.Options) error {
	if opts.GlobalOptions.ConfigFile == "bad-config" {
		return errors.New("bad post-run")
	}
	return nil
}

// ScanHook implementation
func (*testHook) PreScan(_ context.Context, target *types.ScanTarget, _ types.ScanOptions) error {
	if target.Name == "bad-pre" {
		return errors.New("bad pre-scan")
	}
	target.Name += " (pre-scan)"
	return nil
}

func (*testHook) PostScan(_ context.Context, results types.Results) (types.Results, error) {
	for i, r := range results {
		if r.Target == "bad" {
			return nil, errors.New("bad")
		}
		for j := range r.Vulnerabilities {
			results[i].Vulnerabilities[j].References = []string{
				"https://example.com/post-scan",
			}
		}
	}
	return results, nil
}

// ReportHook implementation
func (*testHook) PreReport(_ context.Context, report *types.Report, _ flag.Options) error {
	if report.ArtifactName == "bad-report" {
		return errors.New("bad pre-report")
	}

	// Modify the report
	for i := range report.Results {
		for j := range report.Results[i].Vulnerabilities {
			report.Results[i].Vulnerabilities[j].Title = "Modified by pre-report hook"
		}
	}
	return nil
}

func (*testHook) PostReport(_ context.Context, report *types.Report, _ flag.Options) error {
	if report.ArtifactName == "bad-report" {
		return errors.New("bad post-report")
	}

	// Modify the report
	for i := range report.Results {
		for j := range report.Results[i].Vulnerabilities {
			report.Results[i].Vulnerabilities[j].Description = "Modified by post-report hook"
		}
	}
	return nil
}

func Init(t *testing.T) {
	h := &testHook{}
	extension.RegisterHook(h)
	t.Cleanup(func() {
		extension.DeregisterHook(h.Name())
	})
}