Fixed various yara performance and memory issues

Core/nbproject/project.xml

@@ -561,7 +561,7 @@
             </class-path-extension>
             <class-path-extension>
                 <runtime-relative-path>ext/YaraJNIWrapper.jar</runtime-relative-path>
-                <binary-origin>release/modules/ext/YaraJNIWrapper.jar</binary-origin>
+                <binary-origin>release\modules\ext\YaraJNIWrapper.jar</binary-origin>
             </class-path-extension>
             <class-path-extension>
                 <runtime-relative-path>ext/grpc-context-1.19.0.jar</runtime-relative-path>

Core/src/org/sleuthkit/autopsy/modules/yara/YaraIngestHelper.java

@@ -76,19 +76,43 @@ final class YaraIngestHelper {
      * rule set.
      *
      * @param file                 The file to scan.
-     * @param baseDirectory Base directory for the compiled rule sets.
+     * @param baseRuleSetDirectory Base directory for the compiled rule sets.
      *
      * @throws TskCoreException
      */
-    static List<BlackboardArtifact> scanFileForMatches(AbstractFile file, File baseDirectory) throws TskCoreException, YaraWrapperException {
+    static List<BlackboardArtifact> scanFileForMatches(AbstractFile file, File baseRuleSetDirectory, byte[] fileData, int fileDataSize, int timeout) throws TskCoreException, YaraWrapperException {
         List<BlackboardArtifact> artifacts = new ArrayList<>();
 
-        byte[] fileBytes = new byte[(int) file.getSize()];
+        File[] ruleSetDirectories = baseRuleSetDirectory.listFiles();
-        file.read(fileBytes, 0, fileBytes.length);
-
-        File[] ruleSetDirectories = baseDirectory.listFiles();
         for (File ruleSetDirectory : ruleSetDirectories) {
-            List<String> ruleMatches = YaraIngestHelper.scanFileForMatches(fileBytes, ruleSetDirectory);
+
+            List<String> ruleMatches = YaraIngestHelper.scanFileForMatches(fileData, fileDataSize, ruleSetDirectory, timeout);
+            if (!ruleMatches.isEmpty()) {
+                artifacts.addAll(YaraIngestHelper.createArtifact(file, ruleSetDirectory.getName(), ruleMatches));
+            }
+        }
+
+        return artifacts;
+    }
+
+    /**
+     *
+     * @param file                 The Abstract File being processed.
+     * @param baseRuleSetDirectory Base directory of the compiled rule sets.
+     * @param localFile            Local copy of file.
+     * @param timeout              Yara file scan timeout in seconds.
+     *
+     * @return
+     *
+     * @throws TskCoreException
+     * @throws YaraWrapperException
+     */
+    static List<BlackboardArtifact> scanFileForMatches(AbstractFile file, File baseRuleSetDirectory, File localFile, int timeout) throws TskCoreException, YaraWrapperException {
+        List<BlackboardArtifact> artifacts = new ArrayList<>();
+
+        File[] ruleSetDirectories = baseRuleSetDirectory.listFiles();
+        for (File ruleSetDirectory : ruleSetDirectories) {
+            List<String> ruleMatches = YaraIngestHelper.scanFileForMatch(localFile, ruleSetDirectory, timeout);
             if (!ruleMatches.isEmpty()) {
                 artifacts.addAll(YaraIngestHelper.createArtifact(file, ruleSetDirectory.getName(), ruleMatches));
             }
@@ -109,13 +133,25 @@ final class YaraIngestHelper {
      *
      * @throws TskCoreException
      */
-    private static List<String> scanFileForMatches(byte[] fileBytes, File ruleSetDirectory) throws TskCoreException, YaraWrapperException {
+    private static List<String> scanFileForMatches(byte[] fileBytes, int fileSize, File ruleSetDirectory, int timeout) throws YaraWrapperException {
         List<String> matchingRules = new ArrayList<>();
 
         File[] ruleSetCompiledFileList = ruleSetDirectory.listFiles();
 
         for (File ruleFile : ruleSetCompiledFileList) {
-            matchingRules.addAll(YaraJNIWrapper.findRuleMatch(ruleFile.getAbsolutePath(), fileBytes));
+            matchingRules.addAll(YaraJNIWrapper.findRuleMatch(ruleFile.getAbsolutePath(), fileBytes, fileSize, timeout));
+        }
+
+        return matchingRules;
+    }
+
+    private static List<String> scanFileForMatch(File scanFile, File ruleSetDirectory, int timeout) throws YaraWrapperException {
+        List<String> matchingRules = new ArrayList<>();
+
+        File[] ruleSetCompiledFileList = ruleSetDirectory.listFiles();
+
+        for (File ruleFile : ruleSetCompiledFileList) {
+            matchingRules.addAll(YaraJNIWrapper.findRuleMatchFile(ruleFile.getAbsolutePath(), scanFile.getAbsolutePath(), timeout));
         }
 
         return matchingRules;

Core/src/org/sleuthkit/autopsy/modules/yara/YaraIngestModule.java

@@ -18,8 +18,11 @@
  */
 package org.sleuthkit.autopsy.modules.yara;
 
+import java.io.File;
+import java.io.IOException;
 import java.nio.file.Path;
 import java.nio.file.Paths;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
@@ -30,6 +33,7 @@ import org.sleuthkit.autopsy.casemodule.Case;
 import org.sleuthkit.autopsy.casemodule.NoCurrentCaseException;
 import org.sleuthkit.autopsy.coreutils.Logger;
 import org.sleuthkit.autopsy.coreutils.PlatformUtil;
+import org.sleuthkit.autopsy.datamodel.ContentUtils;
 import org.sleuthkit.autopsy.ingest.FileIngestModuleAdapter;
 import org.sleuthkit.autopsy.ingest.IngestJobContext;
 import org.sleuthkit.autopsy.ingest.IngestModule;
@@ -40,6 +44,7 @@ import org.sleuthkit.datamodel.Blackboard;
 import org.sleuthkit.datamodel.Blackboard.BlackboardException;
 import org.sleuthkit.datamodel.BlackboardArtifact;
 import org.sleuthkit.datamodel.TskCoreException;
+import org.sleuthkit.datamodel.TskData;
 
 /**
  * An ingest module that runs the yara against the given files.
@@ -47,10 +52,16 @@ import org.sleuthkit.datamodel.TskCoreException;
  */
 public class YaraIngestModule extends FileIngestModuleAdapter {
 
+    // 15MB
+    private static final int FILE_SIZE_THRESHOLD_MB = 100;
+    private static final int FILE_SIZE_THRESHOLD_BYTE = FILE_SIZE_THRESHOLD_MB * 1024 * 1024;
+    private static final int YARA_SCAN_TIMEOUT_SEC = 30 * 60 * 60; // 30 minutes.
+    
     private static final IngestModuleReferenceCounter refCounter = new IngestModuleReferenceCounter();
     private final static Logger logger = Logger.getLogger(YaraIngestModule.class.getName());
     private static final String YARA_DIR = "yara";
     private static final Map<Long, Path> pathsByJobId = new ConcurrentHashMap<>();
+    private static final String RULESET_DIR = "RuleSets";
 
     private final YaraIngestJobSettings settings;
 
@@ -81,8 +92,12 @@ public class YaraIngestModule extends FileIngestModuleAdapter {
         if (refCounter.incrementAndGet(jobId) == 1) {
             // compile the selected rules & put into temp folder based on jobID
             Path tempDir = getTempDirectory(jobId);
+            Path tempRuleSetDir = Paths.get(tempDir.toString(), RULESET_DIR);
+            if(!tempRuleSetDir.toFile().exists()) {
+                tempRuleSetDir.toFile().mkdir();
+            }
 
-            YaraIngestHelper.compileRules(settings.getSelectedRuleSetNames(), tempDir);
+            YaraIngestHelper.compileRules(settings.getSelectedRuleSetNames(), tempRuleSetDir);
         }
     }
 
@@ -108,13 +123,31 @@ public class YaraIngestModule extends FileIngestModuleAdapter {
             }
         }
 
-        // Skip the file if its 0 in length.
+        // Skip the file if its 0 in length or a directory.
-        if (file.getSize() == 0) {
+        if (file.getSize() == 0 || 
+                file.isDir() || 
+                file.getType() == TskData.TSK_DB_FILES_TYPE_ENUM.UNALLOC_BLOCKS) {
             return ProcessResult.OK;
         }
 
         try {
-            List<BlackboardArtifact> artifacts = YaraIngestHelper.scanFileForMatches(file, getTempDirectory(jobId).toFile());
+            List<BlackboardArtifact> artifacts = new ArrayList<>();
+            File ruleSetsDir = Paths.get(getTempDirectory(jobId).toString(), RULESET_DIR).toFile();
+            
+            // If the file size is less than FILE_SIZE_THRESHOLD_BYTE read the file
+            // into a buffer, else make a local copy of the file.
+            if(file.getSize() < FILE_SIZE_THRESHOLD_BYTE) {
+                byte[] fileBuffer = new byte[(int)file.getSize()];
+                
+                int dataRead = file.read(fileBuffer, 0, file.getSize());
+                if(dataRead != 0) {
+                    artifacts.addAll( YaraIngestHelper.scanFileForMatches(file, ruleSetsDir, fileBuffer, dataRead, YARA_SCAN_TIMEOUT_SEC));
+                } 
+            } else {
+                File tempCopy = createLocalCopy(file);
+                artifacts.addAll( YaraIngestHelper.scanFileForMatches(file, ruleSetsDir, tempCopy, YARA_SCAN_TIMEOUT_SEC));
+                tempCopy.delete();
+            }
             
             if(!artifacts.isEmpty()) {
                 Blackboard blackboard = Case.getCurrentCaseThrows().getSleuthkitCase().getBlackboard();
@@ -122,9 +155,13 @@ public class YaraIngestModule extends FileIngestModuleAdapter {
             }
             
         } catch (BlackboardException | NoCurrentCaseException | IngestModuleException | TskCoreException | YaraWrapperException ex) {
-            logger.log(Level.SEVERE, "YARA ingest module failed to process file.", ex);
+            logger.log(Level.SEVERE, String.format("YARA ingest module failed to process file id %d", file.getId()), ex);
+            return ProcessResult.ERROR;
+        } catch(IOException ex) {
+            logger.log(Level.SEVERE, String.format("YARA ingest module failed to make a local copy of given file id %d", file.getId()), ex);
             return ProcessResult.ERROR;
         }
+        
         return ProcessResult.OK;
     }
 
@@ -165,4 +202,23 @@ public class YaraIngestModule extends FileIngestModuleAdapter {
         return jobPath;
     }
     
+    /**
+     * Create a local copy of the given AbstractFile.
+     * 
+     * @param file AbstractFile to make a copy of.
+     * 
+     * @return A File object representation of the local copy.
+     * 
+     * @throws org.sleuthkit.autopsy.ingest.IngestModule.IngestModuleException
+     * @throws IOException 
+     */
+    protected File createLocalCopy(AbstractFile file) throws IngestModuleException, IOException {
+        String tempFileName = RandomStringUtils.randomAlphabetic(15) + file.getId() + ".temp";
+        
+        File tempFile = Paths.get(getTempDirectory(context.getJobId()).toString(), tempFileName).toFile();
+        ContentUtils.writeToFile(file, tempFile, context::dataSourceIngestIsCancelled);
+        
+        return tempFile;
+    }
+
 }

thirdparty/yara/YaraJNIWrapper/nbproject/private/private.xml

@@ -1,4 +1,9 @@
 <?xml version="1.0" encoding="UTF-8"?>
 <project-private xmlns="http://www.netbeans.org/ns/project-private/1">
     <editor-bookmarks xmlns="http://www.netbeans.org/ns/editor-bookmarks/2" lastBookmarkId="0"/>
+    <open-files xmlns="http://www.netbeans.org/ns/projectui-open-files/2">
+        <group>
+            <file>file:/C:/Users/kelly/Workspace/autopsy/thirdparty/yara/YaraJNIWrapper/src/org/sleuthkit/autopsy/yara/YaraJNIWrapper.java</file>
+        </group>
+    </open-files>
 </project-private>

thirdparty/yara/YaraJNIWrapper/src/org/sleuthkit/autopsy/yara/YaraJNIWrapper.java

@@ -46,14 +46,30 @@ public class YaraJNIWrapper {
      *
      * The rule path must be to a yara compile rule file.
      *
-     * @param compiledRulesPath
+     * @param compiledRulesPath Absolute path to a compiled YARA rule file.
-     * @param byteBuffer
+     * @param byteBuffer        File buffer.
+     * @param bufferSize        Size of the byte to read in the given buffer
+     * @param timeoutSec        Scan timeout value in seconds.
      *
      * @return List of rules found rules. Null maybe returned if error occurred.
      *
      * @throws YaraWrapperException
      */
-    static public native List<String> findRuleMatch(String compiledRulesPath, byte[] byteBuffer) throws YaraWrapperException;
+    static public native List<String> findRuleMatch(String compiledRulesPath, byte[] byteBuffer, int bufferSize, int timeoutSec) throws YaraWrapperException;
+
+    /**
+     * Returns a list of matching YARA rules found in the given file.
+     *
+     * @param compiledRulePath Absolute path to a compiled YARA rule file.
+     * @param filePath         Absolute path to the file to search.
+     * @param timeoutSec       Scan timeout value in seconds.
+     *
+     * @return List of rules found rules. Null maybe returned if error occurred.
+     *
+     *
+     * @throws YaraWrapperException
+     */
+    static public native List<String> findRuleMatchFile(String compiledRulePath, String filePath, int timeoutSec) throws YaraWrapperException;
 
     /**
      * Copy yarabridge.dll from inside the jar to a temp file that can be loaded

thirdparty/yara/YaraWrapperTest/src/org/sleuthkit/autopsy/yara/YaraWrapperTest.java

@@ -26,8 +26,6 @@ import java.util.List;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
-import org.sleuthkit.autopsy.yara.YaraJNIWrapper;
-import org.sleuthkit.autopsy.yara.YaraWrapperException;
 
 /**
  * Tests the YaraJNIWrapper code.
@@ -43,6 +41,7 @@ public class YaraWrapperTest {
         }
 
         testFileRuleMatch(args[0], args[1]);
+        testFileRuleMatchFile(args[0], args[1]);
     }
 
     /**
@@ -58,7 +57,7 @@ public class YaraWrapperTest {
         try {
             byte[] data = Files.readAllBytes(path);
 
-            List<String> list = YaraJNIWrapper.findRuleMatch(compiledRulePath, data);
+            List<String> list = YaraJNIWrapper.findRuleMatch(compiledRulePath, data, data.length, 100);
 
             if (list != null) {
                 if (list.isEmpty()) {
@@ -78,4 +77,33 @@ public class YaraWrapperTest {
         }
     }
     
+    /**
+     * Test the call to findRuleMatchFile which takes a compiled rule file
+     * path and a path to a file.
+     * 
+     * @param compiledRulePath
+     * @param filePath 
+     */
+    private static void testFileRuleMatchFile(String compiledRulePath, String filePath) {
+        try {
+            List<String> list = YaraJNIWrapper.findRuleMatchFile(compiledRulePath, filePath, 100);
+
+            if (list != null) {
+                if (list.isEmpty()) {
+                    System.out.println("FindRuleMatch return an empty list");
+                } else {
+                    System.out.println("Matching Rules:");
+                    for (String s : list) {
+                        System.out.println(s);
+                    }
+                }
+            } else {
+                logger.log(Level.SEVERE, "FindRuleMatch return a null list");
+            }
+
+        } catch (YaraWrapperException ex) {
+            logger.log(Level.SEVERE, "Error thrown from yarabridge", ex);
+        }
+    }
+
 }

thirdparty/yara/yarabridge/yarabridge/YaraJNIWrapper.cpp

@@ -20,7 +20,6 @@
 using std::string;
 using std::vector;
 
-
 /*
 	Callback method to be passed to yr_rules_scan_mem method.
 	user_data is expected to be a pointer to a string vector.
@@ -79,49 +78,85 @@ jobject createArrayList(JNIEnv *env, std::vector<std::string> vector) {
 	return list;
 }
 
+/*
+	Loads the compiled rules file returning a YARA error code.
+	Throws a java exeception if there are any issues.
+*/
+int loadRuleFile(JNIEnv * env, jstring compiledRulePath, YR_RULES **rules) {
+	char errorMessage[256];
+	const char *nativeString = env->GetStringUTFChars(compiledRulePath, 0);
+	int result = yr_rules_load(nativeString, rules);
+
+	if (result != ERROR_SUCCESS) {
+		sprintf_s(errorMessage, "Failed to load compiled yara rule %s (error code = %d)\n", nativeString, result);
+		throwException(env, errorMessage);
+	}
+
+	env->ReleaseStringUTFChars(compiledRulePath, nativeString);
+
+	return result;
+}
+
+/*
+	Initalize the YARA library, if needed. yr_initialize only needs to be called once.
+*/
+int initalizeYaraLibrary(JNIEnv * env) {
+	static int library_initalized = 0;
+	char errorMessage[256];
+	int result = ERROR_SUCCESS;
+	if (library_initalized == 0) {
+		if ((result = yr_initialize()) != ERROR_SUCCESS) {
+			sprintf_s(errorMessage, "libyara initialization error (%d)\n", result);
+			throwException(env, errorMessage);
+		}
+		library_initalized = 1;
+	}
+
+	return result;
+}
+
 /*
 * Class:     org_sleuthkit_autopsy_yara_YaraJNIWrapper
 * Method:    FindRuleMatch
 * Signature: (Ljava/lang/String;[B)Ljava/util/List;
 */
 JNIEXPORT jobject JNICALL Java_org_sleuthkit_autopsy_yara_YaraJNIWrapper_findRuleMatch
-(JNIEnv * env, jclass cls, jstring compiledRulePath, jbyteArray fileByteArray) {
+(JNIEnv * env, jclass cls, jstring compiledRulePath, jbyteArray fileByteArray, jint byteArrayLength, jint timeoutSec) {
 	
 	char errorMessage[256];
-	const char *nativeString = env->GetStringUTFChars(compiledRulePath, 0);
 	jobject resultList = NULL;
-	
 	int result;
-	if ((result = yr_initialize()) != ERROR_SUCCESS) {
+	YR_RULES *rules = NULL;
-		sprintf_s(errorMessage, "libyara initialization error (%d)\n", result);
+
-		throwException(env, errorMessage);
+	if ((result = initalizeYaraLibrary(env)) != ERROR_SUCCESS) {
 		return resultList;
 	}
 
+	
 	while (1) {
-		YR_RULES *rules = NULL;
+		if((result = loadRuleFile(env, compiledRulePath, &rules)) != ERROR_SUCCESS) {
-		if ((result = yr_rules_load(nativeString, &rules)) != ERROR_SUCCESS) {
-			sprintf_s(errorMessage, "Failed to load compiled yara rules (%d)\n", result);
-			throwException(env, errorMessage);
 			break;
 		}
 
-		boolean isCopy;
-		int byteArrayLength = env->GetArrayLength(fileByteArray);
 		if (byteArrayLength == 0) {
 			throwException(env, "Unable to scan for matches. File byte array size was 0.");
 			break;
 		}
 
+		boolean isCopy;
 		jbyte* nativeByteArray = env->GetByteArrayElements(fileByteArray, &isCopy);
-		int flags = 0;
 		std::vector<std::string> scanResults;
 
-		result = yr_rules_scan_mem(rules, (unsigned char*)nativeByteArray, byteArrayLength, flags, callback, &scanResults, 1000000);
+		result = yr_rules_scan_mem(rules, (unsigned char*)nativeByteArray, byteArrayLength, 0, callback, &scanResults, timeoutSec);
 		env->ReleaseByteArrayElements(fileByteArray, nativeByteArray, 0);
 
 		if (result != ERROR_SUCCESS) {
+			if (result == ERROR_SCAN_TIMEOUT) {
+				sprintf_s(errorMessage, "Yara file scan timed out");
+			}
+			else {
 				sprintf_s(errorMessage, "Yara file scan failed (%d)\n", result);
+			}
 			throwException(env, errorMessage);
 			break;
 		}
@@ -130,9 +165,60 @@ JNIEXPORT jobject JNICALL Java_org_sleuthkit_autopsy_yara_YaraJNIWrapper_findRul
 		break;
 	}
 
-	env->ReleaseStringUTFChars(compiledRulePath, nativeString);
+	if (rules != NULL) {
-	yr_finalize();
+		yr_rules_destroy(rules);
+	}
 
 	return resultList;
 
 }
+
+/*
+* Class:     org_sleuthkit_autopsy_yara_YaraJNIWrapper
+* Method:    findRuleMatchFile
+* Signature: (Ljava/lang/String;Ljava/lang/String;)Ljava/util/List;
+*/
+JNIEXPORT jobject JNICALL Java_org_sleuthkit_autopsy_yara_YaraJNIWrapper_findRuleMatchFile
+(JNIEnv * env, jclass cls, jstring compiledRulePath, jstring filePath, jint timeoutSec) {
+
+	char errorMessage[256];
+	jobject resultList = NULL;
+	int result;
+	YR_RULES *rules = NULL;
+
+	if ((result = initalizeYaraLibrary(env)) != ERROR_SUCCESS) {
+		return resultList;
+	}
+
+
+	while (1) {
+		if ((result = loadRuleFile(env, compiledRulePath, &rules)) != ERROR_SUCCESS) {
+			break;
+		}
+
+		std::vector<std::string> scanResults;
+		const char *nativeString = env->GetStringUTFChars(filePath, 0);
+
+		result = yr_rules_scan_file(rules, nativeString, 0, callback, &scanResults, timeoutSec);
+
+		if (result != ERROR_SUCCESS) {
+			if (result == ERROR_SCAN_TIMEOUT) {
+				sprintf_s(errorMessage, "Yara file scan timed out on file %s", nativeString);
+			}
+			else {
+				sprintf_s(errorMessage, "Yara file scan failed (%d)\n", result);
+			}
+			throwException(env, errorMessage);
+			break;
+		}
+
+		resultList = createArrayList(env, scanResults);
+		break;
+	}
+
+	if (rules != NULL) {
+		yr_rules_destroy(rules);
+	}
+
+	return resultList;
+}

thirdparty/yara/yarabridge/yarabridge/YaraJNIWrapper.h

@@ -13,7 +13,15 @@ extern "C" {
 	* Signature: (Ljava/lang/String;[B)Ljava/util/List;
 	*/
 	JNIEXPORT jobject JNICALL Java_org_sleuthkit_autopsy_yara_YaraJNIWrapper_findRuleMatch
-	(JNIEnv *, jclass, jstring, jbyteArray);
+	(JNIEnv *, jclass, jstring, jbyteArray, jint, jint);
+
+	/*
+	* Class:     org_sleuthkit_autopsy_yara_YaraJNIWrapper
+	* Method:    findRuleMatchFile
+	* Signature: (Ljava/lang/String;Ljava/lang/String;)Ljava/util/List;
+	*/
+	JNIEXPORT jobject JNICALL Java_org_sleuthkit_autopsy_yara_YaraJNIWrapper_findRuleMatchFile
+	(JNIEnv *, jclass, jstring, jstring, jint);
 
 #ifdef __cplusplus
 }