llmgw/internal/guard/detect/detect.go

// Package detect finds what a guard policy acts on in a piece of text:
// personal data and secrets, by where they are, so they can be rewritten in
// place; and signs of prompt injection, in the text as a model would read
// it once disguises are taken off.
//
// It reports kinds and positions, never values: what a caller does with a
// span is the caller's business, and nothing here keeps the text.
package detect

import (
	"regexp"
	"slices"
	"strconv"
	"strings"
)

// Kinds of finding.
const (
	Email  = "email"
	Phone  = "phone"
	IBAN   = "iban"
	Card   = "card"
	TCKN   = "tckn"
	IP     = "ip"
	Secret = "secret"

	// Injection kinds: a phrase telling the model to drop its
	// instructions, a chat template's control tokens in user text, text a
	// person cannot see, and a long encoded blob.
	InjectionPhrase    = "injection_phrase"
	InjectionTemplate  = "injection_template"
	InjectionInvisible = "injection_invisible"
	InjectionEncoded   = "injection_encoded"
)

// PIIKinds are the kinds Spans can find, in the order a policy lists them.
var PIIKinds = []string{Email, Phone, IBAN, Card, TCKN, IP, Secret}

// Span is one finding in a text: its kind and byte range.
type Span struct {
	Kind       string
	Start, End int
}

var (
	emailRe = regexp.MustCompile(`[A-Za-z0-9._%+\-]+@[A-Za-z0-9\-]+(?:\.[A-Za-z0-9\-]+)*\.[A-Za-z]{2,}`)
	// Digits with the separators people write them with; validated after.
	phoneRe = regexp.MustCompile(`(?:\+|\b00)?\(?\d[\d \-().]{8,18}\d\b`)
	ibanRe  = regexp.MustCompile(`\b[A-Z]{2}\d{2}(?: ?[A-Z0-9]){11,30}\b`)
	cardRe  = regexp.MustCompile(`\b\d(?:[ \-]?\d){12,18}\b`)
	tcknRe  = regexp.MustCompile(`\b[1-9]\d{10}\b`)
	ipv4Re  = regexp.MustCompile(`\b(?:\d{1,3}\.){3}\d{1,3}\b`)

	secretRes = []*regexp.Regexp{
		regexp.MustCompile(`\bsk-(?:proj-|ant-)?[A-Za-z0-9_\-]{20,}`),                            // OpenAI, Anthropic
		regexp.MustCompile(`\b(?:AKIA|ASIA)[0-9A-Z]{16}\b`),                                      // AWS access key id
		regexp.MustCompile(`\bgh[pousr]_[A-Za-z0-9]{36,}\b`),                                     // GitHub
		regexp.MustCompile(`\bxox[abprs]-[A-Za-z0-9\-]{10,}`),                                    // Slack
		regexp.MustCompile(`\bAIza[0-9A-Za-z_\-]{35}\b`),                                         // Google API key
		regexp.MustCompile(`\beyJ[A-Za-z0-9_\-]{8,}\.eyJ[A-Za-z0-9_\-]{8,}\.[A-Za-z0-9_\-]{8,}`), // JWT
		regexp.MustCompile(`-----BEGIN [A-Z ]*PRIVATE KEY-----[\s\S]*?(?:-----END [A-Z ]*PRIVATE KEY-----|$)`),
		regexp.MustCompile(`\bllmgw_[A-Za-z0-9]{20,}\b`), // this gateway's own keys
	}
)

// Spans finds personal data and secrets in text, keeping only those that
// pass their validators — a checksum, a range — and the longest where they
// overlap.
func Spans(text string) []Span {
	var out []Span
	add := func(kind string, re *regexp.Regexp, valid func(string) bool) {
		for _, m := range re.FindAllStringIndex(text, -1) {
			if valid == nil || valid(text[m[0]:m[1]]) {
				out = append(out, Span{Kind: kind, Start: m[0], End: m[1]})
			}
		}
	}
	for _, re := range secretRes {
		add(Secret, re, nil)
	}
	add(Email, emailRe, nil)
	add(IBAN, ibanRe, validIBAN)
	add(Card, cardRe, validCard)
	add(TCKN, tcknRe, validTCKN)
	add(IP, ipv4Re, validIPv4)
	add(Phone, phoneRe, validPhone)
	return longest(out)
}

// longest keeps, of overlapping spans, the one found first in priority
// order (secrets, then the rest as added) unless a later one is longer.
func longest(in []Span) []Span {
	var out []Span
	for _, s := range in {
		keep := true
		for i := 0; i < len(out); i++ {
			o := out[i]
			if s.Start >= o.End || o.Start >= s.End {
				continue
			}
			if s.End-s.Start > o.End-o.Start {
				out = slices.Delete(out, i, i+1)
				i--
				continue
			}
			keep = false
			break
		}
		if keep {
			out = append(out, s)
		}
	}
	slices.SortFunc(out, func(a, b Span) int { return a.Start - b.Start })
	return out
}

func digits(s string) string {
	var b strings.Builder
	for _, r := range s {
		if r >= '0' && r <= '9' {
			b.WriteRune(r)
		}
	}
	return b.String()
}

func validCard(s string) bool {
	d := digits(s)
	if len(d) < 13 || len(d) > 19 {
		return false
	}
	// Issuers' prefixes: Visa, Mastercard, Amex, Discover, Troy, UnionPay,
	// JCB, Diners — anything else is a number, not a card.
	switch {
	case d[0] == '4', d[0] == '5', d[0] == '2', d[:2] == "34", d[:2] == "37", d[0] == '6',
		d[:2] == "35", d[:2] == "30", d[:2] == "36", d[:2] == "38", d[:2] == "65", d[:4] == "9792":
	default:
		return false
	}
	return luhn(d)
}

func luhn(d string) bool {
	sum, alt := 0, false
	for i := len(d) - 1; i >= 0; i-- {
		n := int(d[i] - '0')
		if alt {
			n *= 2
			if n > 9 {
				n -= 9
			}
		}
		sum += n
		alt = !alt
	}
	return sum%10 == 0
}

func validIBAN(s string) bool {
	s = strings.ReplaceAll(s, " ", "")
	if len(s) < 15 || len(s) > 34 {
		return false
	}
	r := s[4:] + s[:4]
	mod := 0
	for _, c := range r {
		var v int
		switch {
		case c >= '0' && c <= '9':
			v = int(c - '0')
		case c >= 'A' && c <= 'Z':
			v = int(c-'A') + 10
		default:
			return false
		}
		if v >= 10 {
			mod = (mod*100 + v) % 97
		} else {
			mod = (mod*10 + v) % 97
		}
	}
	return mod == 1
}

// validTCKN checks a Turkish identity number's two check digits.
func validTCKN(s string) bool {
	if len(s) != 11 || s[0] == '0' {
		return false
	}
	d := make([]int, 11)
	for i := range s {
		d[i] = int(s[i] - '0')
	}
	odd := d[0] + d[2] + d[4] + d[6] + d[8]
	even := d[1] + d[3] + d[5] + d[7]
	if ((odd*7-even)%10+10)%10 != d[9] {
		return false
	}
	sum := 0
	for _, v := range d[:10] {
		sum += v
	}
	return sum%10 == d[10]
}

func validIPv4(s string) bool {
	for _, p := range strings.Split(s, ".") {
		n, err := strconv.Atoi(p)
		if err != nil || n > 255 || (len(p) > 1 && p[0] == '0') {
			return false
		}
	}
	// Versions are written like addresses; 1.2.3.4-shaped strings with
	// small parts are far likelier to be one.
	return !strings.HasPrefix(s, "0.")
}

// validPhone accepts numbers written the way they are dialled: +country,
// 00country, or a national 0 and an area code. A bare run of digits is
// anything — an order number, part of an account number — and is not taken.
func validPhone(s string) bool {
	d := digits(s)
	if len(d) < 10 || len(d) > 15 || strings.Count(d, d[:1]) == len(d) {
		return false
	}
	t := strings.TrimLeft(s, "(")
	switch {
	case strings.HasPrefix(t, "+"):
		return d[0] != '0'
	case strings.HasPrefix(d, "00"):
		return len(d) >= 12 && d[2] != '0'
	case d[0] == '0':
		return d[1] != '0' && len(d) <= 11
	}
	return false
}