ONNX
security
malware-detection
Vigil / source /pkg /compactmodel /vectorize.go
turentomer's picture
Publish self-contained Vigil distribution
d2507b5 verified
Raw
History Blame Contribute Delete
5.19 kB
package compactmodel
import (
"encoding/binary"
"math"
"unicode"
"golang.org/x/text/cases"
"golang.org/x/text/language"
)
// Vectorize serializes an already-scanned package and returns the fixed 65,552
// float32 feature vector expected by the compact linear scorer.
func Vectorize(pkg Package) ([]float32, error) {
document, err := Serialize(pkg)
if err != nil {
return nil, err
}
return VectorizeDocument(document)
}
// VectorizeDocument hashes a previously serialized document. The word and
// character blocks are independently L2-normalized, matching scikit-learn's
// HashingVectorizer configuration; structural values are copied unchanged.
func VectorizeDocument(document Document) ([]float32, error) {
if err := ValidateRuntime(); err != nil {
return nil, err
}
result := make([]float32, TotalFeatures)
lowered := cases.Lower(language.Und).String(document.Text)
hashWordFeatures(result[:WordFeatures], wordTokens(lowered))
hashCharFeatures(result[WordFeatures:WordFeatures+CharFeatures], lowered)
copy(result[WordFeatures+CharFeatures:], document.Structured[:])
return result, nil
}
func isWordRune(value rune) bool {
return value == '_' || unicode.IsLetter(value) || unicode.IsNumber(value)
}
func isTokenRune(value rune) bool {
if isWordRune(value) {
return true
}
switch value {
case '.', '/', '$', ':', '@', '-':
return true
default:
return false
}
}
// wordTokens implements (?u)\b[\w./$:@-]{2,}\b. Runs are trimmed to word
// characters because Python's \b assertions only consider \w, even though the
// interior class admits punctuation.
func wordTokens(lowered string) []string {
runes := []rune(lowered)
tokens := make([]string, 0)
for start := 0; start < len(runes); {
for start < len(runes) && !isTokenRune(runes[start]) {
start++
}
end := start
for end < len(runes) && isTokenRune(runes[end]) {
end++
}
wordStart, wordEnd := start, end
for wordStart < wordEnd && !isWordRune(runes[wordStart]) {
wordStart++
}
for wordEnd > wordStart && !isWordRune(runes[wordEnd-1]) {
wordEnd--
}
if wordEnd-wordStart >= 2 {
tokens = append(tokens, string(runes[wordStart:wordEnd]))
}
start = end
}
return tokens
}
func hashWordFeatures(destination []float32, tokens []string) {
for _, token := range tokens {
addHashed(destination, token)
}
for index := 0; index+1 < len(tokens); index++ {
addHashed(destination, tokens[index]+" "+tokens[index+1])
}
normalizeL2(destination)
}
func pythonWhitespace(value rune) bool {
switch value {
case '\t', '\n', '\v', '\f', '\r', ' ',
0x1c, 0x1d, 0x1e, 0x1f, 0x85, 0x2028, 0x2029:
return true
default:
return unicode.Is(unicode.Zs, value)
}
}
func normalizeCharWhitespace(lowered string) []rune {
input := []rune(lowered)
result := make([]rune, 0, len(input))
for start := 0; start < len(input); {
if !pythonWhitespace(input[start]) {
result = append(result, input[start])
start++
continue
}
end := start + 1
for end < len(input) && pythonWhitespace(input[end]) {
end++
}
if end-start >= 2 {
result = append(result, ' ')
} else {
result = append(result, input[start])
}
start = end
}
return result
}
func hashCharFeatures(destination []float32, lowered string) {
runes := normalizeCharWhitespace(lowered)
for index := 0; index+4 <= len(runes); index++ {
addHashed(destination, string(runes[index:index+4]))
}
normalizeL2(destination)
}
func addHashed(destination []float32, feature string) {
hash := int32(murmurHash3X86_32([]byte(feature), 0))
// scikit-learn hashes to abs(signed_hash) % n_features and uses the
// signed hash for alternate_sign. MurmurHash3's sole MinInt32 output is
// handled without overflowing the absolute value.
magnitude := int64(hash)
value := float32(1)
if magnitude < 0 {
magnitude = -magnitude
value = -1
}
index := int(magnitude % int64(len(destination)))
destination[index] += value
}
func normalizeL2(values []float32) {
var sumSquares float64
for _, value := range values {
// scikit's fused float32 kernel multiplies in float32, then adds the
// rounded product to its double accumulator.
sumSquares += float64(value * value)
}
if sumSquares == 0 {
return
}
norm := math.Sqrt(sumSquares)
for index := range values {
values[index] = float32(float64(values[index]) / norm)
}
}
func murmurHash3X86_32(data []byte, seed uint32) uint32 {
const (
c1 = uint32(0xcc9e2d51)
c2 = uint32(0x1b873593)
)
hash := seed
blocks := len(data) / 4
for index := 0; index < blocks; index++ {
key := binary.LittleEndian.Uint32(data[index*4:])
key *= c1
key = key<<15 | key>>(32-15)
key *= c2
hash ^= key
hash = hash<<13 | hash>>(32-13)
hash = hash*5 + 0xe6546b64
}
var tail uint32
remainder := data[blocks*4:]
switch len(remainder) {
case 3:
tail ^= uint32(remainder[2]) << 16
fallthrough
case 2:
tail ^= uint32(remainder[1]) << 8
fallthrough
case 1:
tail ^= uint32(remainder[0])
tail *= c1
tail = tail<<15 | tail>>(32-15)
tail *= c2
hash ^= tail
}
hash ^= uint32(len(data))
hash ^= hash >> 16
hash *= 0x85ebca6b
hash ^= hash >> 13
hash *= 0xc2b2ae35
hash ^= hash >> 16
return hash
}