github.com/khulnasoft-lab/defsec@v1.0.5-0.20230827010352-5e9f46893d95/pkg/scan/highlighting.go (about)

     1  package scan
     2  
     3  import (
     4  	"bytes"
     5  	"fmt"
     6  	"strings"
     7  	"sync"
     8  
     9  	"github.com/alecthomas/chroma"
    10  	"github.com/alecthomas/chroma/formatters"
    11  	"github.com/alecthomas/chroma/lexers"
    12  	"github.com/alecthomas/chroma/styles"
    13  )
    14  
    15  type cache struct {
    16  	sync.RWMutex
    17  	data map[string][]string
    18  }
    19  
    20  func (c *cache) Get(key string) ([]string, bool) {
    21  	c.RLock()
    22  	defer c.RUnlock()
    23  	data, ok := c.data[key]
    24  	return data, ok
    25  }
    26  
    27  func (c *cache) Set(key string, data []string) {
    28  	c.Lock()
    29  	defer c.Unlock()
    30  	c.data[key] = data
    31  }
    32  
    33  var globalCache = &cache{
    34  	data: make(map[string][]string),
    35  }
    36  
    37  func highlight(fsKey string, filename string, input []byte, theme string) []string {
    38  
    39  	key := fmt.Sprintf("%s|%s", fsKey, filename)
    40  	if lines, ok := globalCache.Get(key); ok {
    41  		return lines
    42  	}
    43  
    44  	lexer := lexers.Match(filename)
    45  	if lexer == nil {
    46  		lexer = lexers.Fallback
    47  	}
    48  	lexer = chroma.Coalesce(lexer)
    49  
    50  	style := styles.Get(theme)
    51  	if style == nil {
    52  		style = styles.Fallback
    53  	}
    54  	formatter := formatters.Get("terminal256")
    55  	if formatter == nil {
    56  		formatter = formatters.Fallback
    57  	}
    58  
    59  	// replace windows line endings
    60  	input = bytes.ReplaceAll(input, []byte{0x0d}, []byte{})
    61  	iterator, err := lexer.Tokenise(nil, string(input))
    62  	if err != nil {
    63  		return nil
    64  	}
    65  
    66  	buffer := bytes.NewBuffer([]byte{})
    67  	if err := formatter.Format(buffer, style, iterator); err != nil {
    68  		return nil
    69  	}
    70  
    71  	raw := shiftANSIOverLineEndings(buffer.Bytes())
    72  	lines := strings.Split(string(raw), "\n")
    73  	globalCache.Set(key, lines)
    74  	return lines
    75  }
    76  
    77  func shiftANSIOverLineEndings(input []byte) []byte {
    78  	var output []byte
    79  	prev := byte(0)
    80  	inCSI := false
    81  	csiShouldCarry := false
    82  	var csi []byte
    83  	var skipOutput bool
    84  	for _, r := range input {
    85  		skipOutput = false
    86  		if !inCSI {
    87  			switch {
    88  			case r == '\n':
    89  				if csiShouldCarry && len(csi) > 0 {
    90  					skipOutput = true
    91  					output = append(output, '\n')
    92  					output = append(output, csi...)
    93  					csi = nil
    94  					csiShouldCarry = false
    95  				}
    96  			case r == '[' && prev == 0x1b:
    97  				inCSI = true
    98  				csi = append(csi, 0x1b, '[')
    99  				output = output[:len(output)-1]
   100  				skipOutput = true
   101  			default:
   102  				csiShouldCarry = false
   103  				if len(csi) > 0 {
   104  					output = append(output, csi...)
   105  					csi = nil
   106  				}
   107  			}
   108  		} else {
   109  			csi = append(csi, r)
   110  			skipOutput = true
   111  			switch {
   112  			case r >= 0x40 && r <= 0x7E:
   113  				csiShouldCarry = true
   114  				inCSI = false
   115  			}
   116  		}
   117  		if !skipOutput {
   118  			output = append(output, r)
   119  		}
   120  		prev = r
   121  	}
   122  
   123  	return append(output, csi...)
   124  }