github.com/devseccon/trivy@v0.47.1-0.20231123133102-bd902a0bd996/pkg/report/template.go (about)

     1  package report
     2  
     3  import (
     4  	"bytes"
     5  	"encoding/xml"
     6  	"html"
     7  	"io"
     8  	"os"
     9  	"strings"
    10  	"text/template"
    11  
    12  	"github.com/Masterminds/sprig/v3"
    13  	"golang.org/x/xerrors"
    14  
    15  	dbTypes "github.com/aquasecurity/trivy-db/pkg/types"
    16  	"github.com/devseccon/trivy/pkg/log"
    17  	"github.com/devseccon/trivy/pkg/types"
    18  )
    19  
    20  // CustomTemplateFuncMap is used to overwrite existing functions for testing.
    21  var CustomTemplateFuncMap = make(map[string]interface{})
    22  
    23  // TemplateWriter write result in custom format defined by user's template
    24  type TemplateWriter struct {
    25  	Output   io.Writer
    26  	Template *template.Template
    27  }
    28  
    29  // NewTemplateWriter is the factory method to return TemplateWriter object
    30  func NewTemplateWriter(output io.Writer, outputTemplate string) (*TemplateWriter, error) {
    31  	if strings.HasPrefix(outputTemplate, "@") {
    32  		buf, err := os.ReadFile(strings.TrimPrefix(outputTemplate, "@"))
    33  		if err != nil {
    34  			return nil, xerrors.Errorf("error retrieving template from path: %w", err)
    35  		}
    36  		outputTemplate = string(buf)
    37  	}
    38  	var templateFuncMap template.FuncMap = sprig.GenericFuncMap()
    39  	templateFuncMap["escapeXML"] = func(input string) string {
    40  		escaped := &bytes.Buffer{}
    41  		if err := xml.EscapeText(escaped, []byte(input)); err != nil {
    42  			log.Logger.Error("error while escapeString to XML: %s", err)
    43  			return input
    44  		}
    45  		return escaped.String()
    46  	}
    47  	templateFuncMap["endWithPeriod"] = func(input string) string {
    48  		if !strings.HasSuffix(input, ".") {
    49  			input += "."
    50  		}
    51  		return input
    52  	}
    53  	templateFuncMap["escapeString"] = html.EscapeString
    54  	templateFuncMap["sourceID"] = func(input string) dbTypes.SourceID {
    55  		return dbTypes.SourceID(input)
    56  	}
    57  
    58  	// Overwrite functions
    59  	for k, v := range CustomTemplateFuncMap {
    60  		templateFuncMap[k] = v
    61  	}
    62  
    63  	tmpl, err := template.New("output template").Funcs(templateFuncMap).Parse(outputTemplate)
    64  	if err != nil {
    65  		return nil, xerrors.Errorf("error parsing template: %w", err)
    66  	}
    67  	return &TemplateWriter{
    68  		Output:   output,
    69  		Template: tmpl,
    70  	}, nil
    71  }
    72  
    73  // Write writes result
    74  func (tw TemplateWriter) Write(report types.Report) error {
    75  	err := tw.Template.Execute(tw.Output, report.Results)
    76  	if err != nil {
    77  		return xerrors.Errorf("failed to write with template: %w", err)
    78  	}
    79  	return nil
    80  }