chardet.go

package xtractr

import (
	"archive/zip"
	"encoding/binary"
	"hash/crc32"
	"strings"
	"unicode/utf8"

	"github.com/saintfish/chardet"
	"golang.org/x/text/encoding"
	"golang.org/x/text/encoding/charmap"
	"golang.org/x/text/encoding/japanese"
	"golang.org/x/text/encoding/korean"
	"golang.org/x/text/encoding/simplifiedchinese"
	"golang.org/x/text/encoding/traditionalchinese"
)

// zipNameDecoders stores per-entry filename encodings with an archive-level fallback.
type zipNameDecoders struct {
	defaultEncoding encoding.Encoding
	nameEncodings   map[string]encoding.Encoding
	partNames       map[string]string
}

// Encoding preference scores for tiebreaking. Higher = more commonly seen in non-UTF8 zips.
// GBK/GB-18030 is most common because Chinese Windows is the largest source
// of non-UTF8 zip files. Shift-JIS is next (Japanese Windows). These are
// preferred over less common encodings when all other signals are equal.
const (
	prefGBK     = 5 // Chinese Windows is the largest source of non-UTF8 zip files
	prefShiftJS = 4 // Japanese Windows
	prefBig5    = 3 // Traditional Chinese
	prefEUCKR   = 2 // Korean
	prefEUCJP   = 1 // Japanese (less common than Shift-JIS in zips)
)

// Script scoring constants for disambiguating CJK encodings.
const (
	scoreKanaBonus      = 200 // kana is an unambiguous marker for Japanese
	scoreHangulBonus    = 150 // pure Hangul is a clear marker for Korean
	scoreMixedHangulCJK = 50  // mixed Hangul + CJK is suspicious (likely wrong encoding)
	scorePureCJK        = 100 // pure CJK Unified: consistent Chinese text
	scorePercentDivisor = 100 // used to compute percentage-based scores
	maxASCII            = 0x80
)

const zipExtraUnicodePathID = 0x7075 // Info-ZIP Unicode Path extra field.

// encodingCandidate holds a charset detection result with its validation score.
type encodingCandidate struct {
	charset    string
	confidence int
	enc        encoding.Encoding
	score      int
}

// detectZipEncoding scans all zip file entries for non-UTF8 filenames and
// attempts to detect their character encoding. It uses chardet for initial
// candidates, then validates and scores each one. Returns nil if no non-UTF8
// filenames are found or no suitable decoder can be determined.
func detectZipEncoding(xFile *XFile, entries []*zip.File) *zipNameDecoders {
	var rawNames []string

	for _, f := range entries {
		if f.NonUTF8 || !utf8.ValidString(f.Name) {
			rawNames = append(rawNames, f.Name)
		}
	}

	if len(rawNames) == 0 {
		return nil
	}

	// Concatenate all non-UTF8 names for an archive-level fallback.
	var allBytes []byte
	for _, name := range rawNames {
		allBytes = append(allBytes, []byte(name)...)
	}

	// Get a map of filename->encoding.
	detector, decoders := detectEachFileEncoding(rawNames)

	results, err := detector.DetectAll(allBytes)
	if err != nil {
		xFile.Debugf("Charset detection failed for zip filenames: %v", err)

		if len(decoders.nameEncodings) == 0 {
			return nil
		}

		return decoders
	}

	best := pickBestEncoding(results, rawNames)
	if best != nil {
		decoders.defaultEncoding = best.enc
		xFile.Debugf("Detected zip fallback filename encoding: %s (confidence: %d, score: %d)",
			best.charset, best.confidence, best.score)
	}

	if len(decoders.nameEncodings) == 0 && decoders.defaultEncoding == nil {
		xFile.Debugf("No suitable encoding found for %d non-UTF8 zip filenames", len(rawNames))
		return nil
	}

	xFile.Debugf("Detected zip filename encodings for %d/%d entries and %d path parts",
		len(decoders.nameEncodings), len(rawNames), len(decoders.partNames))

	return decoders
}

func detectEachFileEncoding(rawNames []string) (*chardet.Detector, *zipNameDecoders) {
	detector := chardet.NewTextDetector()
	decoders := &zipNameDecoders{
		nameEncodings: make(map[string]encoding.Encoding, len(rawNames)),
		partNames:     map[string]string{},
	}
	// Detect the encoding of the filenames, per filename.
	for _, name := range rawNames {
		results, err := detector.DetectAll([]byte(name))
		if err != nil {
			continue
		}

		best := pickBestEncoding(results, []string{name})
		if best == nil {
			continue
		}

		decoders.nameEncodings[name] = best.enc

		decodedName, err := best.enc.NewDecoder().String(name)
		if err != nil {
			continue
		}

		rawParts := strings.Split(name, "/")

		decodedParts := strings.Split(decodedName, "/")
		if len(rawParts) != len(decodedParts) {
			continue
		}

		for idx, part := range rawParts {
			if part == "" {
				continue
			}

			if _, ok := decoders.partNames[part]; ok {
				continue
			}

			decoders.partNames[part] = decodedParts[idx]
		}
	}

	return detector, decoders
}

// pickBestEncoding evaluates chardet results against the raw filenames and
// returns the candidate with the highest combined score, or nil if none are valid.
func pickBestEncoding(results []chardet.Result, rawNames []string) *encodingCandidate {
	var candidates []encodingCandidate

	for _, result := range results {
		enc := charsetToEncoding(result.Charset)
		if enc == nil {
			continue
		}

		decoder := enc.NewDecoder()

		decoded, valid := decodeAll(decoder, rawNames)
		if !valid {
			continue
		}

		// Score: chardet confidence + script analysis + encoding preference.
		score := result.Confidence
		score += scriptConsistencyScore(decoded)
		score += encodingPreference(result.Charset)

		candidates = append(candidates, encodingCandidate{
			charset:    result.Charset,
			confidence: result.Confidence,
			enc:        enc,
			score:      score,
		})
	}

	if len(candidates) == 0 {
		return nil
	}

	// Pick the candidate with the highest combined score.
	best := &candidates[0]
	for idx := range candidates[1:] {
		if candidates[idx+1].score > best.score {
			best = &candidates[idx+1]
		}
	}

	return best
}

// decodeZipFilename decodes a zip entry filename if it's non-UTF8 and a decoder is available.
func decodeZipFilename(name string, extra []byte, nonUTF8 bool, decoders *zipNameDecoders) string {
	// Prefer ZIP's Unicode Path extra field when present. This is explicit metadata
	// and avoids heuristic guessing for mixed-language filenames.
	if unicodeName, ok := decodeUnicodePathExtra(name, extra); ok {
		return unicodeName
	}

	if decoders == nil {
		return name
	}

	if !nonUTF8 && utf8.ValidString(name) {
		return name
	}

	parts := strings.Split(name, "/")
	decoded := make([]string, len(parts))

	for idx, part := range parts {
		if part == "" {
			decoded[idx] = part
			continue
		}

		enc := decoders.defaultEncoding
		if value, ok := decoders.partNames[part]; ok {
			decoded[idx] = value
			continue
		} else if specific, ok := decoders.nameEncodings[name]; ok {
			enc = specific
		}

		if enc == nil {
			decoded[idx] = part
			continue
		}

		value, err := enc.NewDecoder().String(part)
		if err != nil {
			decoded[idx] = part // fall back to original on error
			continue
		}

		decoded[idx] = value
	}

	return strings.Join(decoded, "/")
}

func decodeUnicodePathExtra(rawName string, extra []byte) (string, bool) {
	const extraFieldBytes = 4
	for idx := 0; idx+extraFieldBytes <= len(extra); {
		fieldID := binary.LittleEndian.Uint16(extra[idx : idx+2])
		fieldSize := int(binary.LittleEndian.Uint16(extra[idx+2 : idx+extraFieldBytes]))
		start := idx + extraFieldBytes

		end := start + fieldSize
		if end > len(extra) {
			return "", false
		}

		if fieldID == zipExtraUnicodePathID && fieldSize >= 5 {
			fieldData := extra[start:end]

			nameCRC := binary.LittleEndian.Uint32(fieldData[1:5])
			if nameCRC != crc32.ChecksumIEEE([]byte(rawName)) {
				idx = end
				continue
			}

			unicodeName := string(fieldData[5:])
			if unicodeName != "" && utf8.ValidString(unicodeName) {
				return unicodeName, true
			}
		}

		idx = end
	}

	return "", false
}

func encodingPreference(charset string) int {
	switch strings.ToLower(charset) {
	case "gb-2312", "gb2312", "gbk", "gb18030", "gb-18030":
		return prefGBK
	case "shift_jis", "shift-jis":
		return prefShiftJS
	case "big5", "big-5":
		return prefBig5
	case "euc-kr":
		return prefEUCKR
	case "euc-jp":
		return prefEUCJP
	default:
		return 0
	}
}

// decodeAll attempts to decode every name using the given decoder.
// Returns the decoded names and whether all decoded successfully into valid UTF-8.
func decodeAll(decoder *encoding.Decoder, names []string) ([]string, bool) {
	decoded := make([]string, len(names))

	for idx, name := range names {
		d, err := decoder.String(name)
		if err != nil || !utf8.ValidString(d) {
			return nil, false
		}

		decoded[idx] = d
	}

	return decoded, true
}

// scriptConsistencyScore examines decoded filenames and returns a score based
// on how consistent the non-ASCII characters are within known Unicode blocks.
// Higher scores indicate text that looks like a natural writing system rather
// than garbled characters from a wrong encoding.
func scriptConsistencyScore(decoded []string) int {
	counts := countScriptRunes(decoded)

	total := counts.cjk + counts.hiragana + counts.katakana + counts.hangul + counts.latin + counts.other
	if total == 0 {
		return 0
	}

	kana := counts.hiragana + counts.katakana

	// Kana characters are unambiguous markers for Japanese text.
	// When decoded text contains kana, that encoding is almost certainly correct.
	if kana > 0 {
		// Some mojibake patterns produce a small amount of kana mixed into mostly
		// CJK text. That's usually not Japanese; avoid a large kana bonus there.
		if counts.cjk > 0 && kana*4 < counts.cjk {
			return counts.cjk * scorePercentDivisor / total
		}

		return scoreKanaBonus + kana*scorePercentDivisor/total
	}

	// Pure Hangul (no CJK mix) is a clear marker for Korean.
	if counts.hangul > 0 && counts.cjk == 0 {
		return scoreHangulBonus + counts.hangul*scorePercentDivisor/total
	}

	// Mixed Hangul + CJK is suspicious - usually means wrong encoding
	// (e.g., GBK bytes decoded as EUC-KR).
	if counts.hangul > 0 && counts.cjk > 0 {
		return scoreMixedHangulCJK
	}

	// Pure CJK Unified: consistent Chinese text (GBK/Big5).
	if counts.cjk > 0 && counts.other == 0 && counts.latin == 0 {
		return scorePureCJK
	}

	// General consistency score.
	dominant := max(counts.latin, max(counts.hangul, max(kana, counts.cjk)))

	return dominant * scorePercentDivisor / total
}

// scriptCounts holds per-script character counts from decoded filenames.
type scriptCounts struct {
	cjk      int // CJK Unified Ideographs (U+4E00-U+9FFF)
	hiragana int // Japanese Hiragana (U+3040-U+309F)
	katakana int // Japanese Katakana (U+30A0-U+30FF)
	hangul   int // Korean Hangul Syllables (U+AC00-U+D7AF)
	latin    int // Latin Extended (U+00C0-U+024F)
	other    int // everything else non-ASCII
}

// countScriptRunes classifies non-ASCII runes in the decoded names by Unicode block.
func countScriptRunes(decoded []string) scriptCounts {
	var counts scriptCounts

	for _, name := range decoded {
		for _, char := range name {
			switch {
			case char < maxASCII:
				// ASCII - ignore for scoring
			case char >= 0x4E00 && char <= 0x9FFF:
				counts.cjk++
			case char >= 0x3040 && char <= 0x309F:
				counts.hiragana++
			case char >= 0x30A0 && char <= 0x30FF:
				counts.katakana++
			case char >= 0xAC00 && char <= 0xD7AF:
				counts.hangul++
			case char >= 0x00C0 && char <= 0x024F:
				counts.latin++
			default:
				counts.other++
			}
		}
	}

	return counts
}

// charsetToEncoding maps charset names returned by chardet to golang.org/x/text encodings.
func charsetToEncoding(charset string) encoding.Encoding {
	switch strings.ToLower(charset) {
	case "gb-2312", "gb2312", "gbk", "gb18030", "gb-18030":
		return simplifiedchinese.GBK
	case "big5", "big-5":
		return traditionalchinese.Big5
	case "euc-jp":
		return japanese.EUCJP
	case "shift_jis", "shift-jis":
		return japanese.ShiftJIS
	case "iso-2022-jp":
		return japanese.ISO2022JP
	case "euc-kr":
		return korean.EUCKR
	case "iso-8859-1", "windows-1252":
		return charmap.Windows1252
	case "iso-8859-2":
		return charmap.ISO8859_2
	case "iso-8859-5":
		return charmap.ISO8859_5
	case "iso-8859-6":
		return charmap.ISO8859_6
	case "iso-8859-7":
		return charmap.ISO8859_7
	case "iso-8859-8":
		return charmap.ISO8859_8
	case "iso-8859-9":
		return charmap.ISO8859_9
	case "windows-1250":
		return charmap.Windows1250
	case "windows-1251":
		return charmap.Windows1251
	case "windows-1253":
		return charmap.Windows1253
	case "windows-1255":
		return charmap.Windows1255
	case "windows-1256":
		return charmap.Windows1256
	case "koi8-r":
		return charmap.KOI8R
	default:
		return nil
	}
}