diff --git a/Core/src/com/basistech/df/cybertriage/autopsy/malwarescan/MalwareScanIngestModule.java b/Core/src/com/basistech/df/cybertriage/autopsy/malwarescan/MalwareScanIngestModule.java index c982220e8b..8db2ba70e8 100644 --- a/Core/src/com/basistech/df/cybertriage/autopsy/malwarescan/MalwareScanIngestModule.java +++ b/Core/src/com/basistech/df/cybertriage/autopsy/malwarescan/MalwareScanIngestModule.java @@ -23,15 +23,12 @@ import com.basistech.df.cybertriage.autopsy.ctapi.CTCloudException; import com.basistech.df.cybertriage.autopsy.ctapi.json.AuthTokenResponse; import com.basistech.df.cybertriage.autopsy.ctapi.json.AuthenticatedRequestData; import com.basistech.df.cybertriage.autopsy.ctapi.json.CTCloudBean; -import com.basistech.df.cybertriage.autopsy.ctapi.json.DecryptedLicenseResponse; import com.basistech.df.cybertriage.autopsy.ctapi.json.LicenseInfo; -import com.basistech.df.cybertriage.autopsy.ctapi.json.MalwareResultBean; import com.basistech.df.cybertriage.autopsy.ctapi.json.MalwareResultBean.Status; import com.basistech.df.cybertriage.autopsy.ctapi.json.MetadataUploadRequest; import com.basistech.df.cybertriage.autopsy.ctoptions.ctcloud.CTLicensePersistence; import java.text.MessageFormat; import java.util.ArrayList; -import java.util.Collection; import java.util.Collections; import java.util.HashMap; import java.util.List; @@ -42,12 +39,11 @@ import java.util.logging.Level; import java.util.stream.Collectors; import java.util.stream.Stream; import org.apache.commons.collections4.CollectionUtils; +import org.apache.commons.collections4.MapUtils; import org.apache.commons.lang3.StringUtils; -import org.apache.commons.lang3.tuple.Pair; import org.apache.curator.shaded.com.google.common.collect.Lists; import org.openide.util.NbBundle.Messages; import org.sleuthkit.autopsy.casemodule.Case; -import org.sleuthkit.autopsy.casemodule.NoCurrentCaseException; import org.sleuthkit.autopsy.coreutils.Logger; import org.sleuthkit.autopsy.coreutils.MessageNotifyUtil; import org.sleuthkit.autopsy.ingest.FileIngestModule; @@ -62,7 +58,6 @@ import org.sleuthkit.datamodel.ReadContentInputStream; import org.sleuthkit.datamodel.HashUtility; import org.sleuthkit.datamodel.HashUtility.HashResult; import org.sleuthkit.datamodel.HashUtility.HashType; -import org.sleuthkit.datamodel.IngestJobInfo; import org.sleuthkit.datamodel.Score; import org.sleuthkit.datamodel.SleuthkitCase; import org.sleuthkit.datamodel.TskCoreException; @@ -109,7 +104,7 @@ public class MalwareScanIngestModule implements FileIngestModule { private static final long MIN_UPLOAD_SIZE = 1; private static final long MAX_UPLOAD_SIZE = 1_000_000_000; - private static final int NUM_FILE_UPLOAD_RETRIES = 60 * 5; + private static final int NUM_FILE_UPLOAD_RETRIES = 7; private static final long FILE_UPLOAD_RETRY_SLEEP_MILLIS = 60 * 1000; private static final Set EXECUTABLE_MIME_TYPES = Stream.of( @@ -131,13 +126,16 @@ public class MalwareScanIngestModule implements FileIngestModule { private static final Logger logger = Logger.getLogger(MalwareScanIngestModule.class.getName()); - private final BatchProcessor batchProcessor = new BatchProcessor(BATCH_SIZE, FLUSH_SECS_TIMEOUT, this::handleBatch); + private final BatchProcessor batchProcessor = new BatchProcessor( + BATCH_SIZE, + FLUSH_SECS_TIMEOUT, + (lst) -> SharedProcessing.this.handleBatch(SharedProcessing.this.ingestJobState, lst)); + private final CTLicensePersistence ctSettingsPersistence = CTLicensePersistence.getInstance(); private final CTApiDAO ctApiDAO = CTApiDAO.getInstance(); - // TODO minimize state -// private RunState runState = null; -// private IngestJobState ingestJobState = null; + private IngestJobState ingestJobState = null; + @Messages({ "MalwareScanIngestModule_malwareTypeDisplayName=Malware", "MalwareScanIngestModule_ShareProcessing_noLicense_title=No Cyber Triage License", @@ -154,16 +152,14 @@ public class MalwareScanIngestModule implements FileIngestModule { "MalwareScanIngestModule_ShareProcessing_lowUploadsLimitWarning_desc=This license only has {0} file uploads remaining.",}) synchronized void startUp(IngestJobContext context) throws IngestModuleException { // only run this code once per startup - if (runState == RunState.STARTED_UP || runState == RunState.DISABLED) { + if (ingestJobState != null) { return; } try { - Pair jobStateResult = getNewJobState(context); - runState = jobStateResult.getLeft(); - ingestJobState = jobStateResult.getRight(); + ingestJobState = getNewJobState(context); } catch (Exception ex) { - runState = RunState.DISABLED; + ingestJobState = IngestJobState.DISABLED; throw new IngestModuleException("An exception occurred on MalwareScanIngestModule startup", ex); } } @@ -176,7 +172,7 @@ public class MalwareScanIngestModule implements FileIngestModule { * parameters required for the job. * @throws Exception */ - private Pair getNewJobState(IngestJobContext context) throws Exception { + private IngestJobState getNewJobState(IngestJobContext context) throws Exception { // get saved license Optional licenseInfoOpt = ctSettingsPersistence.loadLicenseInfo(); if (licenseInfoOpt.isEmpty() || licenseInfoOpt.get().getDecryptedLicense() == null) { @@ -184,7 +180,8 @@ public class MalwareScanIngestModule implements FileIngestModule { Bundle.MalwareScanIngestModule_ShareProcessing_noLicense_title(), Bundle.MalwareScanIngestModule_ShareProcessing_noLicense_desc(), null); - return Pair.of(RunState.DISABLED, null); + + return IngestJobState.DISABLED; } AuthTokenResponse authTokenResponse = ctApiDAO.getAuthToken(licenseInfoOpt.get().getDecryptedLicense()); @@ -197,7 +194,8 @@ public class MalwareScanIngestModule implements FileIngestModule { Bundle.MalwareScanIngestModule_ShareProcessing_noLookupsRemaining_title(), Bundle.MalwareScanIngestModule_ShareProcessing_noLookupsRemaining_desc(), null); - return Pair.of(RunState.DISABLED, null); + + return IngestJobState.DISABLED; } else if (lookupsRemaining < LOW_LOOKUPS_REMAINING) { notifyWarning( Bundle.MalwareScanIngestModule_ShareProcessing_lowLookupsLimitWarning_title(), @@ -230,16 +228,16 @@ public class MalwareScanIngestModule implements FileIngestModule { Bundle.MalwareScanIngestModule_malwareTypeDisplayName(), BlackboardArtifact.Category.ANALYSIS_RESULT); - IngestJobState ingestJobState = new IngestJobState( + return new IngestJobState( tskCase, new FileTypeDetector(), licenseInfoOpt.get(), malwareType, context.getDataSource().getId(), context.getJobId(), - uploadFiles); - - return Pair.of(RunState.STARTED_UP, ingestJobState); + uploadFiles, + true + ); } /** @@ -286,17 +284,26 @@ public class MalwareScanIngestModule implements FileIngestModule { return null; } - // TODO return to this + /** + * Processes a file. The file goes through the lookup process if the + * file meets acceptable criteria: 1) not FileKnown.KNOWN 2) is + * executable 3) does not have any pre-existing TSK_MALWARE results 4) + * file lookup has not been disabled. + * + * @param af The file. + * @return OK or ERROR. + */ @Messages({ "MalwareScanIngestModule_ShareProcessing_batchTimeout_title=Batch Processing Timeout", "MalwareScanIngestModule_ShareProcessing_batchTimeout_desc=Batch processing timed out" }) IngestModule.ProcessResult process(AbstractFile af) { try { - if (runState == RunState.STARTED_UP + if (ingestJobState != null + && ingestJobState.isDoFileLookups() && af.getKnown() != TskData.FileKnown.KNOWN - && EXECUTABLE_MIME_TYPES.contains(StringUtils.defaultString(fileTypeDetector.getMIMEType(af)).trim().toLowerCase()) - && CollectionUtils.isEmpty(af.getAnalysisResults(malwareType))) { + && EXECUTABLE_MIME_TYPES.contains(StringUtils.defaultString(ingestJobState.getFileTypeDetector().getMIMEType(af)).trim().toLowerCase()) + && CollectionUtils.isEmpty(af.getAnalysisResults(ingestJobState.getMalwareType()))) { String md5 = getOrCalcHash(af); if (StringUtils.isNotBlank(md5)) { @@ -319,7 +326,13 @@ public class MalwareScanIngestModule implements FileIngestModule { } } - // TODO return to this + /** + * Handles a batch of files to be sent to CT file lookup for results. + * + * @param ingestJobState The current state of operation for the ingest + * job. + * @param fileRecords The file records to be uploaded. + */ @Messages({ "MalwareScanIngestModule_SharedProcessing_authTokenResponseError_title=Authentication API error", "# {0} - errorResponse", @@ -331,8 +344,8 @@ public class MalwareScanIngestModule implements FileIngestModule { "MalwareScanIngestModule_SharedProcessing_exhaustedHashLookups_desc=The remaining hash lookups for this license have been exhausted", "MalwareScanIngestModule_SharedProcessing_generalProcessingError_title=Hash Lookup Error", "MalwareScanIngestModule_SharedProcessing_generalProcessingError_desc=An error occurred while processing hash lookup results",}) - private void handleBatch(List fileRecords) { - if (runState != RunState.STARTED_UP || fileRecords == null || fileRecords.isEmpty()) { + private void handleBatch(IngestJobState ingestJobState, List fileRecords) { + if (ingestJobState == null || !ingestJobState.isDoFileLookups() || fileRecords == null || fileRecords.isEmpty()) { return; } @@ -348,7 +361,6 @@ public class MalwareScanIngestModule implements FileIngestModule { md5ToObjId .computeIfAbsent(sanitizedMd5, (k) -> new ArrayList<>()) .add(fr.getObjId()); - } List md5Hashes = new ArrayList<>(md5ToObjId.keySet()); @@ -358,36 +370,8 @@ public class MalwareScanIngestModule implements FileIngestModule { } try { - List repResult = getHashLookupResults(md5Hashes); - Map> statusGroupings = repResult.stream() - .filter(bean -> bean.getMalwareResult() != null) - .collect(Collectors.groupingBy(bean -> bean.getMalwareResult().getStatus())); - - List found = statusGroupings.get(Status.FOUND); - createArtifacts(found, md5ToObjId); - - // if being scanned, check list to run later - List beingScannedList = statusGroupings.get(Status.BEING_SCANNED); - processMissing(beingScannedList, md5ToObjId, false); - - // if not found, try upload - List notFound = statusGroupings.get(Status.NOT_FOUND); - processMissing(notFound, md5ToObjId, true); - - if (CollectionUtils.isNotEmpty(statusGroupings.get(Status.ERROR))) { - notifyWarning( - Bundle.MalwareScanIngestModule_SharedProcessing_generalProcessingError_title(), - Bundle.MalwareScanIngestModule_SharedProcessing_generalProcessingError_desc(), - null); - } - - if (CollectionUtils.isNotEmpty(statusGroupings.get(Status.LIMITS_EXCEEDED))) { - notifyWarning( - Bundle.MalwareScanIngestModule_SharedProcessing_exhaustedHashLookups_title(), - Bundle.MalwareScanIngestModule_SharedProcessing_exhaustedHashLookups_desc(), - null); - } - + List repResult = getHashLookupResults(ingestJobState, md5Hashes); + handleLookupResults(ingestJobState, md5ToObjId, repResult); } catch (Exception ex) { notifyWarning( Bundle.MalwareScanIngestModule_SharedProcessing_generalProcessingError_title(), @@ -396,35 +380,110 @@ public class MalwareScanIngestModule implements FileIngestModule { } } - // TODO return to this - private void processMissing(Collection results, Map> md5ToObjId, boolean doFileUpload) throws CTCloudException, TskCoreException { - for (CTCloudBean beingScanned : CollectionUtils.emptyIfNull(results)) { + /** + * Handles results received from CT Cloud. + * + * @param ingestJobState The current state of operations of the ingest + * module. + * @param md5ToObjId The mapping of md5 to a list of object ids. + * @param repResult The ct cloud results. + * @throws org.sleuthkit.datamodel.Blackboard.BlackboardException + * @throws TskCoreException + * @throws TskCoreException + */ + @Messages({ + "MalwareScanIngestModule_SharedProcessing_exhaustedResultsHashLookups_title=Some Lookup Results Not Processed", + "MalwareScanIngestModule_SharedProcessing_exhaustedResultsHashLookups_desc=Some lookup results were not processed due to exceeding limits. Please try again later.",}) + private void handleLookupResults(IngestJobState ingestJobState, Map> md5ToObjId, List repResult) throws Blackboard.BlackboardException, TskCoreException, TskCoreException, CTCloudException { + if (CollectionUtils.isEmpty(repResult)) { + return; + } - String sanitizedMd5 = normalizedMd5(beingScanned.getMd5HashValue()); - if (StringUtils.isBlank(sanitizedMd5)) { - continue; - } - List correspondingObjIds = md5ToObjId.get(sanitizedMd5); - if (CollectionUtils.isEmpty(correspondingObjIds)) { - continue; - } + Map> statusGroupings = repResult.stream() + .filter(bean -> bean.getMalwareResult() != null) + .collect(Collectors.groupingBy(bean -> bean.getMalwareResult().getStatus())); - if (doFileUpload) { - uploadFile(beingScanned, correspondingObjIds.get(0)); - } - this.unidentifiedHashes.put(sanitizedMd5, correspondingObjIds); + // for all found items, create analysis results + List found = statusGroupings.get(Status.FOUND); + createAnalysisResults(ingestJobState, found, md5ToObjId); + + // if being scanned, check list to run later + handleNonFoundResults(ingestJobState, md5ToObjId, statusGroupings.get(Status.BEING_SCANNED), false); + + // if not found, try upload + handleNonFoundResults(ingestJobState, md5ToObjId, statusGroupings.get(Status.NOT_FOUND), true); + + // indicate a general error if some result in an error + if (CollectionUtils.isNotEmpty(statusGroupings.get(Status.ERROR))) { + notifyWarning( + Bundle.MalwareScanIngestModule_SharedProcessing_generalProcessingError_title(), + Bundle.MalwareScanIngestModule_SharedProcessing_generalProcessingError_desc(), + null); + } + + // indicate some results were not processed if limits exceeded in results + if (CollectionUtils.isNotEmpty(statusGroupings.get(Status.LIMITS_EXCEEDED))) { + notifyWarning( + Bundle.MalwareScanIngestModule_SharedProcessing_exhaustedResultsHashLookups_title(), + Bundle.MalwareScanIngestModule_SharedProcessing_exhaustedResultsHashLookups_desc(), + null); } } - // TODO return to this - private List getHashLookupResults(List md5Hashes) throws CTCloudException { + /** + * Handles a CT cloud response objects that have a status that isn't + * FOUND but still are queryable (i.e. NOT_FOUND, BEING_SCANNED). + * + * @param ingestJobState The current state of operations of the ingest + * module. + * @param md5ToObjId The mapping of md5 to a list of object ids. + * @param results The ct cloud results. + * @param performFileUpload True if the class of results warrants file + * upload (i.e. NOT_FOUND) + */ + private void handleNonFoundResults(IngestJobState ingestJobState, Map> md5ToObjId, List results, boolean performFileUpload) throws CTCloudException, TskCoreException { + if (CollectionUtils.isNotEmpty(results) + && ingestJobState.isDoFileLookups() + && ((performFileUpload && ingestJobState.isUploadUnknownFiles()) || (!performFileUpload && ingestJobState.isQueryForMissing()))) { + + for (CTCloudBean beingScanned : CollectionUtils.emptyIfNull(results)) { + + String sanitizedMd5 = normalizedMd5(beingScanned.getMd5HashValue()); + if (StringUtils.isBlank(sanitizedMd5)) { + continue; + } + List correspondingObjIds = md5ToObjId.get(sanitizedMd5); + if (CollectionUtils.isEmpty(correspondingObjIds)) { + continue; + } + + if (performFileUpload) { + uploadFile(ingestJobState, correspondingObjIds.get(0)); + } + + ingestJobState.getUnidentifiedHashes().put(sanitizedMd5, correspondingObjIds); + } + } + } + + /** + * Makes CT Cloud REST API query for results regarding the status of a + * list of md5 hashes for executables. + * + * @param ingestJobState The current state of operations of the ingest + * module. + * @param md5Hashes The md5 hashes to check. + * @return The results from CT Cloud. + * @throws CTCloudException + */ + private List getHashLookupResults(IngestJobState ingestJobState, List md5Hashes) throws CTCloudException { // get an auth token with the license - AuthTokenResponse authTokenResponse = ctApiDAO.getAuthToken(licenseInfo.getDecryptedLicense()); + AuthTokenResponse authTokenResponse = ctApiDAO.getAuthToken(ingestJobState.getLicenseInfo().getDecryptedLicense()); // make sure we are in bounds for the remaining scans long remainingScans = remaining(authTokenResponse.getHashLookupLimit(), authTokenResponse.getHashLookupCount()); if (remainingScans <= 0) { - runState = RunState.DISABLED; + ingestJobState.disableDoFileLookups(); notifyWarning( Bundle.MalwareScanIngestModule_SharedProcessing_exhaustedHashLookups_title(), Bundle.MalwareScanIngestModule_SharedProcessing_exhaustedHashLookups_desc(), @@ -434,13 +493,14 @@ public class MalwareScanIngestModule implements FileIngestModule { // using auth token, get results return ctApiDAO.getReputationResults( - new AuthenticatedRequestData(licenseInfo.getDecryptedLicense(), authTokenResponse), + new AuthenticatedRequestData(ingestJobState.getLicenseInfo().getDecryptedLicense(), authTokenResponse), md5Hashes ); } /** * Normalizes an md5 string for the purposes of lookup in a map. + * * @param orig The original value. * @return The normalized value */ @@ -448,19 +508,6 @@ public class MalwareScanIngestModule implements FileIngestModule { return StringUtils.defaultString(orig).trim().toLowerCase(); } - /** - * Returns true if the cloud bean indicates that the file is not - * currently known by CT cloud. - * - * @param cloudBean The cloud bean. - * @return True if not known by CT cloud. - */ - private static boolean isUnknown(CTCloudBean cloudBean) { - return cloudBean != null - && cloudBean.getMalwareResult() != null - && cloudBean.getMalwareResult().getStatus() == MalwareResultBean.Status.NOT_FOUND; - } - /** * Whether or not an abstract file meets the requirements to be * uploaded. @@ -473,22 +520,22 @@ public class MalwareScanIngestModule implements FileIngestModule { return size >= MIN_UPLOAD_SIZE && size <= MAX_UPLOAD_SIZE; } - // TODO return to this - private boolean uploadFile(IngestJobState ingestJobState, CTCloudBean cloudBean, long objId) throws CTCloudException, TskCoreException { - if (!ingestJobState.uploadUnknownFiles()) { - return false; - } - - if (!isUnknown(cloudBean)) { + /** + * Uploads a file to CT Cloud if the file is valid for upload. + * + * @param ingestJobState The current state of the ingest job. + * @param objId The object id of the file to upload to CT cloud. + * @return True if successfully uploaded. + * @throws CTCloudException + * @throws TskCoreException + */ + private boolean uploadFile(IngestJobState ingestJobState, long objId) throws CTCloudException, TskCoreException { + if (!ingestJobState.isUploadUnknownFiles()) { return false; } AbstractFile af = ingestJobState.getTskCase().getAbstractFileById(objId); - if (af == null) { - return false; - } - - if (!isUploadable(af)) { + if (af == null || !isUploadable(af)) { return false; } @@ -498,7 +545,7 @@ public class MalwareScanIngestModule implements FileIngestModule { throw new CTCloudException(CTCloudException.ErrorCode.NETWORK_ERROR); } else if (remaining(authTokenResponse.getFileUploadLimit(), authTokenResponse.getFileUploadCount()) <= 0) { // don't proceed with upload if reached limit - ingestJobState.setUploadUnknownFiles(false); + ingestJobState.disableUploadUnknownFiles(); return false; } @@ -520,25 +567,58 @@ public class MalwareScanIngestModule implements FileIngestModule { return true; } - // TODO return to this - private boolean getUploadedFileResults(Map> md5objIdMapping) throws InterruptedException, CTCloudException, Blackboard.BlackboardException, TskCoreException { - Map> remaining = new HashMap<>(md5objIdMapping); + /** + * + * @param ingestJobState + * @param md5objIdMapping + * @return + * @throws InterruptedException + * @throws CTCloudException + * @throws org.sleuthkit.datamodel.Blackboard.BlackboardException + * @throws TskCoreException + */ + private void longPollForNotFound(IngestJobState ingestJobState) throws InterruptedException, CTCloudException, Blackboard.BlackboardException, TskCoreException { + if (!ingestJobState.isDoFileLookups() || !ingestJobState.isQueryForMissing() || MapUtils.isEmpty(ingestJobState.getUnidentifiedHashes())) { + return; + } + + // TODO notify as a part of status that we are long polling for any missing + Map> remaining = new HashMap<>(ingestJobState.getUnidentifiedHashes()); for (int retry = 0; retry < NUM_FILE_UPLOAD_RETRIES; retry++) { List> md5Batches = Lists.partition(new ArrayList<>(remaining.keySet()), BATCH_SIZE); for (List batch : md5Batches) { - List repResult = getHashLookupResults(batch); - createArtifacts(repResult, remaining); + // if we have exceeded limits, then we're done. + if (!ingestJobState.isDoFileLookups()) { + return; + } + + List repResult = getHashLookupResults(ingestJobState, batch); + + Map> statusGroupings = repResult.stream() + .filter(bean -> bean.getMalwareResult() != null) + .collect(Collectors.groupingBy(bean -> bean.getMalwareResult().getStatus())); + + // for all found items, create analysis results + List found = statusGroupings.get(Status.FOUND); + + createAnalysisResults(ingestJobState, found, remaining); + + // remove any found items from the list of items to long poll for + for (CTCloudBean foundItem : found) { + String normalizedMd5 = normalizedMd5(foundItem.getMd5HashValue()); + remaining.remove(normalizedMd5); + } } if (remaining.isEmpty()) { - return true; + return; } - Thread.sleep(FILE_UPLOAD_RETRY_SLEEP_MILLIS); + Thread.sleep(FILE_UPLOAD_RETRY_SLEEP_MILLIS * ((long) Math.pow(2, retry))); } - return false; + // TODO if we get here, notify that not all lookups performed } /** @@ -548,49 +628,52 @@ public class MalwareScanIngestModule implements FileIngestModule { * @param ingestJobState The ingest job state. * @param repResult The list of cloud beans. Only cloud beans with a * malware status - * @param md5ToObjId + * @param md5ToObjId The mapping of md5 * @throws org.sleuthkit.datamodel.Blackboard.BlackboardException * @throws TskCoreException */ private void createAnalysisResults(IngestJobState ingestJobState, List repResult, Map> md5ToObjId) throws Blackboard.BlackboardException, TskCoreException { - List createdArtifacts = new ArrayList<>(); - if (!CollectionUtils.isEmpty(repResult)) { - SleuthkitCase.CaseDbTransaction trans = null; - try { - trans = ingestJobState.getTskCase().beginTransaction(); - for (CTCloudBean result : repResult) { - String sanitizedMd5 = normalizedMd5(result.getMd5HashValue()); - List objIds = md5ToObjId.remove(sanitizedMd5); - if (objIds == null || objIds.isEmpty()) { - continue; - } + if (CollectionUtils.isEmpty(repResult)) { + return; + } - for (Long objId : objIds) { - AnalysisResult res = createAnalysisResult(ingestJobState, trans, result, objId); - if (res != null) { - createdArtifacts.add(res); - } - } + List createdArtifacts = new ArrayList<>(); + SleuthkitCase.CaseDbTransaction trans = null; + try { + trans = ingestJobState.getTskCase().beginTransaction(); + for (CTCloudBean result : repResult) { + String sanitizedMd5 = normalizedMd5(result.getMd5HashValue()); + List objIds = md5ToObjId.remove(sanitizedMd5); + if (CollectionUtils.isEmpty(objIds)) { + continue; } - trans.commit(); - trans = null; - } finally { - if (trans != null) { - trans.rollback(); - createdArtifacts.clear(); - trans = null; + for (Long objId : objIds) { + AnalysisResult res = createAnalysisResult(ingestJobState, trans, result, objId); + if (res != null) { + createdArtifacts.add(res); + } } } - if (!CollectionUtils.isEmpty(createdArtifacts)) { - ingestJobState.getTskCase().getBlackboard().postArtifacts( - createdArtifacts, - Bundle.MalwareScanIngestModuleFactory_displayName(), - ingestJobState.getIngestJobId() - ); + trans.commit(); + trans = null; + } finally { + if (trans != null) { + trans.rollback(); + createdArtifacts.clear(); + trans = null; } } + + if (!CollectionUtils.isEmpty(createdArtifacts)) { + ingestJobState.getTskCase().getBlackboard().postArtifacts( + createdArtifacts, + Bundle.MalwareScanIngestModuleFactory_displayName(), + ingestJobState.getIngestJobId() + ); + } + } /** @@ -652,15 +735,14 @@ public class MalwareScanIngestModule implements FileIngestModule { }) synchronized void shutDown() { // if already shut down, return - if (runState == RunState.SHUT_DOWN) { + if (ingestJobState == null) { return; } // flush any remaining items try { + longPollForNotFound(ingestJobState); batchProcessor.flushAndReset(); - - getUploadedFileResults(this.unidentifiedHashes); } catch (InterruptedException ex) { notifyWarning( Bundle.MalwareScanIngestModule_SharedProcessing_flushTimeout_title(), @@ -673,7 +755,7 @@ public class MalwareScanIngestModule implements FileIngestModule { ex); } finally { // set state to shut down and clear any remaining - runState = RunState.SHUT_DOWN; + ingestJobState = null; } } @@ -690,10 +772,6 @@ public class MalwareScanIngestModule implements FileIngestModule { logger.log(Level.WARNING, message, ex); } - private enum RunState { - STARTED_UP, DISABLED, SHUT_DOWN - } - class FileRecord { private final long objId; @@ -714,7 +792,28 @@ public class MalwareScanIngestModule implements FileIngestModule { } - class IngestJobState { + /** + * Represents the state of the current ingest job. + * + * NOTE: if doFileLookups is false, most variables will likely be null + * (TSK case, file type detector, etc.) and should not be used. The + * contract for this class should be that if doFileLookups is true or + * uploadUnknownFiles is true, the remaining variables should be non + * null, if doFileLookups is false and uploadUnknownFiles is false, no + * other access to this class can be made reliably. + */ + static class IngestJobState { + + static final IngestJobState DISABLED = new IngestJobState( + null, + null, + null, + null, + 0L, + 0L, + false, + false + ); private final SleuthkitCase tskCase; private final FileTypeDetector fileTypeDetector; @@ -722,19 +821,24 @@ public class MalwareScanIngestModule implements FileIngestModule { private final BlackboardArtifact.Type malwareType; private final long dsId; private final long ingestJobId; + private final boolean queryForMissing; private final Map> unidentifiedHashes = new HashMap<>(); // this can change mid run private boolean uploadUnknownFiles; + private boolean doFileLookups; - IngestJobState(SleuthkitCase tskCase, FileTypeDetector fileTypeDetector, LicenseInfo licenseInfo, BlackboardArtifact.Type malwareType, long dsId, long ingestJobId, boolean uploadUnknownFiles) { + IngestJobState(SleuthkitCase tskCase, FileTypeDetector fileTypeDetector, LicenseInfo licenseInfo, BlackboardArtifact.Type malwareType, long dsId, long ingestJobId, boolean uploadUnknownFiles, boolean doFileLookups) { this.tskCase = tskCase; this.fileTypeDetector = fileTypeDetector; this.licenseInfo = licenseInfo; this.malwareType = malwareType; this.dsId = dsId; this.ingestJobId = ingestJobId; + // for now, querying for any missing files will be tied to whether initially we should upload files and do lookups at all + this.queryForMissing = uploadUnknownFiles && doFileLookups; this.uploadUnknownFiles = uploadUnknownFiles; + this.doFileLookups = doFileLookups; } SleuthkitCase getTskCase() { @@ -765,12 +869,24 @@ public class MalwareScanIngestModule implements FileIngestModule { return unidentifiedHashes; } - boolean uploadUnknownFiles() { + boolean isQueryForMissing() { + return queryForMissing; + } + + boolean isUploadUnknownFiles() { return uploadUnknownFiles; } - void setUploadUnknownFiles(boolean uploadUnknownFiles) { - this.uploadUnknownFiles = uploadUnknownFiles; + void disableUploadUnknownFiles() { + this.uploadUnknownFiles = false; + } + + boolean isDoFileLookups() { + return doFileLookups; + } + + void disableDoFileLookups() { + this.doFileLookups = false; } } }