refactoring

Core/src/com/basistech/df/cybertriage/autopsy/malwarescan/MalwareScanIngestModule.java

@@ -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 RunState runState = null;
-
+//        private IngestJobState ingestJobState = 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;
 
         @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 Pair.of(RunState.DISABLED, null);
-                    return;
                 }
 
                 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 Pair.of(RunState.DISABLED, null);
-                    return;
                 } 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();
+                SleuthkitCase tskCase = Case.getCurrentCaseThrows().getSleuthkitCase();
-                malwareType = tskCase.getBlackboard().getOrAddArtifactType(
+                BlackboardArtifact.Type malwareType = tskCase.getBlackboard().getOrAddArtifactType(
                         MALWARE_TYPE_NAME,
                         Bundle.MalwareScanIngestModule_malwareTypeDisplayName(),
                         BlackboardArtifact.Category.ANALYSIS_RESULT);
-                fileTypeDetector = new FileTypeDetector();
+                
-                dsId = context.getDataSource().getId();
+                IngestJobState ingestJobState = new IngestJobState(
-                ingestJobId = context.getJobId();
+                        tskCase,
-                licenseInfo = licenseInfoOpt.get();
+                        new FileTypeDetector(), 
-                uploadUnknownFiles = uploadFiles;
+                        licenseInfoOpt.get(), 
-                unidentifiedHashes = new HashMap<>();
+                        malwareType, 
-
+                        context.getDataSource().getId(), 
-                // set run state to initialized
+                        context.getJobId(), 
-                runState = RunState.STARTED_UP;
+                        uploadFiles);
-            } catch (Exception ex) {
+                
-                runState = RunState.DISABLED;
+                return Pair.of(RunState.STARTED_UP, ingestJobState);
-                throw new IngestModuleException("An exception occurred on MalwareScanIngestModule startup", ex);
-            }
         }
 
+        /**
+         * 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;
+            }
+        }
     }
 }