Script to clear Yara rules (remove duplicates, move files with syntax errors, etc.)

cleaner.go

package main

import (
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"
	"strings"

	log "github.com/Sirupsen/logrus"
	yara "github.com/hillu/go-yara"
)

var (
	compiler    *yara.Compiler
	blacklisted []string
)

func main() {
	if err := compileRules("./"); err != nil {
		log.Fatal(err)
	}
	log.Info("[*] Cleansing finished")
}

// StringInSlice returns whether or not a string exists in a slice
func StringInSlice(a string, list []string) bool {
	for _, b := range list {
		if b == a {
			return true
		}
	}
	return false
}

type infoStruct struct {
	rule   string
	lineNo int
	file   string
}

func compileRules(RulesDir string) error {
	fileList := []string{}

	// walk rules directory
	err := filepath.Walk(RulesDir, func(path string, f os.FileInfo, err error) error {
		if !f.IsDir() {
			fileList = append(fileList, path)
		}
		return nil
	})
	if err != nil {
		return err
	}

	// new yara compiler
	compiler, err = yara.NewCompiler()
	if err != nil {
		return err
	}
	// compile all yara rules
	for _, file := range fileList {
		if file == "main.go" {
			continue
		}
		// log.Infof("[*] Adding rule %s", file)
		if StringInSlice(file, blacklisted) {
			continue
		}

		f, err := os.Open(file)
		if err != nil {
			return err
		}

		// log.Debug("Adding rule: ", file)
		err = compiler.AddFile(f, "webtotem")
		if err != nil {
			var process bool
			var replace bool
			var syntaxErrReplace bool
			var identifierReplace bool

			var toProcess []infoStruct
			blacklisted = append(blacklisted, file)
			for _, er := range compiler.Errors {
				if replace || syntaxErrReplace || identifierReplace {
					continue
				}
				if strings.Contains(er.Text, `duplicated identifier`) {
					process = true
					toProcess = append(toProcess, infoStruct{er.Text, er.Line, er.Filename})
				}
				if strings.Contains(er.Text, `unknown module`) {
					replace = true
				}
				if strings.Contains(er.Text, `can't open include`) || strings.Contains(er.Text, `syntax error`) || strings.Contains(er.Text, `duplicated string`) || strings.Contains(er.Text, `invalid hex string`) || strings.Contains(er.Text, `non-ascii`) || strings.Contains(er.Text, `empty string`) || strings.Contains(er.Text, `unterminated regular expression`) {
					syntaxErrReplace = true
				}
				if strings.Contains(er.Text, `undefined identifier`) {
					identifierReplace = true
				}

				log.WithFields(log.Fields{
					"rule":    er.Filename,
					"line_no": er.Line,
				}).Error(er.Text)
			}
			// for _, wr := range yaraCompiler.Warnings {
			// 	log.Warn(wr)
			// }
			f.Close()
			if process {
				if err := processDuplicateRule(file, toProcess); err != nil {
					log.Fatal(err)
				} else {
					log.Infof("removed duplicated rules from file %s", file)
				}
			}
			if replace {

				err := os.Rename(file, "../sortedOut/unsupported/"+file)
				if err != nil {
					panic(err)
				} else {
					log.Infof("moved file %s", file)
				}
			}
			if syntaxErrReplace {
				err := os.Rename(file, "../sortedOut/syntaxErr/"+file)
				if err != nil {
					panic(err)
				} else {
					log.Infof("moved file %s", file)
				}
			}
			if identifierReplace {
				err := os.Rename(file, "../sortedOut/undefinedIdentifier/"+file)
				if err != nil {
					panic(err)
				} else {
					log.Infof("moved file %s", file)
				}
			}
			// destroy unstable YR_COMPILER
			// (see https://github.com/hillu/go-yara/issues/32#issuecomment-416040753)
			log.Debug("destroying unstable yara compiler")
			compiler.Destroy()
			log.Debug("recreating yara compiler")
			// os.Exit(1)
			if err := compileRules(RulesDir); err != nil {
				log.Fatal(err)
			}
		}
		f.Close()
	}

	return nil
}

func processDuplicateRule(file string, rulesForRemoval []infoStruct) (err error) {
	// start := regexp.MustCompile("{\n")
	// end := regexp.MustCompile("}\n")
	regRuleName := regexp.MustCompile(`"(.*?)"`)
	regRuleEnd := regexp.MustCompile(`^\s*}`)

	for _, r := range rulesForRemoval {
		ruleName := regRuleName.FindStringSubmatch(r.rule)[1]
		log.Infof("[*] Rule for removal: %s", ruleName)
		regRuleNameFinder := regexp.MustCompile(fmt.Sprintf(`rule\s%s\b`, ruleName))

		f, err := os.Open(r.file)
		if err != nil {
			return err
		}
		scanner := bufio.NewScanner(f)
		lineCount := 0

		// Search for the start of the rule
		var start int
	firstLoop:
		for scanner.Scan() {
			lineCount++
			line := scanner.Text()
			// log.Infof("[*] line %d: %s", lineCount, line)
			if regRuleNameFinder.MatchString(line) {
				log.Infof("[*] Found rule: %s, line %d", line, lineCount)
				start = lineCount
				break firstLoop
			}
		}
		// Search for the end of the rule
		var end int
	secondLoop:
		for scanner.Scan() {
			lineCount++
			line := scanner.Text()
			if regRuleEnd.MatchString(line) {
				log.Infof("[*] Rule ends at line %d", lineCount)
				end = lineCount
				break secondLoop
			}
		}
		f.Close()
		// Check if end of rule has been found
		if end == 0 {
			return fmt.Errorf("Rule end hasn't been found [start: %d], file: %s", start, file)
		}
		// Remove rule
		numberOfLinesToRemove := end - start + 1
		if err := removeLines(r.file, start, numberOfLinesToRemove); err != nil {
			return err
		}
	}
	return nil
}

func removeLines(fn string, start, n int) (err error) {
	if start < 1 {
		return errors.New("invalid request.  line numbers start at 1")
	}
	if n < 0 {
		return errors.New("invalid request.  negative number to remove")
	}
	var f *os.File
	if f, err = os.OpenFile(fn, os.O_RDWR, 0); err != nil {
		return
	}
	defer func() {
		if cErr := f.Close(); err == nil {
			err = cErr
		}
	}()
	var b []byte
	if b, err = ioutil.ReadAll(f); err != nil {
		return
	}
	cut, ok := skip(b, start-1)
	if !ok {
		return fmt.Errorf("less than %d lines", start)
	}
	if n == 0 {
		return nil
	}
	tail, ok := skip(cut, n)
	if !ok {
		return fmt.Errorf("less than %d lines after line %d", n, start)
	}
	t := int64(len(b) - len(cut))
	if err = f.Truncate(t); err != nil {
		return
	}
	if len(tail) > 0 {
		_, err = f.WriteAt(tail, t)
	}
	return
}

func skip(b []byte, n int) ([]byte, bool) {
	for ; n > 0; n-- {
		if len(b) == 0 {
			return nil, false
		}
		x := bytes.IndexByte(b, '\n')
		if x < 0 {
			x = len(b)
		} else {
			x++
		}
		b = b[x:]
	}
	return b, true
}