sm-seo-plugin/classes/csv.go

package classes

import (
	"bytes"
	"encoding/csv"
	"errors"
	"io"
	"strings"
)

// BOM is the UTF-8 byte order mark export files start with, so Excel
// opens Polish characters correctly; import strips it.
const BOM = "\ufeff"

// ErrUnreadable refuses a file that is not CSV with a header row
// (transfer.unreadable).
var ErrUnreadable = errors.New("golem15.seo: the file could not be read as CSV")

// CSVDelimiter returns the delimiter rune of a setting: ";" or ",", comma
// by default.
func CSVDelimiter(d string) rune {
	if strings.TrimSpace(d) == ";" {
		return ';'
	}
	return ','
}

// formulaStart reports whether a cell would be read as a formula by a
// spreadsheet: it starts with = + - @, a tab or a carriage return.
func formulaStart(s string) bool {
	if s == "" {
		return false
	}
	switch s[0] {
	case '=', '+', '-', '@', '\t', '\r':
		return true
	}
	return false
}

// escapeCell prefixes a formula cell with an apostrophe (CSV formula
// injection, T-32-24).
func escapeCell(s string) string {
	if formulaStart(s) {
		return "'" + s
	}
	return s
}

// unescapeCell drops the apostrophe escapeCell added, so an exported file
// imports unchanged.
func unescapeCell(s string) string {
	if strings.HasPrefix(s, "'") && formulaStart(s[1:]) {
		return s[1:]
	}
	return s
}

// WriteCSV writes header and rows as UTF-8 with a BOM and the delimiter,
// every formula cell escaped.
func WriteCSV(header []string, rows [][]string, delimiter string) (string, error) {
	var b bytes.Buffer
	b.WriteString(BOM)
	w := csv.NewWriter(&b)
	w.Comma = CSVDelimiter(delimiter)
	if err := w.Write(header); err != nil {
		return "", err
	}
	for _, row := range rows {
		cells := make([]string, len(row))
		for i, c := range row {
			cells[i] = escapeCell(c)
		}
		if err := w.Write(cells); err != nil {
			return "", err
		}
	}
	w.Flush()
	return b.String(), w.Error()
}

// CSVRow is one data row of a file: its line number and its cells by
// column name.
type CSVRow struct {
	Line   int
	Fields map[string]string
}

// ReadCSV reads a file with a header row: the BOM is stripped, the
// delimiter is the one given (detected from the header when empty), the
// column names are trimmed and lowercased and the apostrophe of escaped
// formula cells is dropped. Rows whose cells are all empty are skipped.
// A file that is not CSV, has no header or was read with the wrong
// delimiter answers ErrUnreadable.
func ReadCSV(data, delimiter string) ([]string, []CSVRow, error) {
	data = strings.TrimPrefix(data, BOM)
	if strings.TrimSpace(data) == "" {
		return nil, nil, ErrUnreadable
	}
	first, _, _ := strings.Cut(data, "\n")
	comma := CSVDelimiter(delimiter)
	if strings.TrimSpace(delimiter) == "" && strings.Count(first, ";") > strings.Count(first, ",") {
		comma = ';'
	}
	r := csv.NewReader(strings.NewReader(data))
	r.Comma = comma
	r.FieldsPerRecord = -1
	header, err := r.Read()
	if err != nil {
		return nil, nil, ErrUnreadable
	}
	for i, h := range header {
		header[i] = strings.ToLower(strings.TrimSpace(h))
	}
	if len(header) == 1 && strings.ContainsAny(header[0], ",;") {
		return nil, nil, ErrUnreadable
	}
	var rows []CSVRow
	for {
		rec, err := r.Read()
		if errors.Is(err, io.EOF) {
			break
		}
		if err != nil {
			return nil, nil, ErrUnreadable
		}
		line, _ := r.FieldPos(0)
		row := CSVRow{Line: line, Fields: map[string]string{}}
		empty := true
		for i, cell := range rec {
			if i >= len(header) || header[i] == "" {
				continue
			}
			cell = unescapeCell(strings.TrimSpace(cell))
			if cell != "" {
				empty = false
			}
			row.Fields[header[i]] = cell
		}
		if !empty {
			rows = append(rows, row)
		}
	}
	return header, rows, nil
}