MalwareScanIngestModule.java

Go to the documentation of this file.

/*
 * Autopsy Forensic Browser
 *
 * Copyright 2023 Basis Technology Corp.
 * Contact: carrier <at> sleuthkit <dot> org
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.basistech.df.cybertriage.autopsy.malwarescan;

import com.basistech.df.cybertriage.autopsy.ctapi.CTApiDAO;
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.LicenseInfo;
import com.basistech.df.cybertriage.autopsy.ctoptions.ctcloud.CTLicensePersistence;
import java.text.MessageFormat;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
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.lang3.StringUtils;
import org.openide.util.NbBundle.Messages;
import org.sleuthkit.autopsy.casemodule.Case;
import org.sleuthkit.autopsy.coreutils.Logger;
import org.sleuthkit.autopsy.coreutils.MessageNotifyUtil;
import org.sleuthkit.autopsy.ingest.FileIngestModule;
import org.sleuthkit.autopsy.ingest.IngestJobContext;
import org.sleuthkit.autopsy.ingest.IngestModule;
import org.sleuthkit.autopsy.modules.filetypeid.FileTypeDetector;
import org.sleuthkit.datamodel.AbstractFile;
import org.sleuthkit.datamodel.AnalysisResult;
import org.sleuthkit.datamodel.Blackboard;
import org.sleuthkit.datamodel.BlackboardArtifact;
import org.sleuthkit.datamodel.HashUtility;
import org.sleuthkit.datamodel.HashUtility.HashResult;
import org.sleuthkit.datamodel.HashUtility.HashType;
import org.sleuthkit.datamodel.Score;
import org.sleuthkit.datamodel.SleuthkitCase;
import org.sleuthkit.datamodel.TskCoreException;
import org.sleuthkit.datamodel.TskData;

public class MalwareScanIngestModule implements FileIngestModule {

    private static final SharedProcessing sharedProcessing = new SharedProcessing();

    @Override
    public void startUp(IngestJobContext context) throws IngestModuleException {
        sharedProcessing.startUp(context);
    }

    @Override
    public ProcessResult process(AbstractFile af) {
        return sharedProcessing.process(af);
    }

    @Override
    public void shutDown() {
        sharedProcessing.shutDown();
    }

    private static class SharedProcessing {

        // batch size of 200 files max
        private static final int BATCH_SIZE = 200;
        // 1 day timeout for all API requests
        private static final long FLUSH_SECS_TIMEOUT = 24 * 60 * 60;

        //minimum lookups left before issuing warning
        private static final long LOW_LOOKUPS_REMAINING = 250;

        private static final Set<String> EXECUTABLE_MIME_TYPES = Stream.of(
                "application/x-bat",//NON-NLS
                "application/x-dosexec",//NON-NLS
                "application/vnd.microsoft.portable-executable",//NON-NLS
                "application/x-msdownload",//NON-NLS
                "application/exe",//NON-NLS
                "application/x-exe",//NON-NLS
                "application/dos-exe",//NON-NLS
                "vms/exe",//NON-NLS
                "application/x-winexe",//NON-NLS
                "application/msdos-windows",//NON-NLS
                "application/x-msdos-program"//NON-NLS
        ).collect(Collectors.toSet());

        private static final String MALWARE_TYPE_NAME = "TSK_MALWARE";
        private static final String MALWARE_CONFIG = "Cyber Triage Cloud";

        private static final Logger logger = Logger.getLogger(MalwareScanIngestModule.class.getName());
        private final BatchProcessor<FileRecord> batchProcessor = new BatchProcessor<FileRecord>(BATCH_SIZE, FLUSH_SECS_TIMEOUT, this::handleBatch);

        private final CTLicensePersistence ctSettingsPersistence = CTLicensePersistence.getInstance();
        private final CTApiDAO ctApiDAO = CTApiDAO.getInstance();

        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;

        @Messages({
            "MalwareScanIngestModule_ShareProcessing_lowLimitWarning_title=Hash Lookups Low",
            "# {0} - remainingLookups",
            "MalwareScanIngestModule_ShareProcessing_lowLimitWarning_desc=This license only has {0} lookups remaining",
            "MalwareScanIngestModule_malwareTypeDisplayName=Malware",
            "MalwareScanIngestModule_ShareProcessing_noLicense_title=No Cyber Triage License",
            "MalwareScanIngestModule_ShareProcessing_noLicense_desc=No Cyber Triage license could be loaded.  Cyber Triage processing will be disabled.",
            "MalwareScanIngestModule_ShareProcessing_noRemaining_title=No remaining lookups",
            "MalwareScanIngestModule_ShareProcessing_noRemaining_desc=There are no more remaining hash lookups for this license at this time.  Cyber Triage processing will be disabled."
        })
        synchronized void startUp(IngestJobContext context) throws IngestModuleException {
            // only run this code once per startup
            if (runState == RunState.STARTED_UP || runState == RunState.DISABLED) {
                return;
            }

            try {
                // 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;
                }

                AuthTokenResponse authTokenResponse = ctApiDAO.getAuthToken(licenseInfoOpt.get().getDecryptedLicense());
                // syncronously fetch malware scans info

                // determine lookups remaining
                long lookupsRemaining = remaining(authTokenResponse.getHashLookupLimit(), authTokenResponse.getHashLookupCount());
                if (lookupsRemaining <= 0) {
                    notifyWarning(
                            Bundle.MalwareScanIngestModule_ShareProcessing_noRemaining_title(),
                            Bundle.MalwareScanIngestModule_ShareProcessing_noRemaining_desc(),
                            null);
                    runState = RunState.DISABLED;
                    return;
                } else if (lookupsRemaining < LOW_LOOKUPS_REMAINING) {
                    notifyWarning(
                            Bundle.MalwareScanIngestModule_ShareProcessing_lowLimitWarning_title(),
                            Bundle.MalwareScanIngestModule_ShareProcessing_lowLimitWarning_desc(lookupsRemaining),
                            null);
                }

                // setup necessary variables for processing
                tskCase = Case.getCurrentCaseThrows().getSleuthkitCase();
                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();

                // 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);
            }
        }

        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) {
            if (StringUtils.isNotBlank(af.getMd5Hash())) {
                return af.getMd5Hash();
            }

            try {
                List<HashResult> hashResults = HashUtility.calculateHashes(af, Collections.singletonList(HashType.MD5));
                if (CollectionUtils.isNotEmpty(hashResults)) {
                    for (HashResult hashResult : hashResults) {
                        if (hashResult.getType() == HashType.MD5) {
                            return hashResult.getValue();
                        }
                    }
                }
            } catch (TskCoreException ex) {
                logger.log(Level.WARNING,
                        MessageFormat.format("An error occurred while processing file name: {0} and obj id: {1}.",
                                af.getName(),
                                af.getId()),
                        ex);
            }

            return null;
        }

        @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
                        && af.getKnown() != TskData.FileKnown.KNOWN
                        && EXECUTABLE_MIME_TYPES.contains(StringUtils.defaultString(fileTypeDetector.getMIMEType(af)).trim().toLowerCase())
                        && CollectionUtils.isEmpty(af.getAnalysisResults(malwareType))) {

                    String md5 = getOrCalcHash(af);
                    if (StringUtils.isNotBlank(md5)) {
                        batchProcessor.add(new FileRecord(af.getId(), md5));
                    }
                }
                return ProcessResult.OK;
            } catch (TskCoreException ex) {
                notifyWarning(
                        Bundle.MalwareScanIngestModule_SharedProcessing_generalProcessingError_title(),
                        Bundle.MalwareScanIngestModule_SharedProcessing_generalProcessingError_desc(),
                        ex);
                return IngestModule.ProcessResult.ERROR;
            } catch (InterruptedException ex) {
                notifyWarning(
                        Bundle.MalwareScanIngestModule_ShareProcessing_batchTimeout_title(),
                        Bundle.MalwareScanIngestModule_ShareProcessing_batchTimeout_desc(),
                        ex);
                return IngestModule.ProcessResult.ERROR;
            }
        }

        @Messages({
            "MalwareScanIngestModule_SharedProcessing_authTokenResponseError_title=Authentication API error",
            "# {0} - errorResponse",
            "MalwareScanIngestModule_SharedProcessing_authTokenResponseError_desc=Received error: ''{0}'' when fetching the API authentication token for the license",
            "MalwareScanIngestModule_SharedProcessing_repServicenResponseError_title=Lookup API error",
            "# {0} - errorResponse",
            "MalwareScanIngestModule_SharedProcessing_repServicenResponseError_desc=Received error: ''{0}'' when fetching hash lookup results",
            "MalwareScanIngestModule_SharedProcessing_exhaustedHashLookups_title=Hash Lookups Exhausted",
            "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<FileRecord> fileRecords) {
            if (runState != RunState.STARTED_UP || fileRecords == null || fileRecords.isEmpty()) {
                return;
            }

            // create mapping of md5 to corresponding object ids as well as just the list of md5's
            Map<String, List<Long>> md5ToObjId = new HashMap<>();

            for (FileRecord fr : fileRecords) {
                if (fr == null || StringUtils.isBlank(fr.getMd5hash()) || fr.getObjId() <= 0) {
                    continue;
                }

                String sanitizedMd5 = sanitizedMd5(fr.getMd5hash());
                md5ToObjId
                        .computeIfAbsent(sanitizedMd5, (k) -> new ArrayList<>())
                        .add(fr.getObjId());

            }

            List<String> md5Hashes = new ArrayList<>(md5ToObjId.keySet());

            if (md5Hashes.isEmpty()) {
                return;
            }

            try {
                // get an auth token with the license
                AuthTokenResponse authTokenResponse = ctApiDAO.getAuthToken(licenseInfo.getDecryptedLicense());

                // make sure we are in bounds for the remaining scans
                long remainingScans = remaining(authTokenResponse.getHashLookupLimit(), authTokenResponse.getHashLookupCount());
                if (remainingScans <= 0) {
                    runState = RunState.DISABLED;
                    notifyWarning(
                            Bundle.MalwareScanIngestModule_SharedProcessing_exhaustedHashLookups_title(),
                            Bundle.MalwareScanIngestModule_SharedProcessing_exhaustedHashLookups_desc(),
                            null);
                    return;
                }

                // using auth token, get results
                List<CTCloudBean> repResult = ctApiDAO.getReputationResults(
                        new AuthenticatedRequestData(licenseInfo.getDecryptedLicense(), authTokenResponse),
                        md5Hashes
                );

                List<BlackboardArtifact> createdArtifacts = new ArrayList<>();
                if (!CollectionUtils.isEmpty(repResult)) {
                    SleuthkitCase.CaseDbTransaction trans = null;
                    try {
                        trans = tskCase.beginTransaction();
                        for (CTCloudBean result : repResult) {
                            String sanitizedMd5 = sanitizedMd5(result.getMd5HashValue());
                            List<Long> objIds = md5ToObjId.remove(sanitizedMd5);
                            if (objIds == null || objIds.isEmpty()) {
                                continue;
                            }

                            for (Long objId : objIds) {
                                AnalysisResult res = createAnalysisResult(objId, result, trans);
                                if (res != null) {
                                    createdArtifacts.add(res);
                                }
                            }
                        }

                        trans.commit();
                        trans = null;
                    } finally {
                        if (trans != null) {
                            trans.rollback();
                            createdArtifacts.clear();
                            trans = null;
                        }
                    }

                    if (!CollectionUtils.isEmpty(createdArtifacts)) {
                        tskCase.getBlackboard().postArtifacts(createdArtifacts, Bundle.MalwareScanIngestModuleFactory_displayName(), ingestJobId);
                    }
                }
            } catch (Exception ex) {
                notifyWarning(
                        Bundle.MalwareScanIngestModule_SharedProcessing_generalProcessingError_title(),
                        Bundle.MalwareScanIngestModule_SharedProcessing_generalProcessingError_desc(),
                        ex);
            }
        }

        private String sanitizedMd5(String orig) {
            return StringUtils.defaultString(orig).trim().toLowerCase();
        }

        @Messages({
            "MalwareScanIngestModule_SharedProcessing_createAnalysisResult_Yes=YES",
            "MalwareScanIngestModule_SharedProcessing_createAnalysisResult_No=NO"
        })
        private AnalysisResult createAnalysisResult(Long objId, CTCloudBean cloudBean, SleuthkitCase.CaseDbTransaction trans) throws Blackboard.BlackboardException {
            if (objId == null || cloudBean == null || cloudBean.getMalwareResult() == null) {
                return null;
            }

            Score score = cloudBean.getMalwareResult().getCTScore() == null
                    ? Score.SCORE_UNKNOWN
                    : cloudBean.getMalwareResult().getCTScore().getTskCore();

            String conclusion = score.getSignificance() == Score.Significance.NOTABLE || score.getSignificance() == Score.Significance.LIKELY_NOTABLE
                    ? Bundle.MalwareScanIngestModule_SharedProcessing_createAnalysisResult_Yes()
                    : Bundle.MalwareScanIngestModule_SharedProcessing_createAnalysisResult_No();

            String justification = cloudBean.getMalwareResult().getStatusDescription();

            return tskCase.getBlackboard().newAnalysisResult(
                    malwareType,
                    objId,
                    dsId,
                    score,
                    conclusion,
                    MALWARE_CONFIG,
                    justification,
                    Collections.emptyList(),
                    trans).getAnalysisResult();
        }

        @Messages({
            "MalwareScanIngestModule_SharedProcessing_flushTimeout_title=Processing Timeout",
            "MalwareScanIngestModule_SharedProcessing_flushTimeout_desc=A timeout occurred while finishing processing"
        })
        synchronized void shutDown() {
            // if already shut down, return
            if (runState == RunState.SHUT_DOWN) {
                return;
            }

            // flush any remaining items
            try {
                batchProcessor.flushAndReset();
            } catch (InterruptedException ex) {
                notifyWarning(
                        Bundle.MalwareScanIngestModule_SharedProcessing_flushTimeout_title(),
                        Bundle.MalwareScanIngestModule_SharedProcessing_flushTimeout_desc(),
                        ex);
            } finally {
                // set state to shut down and clear any remaining
                runState = RunState.SHUT_DOWN;
            }
        }

        private void notifyWarning(String title, String message, Exception ex) {
            MessageNotifyUtil.Notify.warn(title, message);
            logger.log(Level.WARNING, message, ex);
        }

        private enum RunState {
            STARTED_UP, DISABLED, SHUT_DOWN
        }

        class FileRecord {

            private final long objId;
            private final String md5hash;

            FileRecord(long objId, String md5hash) {
                this.objId = objId;
                this.md5hash = md5hash;
            }

            long getObjId() {
                return objId;
            }

            String getMd5hash() {
                return md5hash;
            }

        }
    }
}