- Removed messy writer implementation

- Fixed reader implementation
- Fixed memory leaks

Author: gboddin

README.md

@@ -1,87 +1,99 @@
-# Yara/ClamAV/Other Combo
+# Yara scanner for Golang
 
-Yara scanner library compatible with `io.TeeReader` for streaming
+Yara scanner library compatible with `io.Reader` for streaming
 
-## Requirements
-
-All requirements from [github.com/hillu/go-yara](https://github.com/hillu/go-yara) apply
-
-## Implementation
+## Features
 
-### Must implement WriteCloser
+- Read multiple rules for multiple directories
+- Scan inside archives with maximum depth
 
-When using a Tee reader, the `TeeReaderAutoClose` should be used to clean resources.
+## Requirements
 
-Make sure to call `Close()` if you're doing a custom implementation.
+All requirements from [github.com/hillu/go-yara](https://github.com/hillu/go-yara) apply
 
 ## Example
 
 ```golang
 package main
 
 import (
-	"context"
 	"crypto/sha1"
 	"crypto/sha256"
-	"encoding/hex"
 	"flag"
 	"github.com/LeakIX/YaraStream"
-	"github.com/elvinchan/clamd"
 	"io"
+	"io/fs"
 	"log"
 	"os"
+	"path/filepath"
+	"runtime"
+	"sync"
 )
 
-var (
-	clamdSock = flag.String("clamd-sock", "", "ClamD socket")
-)
+var yaraScanner *YaraStream.YaraScanner
+var fileChan = make(chan string)
+var wg sync.WaitGroup
 
 func main() {
+	var err error
 	flag.Parse()
 	if len(flag.Args()) != 1 {
 		log.Fatal("You must provide a file to scan")
 	}
-	clamClient, err := clamd.NewClient("unix", *clamdSock)
-	if err != nil {
-		panic(err)
-	}
-	scanner, err := YaraStream.NewYaraScanner(
+	yaraScanner, err = YaraStream.NewYaraScanner(
 		YaraStream.RuleDirectory{Namespace: "AbuseCH", Path: "./rules/abusech"},
 		YaraStream.RuleDirectory{Namespace: "ReversingLabs", Path: "./rules/reversinglabs"},
 		YaraStream.RuleDirectory{Namespace: "ESET", Path: "./rules/eset"},
+		YaraStream.RuleDirectory{Namespace: "AlienVault", Path: "./rules/otx"},
 	)
 	if err != nil {
 		panic(err)
 	}
-	file, err := os.Open(flag.Arg(0))
+
+	for w := 1; w <= runtime.NumCPU(); w++ {
+		go worker()
+	}
+	err = filepath.Walk(flag.Arg(0), func(path string, info fs.FileInfo, err error) error {
+		if err != nil {
+			return nil
+		}
+		if info.IsDir() || !info.Mode().IsRegular() {
+			return nil
+		}
+		wg.Add(1)
+		fileChan <- path
+		return nil
+	})
+
 	if err != nil {
 		panic(err)
 	}
-	yaraWriter := scanner.NewYaraWriter()
+	wg.Wait()
+}
+
+func worker() {
+	for filePath := range fileChan {
+		scanFile(filePath)
+		wg.Done()
+	}
+}
+
+func scanFile(path string) error {
+	log.Printf("Scanning %s", path)
+	file, err := os.Open(path)
+	if err != nil {
+		log.Println(err)
+		return nil
+	}
 	sha256Hasher := sha256.New()
 	sha1Hasher := sha1.New()
 	sha256Tee := io.TeeReader(file, sha256Hasher)
 	sha1Tee := io.TeeReader(sha256Tee, sha1Hasher)
-	yaraTee := YaraStream.TeeReaderAutoClose(sha1Tee, yaraWriter)
-	scanResults, err := clamClient.ScanReader(context.Background(), yaraTee)
-	if err != nil {
-		panic(err)
-	}
-	infected := false
-	for _, scanResult := range scanResults {
-		if scanResult.Status == "FOUND" {
-			infected = true
-			log.Printf("[ClamAV] Infection found: %s", scanResult.Signature)
-		}
-	}
-	for _, scanResult := range yaraWriter.MatchedRules {
-		infected = true
-		log.Printf("[YARA] Infection found: %s/%s", scanResult.Namespace(), scanResult.Identifier())
-	}
-	log.Printf("SHA256: %s", hex.EncodeToString(sha256Hasher.Sum(nil)))
-	log.Printf("SHA1: %s", hex.EncodeToString(sha1Hasher.Sum(nil)))
-	if infected {
-		os.Exit(1)
+	matches, err := yaraScanner.ScanReader(sha1Tee, YaraStream.WithFilenameTip(path), YaraStream.WithMaxLevel(3))
+	for _, scanResult := range matches {
+		log.Printf("[YARA] Infection found: %s/%s in %s\n", scanResult.Namespace(), scanResult.Identifier(), path)
 	}
+	return nil
 }
+
 ```

TeeAutoCloser.go

@@ -20,18 +20,20 @@ type YaraReader struct {
 	Infected       bool
 	filename       string
 	level          int
+	maxBlockSize   int
 }
 
+// ScanReader Will scan a given reader until EOF or an error happens. It will scan archives.
 func (s *YaraScanner) ScanReader(reader io.Reader, opts ...func(writer *YaraReader)) ([]*yara.Rule, error) {
-	bufferedReader := bufio.NewReaderSize(reader, 16*1024)
 	testReader := &YaraReader{
-		r:     bufferedReader,
-		level: 10,
+		level:        10,
+		maxBlockSize: 16 * 1024,
 	}
 	for _, option := range opts {
 		option(testReader)
 	}
-	dec, err := decoder.GetDecoder("", testReader.r)
+	testReader.r = bufio.NewReaderSize(reader, testReader.maxBlockSize)
+	dec, err := decoder.GetDecoder(testReader.filename, testReader.r)
 	if err == nil {
 		if testReader.level < 1 {
 			return nil, nil
@@ -42,7 +44,11 @@ func (s *YaraScanner) ScanReader(reader io.Reader, opts ...func(writer *YaraRead
 				return testReader.mrs, nil
 			}
 			if entry.IsFile() {
-				partResults, _ := s.ScanReader(dec, ReaderWithFilenameTip(entry.Filename), ReaderWithCurrentLevel(testReader.level-1))
+				partResults, _ := s.ScanReader(dec,
+					WithFilenameTip(entry.Filename),
+					WithMaxLevel(testReader.level-1),
+					WithBlockSize(testReader.maxBlockSize),
+				)
 				testReader.mrs = append(testReader.mrs, partResults...)
 			}
 		}
@@ -51,13 +57,14 @@ func (s *YaraScanner) ScanReader(reader io.Reader, opts ...func(writer *YaraRead
 	defer testReader.scanner.Destroy()
 	testReader.scanner.SetFlags(yara.ScanFlagsProcessMemory)
 	testReader.scanner.SetCallback(testReader)
-	testReader.buf = make([]byte, 16*1024)
-	testReader.firstBlockData = make([]byte, 16*1024)
+	testReader.buf = make([]byte, testReader.maxBlockSize)
+	testReader.firstBlockData = make([]byte, testReader.maxBlockSize)
 	err = testReader.scanner.ScanMemBlocks(testReader)
 
 	return testReader.mrs, err
 }
 
+// First Will fetch the first block and cache it in our reader for further calls
 func (s *YaraReader) First() *yara.MemoryBlock {
 	if s.firstBlock == nil {
 		s.firstBlock = s.Next()
@@ -71,6 +78,7 @@ func (s *YaraReader) First() *yara.MemoryBlock {
 	return s.firstBlock
 }
 
+// Next Will fetch the next block for scanning
 func (s *YaraReader) Next() *yara.MemoryBlock {
 	n, err := s.r.Read(s.buf)
 	if err != nil && n == 0 {
@@ -85,29 +93,40 @@ func (s *YaraReader) Next() *yara.MemoryBlock {
 	}
 }
 
+// RuleMatching will be called by the engine when a rule is matched
 func (y *YaraReader) RuleMatching(_ *yara.ScanContext, rule *yara.Rule) (bool, error) {
 	y.Infected = true
 	y.mrs = append(y.mrs, rule)
 	return true, nil
 }
 
+// copy Helper for Next()
 func (s *YaraReader) copy(buf []byte) {
 	copy(buf, s.buf[:s.length])
 }
+
+// first Helper for First()
 func (s *YaraReader) first(buf []byte) {
 	copy(buf, s.firstBlockData[:s.firstBlock.Size])
-	//log.Println(s.filename, s.firstBlock.Base, s.firstBlock.Size, string(buf))
-
 }
 
-func ReaderWithFilenameTip(filename string) func(writer *YaraReader) {
+// WithFilenameTip Will tip the Decoder on possible archive types
+func WithFilenameTip(filename string) func(writer *YaraReader) {
 	return func(writer *YaraReader) {
 		writer.filename = filename
 	}
 }
 
-func ReaderWithCurrentLevel(level int) func(writer *YaraReader) {
+// WithMaxLevel Will prevent the Reader to inspect archives under and given level
+func WithMaxLevel(level int) func(writer *YaraReader) {
 	return func(writer *YaraReader) {
 		writer.level = level
 	}
 }
+
+// WithBlockSize Sets the default buffer and block size for in-memory scanning
+func WithBlockSize(size int) func(writer *YaraReader) {
+	return func(writer *YaraReader) {
+		writer.maxBlockSize = size
+	}
+}

YaraReader.go

@@ -17,6 +17,7 @@ type RuleDirectory struct {
 	Path      string
 }
 
+// NewYaraScanner Will return a new scanner with rules compiled
 func NewYaraScanner(ruleDirectories ...RuleDirectory) (*YaraScanner, error) {
 	scanner := &YaraScanner{}
 	compiler, err := yara.NewCompiler()