github.com/raphaelreyna/latte@v0.11.2-0.20220317193248-98e2fcef4eef/internal/job/request.go (about)

     1  package job
     2  
     3  import (
     4  	"github.com/raphaelreyna/go-recon"
     5  	"text/template"
     6  	"encoding/hex"
     7  	"crypto/md5"
     8  	"encoding/base64"
     9  	"path/filepath"
    10  	"io/ioutil"
    11  	"os"
    12  )
    13  
    14  type Request struct {
    15  	Template string `json:"template"`
    16  
    17  	Details map[string]interface{} `json:"details"`
    18  
    19  	Resources map[string]string `json:"resources"`
    20  
    21  	Delimiters Delimiters `json:"delimiters"`
    22  	OnMissingKey MissingKeyOpt `json:"onMissingKey"`
    23  	Compiler Compiler `json:"compiler"`
    24  	Count uint `json:"count"`
    25  }
    26  
    27  func (r *Request) NewJob(root string, sc recon.SourceChain, cache *TemplateCache) (*Job, error) {
    28  	var err error
    29  	j := &Job{Opts: DefaultOptions}
    30  	j.Root = root
    31  	j.SourceChain = sc
    32  
    33  	opts := DefaultOptions
    34  
    35  	if r.Delimiters != EmptyDelimiters && r.Delimiters != BadDefaultDelimiters {
    36  		opts.Delims = r.Delimiters
    37  	} else {
    38  		opts.Delims = DefaultDelimiters
    39  		r.Delimiters = DefaultDelimiters
    40  	}
    41  
    42  	if x := r.OnMissingKey; x != "" {
    43  		opts.OnMissingKey = x
    44  	}
    45  	if x := r.Compiler; x != "" {
    46  		opts.CC = x
    47  	}
    48  	if x := r.Count; x > 0 {
    49  		opts.N = x
    50  	}
    51  
    52  	j.Opts = opts
    53  	j.Details = r.Details
    54  
    55  	if j.Template, err = r.parseTemplate(cache); err != nil {
    56  		return nil, err
    57  	}
    58  
    59  
    60  	// Write resources files into working directory
    61  	for name, data := range r.Resources {
    62  		fname := filepath.Join(root, name)
    63  		bytes, err := base64.StdEncoding.DecodeString(data)
    64  		if err != nil {
    65  			return nil, err
    66  		}
    67  		err = ioutil.WriteFile(fname, bytes, os.ModePerm)
    68  		if err != nil {
    69  			return nil, err
    70  		}
    71  	}
    72  
    73  	return j, nil
    74  }
    75  
    76  func (r *Request) parseTemplate(cache *TemplateCache) (*template.Template, error) {
    77  	if r.Template == "" {
    78  		return nil, nil
    79  	}
    80  
    81  	// Check if we've already parsed this template; if not, parse it and cache the results
    82  	tHash := md5.Sum([]byte(r.Template))
    83  	// We append template delimiters to account for the same file being uploaded with different delimiters.
    84  	// This would really only happen on accident but not taking it into account leads to unexpected caching behavior.
    85  	cid := hex.EncodeToString(tHash[:]) + r.Delimiters.Left + r.Delimiters.Right
    86  	cache.Lock()
    87  	defer cache.Unlock()
    88  	ti, exists := cache.Get(cid)
    89  	var t *template.Template
    90  	if !exists {
    91  		tBytes, err := base64.StdEncoding.DecodeString(r.Template)
    92  		if err != nil {
    93  			cache.Unlock()
    94  			return nil, err
    95  		}
    96  		t = template.New(cid).Delims(r.Delimiters.Left, r.Delimiters.Right)
    97  		t, err = t.Parse(string(tBytes))
    98  		if err != nil {
    99  			cache.Unlock()
   100  			return nil, err
   101  		}
   102  
   103  		cache.Add(cid, t)
   104  	} else {
   105  		t = ti.(*template.Template)
   106  	}
   107  
   108  	return t.Option("missingkey=" + r.OnMissingKey.Val()), nil
   109  }