refactoring

This commit is contained in:
Greg DiCristofaro 2023-07-26 16:02:43 -04:00
parent 673f623c7b
commit 10ae3411ac

View File

@ -43,9 +43,11 @@ import java.util.stream.Collectors;
import java.util.stream.Stream;
import org.apache.commons.collections4.CollectionUtils;
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;
@ -133,17 +135,8 @@ public class MalwareScanIngestModule implements FileIngestModule {
private final CTApiDAO ctApiDAO = CTApiDAO.getInstance();
// TODO minimize state
private RunState runState = null;
private SleuthkitCase tskCase = null;
private FileTypeDetector fileTypeDetector = null;
private LicenseInfo licenseInfo = null;
private BlackboardArtifact.Type malwareType = null;
private long dsId = 0;
private long ingestJobId = 0;
private boolean uploadUnknownFiles = false;
private Map<String, List<Long>> unidentifiedHashes = null;
// private RunState runState = null;
// private IngestJobState ingestJobState = null;
@Messages({
"MalwareScanIngestModule_malwareTypeDisplayName=Malware",
@ -166,15 +159,30 @@ public class MalwareScanIngestModule implements FileIngestModule {
}
try {
// get saved license
Pair<RunState, IngestJobState> jobStateResult = getNewJobState(context);
runState = jobStateResult.getLeft();
ingestJobState = jobStateResult.getRight();
} catch (Exception ex) {
runState = RunState.DISABLED;
throw new IngestModuleException("An exception occurred on MalwareScanIngestModule startup", ex);
}
}
/**
* Sets up the state necessary for a new ingest job.
* @param context The ingest job context.
* @return A pair of the runtime state (i.e. started up, disabled) and parameters required for the job.
* @throws Exception
*/
private Pair<RunState, IngestJobState> getNewJobState(IngestJobContext context) throws Exception {
// get saved license
Optional<LicenseInfo> licenseInfoOpt = ctSettingsPersistence.loadLicenseInfo();
if (licenseInfoOpt.isEmpty() || licenseInfoOpt.get().getDecryptedLicense() == null) {
notifyWarning(
Bundle.MalwareScanIngestModule_ShareProcessing_noLicense_title(),
Bundle.MalwareScanIngestModule_ShareProcessing_noLicense_desc(),
null);
runState = RunState.DISABLED;
return;
return Pair.of(RunState.DISABLED, null);
}
AuthTokenResponse authTokenResponse = ctApiDAO.getAuthToken(licenseInfoOpt.get().getDecryptedLicense());
@ -187,8 +195,7 @@ public class MalwareScanIngestModule implements FileIngestModule {
Bundle.MalwareScanIngestModule_ShareProcessing_noLookupsRemaining_title(),
Bundle.MalwareScanIngestModule_ShareProcessing_noLookupsRemaining_desc(),
null);
runState = RunState.DISABLED;
return;
return Pair.of(RunState.DISABLED, null);
} else if (lookupsRemaining < LOW_LOOKUPS_REMAINING) {
notifyWarning(
Bundle.MalwareScanIngestModule_ShareProcessing_lowLookupsLimitWarning_title(),
@ -215,33 +222,42 @@ public class MalwareScanIngestModule implements FileIngestModule {
}
// setup necessary variables for processing
tskCase = Case.getCurrentCaseThrows().getSleuthkitCase();
malwareType = tskCase.getBlackboard().getOrAddArtifactType(
SleuthkitCase tskCase = Case.getCurrentCaseThrows().getSleuthkitCase();
BlackboardArtifact.Type malwareType = tskCase.getBlackboard().getOrAddArtifactType(
MALWARE_TYPE_NAME,
Bundle.MalwareScanIngestModule_malwareTypeDisplayName(),
BlackboardArtifact.Category.ANALYSIS_RESULT);
fileTypeDetector = new FileTypeDetector();
dsId = context.getDataSource().getId();
ingestJobId = context.getJobId();
licenseInfo = licenseInfoOpt.get();
uploadUnknownFiles = uploadFiles;
unidentifiedHashes = new HashMap<>();
// set run state to initialized
runState = RunState.STARTED_UP;
} catch (Exception ex) {
runState = RunState.DISABLED;
throw new IngestModuleException("An exception occurred on MalwareScanIngestModule startup", ex);
}
IngestJobState ingestJobState = new IngestJobState(
tskCase,
new FileTypeDetector(),
licenseInfoOpt.get(),
malwareType,
context.getDataSource().getId(),
context.getJobId(),
uploadFiles);
return Pair.of(RunState.STARTED_UP, ingestJobState);
}
/**
* Determines remaining given a possibly null limit and used count.
* @param limit The limit (can be null).
* @param used The number used (can be null).
* @return The remaining amount.
*/
private static long remaining(Long limit, Long used) {
limit = limit == null ? 0 : limit;
used = used == null ? 0 : used;
return limit - used;
}
private String getOrCalcHash(AbstractFile af) {
/**
* Gets the md5 hash from the abstract file or calculates it.
* @param af The abstract file.
* @return The md5 hash (or null if could not be determined).
*/
private static String getOrCalcHash(AbstractFile af) {
if (StringUtils.isNotBlank(af.getMd5Hash())) {
return af.getMd5Hash();
}
@ -342,22 +358,29 @@ public class MalwareScanIngestModule implements FileIngestModule {
.collect(Collectors.groupingBy(bean -> bean.getMalwareResult().getStatus()));
List<CTCloudBean> found = statusGroupings.get(Status.FOUND);
createArtifacts(repResult, md5ToObjId);
createArtifacts(found, md5ToObjId);
// if being scanned, check list to run later
List<CTCloudBean> beingScannedList = statusGroupings.get(Status.BEING_SCANNED);
processMissing(beingScannedList, md5ToObjId, false);
// if not found, try upload
List<CTCloudBean> 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);
}
// TODO handle caching list and creating new items
} catch (Exception ex) {
notifyWarning(
@ -367,7 +390,7 @@ public class MalwareScanIngestModule implements FileIngestModule {
}
}
private void processMissing(Collection<CTCloudBean> results, Map<String, List<Long>> md5ToObjId, boolean doFileUpload) throws CTCloudException {
private static void processMissing(Collection<CTCloudBean> results, Map<String, List<Long>> md5ToObjId, boolean doFileUpload) throws CTCloudException, TskCoreException {
for (CTCloudBean beingScanned : CollectionUtils.emptyIfNull(results)) {
String sanitizedMd5 = sanitizedMd5(beingScanned.getMd5HashValue());
@ -507,7 +530,6 @@ public class MalwareScanIngestModule implements FileIngestModule {
}
private boolean getUploadedFileResults(Map<String, List<Long>> md5objIdMapping) throws InterruptedException, CTCloudException, Blackboard.BlackboardException, TskCoreException {
// TODO integrate this
Map<String, List<Long>> remaining = new HashMap<>(md5objIdMapping);
for (int retry = 0; retry < NUM_FILE_UPLOAD_RETRIES; retry++) {
@ -571,11 +593,18 @@ public class MalwareScanIngestModule implements FileIngestModule {
// flush any remaining items
try {
batchProcessor.flushAndReset();
getUploadedFileResults(this.unidentifiedHashes);
} catch (InterruptedException ex) {
notifyWarning(
Bundle.MalwareScanIngestModule_SharedProcessing_flushTimeout_title(),
Bundle.MalwareScanIngestModule_SharedProcessing_flushTimeout_desc(),
ex);
} catch (Exception ex) {
notifyWarning(
Bundle.MalwareScanIngestModule_SharedProcessing_generalProcessingError_title(),
Bundle.MalwareScanIngestModule_SharedProcessing_generalProcessingError_desc(),
ex);
} finally {
// set state to shut down and clear any remaining
runState = RunState.SHUT_DOWN;
@ -610,5 +639,64 @@ public class MalwareScanIngestModule implements FileIngestModule {
}
}
class IngestJobState {
private final SleuthkitCase tskCase;
private final FileTypeDetector fileTypeDetector;
private final LicenseInfo licenseInfo;
private final BlackboardArtifact.Type malwareType;
private final long dsId;
private final long ingestJobId;
private final Map<String, List<Long>> unidentifiedHashes = new HashMap<>();
// this can change mid run
private boolean uploadUnknownFiles;
IngestJobState(SleuthkitCase tskCase, FileTypeDetector fileTypeDetector, LicenseInfo licenseInfo, BlackboardArtifact.Type malwareType, long dsId, long ingestJobId, boolean uploadUnknownFiles) {
this.tskCase = tskCase;
this.fileTypeDetector = fileTypeDetector;
this.licenseInfo = licenseInfo;
this.malwareType = malwareType;
this.dsId = dsId;
this.ingestJobId = ingestJobId;
this.uploadUnknownFiles = uploadUnknownFiles;
}
SleuthkitCase getTskCase() {
return tskCase;
}
FileTypeDetector getFileTypeDetector() {
return fileTypeDetector;
}
LicenseInfo getLicenseInfo() {
return licenseInfo;
}
BlackboardArtifact.Type getMalwareType() {
return malwareType;
}
long getDsId() {
return dsId;
}
long getIngestJobId() {
return ingestJobId;
}
Map<String, List<Long>> getUnidentifiedHashes() {
return unidentifiedHashes;
}
boolean uploadUnknownFiles() {
return uploadUnknownFiles;
}
void setUploadUnknownFiles(boolean uploadUnknownFiles) {
this.uploadUnknownFiles = uploadUnknownFiles;
}
}
}
}