Skip to content

Commit 387fa4e

Browse files
codeyourwebcodeyourweb
committed
performance improvements
1 parent 89fd139 commit 387fa4e

File tree

4 files changed

+99
-77
lines changed

4 files changed

+99
-77
lines changed

‎analysis.go‎

Lines changed: 38 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,14 @@
11
package main
22

33
import (
4-
"crypto/md5"
54
"crypto/rc4"
65
b64 "encoding/base64"
76
"fmt"
87
"io/ioutil"
98
"log"
109
"os"
1110
"path/filepath"
11+
"runtime/debug"
1212
"time"
1313

1414
"github.com/hillu/go-yara"
@@ -28,8 +28,10 @@ func FileAnalysis(path string, pQuarantine string, pKill bool, pAggressive bool,
2828
var content []byte
2929
var result yara.MatchRules
3030

31-
if f, err = os.Stat(path); err != nil && pVerbose {
32-
log.Println("[ERROR]", path, err)
31+
if f, err = os.Stat(path); err != nil {
32+
if pVerbose {
33+
log.Println("[ERROR]", path, err)
34+
}
3335
} else {
3436
if RegisterFileInHistory(f, path, &filescanHistory) {
3537
content, err = ioutil.ReadFile(path)
@@ -41,7 +43,12 @@ func FileAnalysis(path string, pQuarantine string, pKill bool, pAggressive bool,
4143
log.Println("[INFO] ["+sourceIndex+"] Analyzing", path)
4244
}
4345

44-
result = PerformYaraScan(content, rules, pVerbose)
46+
// cleaning memory if file size is greater than 1Gb
47+
if len(content) > 1024*1024*1024 {
48+
defer debug.FreeOSMemory()
49+
}
50+
51+
result = PerformYaraScan(&content, rules, pVerbose)
4552

4653
if len(result) > 0 {
4754
// windows notifications
@@ -73,51 +80,45 @@ func FileAnalysis(path string, pQuarantine string, pKill bool, pAggressive bool,
7380
}
7481

7582
// MemoryAnalysis sub-routine for running processes analysis
76-
func MemoryAnalysis(proc ProcessInformation, pQuarantine string, pKill bool, pAggressive bool, pNotifications bool, pVerbose bool, rules *yara.Rules) {
77-
memoryHash := fmt.Sprintf("%x", md5.Sum(proc.ProcessMemory))
83+
func MemoryAnalysis(proc *ProcessInformation, pQuarantine string, pKill bool, pAggressive bool, pNotifications bool, pVerbose bool, rules *yara.Rules) {
84+
if pVerbose {
85+
log.Println("[INFO] [MEMORY] Analyzing", proc.ProcessName, "PID:", proc.PID)
86+
}
7887

79-
// if hash isn't already whitelisted, yara scan it
80-
if !StringInSlice(memoryHash, memoryscanHistory) {
81-
if pVerbose {
82-
log.Println("[INFO] [MEMORY] Analyzing", proc.ProcessName, "PID:", proc.PID)
88+
result := PerformYaraScan(&proc.MemoryDump, rules, pVerbose)
89+
if len(result) > 0 {
90+
// windows notifications
91+
if pNotifications {
92+
NotifyUser("YARA match", proc.ProcessName+" - PID:"+fmt.Sprint(proc.PID)+" match "+fmt.Sprint(len(result))+" rules")
8393
}
8494

85-
result := PerformYaraScan(proc.ProcessMemory, rules, pVerbose)
86-
if len(result) > 0 {
87-
// windows notifications
88-
if pNotifications {
89-
NotifyUser("YARA match", proc.ProcessName+" - PID:"+fmt.Sprint(proc.PID)+" match "+fmt.Sprint(len(result))+" rules")
90-
}
91-
92-
// logging
93-
for _, match := range result {
94-
log.Println("[ALERT]", "[MEMORY] YARA match", proc.ProcessName, "PID:", fmt.Sprint(proc.PID), match.Namespace, match.Rule)
95-
}
95+
// logging
96+
for _, match := range result {
97+
log.Println("[ALERT]", "[MEMORY] YARA match", proc.ProcessName, "PID:", fmt.Sprint(proc.PID), match.Namespace, match.Rule)
98+
}
9699

97-
// dump matching process to quarantine
98-
if len(pQuarantine) > 0 {
99-
log.Println("[INFO]", "DUMPING PID", proc.PID)
100-
err := QuarantineProcess(proc, pQuarantine)
101-
if err != nil {
102-
log.Println("[ERROR]", "Cannot quarantine PID", proc.PID, err)
103-
}
100+
// dump matching process to quarantine
101+
if len(pQuarantine) > 0 {
102+
log.Println("[INFO]", "DUMPING PID", proc.PID)
103+
err := QuarantineProcess(proc, pQuarantine)
104+
if err != nil {
105+
log.Println("[ERROR]", "Cannot quarantine PID", proc.PID, err)
104106
}
107+
}
105108

106-
// killing process
107-
if pKill {
108-
log.Println("[INFO]", "KILLING PID", proc.PID)
109-
KillProcessByID(proc.PID, pVerbose)
110-
}
111-
} else {
112-
memoryscanHistory = append(memoryscanHistory, memoryHash)
109+
// killing process
110+
if pKill {
111+
log.Println("[INFO]", "KILLING PID", proc.PID)
112+
KillProcessByID(proc.PID, pVerbose)
113113
}
114114
}
115+
115116
}
116117

117118
// QuarantineProcess dump process memory and cipher them in quarantine folder
118-
func QuarantineProcess(proc ProcessInformation, quarantinePath string) (err error) {
119+
func QuarantineProcess(proc *ProcessInformation, quarantinePath string) (err error) {
119120

120-
err = quarantineContent(proc.ProcessMemory, proc.ProcessName+fmt.Sprint(proc.PID)+".mem", quarantinePath)
121+
err = quarantineContent(proc.MemoryDump, proc.ProcessName+fmt.Sprint(proc.PID)+".mem", quarantinePath)
121122
if err != nil {
122123
return err
123124
}

‎main.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,8 @@ import (
1717

1818
var (
1919
notificationsHistory []string
20-
filescanHistory []string
21-
memoryscanHistory []string
20+
filescanHistory []FileDescriptor
21+
memoryHashHistory []string
2222
killQueue []string
2323
exit = make(chan bool)
2424
)

‎procsmemory.go‎

Lines changed: 57 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -17,10 +17,10 @@ import (
1717

1818
// ProcessInformation wrap basic process information and memory dump in a structure
1919
type ProcessInformation struct {
20-
PID uint32
21-
ProcessName string
22-
ProcessPath string
23-
ProcessMemory []byte
20+
PID uint32
21+
ProcessName string
22+
ProcessPath string
23+
MemoryDump []byte
2424
}
2525

2626
// MemoryAnalysisRoutine analyse processes memory every 5 seconds
@@ -29,28 +29,34 @@ func MemoryAnalysisRoutine(pDump string, pQuarantine string, pKill bool, pAggres
2929
// list process information and memory
3030
procs := ListProcess(pVerbose)
3131

32-
// dump process memory and quit the program
33-
if len(pDump) > 0 {
34-
for _, proc := range procs {
35-
if err := WriteProcessMemoryToFile(pDump, proc.ProcessName+fmt.Sprint(proc.PID)+".dmp", proc.ProcessMemory); err != nil && pVerbose {
32+
// analyze process memory and executable
33+
for _, proc := range procs {
34+
// dump processes memory and quit the program
35+
if len(pDump) > 0 {
36+
if err := WriteProcessMemoryToFile(pDump, proc.ProcessName+fmt.Sprint(proc.PID)+".dmp", proc.MemoryDump); err != nil && pVerbose {
3637
log.Println("[ERROR]", err)
3738
}
3839
}
39-
os.Exit(0)
40-
}
41-
42-
// analyze process memory and executable
43-
for _, proc := range procs {
4440

4541
// parsing kill queue
4642
if StringInSlice(proc.ProcessPath, killQueue) && pKill {
4743
log.Println("[INFO]", "KILLING PID", proc.PID)
4844
KillProcessByID(proc.PID, pVerbose)
4945
} else {
50-
MemoryAnalysis(proc, pQuarantine, pKill, pAggressive, pNotifications, pVerbose, rules)
46+
// analyzing process memory and cleaning memory buffer
47+
MemoryAnalysis(&proc, pQuarantine, pKill, pAggressive, pNotifications, pVerbose, rules)
48+
proc.MemoryDump = nil
49+
50+
// analyzing process executable
5151
FileAnalysis(proc.ProcessPath, pQuarantine, pKill, pAggressive, pNotifications, pVerbose, rules, "MEMORY")
5252
}
5353
}
54+
55+
if len(pDump) > 0 {
56+
log.Println("[INFO] Processes memory dump completed - Exiting program.")
57+
os.Exit(0)
58+
}
59+
5460
killQueue = nil
5561

5662
time.Sleep(5 * time.Second)
@@ -73,38 +79,53 @@ func ListProcess(verbose bool) (procsInfo []ProcessInformation) {
7379
}
7480

7581
if err == nil && procHandle > 0 {
76-
procFilename, modules, err := GetProcessModulesHandles(procHandle)
77-
if err == nil {
78-
for _, moduleHandle := range modules {
79-
if moduleHandle != 0 {
80-
moduleRawName, err := GetModuleFileNameEx(procHandle, moduleHandle, 512)
81-
if err != nil && verbose {
82-
log.Println("[ERROR]", err)
83-
}
84-
moduleRawName = bytes.Trim(moduleRawName, "\x00")
85-
modulePath := strings.Split(string(moduleRawName), "\\")
86-
moduleFileName := modulePath[len(modulePath)-1]
87-
88-
if procFilename == moduleFileName {
89-
memdump := DumpModuleMemory(procHandle, moduleHandle, verbose)
90-
if len(memdump) > 0 {
91-
proc := ProcessInformation{PID: procsIds[i], ProcessName: procFilename, ProcessPath: string(moduleRawName), ProcessMemory: memdump}
92-
if !StringInSlice(fmt.Sprintf("%x", md5.Sum(memdump)), memoryscanHistory) {
93-
procsInfo = append(procsInfo, proc)
94-
memoryscanHistory = append(memoryscanHistory, fmt.Sprintf("%x", md5.Sum(memdump)))
95-
}
96-
}
97-
}
82+
if proc, memdump, err := GetProcessMemory(procsIds[i], procHandle, verbose); err != nil && verbose {
83+
log.Println("[ERROR]", err)
84+
} else {
85+
if len(memdump) > 0 {
86+
// return process memory only if it has changed since the last process scan
87+
if !StringInSlice(fmt.Sprintf("%x", md5.Sum(memdump)), memoryHashHistory) {
88+
proc.MemoryDump = memdump
89+
procsInfo = append(procsInfo, proc)
90+
memoryHashHistory = append(memoryHashHistory, fmt.Sprintf("%x", md5.Sum(memdump)))
9891
}
9992
}
10093
}
94+
10195
}
10296
windows.CloseHandle(procHandle)
10397
}
10498
}
10599
return procsInfo
106100
}
107101

102+
// GetProcessMemory return a process memory dump based on its handle
103+
func GetProcessMemory(pid uint32, handle windows.Handle, verbose bool) (ProcessInformation, []byte, error) {
104+
105+
procFilename, modules, err := GetProcessModulesHandles(handle)
106+
if err != nil {
107+
return ProcessInformation{}, nil, fmt.Errorf("Unable to get PID %d memory: %s", pid, err.Error())
108+
}
109+
110+
for _, moduleHandle := range modules {
111+
if moduleHandle != 0 {
112+
moduleRawName, err := GetModuleFileNameEx(handle, moduleHandle, 512)
113+
if err != nil {
114+
return ProcessInformation{}, nil, err
115+
}
116+
moduleRawName = bytes.Trim(moduleRawName, "\x00")
117+
modulePath := strings.Split(string(moduleRawName), "\\")
118+
moduleFileName := modulePath[len(modulePath)-1]
119+
120+
if procFilename == moduleFileName {
121+
return ProcessInformation{PID: pid, ProcessName: procFilename, ProcessPath: string(moduleRawName)}, DumpModuleMemory(handle, moduleHandle, verbose), nil
122+
}
123+
}
124+
}
125+
126+
return ProcessInformation{}, nil, fmt.Errorf("Unable to get PID %d memory: no module corresponding to process name", pid)
127+
}
128+
108129
// KillProcessByID try to kill the specified PID
109130
func KillProcessByID(procID uint32, verbose bool) (err error) {
110131
hProc, err := GetProcessHandle(procID, windows.PROCESS_QUERY_INFORMATION|windows.PROCESS_TERMINATE)

‎yaraProcessing.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,8 +10,8 @@ import (
1010
)
1111

1212
// PerformYaraScan use provided YARA rules and search for match in the given byte slice
13-
func PerformYaraScan(data []byte, rules *yara.Rules, verbose bool) yara.MatchRules {
14-
result, err := YaraScan(data, rules)
13+
func PerformYaraScan(data *[]byte, rules *yara.Rules, verbose bool) yara.MatchRules {
14+
result, err := YaraScan(*data, rules)
1515
if err != nil && verbose {
1616
log.Println("[ERROR]", err)
1717
}

0 commit comments

Comments
 (0)