nat/lib/styler/styler.go

package styler

import (
	"bufio"
	"bytes"
	"fmt"
	"log"
	"os"
	"os/exec"
	"strings"
	"unicode/utf8"
)

type Highlighter string

const (
	HIGHLIGHTER_NONE   Highlighter = "none"
	HIGHLIGHTER_PYTHON Highlighter = "python"
	HIGHLIGHTER_GO     Highlighter = "go"
)

func (h Highlighter) Cmd() string {
	switch h {
	case HIGHLIGHTER_PYTHON:
		return ".highlighers/highlighter-wrapper.py"
	case HIGHLIGHTER_GO:
		return ".highlighers/highligher.gobin"
	case HIGHLIGHTER_NONE:
		fallthrough
	default:
		return ""
	}
}

type Styler struct {
	listOfLangsFirst map[string]string
	listOfLangsLast  map[string]string
	listOfStyles     map[string]string

	config StylerConfig
}

func (s *Styler) Legacy() [3]map[string]string {
	return [3]map[string]string{
		s.listOfLangsFirst,
		s.listOfLangsLast,
		s.listOfStyles,
	}
}

type StylerConfig struct {
	Highlighter Highlighter `json:"highlighter"` // The name of the highlighter.
}

func New(config ...StylerConfig) *Styler {
	s := &Styler{
		listOfLangsFirst: make(map[string]string),
		listOfLangsLast:  make(map[string]string),
		listOfStyles:     make(map[string]string),
	}
	// default settings
	s.config.Highlighter = "none"
	if len(config) > 0 {
		s.config = config[0]
	}
	return s
}
func isBad(s string) bool {
	for i := 0; i < len(s); i++ {
		b := s[i]
		if ('a' <= b && b <= 'z') ||
			('A' <= b && b <= 'Z') ||
			('0' <= b && b <= '9') ||
			b == ' ' {
			continue
		} else {
			return false
		}
	}
	return true
}

func sanitize(s string) string {
	if !utf8.ValidString(s) {
		return s
	}
	if len(s) < 30 {
		return s
	}
	if !isBad(s) {
		return s
	}
	return ""
}

func (s *Styler) Highlight(paste string, lang string, style string) (string, string, string, string, error) {
	var supported_lang, supported_styles bool
	lang = sanitize(lang)
	style = sanitize(style)
	lang, supported_lang = s.listOfLangsFirst[lang]
	style, supported_styles = s.listOfStyles[style]
	lang = sanitize(lang)
	style = sanitize(style)

	if lang == "" {
		lang = "autodetect"
	}

	if !supported_lang && lang != "autodetect" {
		lang = "text"
	}

	// Same with the styles,
	if !supported_styles {
		style = "autodetect"
	}
	switch s.config.Highlighter {
	case HIGHLIGHTER_PYTHON:
		if _, err := os.Stat(s.config.Highlighter.Cmd()); os.IsNotExist(err) {
			return "", "", "", "", err
		}
		cmd := exec.Command(s.config.Highlighter.Cmd(), lang, style)
		cmd.Stdin = strings.NewReader(paste)

		var stdout bytes.Buffer
		cmd.Stdout = &stdout

		var stderr bytes.Buffer
		cmd.Stderr = &stderr
		err := cmd.Run()
		if err != nil {
			return "", "", "", "", err
		}
		return stdout.String(), stderr.String(), lang, style, nil
	case "none":
		fallthrough
	default:
		return paste, "", lang, style, nil
	}
}

func (s *Styler) loadSupportedStyles() {
	switch s.config.Highlighter {
	case HIGHLIGHTER_PYTHON:
		arg := "getstyles"
		out, err := exec.Command(s.config.Highlighter.Cmd(), arg).Output()
		if err != nil {
			log.Fatal(err)
		}

		// Loop lexers and add them to respectively map,
		for _, line := range strings.Split(string(out), "\n") {
			if line == "" {
				continue
			}
			s.listOfStyles[line] = strings.Title(line)
		}
	case "none":
		fallthrough
	default:
		return
	}
}

func (st *Styler) loadSupportedLangs() error {
	switch st.config.Highlighter {
	case HIGHLIGHTER_PYTHON:
		var prioLexers map[string]string

		// Initialize maps,
		prioLexers = make(map[string]string)

		// Get prioritized lexers and put them in a separate map,
		file, err := os.Open("static/prio-lexers")
		if err != nil {
			return err
		}

		scanner := bufio.NewScanner(file)
		for scanner.Scan() {
			prioLexers[scanner.Text()] = "1"
		}

		if err := scanner.Err(); err != nil {
			log.Fatal(err)
		}
		file.Close()

		arg := "getlexers"
		out, err := exec.Command(st.config.Highlighter.Cmd(), arg).Output()
		if err != nil {
			return fmt.Errorf("%w: %s", err, string(out))
		}

		// Loop lexers and add them to respectively map,
		for _, line := range strings.Split(string(out), "\n") {
			if line == "" {
				continue
			}

			s := strings.Split(line, ";")
			if len(s) != 2 {
				os.Exit(1)
			}
			s[0] = strings.Title(s[0])
			if prioLexers[s[0]] == "1" {
				st.listOfLangsFirst[s[0]] = s[1]
			} else {
				st.listOfLangsLast[s[0]] = s[1]
			}
		}
		return nil
	case "none":
		fallthrough
	default:
		return nil
	}
}