Rewrite to support parallel processing.

Author: averagesecurityguy

config.go

@@ -23,17 +23,18 @@ type Regex struct {
 }
 
 type Config struct {
-	keys      map[string]time.Time
-	db        *bolt.DB
-	Keywords  []*Keyword // A list of keywords to search for in the data.
-	Regexes   []*Regex   // A list of regular expressions to test against data.
-	DbFile    string     `json:"database_file"` // File to use for the Store database.
-	MaxSize   int        `json:"max_size"`      // Do not save files larger than this many bytes.
-	MaxTime   int        `json:"max_time"`      // Max time, in seconds, to store previously downloaded keys.
-	Sleep     int        // Time, in seconds, to wait between each run.
-	GhToken   string     `json:"github_token"` // Github API token
-	Save      bool
-	LocalPath string `json:"local_path"`
+	keys          map[string]time.Time
+	db            *bolt.DB
+	Keywords      []*Keyword // A list of keywords to search for in the data.
+	Regexes       []*Regex   // A list of regular expressions to test against data.
+	DbFile        string     `json:"database_file"` // File to use for the Store database.
+	MaxSize       int        `json:"max_size"`      // Do not save files larger than this many bytes.
+	MaxTime       int        `json:"max_time"`      // Max time, in seconds, to store previously downloaded keys.
+	Sleep         int        // Time, in seconds, to wait between each run.
+	GhToken       string     `json:"github_token"` // Github API token
+	Save          bool
+	LocalPath     string `json:"local_path"`
+	FileBatchSize int    `json:"file_batch_size"`
 }
 
 func newConfig() Config {
@@ -57,5 +58,9 @@ func newConfig() Config {
 		r.compiled = regexp.MustCompile(r.Regex)
 	}
 
+	if c.FileBatchSize == 0 {
+		c.FileBatchSize = 100
+	}
+
 	return c
 }

files.go

@@ -7,63 +7,36 @@ import (
 	"path/filepath"
 )
 
-type File struct {
-	Path    string
-	Key     string
-	Content string
-}
-
-func (f *File) Read() {
-	content, err := ioutil.ReadFile(f.Path)
-	if err != nil {
-		log.Printf("[-] Could not read file %s\n", f.Path)
-		f.Content = ""
-	} else {
-		log.Printf("[+] Reading file: %s\n", f.Key)
-		f.Content = string(content)
-	}
-}
-
-func (f *File) Process() {
-	processContent(f.Key, f.Content)
-}
-
-func (f *File) Delete() {
-	os.Remove(f.Path)
-}
-
-func newFile(path string) *File {
-	f := new(File)
-
-	f.Path = path
-	f.Key = filepath.Base(path)
-
-	return f
-}
-
-func scrapeFiles() {
+func scrapeFiles(c chan<- *ProcessItem) {
 	if conf.LocalPath == "" {
 		return
 	}
 
 	log.Println("[+] Checking for local pastes.")
 
-	filepath.Walk(conf.LocalPath, func(path string, info os.FileInfo, err error) error {
-		if err != nil {
-			log.Printf("[-] Error reading %s: %s\n", conf.LocalPath, err.Error())
-			return nil
+	files, err := ioutil.ReadDir(conf.LocalPath)
+	if err != nil {
+		log.Printf("[-] Error reading %s: %s\n", conf.LocalPath, err)
+		return
+	}
+
+	// Process files in batches
+	for _, file := range files[:conf.FileBatchSize] {
+		if file.IsDir() {
+			log.Printf("[+] Skipping directory %s\n", conf.LocalPath)
+			continue
 		}
 
-		if info.IsDir() {
-			log.Printf("[+] Skipping directory %s\n", path)
-			return nil
+		path := filepath.Join(conf.LocalPath, file.Name())
+		content, err := ioutil.ReadFile(path)
+		if err != nil {
+			log.Printf("[-] Could not read file %s\n", path)
+			continue
 		}
 
-		f := newFile(path)
-		f.Read()
-		f.Process()
-		f.Delete()
+		item := &ProcessItem{Source: "Local", Key: file.Name(), Content: string(content)}
+		c <- item
 
-		return nil
-	})
+		os.Remove(path)
+	}
 }

gist.go

@@ -53,13 +53,7 @@ func (g *Gist) Download() {
 	conf.keys[g.Key] = time.Now()
 }
 
-func (g *Gist) Process() {
-	for _, gist := range g.files {
-		processContent(g.Key, gist.Content)
-	}
-}
-
-func scrapeGists() {
+func scrapeGists(c chan<- *ProcessItem) {
 	var gists []*Gist
 
 	log.Println("[+] Checking for new gists.")
@@ -76,6 +70,10 @@ func scrapeGists() {
 	for i, _ := range gists {
 		g := gists[i]
 		g.Download()
-		g.Process()
+
+		for _, gist := range g.files {
+			item := &ProcessItem{Source: "Gist", Key: g.Key, Content: gist.Content}
+			c <- item
+		}
 	}
 }

paste.go

@@ -9,14 +9,7 @@ import (
 type Paste struct {
 	ScrapeUrl string `json:"scrape_url"`
 	Url       string `json:"full_url"`
-	Date      string
 	Key       string
-	Size      int `json:",string"`
-	Expire    int `json:",string"`
-	Title     string
-	Syntax    string
-	User      string
-	Error     string
 	Content   string
 }
 
@@ -33,11 +26,7 @@ func (p *Paste) Download() {
 	conf.keys[p.Key] = time.Now()
 }
 
-func (p *Paste) Process() {
-	processContent(p.Key, p.Content)
-}
-
-func scrapePastes() {
+func scrapePastes(c chan<- *ProcessItem) {
 	var pastes []*Paste
 
 	log.Println("[+] Checking for new pastes.")
@@ -54,6 +43,8 @@ func scrapePastes() {
 	for i, _ := range pastes {
 		p := pastes[i]
 		p.Download()
-		p.Process()
+
+		item := &ProcessItem{Source: "Pastebin", Key: p.Key, Content: p.Content}
+		c <- item
 	}
 }

process.go

@@ -2,74 +2,75 @@ package main
 
 import (
 	"fmt"
+	"log"
 	"strings"
 )
 
-func savePaste(key, content string) {
-	if conf.Save == false {
+type ProcessItem struct {
+	Source  string
+	Key     string
+	Content string
+	Save    bool
+}
+
+func (p *ProcessItem) Write() {
+	if (conf.Save == false) || (p.Save == false) {
 		return
 	}
 
-	if len(content) > conf.MaxSize {
+	if len(p.Content) > conf.MaxSize {
 		return
 	}
 
-	writeDB(conf.db, "pastes", key, []byte(content))
+	db.Write("pastes", p.Key, []byte(p.Content))
 }
 
-func processRegexes(key, content string) {
-	save := false
+func (p *ProcessItem) Regexes() {
 	for i, _ := range conf.Regexes {
 		r := conf.Regexes[i]
 
 		switch r.Match {
 		case "all":
-			items := r.compiled.FindAllString(content, -1)
+			items := r.compiled.FindAllString(p.Content, -1)
 
 			if items != nil {
-				save = true
+				p.Save = true
 			}
 
 			for k := range items {
-				rKey := fmt.Sprintf("%s-%d", key, k)
-				writeDB(conf.db, r.Bucket, rKey, []byte(items[k]))
+				rKey := fmt.Sprintf("%s-%d", p.Key, k)
+				db.Write(r.Bucket, rKey, []byte(items[k]))
 			}
 		case "one":
-			match := r.compiled.FindString(content)
+			match := r.compiled.FindString(p.Content)
 
 			if match != "" {
-				save = true
-				writeDB(conf.db, r.Bucket, key, []byte(match))
+				p.Save = true
+				db.Write(r.Bucket, p.Key, []byte(match))
 			}
 		default:
 		}
 	}
-
-	if save {
-		savePaste(key, content)
-	}
 }
 
-func processKeywords(key, content string) {
-	save := false
+func (p *ProcessItem) Keywords() {
 	for i, _ := range conf.Keywords {
 		kwd := conf.Keywords[i]
 
-		if strings.Contains(strings.ToLower(content), strings.ToLower(kwd.Keyword)) {
-			save = true
-			writeDB(conf.db, kwd.Bucket, key, nil)
+		if strings.Contains(strings.ToLower(p.Content), strings.ToLower(kwd.Keyword)) {
+			p.Save = true
+			db.Write(kwd.Bucket, p.Key, nil)
 		}
 	}
-
-	if save {
-		savePaste(key, content)
-	}
 }
 
-func processContent(key, content string) {
-	conf.db = getDBConn()
-	defer conf.db.Close()
+func process(sem chan struct{}, pi *ProcessItem) {
+	log.Printf("[+] Processing %s:%s.\n", pi.Source, pi.Key)
+	sem <- struct{}{}
+
+	pi.Regexes()
+	pi.Keywords()
+	pi.Write()
 
-	processRegexes(key, content)
-	processKeywords(key, content)
+	<-sem
 }

scrape.go

@@ -5,6 +5,7 @@ import (
 )
 
 var conf Config
+var db *Database
 
 func cleanKeys() {
 	now := time.Now()
@@ -17,19 +18,43 @@ func cleanKeys() {
 	}
 }
 
-func scrape() {
-	scrapePastes()
-	scrapeGists()
-	scrapeFiles()
-}
+func initDatabase() {
+	db.CreateBucket("pastes")
 
-func main() {
-	conf = newConfig()
-	initDB()
+	for _, kw := range conf.Keywords {
+		db.CreateBucket(kw.Bucket)
+	}
 
+	for _, re := range conf.Regexes {
+		db.CreateBucket(re.Bucket)
+	}
+}
+
+func scrape(piChan chan<- *ProcessItem) {
 	for {
-		scrape()
+		go scrapePastes(piChan)
+		go scrapeGists(piChan)
+		go scrapeFiles(piChan)
+
 		time.Sleep(time.Duration(conf.Sleep) * time.Second)
 		cleanKeys()
 	}
 }
+
+func main() {
+	conf = newConfig()
+	db = newDatabase(conf.DbFile)
+
+	if db != nil {
+		initDatabase()
+
+		processItemChan := make(chan *ProcessItem, 100)
+		processSemaphore := make(chan struct{}, 10)
+
+		go scrape(processItemChan)
+
+		for item := range processItemChan {
+			go process(processSemaphore, item)
+		}
+	}
+}