diff --git a/iped-app/resources/config/IPEDConfig.txt b/iped-app/resources/config/IPEDConfig.txt index b385a8632f..e619943546 100644 --- a/iped-app/resources/config/IPEDConfig.txt +++ b/iped-app/resources/config/IPEDConfig.txt @@ -171,3 +171,8 @@ enableHTMLReport = true # It performs classification of image and video files. # Advanced settings can be modified in file "conf/RemoteImageClassifierConfig.txt" enableRemoteImageClassifier = false + +# Enables remote external image classifier. +# It performs classification of image and video files using review priority and embeddings output. +# Advanced settings can be modified in file "conf/RemoteImageExternalClassifierConfig.txt" +enableRemoteImageExternalClassifier = true diff --git a/iped-app/resources/config/conf/AIFiltersConfig.json b/iped-app/resources/config/conf/AIFiltersConfig.json index d06e0b1340..fae940b2ff 100644 --- a/iped-app/resources/config/conf/AIFiltersConfig.json +++ b/iped-app/resources/config/conf/AIFiltersConfig.json @@ -19,6 +19,14 @@ {"name": "Other"} ]}, + {"name": "External CSAM Detector", "prefix": "ExternalCSAM", "property": "ai\\:externalClassifier\\:reviewPriority", "value": "*", "children":[ + {"name": "Very High", "value": "VERY_HIGH"}, + {"name": "High", "value": "HIGH"}, + {"name": "Medium", "value": "MEDIUM"}, + {"name": "Low", "value": "LOW"}, + {"name": "Very Low", "value": "VERY_LOW"} + ]}, + {"name": "Detected Faces", "prefix": "Face", "property": "face\\_count", "value": "[1 TO *]", "children":[ {"name": "1"}, {"name": "2"}, diff --git a/iped-app/resources/config/conf/RemoteImageExternalClassifierConfig.txt b/iped-app/resources/config/conf/RemoteImageExternalClassifierConfig.txt new file mode 100644 index 0000000000..f96d154f55 --- /dev/null +++ b/iped-app/resources/config/conf/RemoteImageExternalClassifierConfig.txt @@ -0,0 +1,25 @@ +########################################## +# RemoteImageExternalClassifier Task configuration +########################################## + +# URL of the service/central node used by the RemoteImageExternalClassifier implementation +url = http://localhost:8000 + +# Maximum number of thumbs to be included in the zip file to send to the server +batchSize = 50 + +# Optional: used only if backend also returns legacy "class" probabilities. +# Minimum confidence threshold for assigning a class label to an image/video. Values are between 0 and 100. +labelingThreshold = 50 + +# Skip classification of images/videos smaller than a given file size (in bytes; '0' = do not skip) +skipSize = 2048 + +# Skip classification of images/videos smaller than a given dimension, i.e. height or width (in pixels; '0' = do not skip) +skipDimension = 48 + +# Skip classification of images/videos with hits on IPED hashesDB database (if 'hashesDB' is not configured in 'LocalConfig.txt' or 'false', do not skip) +skipHashDBFiles = true + +# Validate server SSL certificate +validateSSL = false diff --git a/iped-app/resources/config/conf/TaskInstaller.xml b/iped-app/resources/config/conf/TaskInstaller.xml index a17823f825..4efe65aa7d 100644 --- a/iped-app/resources/config/conf/TaskInstaller.xml +++ b/iped-app/resources/config/conf/TaskInstaller.xml @@ -40,6 +40,7 @@ + diff --git a/iped-engine/src/main/java/iped/engine/config/RemoteImageExternalClassifierConfig.java b/iped-engine/src/main/java/iped/engine/config/RemoteImageExternalClassifierConfig.java new file mode 100644 index 0000000000..ac129b505d --- /dev/null +++ b/iped-engine/src/main/java/iped/engine/config/RemoteImageExternalClassifierConfig.java @@ -0,0 +1,137 @@ +package iped.engine.config; + +import iped.utils.UTF8Properties; + +public class RemoteImageExternalClassifierConfig extends AbstractTaskPropertiesConfig { + + private static final long serialVersionUID = 1L; + + /** + * Config file name and enable/disable property. + */ + private static final String CONFIG_FILE = "RemoteImageExternalClassifierConfig.txt"; + private static final String ENABLE_PROP = "enableRemoteImageExternalClassifier"; + + /** + * Constants mapping to properties name in config file. + */ + private static final String URL = "url"; + private static final String BATCH_SIZE = "batchSize"; + private static final String LABELING_THRESHOLD = "labelingThreshold"; + private static final String SKIP_SIZE = "skipSize"; + private static final String SKIP_DIMENSION = "skipDimension"; + private static final String SKIP_HASH_DB_FILES = "skipHashDBFiles"; + private static final String VALIDATE_SSL = "validateSSL"; + + // URL of the service/central node used by the RemoteImageClassifier implementation + private String url; + + // Maximum number of thumbs to be included in the zip file to send to the server + private int batchSize = 50; + + // Threshold used to decide if an image is labeled with one class + private double labelingThreshold = 60; + + // Skip classification of images/videos smaller than a given file size (in bytes; '0' = do not skip) + private int skipSize = 0; + + // Skip classification of images/videos smaller than a given dimension, i.e. height or width (in pixels; '0' = do not skip) + private int skipDimension = 0; + + // Skip classification of images/videos with hits on IPED hashesDB database (if 'hashesDB' is not configured in 'LocalConfig.txt' or 'false', do not skip) + private boolean skipHashDBFiles = true; + + // Validate server SSL certificate + private boolean validateSSL = false; + + @Override + public String getTaskEnableProperty() { + return ENABLE_PROP; + } + + @Override + public String getTaskConfigFileName() { + return CONFIG_FILE; + } + + public String getUrl() { + return url; + } + + public int getBatchSize() { + return batchSize; + } + + public double getLabelingThreshold() { + return labelingThreshold; + } + + public int getSkipSize() { + return skipSize; + } + + public int getSkipDimension() { + return skipDimension; + } + + public boolean isSkipHashDBFiles() { + return skipHashDBFiles; + } + + public boolean isValidateSSL() { + return validateSSL; + } + + @Override + void processProperties(UTF8Properties properties) { + + String value = properties.getProperty(URL); + if (value != null && !value.trim().isEmpty()) + url = value.trim(); + + value = properties.getProperty(BATCH_SIZE); + if (value != null && !value.trim().isEmpty()) { + batchSize = Integer.valueOf(value.trim()); + // enforce minimum value + if (batchSize < 1) + batchSize = 1; + } + + value = properties.getProperty(LABELING_THRESHOLD); + if (value != null && !value.trim().isEmpty()) { + labelingThreshold = Double.parseDouble(value.trim()); + // enforce range [0, 100] + if (labelingThreshold < 0) { + labelingThreshold = 0; + } else if (labelingThreshold > 100) { + labelingThreshold = 100; + } + } + + value = properties.getProperty(SKIP_SIZE); + if (value != null && !value.trim().isEmpty()) { + skipSize = Integer.valueOf(value.trim()); + // enforce minimum value + if (skipSize < 0) + skipSize = 0; + } + + value = properties.getProperty(SKIP_DIMENSION); + if (value != null && !value.trim().isEmpty()) { + skipDimension = Integer.valueOf(value.trim()); + // enforce minimum value + if (skipDimension < 0) + skipDimension = 0; + } + + value = properties.getProperty(SKIP_HASH_DB_FILES); + if (value != null && !value.trim().isEmpty()) + skipHashDBFiles = Boolean.valueOf(value.trim()); + + value = properties.getProperty(VALIDATE_SSL); + if (value != null && !value.trim().isEmpty()) + validateSSL = Boolean.valueOf(value.trim()); + + } + +} diff --git a/iped-engine/src/main/java/iped/engine/lucene/analysis/AppAnalyzer.java b/iped-engine/src/main/java/iped/engine/lucene/analysis/AppAnalyzer.java index ce58b52585..e541c8872a 100644 --- a/iped-engine/src/main/java/iped/engine/lucene/analysis/AppAnalyzer.java +++ b/iped-engine/src/main/java/iped/engine/lucene/analysis/AppAnalyzer.java @@ -29,6 +29,7 @@ import iped.engine.config.IndexTaskConfig; import iped.engine.task.HashTask; import iped.engine.task.PhotoDNATask; +import iped.engine.task.RemoteImageExternalClassifierTask; import iped.engine.task.index.IndexItem; import iped.localization.LocalizedProperties; import iped.properties.ExtraProperties; @@ -57,6 +58,8 @@ public static Analyzer get() { analyzerPerField.put(IndexItem.CHANGED, new KeywordAnalyzer()); analyzerPerField.put(IndexItem.TIMESTAMP, new KeywordAnalyzer()); + analyzerPerField.put(RemoteImageExternalClassifierTask.AI_REVIEW_PRIORITY_ATTR, new KeywordAnalyzer()); + IndexTaskConfig indexConfig = ConfigurationManager.get().findObject(IndexTaskConfig.class); StandardASCIIAnalyzer hashAnalyzer = new StandardASCIIAnalyzer(); hashAnalyzer.setMaxTokenLength(Integer.MAX_VALUE); diff --git a/iped-engine/src/main/java/iped/engine/task/RemoteImageExternalClassifierTask.java b/iped-engine/src/main/java/iped/engine/task/RemoteImageExternalClassifierTask.java new file mode 100644 index 0000000000..6acc807ccf --- /dev/null +++ b/iped-engine/src/main/java/iped/engine/task/RemoteImageExternalClassifierTask.java @@ -0,0 +1,876 @@ +package iped.engine.task; + +import java.awt.image.BufferedImage; +import java.io.ByteArrayOutputStream; +import java.io.File; +import java.io.FileOutputStream; +import java.io.IOException; +import java.io.InputStream; +import java.net.SocketTimeoutException; +import java.net.UnknownHostException; +import java.nio.charset.StandardCharsets; +import java.security.cert.CertificateException; +import java.security.cert.X509Certificate; +import java.text.DecimalFormat; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashSet; +import java.util.LinkedList; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.TreeMap; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; +import java.util.zip.CRC32; +import java.util.zip.ZipEntry; +import java.util.zip.ZipOutputStream; + +import javax.imageio.ImageIO; +import javax.net.ssl.SSLContext; + +import org.apache.http.HttpEntity; +import org.apache.http.HttpStatus; +import org.apache.http.client.ClientProtocolException; +import org.apache.http.client.methods.CloseableHttpResponse; +import org.apache.http.client.methods.HttpGet; +import org.apache.http.client.methods.HttpPost; +import org.apache.http.conn.HttpHostConnectException; +import org.apache.http.conn.ssl.NoopHostnameVerifier; +import org.apache.http.conn.ssl.SSLConnectionSocketFactory; +import org.apache.http.conn.ssl.TrustStrategy; +import org.apache.http.entity.ContentType; +import org.apache.http.entity.mime.HttpMultipartMode; +import org.apache.http.entity.mime.MultipartEntityBuilder; +import org.apache.http.impl.client.CloseableHttpClient; +import org.apache.http.impl.client.HttpClients; +import org.apache.http.ssl.SSLContextBuilder; +import org.apache.tika.io.TemporaryResources; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; + +import iped.configuration.Configurable; +import iped.data.IHashValue; +import iped.data.IItem; +import iped.engine.config.ConfigurationManager; +import iped.engine.config.ImageThumbTaskConfig; +import iped.engine.config.RemoteImageExternalClassifierConfig; +import iped.engine.config.VideoThumbsConfig; +import iped.engine.preview.PreviewRepositoryManager; +import iped.engine.task.index.IndexItem; +import iped.parsers.util.MetadataUtil; +import iped.properties.ExtraProperties; +import iped.utils.ImageUtil; + +/** + * Performs remote classification of image and video files. + * + * @implNote Sends files in batches for improved performance (see 'batchSize' config property). + * Stores classification results in evidence's extra attributes. + * Attributes' names are controlled by the remote classifier (usually prefixed by 'AI'). + */ +public class RemoteImageExternalClassifierTask extends AbstractTask { + + private static Logger logger = LoggerFactory.getLogger(RemoteImageExternalClassifierTask.class); + + private RemoteImageExternalClassifierConfig config; + + // Task enabled status + private static boolean enabled = true; + + // Configuration parameters + private String urlZip; + private String urlVersion; + private int batchSize; + private int skipSize; + private int skipDimension; + private boolean skipHashDBFiles; + private boolean validateSSL; + private int thumbSize = 0; + + // AI-related attributes prefix + private static final String aiPrefix = "ai:externalClassifier:"; + + // AI classification extra attributes and values + private static final String AI_CLASSIFICATION_STATUS_ATTR = aiPrefix + "classificationStatus"; + private static final String AI_CLASSIFICATION_SUCCESS = "success"; + private static final String AI_CLASSIFICATION_FAIL_NO_CLASS = "failNoClass"; + private static final String AI_CLASSIFICATION_FAIL_NO_RESULTS = "failNoResults"; + private static final String AI_CLASSIFICATION_SKIP_SIZE = "skippedSize"; + private static final String AI_CLASSIFICATION_SKIP_DIMENSION = "skippedDimension"; + private static final String AI_CLASSIFICATION_SKIP_HASHDB = "skippedHashDB"; + public static final String AI_REVIEW_PRIORITY_ATTR = aiPrefix + "reviewPriority"; + private static final String AI_EXTERNAL_EMBEDDINGS_ATTR = aiPrefix + "externalEmbeddings"; + + // Classifications cache (avoids classification of duplicates) + private static ConcurrentHashMap classifications = new ConcurrentHashMap<>(); + + // Queue to store 'name' to 'evidence' mapping + private Map queue = new TreeMap<>(); + private LinkedList sendToNext = new LinkedList<>(); + + // Zip archive holding files to be sent for classification + private ZipFile zip; + + // Variables related to execution control and information + private static final AtomicBoolean abortNow = new AtomicBoolean(); + private static final AtomicBoolean finishNow = new AtomicBoolean(); + private static final AtomicInteger lastBatch = new AtomicInteger(); + + // Variables related to statistics + private static final AtomicLong classificationTime = new AtomicLong(); + private static final AtomicInteger classificationSuccess = new AtomicInteger(); + private static final AtomicInteger classificationFail = new AtomicInteger(); + private static final AtomicLong sendImageBytes = new AtomicLong(); + private static final AtomicLong sendVideoBytes = new AtomicLong(); + private static final AtomicInteger countImageThumbs = new AtomicInteger(); // also represents 'countImageFiles' for count purposes + private static final AtomicInteger countVideoThumbs = new AtomicInteger(); + private static final AtomicInteger countVideoFiles = new AtomicInteger(); + private static final AtomicInteger skipDuplicatesCount = new AtomicInteger(); + private static final AtomicInteger skipSizeCount = new AtomicInteger(); + private static final AtomicInteger skipDimensionCount = new AtomicInteger(); + private static final AtomicInteger skipHashDBFilesCount = new AtomicInteger(); + private static final AtomicInteger activeInstances = new AtomicInteger(0); + + // Number of the current batch of files being processed + private int currentBatch; + + // Number of retry failed classifications performed + private int retryCount; + + // Maximum number of attempts to retry failed classifications + private static final int MAX_RETRY = 10; + + // Wait time (ms) before retry attempts + private static final int WAIT_BEFORE_RETRY = 100; + + /** + * Represents the result of a classification. + */ + private static class ResultItem { + private String reviewPriority; + private List embeddings; + + public void setReviewPriority(String reviewPriority) { + this.reviewPriority = reviewPriority; + } + + public String getReviewPriority() { + return reviewPriority; + } + + public void setEmbeddings(List embeddings) { + this.embeddings = embeddings; + } + + public List getEmbeddings() { + return embeddings; + } + + public boolean hasAnyResult() { + return reviewPriority != null || embeddings != null; + } + } + + static boolean shouldReplaceReviewPriority(String currentPriority, String newPriority) { + if (newPriority == null || newPriority.isBlank()) { + return false; + } + if (currentPriority == null || currentPriority.isBlank()) { + return true; + } + return getReviewPriorityRank(newPriority) > getReviewPriorityRank(currentPriority); + } + + static int getReviewPriorityRank(String reviewPriority) { + if (reviewPriority == null) { + return Integer.MIN_VALUE; + } + switch (reviewPriority) { + case "VERY_HIGH": + return 5; + case "HIGH": + return 4; + case "MEDIUM": + return 3; + case "LOW": + return 2; + case "VERY_LOW": + return 1; + default: + return Integer.MIN_VALUE; + } + } + + /** + * Represents a Zip archive holding files to be sent for classification. + */ + private static class ZipFile { + TemporaryResources tmp; + File zipFile; + FileOutputStream fos; + ZipOutputStream zos; + int fileCount; + + public ZipFile() throws IOException { + tmp = new TemporaryResources(); + zipFile = tmp.createTemporaryFile(); + fos = new FileOutputStream(zipFile); + zos = new ZipOutputStream(fos); + fileCount = 0; + } + + public int getFileCount() { + return fileCount; + } + + public void addFileToZip(String filename, byte[] data) throws IOException { + ZipEntry zipEntry = new ZipEntry(filename); + zipEntry.setMethod(ZipEntry.STORED); + zipEntry.setSize(data.length); + zipEntry.setCompressedSize(data.length); + CRC32 crc = new CRC32(); + crc.update(data); + zipEntry.setCrc(crc.getValue()); + + zos.putNextEntry(zipEntry); + zos.write(data, 0, data.length); + zos.closeEntry(); + fileCount++; + } + + public File closeAndGetZip() throws IOException { + zos.close(); + fos.close(); + return zipFile; + } + + public void clean() throws IOException { + tmp.close(); + } + } + + @Override + public List> getConfigurables() { + return Arrays.asList(new RemoteImageExternalClassifierConfig()); + } + + public boolean isEnabled() { + return config.isEnabled() && enabled; + } + + @Override + public void init(ConfigurationManager configurationManager) throws Exception { + // Disable task during report generation + if (caseData.isIpedReport()) { + enabled = false; + } + + config = configurationManager.findObject(RemoteImageExternalClassifierConfig.class); + + if (!isEnabled()) { + return; + } + activeInstances.incrementAndGet(); + + urlZip = config.getUrl() + "/zip"; // enforces secure communication (required for sensitive data + // transfer) + urlVersion = config.getUrl() + "/version"; + batchSize = config.getBatchSize(); + skipSize = config.getSkipSize(); + skipDimension = config.getSkipDimension(); + skipHashDBFiles = config.isSkipHashDBFiles(); + validateSSL = config.isValidateSSL(); + + try (CloseableHttpClient client = getClient()) { + HttpGet get = new HttpGet(urlVersion); + + client.execute(get, response -> { + int statusCode = response.getStatusLine().getStatusCode(); + if (statusCode == HttpStatus.SC_OK) { + try (InputStream responseStream = response.getEntity().getContent()) { + ObjectMapper objectMapper = new ObjectMapper(); + JsonNode jsonResponse = objectMapper.readTree(responseStream); + String protocol_version = jsonResponse.get("protocol_version").asText(); + String model_version = jsonResponse.get("model_version").asText(); + JsonNode thumb_size = jsonResponse.get("thumbnail_size"); + if (thumb_size != null) { + thumbSize = thumb_size.asInt(); + } + logger.info( + "RemoteImageExternalClassifierTask: Connected to remote image classifier at '{}' - protocol_version: {} model_version: {} thumbnail_size: {}", urlVersion, protocol_version, model_version, thumbSize); + if (!protocol_version.startsWith("v1")) { + throw new RuntimeException("Incompatible protocol version: " + protocol_version); + } + if (!model_version.startsWith("external_v1")) { + throw new RuntimeException("Incompatible model version: " + model_version); + } + } + } else { + throw new HttpResponseStatusException( + "Failed to connect to remote image classifier at '" + urlVersion + "'", statusCode, null); + } + return null; + }); + } + catch (UnknownHostException | HttpHostConnectException e) { + // Disable task in case of failure to connect to remote image classifier + enabled = false; + logger.error("Task disabled. Failed to connect to remote image classifier at '" + urlVersion + "': " + e.getMessage()); + } + + if (isEnabled() && thumbSize != 0) { + ImageThumbTaskConfig imgThumbConfig = configurationManager.findObject(ImageThumbTaskConfig.class); + VideoThumbsConfig videoConfig = configurationManager.findObject(VideoThumbsConfig.class); + imgThumbConfig.setThumbSize(thumbSize); + videoConfig.setSize(thumbSize); + } + + if (zip == null) { + zip = new ZipFile(); + } + } + + public static String formatNumberToKMGUnits(float bytes) { + if (bytes < 1024 * 1024) { + return new DecimalFormat("#,##0.##").format((float) bytes / 1024.0) + " KB"; + } else if (bytes < 1024L * 1024 * 1024) { + return new DecimalFormat("#,##0.##").format((float) bytes / (1024.0 * 1024)) + " MB"; + } else + return new DecimalFormat("#,##0.##").format((float) bytes / (1024.0 * 1024 * 1024)) + " GB"; + } + + @Override + public void finish() throws Exception { + // "Close" classifications cache + if (activeInstances.decrementAndGet() == 0){ + classifications.clear(); + } + // Summary statistics + if (!finishNow.getAndSet(true)) { + long totClassifications = classificationSuccess.longValue() + classificationFail.longValue(); + long totSkipCount = skipSizeCount.longValue() + skipDimensionCount.longValue() + skipHashDBFilesCount.longValue() + skipDuplicatesCount.longValue(); + logger.info("Total count of files processed: {}", (totClassifications + totSkipCount)); + + // Statistics for files sent for classification + if (totClassifications != 0) { + logger.info(" Files sent for classification: {}", totClassifications); + logger.info(" Successful classifications: {}", classificationSuccess.intValue()); + logger.info(" Failed classifications: {}", classificationFail.intValue()); + if (countImageThumbs.intValue() > 0) + logger.info(" Image files: {} (1 thumb/image)", countImageThumbs.intValue()); + if (countVideoFiles.intValue() > 0) + logger.info(" Video files: {} ({} thumbs/video)", countVideoFiles.intValue(), (countVideoThumbs.intValue() / countVideoFiles.intValue())); + long totThumbs = countImageThumbs.intValue() + countVideoThumbs.intValue(); + if (totThumbs != 0) { + long totSendBytes = sendImageBytes.longValue() + sendVideoBytes.longValue(); + float sendThroughput = (float) totSendBytes / ((float) (classificationTime.longValue()) / 1000); + logger.info(" Thumbs sent: {} ({}); average throughput: {}/s", totThumbs, formatNumberToKMGUnits(totSendBytes), formatNumberToKMGUnits(sendThroughput)); + if (countImageThumbs.intValue() > 0) + logger.info(" Image thumbs: {} ({}); average thumb size: {}", countImageThumbs.intValue(), formatNumberToKMGUnits(sendImageBytes.longValue()), formatNumberToKMGUnits(sendImageBytes.longValue() / countImageThumbs.intValue())); + if (countVideoThumbs.intValue() > 0) + logger.info(" Videos thumbs: {} ({}); average thumb size: {}", countVideoThumbs.intValue(), formatNumberToKMGUnits(sendVideoBytes.longValue()), formatNumberToKMGUnits(sendVideoBytes.longValue() / countVideoThumbs.intValue())); + logger.info(" Average thumb classification time (ms/thumb): {}", String.format("%.3f", (((float) classificationTime.longValue() / this.worker.manager.getNumWorkers()) / totThumbs))); + logger.info(" Average thumb classification throughput (thumbs/s): {}", ((int) (totThumbs / ((float) classificationTime.longValue() / this.worker.manager.getNumWorkers()) * 1000))); + } + } + + // Statistics for files with skipped classification + if (totSkipCount != 0) { + logger.info(" Files with skipped classification: {}", totSkipCount); + logger.info(" Skipped classification by size: {}", skipSizeCount.intValue()); + logger.info(" Skipped classification by dimension: {}", skipDimensionCount.intValue()); + logger.info(" Skipped classification by hashDBFiles: {}", skipHashDBFilesCount.intValue()); + logger.info(" Skipped classification by duplicates: {}", skipDuplicatesCount.intValue()); + } + } + } + + private void sendItemsToNextTask() throws Exception { + sendToNext.addAll(queue.values()); + queue.clear(); + } + + protected void sendToNextTask(IItem item) throws Exception { + if (!isEnabled()) { + super.sendToNextTask(item); + return; + } + + LinkedList localList = new LinkedList(sendToNext); + sendToNext.clear(); + for (IItem it : localList) { + super.sendToNextTask(it); + } + if (!queue.containsValue(item) || item.isQueueEnd()) { + super.sendToNextTask(item); + } + } + + private void processResult(InputStream responseStream) throws IOException { + ObjectMapper objectMapper = new ObjectMapper(); + JsonNode jsonResponse = objectMapper.readTree(responseStream); + logger.debug("Server Response: {}", jsonResponse.toPrettyString()); + + // Queue to store 'name' of failed evidences + Set queueFail = new HashSet<>(); + + // Get 'results' from response + JsonNode resultArray = jsonResponse.get("results"); + Map results = new TreeMap<>(); + if (resultArray != null && resultArray.isArray()) { + // Process 'results' + for (JsonNode item : resultArray) { + // Checks result item for 'filename' field + JsonNode filenameNode = item.get("filename"); + if (filenameNode == null || filenameNode.isNull() || filenameNode.asText().isBlank()) { + logger.warn("Invalid/missing 'filename' field"); + continue; + } + String name = filenameNode.asText(); + ResultItem res = null; + int pathIdx = Math.max(name.lastIndexOf('/'), name.lastIndexOf('\\')); + if (pathIdx >= 0 && pathIdx + 1 < name.length()) { + name = name.substring(pathIdx + 1); + } + // The two 'if' conditions below allow for: + // - storing classification for images + // - storing grouped classification data for video frames and animation image frames + int frameIdxSep = name.indexOf('_'); + if (frameIdxSep > 0 && frameIdxSep < name.length() - 1) { + boolean hasNumericPrefix = true; + for (int i = 0; i < frameIdxSep; i++) { + if (!Character.isDigit(name.charAt(i))) { + hasNumericPrefix = false; + break; + } + } + if (hasNumericPrefix) { + name = name.substring(frameIdxSep + 1); + } + res = results.get(name); + } + if (res == null) { + res = new ResultItem(); + results.put(name, res); + } + + // New API payload + JsonNode reviewPriority = item.get("reviewPriority"); + if (reviewPriority != null && !reviewPriority.isNull() && !reviewPriority.asText().isBlank() + && shouldReplaceReviewPriority(res.getReviewPriority(), reviewPriority.asText())) { + res.setReviewPriority(reviewPriority.asText()); + } + JsonNode embeddings = item.get("embeddings"); + if (embeddings != null && embeddings.isArray()) { + List vector = new ArrayList<>(); + for (JsonNode emb : embeddings) { + vector.add(emb.asDouble()); + } + if (!vector.isEmpty()) { + res.setEmbeddings(vector); + } + } + + if (res.getReviewPriority() == null && res.getEmbeddings() == null) { + // Invalid/missing expected fields from external classifier response + // Server may have encountered a problem during evidence classification (e.g., corrupted image) + // Flag classification fail status in evidence attributes + // Store failed evidence in case cache + + // Check for matching evidence + IItem evidence = queue.get(name); + if (evidence != null) { + // Matching evidence found + // Add classification fail info, if not already exists + // Avoids duplicate counting of failed videos, in case of classification problem with more than one thumb + if (!queueFail.contains(name)) { + // Add evidence to the fail queue + // Avoids signaling evidence classification success later on for videos with classification problem with some, but not all, thumbs + queueFail.add(name); + // Classification fail + classificationFail.incrementAndGet(); + // Add classification status + evidence.setExtraAttribute(AI_CLASSIFICATION_STATUS_ATTR, AI_CLASSIFICATION_FAIL_NO_CLASS); + logger.warn("ClassificationFail::EvidenceNoResult: Invalid/missing 'reviewPriority' and 'embeddings' fields for filename: {}", name); + } + } + else { + // No matching evidence found + logger.warn("ClassificationFail::EvidenceNotFound: Invalid/missing 'reviewPriority' and 'embeddings' fields for filename: {}. No matching evidence found.", name); + } + } + } + + // Add classification to evidences + // Store classification in case cache + for (String name : results.keySet()) { + // Find the evidence to which the classification belongs + IItem evidence; + if (name == null || (evidence = queue.get(name)) == null) { + // No matching evidence found + logger.warn("No matching evidence found"); + continue; + } + // Check if evidence is in fail queue + if (name != null && queueFail.contains(name)) { + // Classification fail + continue; + } + ResultItem res = results.get(name); + if (res.hasAnyResult()) { + // Classification success + classificationSuccess.incrementAndGet(); + // Add classification status + evidence.setExtraAttribute(AI_CLASSIFICATION_STATUS_ATTR, AI_CLASSIFICATION_SUCCESS); + + StringBuilder classes = new StringBuilder(); + if (res.getReviewPriority() != null) { + evidence.setExtraAttribute(AI_REVIEW_PRIORITY_ATTR, res.getReviewPriority()); + classes.append(AI_REVIEW_PRIORITY_ATTR).append('=').append(res.getReviewPriority()); + } + + if (res.getEmbeddings() != null) { + String embeddings = embeddingsToString(res.getEmbeddings()); + evidence.setExtraAttribute(AI_EXTERNAL_EMBEDDINGS_ATTR, embeddings); + if (classes.length() > 0) { + classes.append(';'); + } + classes.append(AI_EXTERNAL_EMBEDDINGS_ATTR).append('=').append(embeddings); + } + + // Store classification in classifications cache + classifications.put(evidence.getHashValue(), classes.toString()); + } + } + } else { + // 'results' array is missing in JSON response + // Server malformed response + // Flag classification fail status in evidence attributes + // Do not store failed evidences in case cache + + // Store fail information for each evidence in queue + for (IItem evidence : queue.values()) { + // Classification fail + classificationFail.incrementAndGet(); + // Add classification status + evidence.setExtraAttribute(AI_CLASSIFICATION_STATUS_ATTR, AI_CLASSIFICATION_FAIL_NO_RESULTS); + } + logger.error("ClassificationFail::NoResults: 'results' array is missing in JSON response. Classification fail for a batch of {} files", queue.size()); + } + } + + private CloseableHttpClient getClient() { + if (!validateSSL) { + SSLContext sslContext; + try { + sslContext = SSLContextBuilder.create().loadTrustMaterial(null, new TrustStrategy() { + @Override + public boolean isTrusted(X509Certificate[] chain, String authType) throws CertificateException { + // Trust all certs + return true; + } + }).build(); + } catch (Exception e) { + throw new RuntimeException("Failed to create a trusting SSL context", e); + } + + // Use NoopHostnameVerifier to ignore host name verification + SSLConnectionSocketFactory sslSocketFactory = new SSLConnectionSocketFactory(sslContext, + NoopHostnameVerifier.INSTANCE); + return HttpClients.custom().setSSLSocketFactory(sslSocketFactory).build(); + } else { + return HttpClients.createDefault(); + } + } + + // Custom HttpResponseStatusException exception + private static class HttpResponseStatusException extends IOException { + private static final long serialVersionUID = -7955212330143589634L; + private final int statusCode; + private final String responseBody; + + public HttpResponseStatusException(String message, int statusCode, String responseBody) { + super(message); + this.statusCode = statusCode; + this.responseBody = responseBody; + } + + public int getStatusCode() { + return statusCode; + } + + @Override + public String toString() { + String str = String.format("%s HTTP Status Code: %d", getMessage(), statusCode); + if (responseBody != null && !responseBody.isEmpty()) { + str += String.format(" - Response: %s", responseBody); + } + return str; + } + } + + private void sendZipFile(File zipFile) throws IOException { + currentBatch = lastBatch.incrementAndGet(); + logger.info("Send ZIP file #{} (files: {})", currentBatch, zip.getFileCount()); + + try (CloseableHttpClient client = getClient()) { + HttpPost post = new HttpPost(urlZip); + + HttpEntity entity = MultipartEntityBuilder.create().setMode(HttpMultipartMode.BROWSER_COMPATIBLE) + .addBinaryBody("file", zipFile, ContentType.APPLICATION_OCTET_STREAM, zipFile.getName()).build(); + + post.setEntity(entity); + + long t = System.currentTimeMillis(); + try (CloseableHttpResponse response = client.execute(post)) { + classificationTime.addAndGet(System.currentTimeMillis() - t); + int statusCode = response.getStatusLine().getStatusCode(); + if (statusCode == HttpStatus.SC_OK) { + try (InputStream responseStream = response.getEntity().getContent()) { + processResult(responseStream); + } + } else { + // Server response with HTTP related issues (e.g., server is busy, overloaded, or down) + // HTTP response might have useful info + String errorMessage; + try (InputStream errorStream = response.getEntity().getContent()) { + errorMessage = new String(errorStream.readAllBytes(), StandardCharsets.UTF_8); + } catch (IOException e) { + errorMessage = "Failed to read error message from response."; + } + String errMsg = String.format("HTTP Status Code: %d - Response:%n%s", statusCode, errorMessage); + throw new HttpResponseStatusException(errMsg, statusCode, errorMessage); + } + } + } + } + + private void sendBatchedFiles() throws Exception { + // In case of errors, retry, if possible + while (retryCount <= MAX_RETRY) { + try { + sendZipFile(zip.closeAndGetZip()); + break; + } + catch (IOException e) { + String baseMsg = String.format(" Failed to upload ZIP file."); + + // In case of non recoverable errors, abort case processing + // Abort if unknown host on URL + if (e instanceof UnknownHostException) { + if (!abortNow.getAndSet(true)) { + logger.error("ClassificationFail::UnknownHost:{} Unknown host on '{}': {}", baseMsg, urlZip, e.getClass().getName()); + throw new RuntimeException("Unknown host on '" + urlZip + "': " + e.getMessage(), e); + } + } + // Abort if HTTP response status code is different from SC_SERVICE_UNAVAILABLE and SC_GATEWAY_TIMEOUT + if (e instanceof HttpResponseStatusException) { + HttpResponseStatusException eHTTP = (HttpResponseStatusException) e; + if (eHTTP.getStatusCode() != HttpStatus.SC_SERVICE_UNAVAILABLE && eHTTP.getStatusCode() != HttpStatus.SC_GATEWAY_TIMEOUT) { + if (!abortNow.getAndSet(true)) { + logger.error("ClassificationFail::HttpStatusNotOK: {}", e.getMessage()); + throw new RuntimeException("HTTP Status Not OK on '" + urlZip + "': " + e.getMessage(), e); + } + } + } + + // In case of recoverable errors, retry sending batched files + // Log warning/error messages + baseMsg = String.format(" Failed to upload ZIP file #%d (files: %d).", currentBatch, zip.getFileCount()); + if (retryCount == MAX_RETRY) + baseMsg = String.format(" Failed to upload ZIP file."); + if (retryCount > 0) + baseMsg = String.format("retry#%d:", retryCount) + baseMsg; + String msg = ""; + if (e instanceof HttpResponseStatusException) { + HttpResponseStatusException eHTTP = (HttpResponseStatusException) e; + if (eHTTP.getStatusCode() == HttpStatus.SC_SERVICE_UNAVAILABLE) + msg += String.format("ClassificationFail::HttpStatusNotOK:%s Service unavailable for '%s': %s: %s", baseMsg, urlZip, e.getClass().getName(), e.getMessage()); + if (eHTTP.getStatusCode() == HttpStatus.SC_GATEWAY_TIMEOUT) + msg += String.format("ClassificationFail::HttpStatusNotOK:%s Gateway timeout for '%s': %s: %s", baseMsg, urlZip, e.getClass().getName(), e.getMessage()); + } + else if (e instanceof HttpHostConnectException) + msg += String.format("ClassificationFail::ConnectionProblem:%s Could not connect to host on '%s': %s: %s", baseMsg, urlZip, e.getClass().getName(), e.getMessage()); + else if (e instanceof SocketTimeoutException) + msg += String.format("ClassificationFail::ConnectionProblem:%s Socket timeout occurred while connecting or reading from '%s': %s: %s", baseMsg, urlZip, e.getClass().getName(), e.getMessage()); + else if (e instanceof ClientProtocolException) + msg += String.format("ClassificationFail::ConnectionProblem:%s HTTP protocol error while communicating with '%s': %s: %s", baseMsg, urlZip, e.getClass().getName(), e.getMessage()); + else + msg += String.format("ClassificationFail::IOProblem:%s I/O error occurred during HTTP request to '%s': %s: %s", baseMsg, urlZip, e.getClass().getName(), e.getMessage()); + if (retryCount < MAX_RETRY) { + // Log warning and retry + logger.warn(msg); + } else { + // Log error and abort case processing + if (!abortNow.getAndSet(true)) { + logger.error(msg); + logger.error("ClassificationFail::TooManyErrors: Aborting case processing"); + throw new RuntimeException( + "Too many errors while communicating with '" + urlZip + "': " + e.getMessage(), e); + } + } + + // Increment 'retryCount' + retryCount++; + + // Wait time before retrying + Thread.sleep(WAIT_BEFORE_RETRY); + } + } + } + + @Override + protected boolean processQueueEnd() { + return true; + } + + @Override + protected void process(IItem evidence) throws Exception { + // Send files to the remote classifier if any condition is met + // (if maximum number of thumbs in current batch has been reached or current batch is the last one) + if (zip.getFileCount() >= batchSize || (evidence.isQueueEnd() && zip.getFileCount() > 0)) { + // Send batched files + sendBatchedFiles(); + zip.clean(); + zip = new ZipFile(); + sendItemsToNextTask(); + } + + // Does not process evidence if any condition is met + if (!isEnabled() || !evidence.isToAddToCase() || evidence.getHashValue() == null || evidence.getThumb() == null + || evidence.getThumb().length < 10 || evidence.isQueueEnd()) { + return; + } + + // Skip classification of images/videos smaller than a given file size (see 'skipSize' config property) + if (skipSize > 0 && skipSize > evidence.getLength()) { + // Add skip classification info + evidence.setExtraAttribute(AI_CLASSIFICATION_STATUS_ATTR, AI_CLASSIFICATION_SKIP_SIZE); + skipSizeCount.incrementAndGet(); + return; + } + + // Skip classification of images/videos smaller than a given dimension, i.e. height or width (see 'skipDimension' config property) + if (skipDimension > 0) { + int width = 0; + int height = 0; + String mediaType = evidence.getMediaType().toString(); + if (mediaType.startsWith("image")) { + try { + if (evidence.getMetadata().get("image:Width") != null) + width = Integer.parseInt(evidence.getMetadata().get("image:Width")); + if (evidence.getMetadata().get("image:Height") != null) + height = Integer.parseInt(evidence.getMetadata().get("image:Height")); + } + catch (NumberFormatException e) { + logger.warn("Invalid dimension for image file '{}': width:{}; height:{}", evidence.getName(), evidence.getMetadata().get("image:Width"), evidence.getMetadata().get("image:Height")); + } + } + else if (mediaType.startsWith("video")) { + try { + if (evidence.getMetadata().get("video:Width") != null) + width = Integer.parseInt(evidence.getMetadata().get("video:Width")); + if (evidence.getMetadata().get("video:Height") != null) + height = Integer.parseInt(evidence.getMetadata().get("video:Height")); + } + catch (NumberFormatException e) { + logger.warn("Invalid dimension for video file '{}': width:{}; height:{}", evidence.getName(), evidence.getMetadata().get("video:Width"), evidence.getMetadata().get("video:Height")); + } + } + if ((skipDimension > width || skipDimension > height) && width > 0 && height > 0) { + // Add skip classification info + evidence.setExtraAttribute(AI_CLASSIFICATION_STATUS_ATTR, AI_CLASSIFICATION_SKIP_DIMENSION); + skipDimensionCount.incrementAndGet(); + return; + } + } + + // Skip classification of images/videos with hits on IPED hashesDB database (see 'skipHashDBFiles' config property) + if (skipHashDBFiles && evidence.getExtraAttribute(ExtraProperties.HASHDB_STATUS) != null) { + // Add skip classification info + evidence.setExtraAttribute(AI_CLASSIFICATION_STATUS_ATTR, AI_CLASSIFICATION_SKIP_HASHDB); + skipHashDBFilesCount.incrementAndGet(); + return; + } + + // Skip classification of images/videos duplicates if classification exists in classifications cache + // Retrieve classification from classifications cache + String classes = classifications.get(evidence.getHashValue()); + if (classes != null) { + // Classification exists in classifications cache + // Add classification status + evidence.setExtraAttribute(AI_CLASSIFICATION_STATUS_ATTR, AI_CLASSIFICATION_SUCCESS); + // Add classification classes to the evidence + String[] classesArray = classes.split(";"); + for (int i = 0; i < classesArray.length; i++) { + // classParts[0] will hold className and classParts[1] will hold classProb + String[] classParts = classesArray[i].split("=", 2); + if (classParts.length != 2) { + continue; + } + // Add classification class to the evidence + String key = classParts[0]; + if (key.equals(AI_REVIEW_PRIORITY_ATTR) || key.equals(AI_EXTERNAL_EMBEDDINGS_ATTR)) { + evidence.setExtraAttribute(key, classParts[1]); + } + } + skipDuplicatesCount.incrementAndGet(); + return; + } + + // 'name' is a key to map to the evidence + String name = evidence.getExtraAttribute(IndexItem.TRACK_ID).toString() + ".jpg"; + if (MetadataUtil.isVideoType(evidence.getMediaType()) || MetadataUtil.isAnimationImage(evidence)) { + // For videos, call the detection method for each extracted frame image (VideoThumbsTask must be enabled) + File viewFile = evidence.getViewFile(); + List frames = null; + if (viewFile != null && viewFile.exists()) { + frames = ImageUtil.getFrames(viewFile); + } else if (evidence.hasPreview()) { + try (InputStream is = PreviewRepositoryManager.get(output).readPreview(evidence, false)) { + frames = ImageUtil.getFrames(is); + } + } + if (frames != null) { + int i = 0; + for (BufferedImage frame : frames) { + String iName = (++i) + "_" + name; + ByteArrayOutputStream baos = new ByteArrayOutputStream(); + ImageIO.write(frame, "jpeg", baos); + zip.addFileToZip(iName, baos.toByteArray()); + sendVideoBytes.addAndGet(baos.toByteArray().length); + } + countVideoThumbs.addAndGet(i); + } else { + countVideoThumbs.incrementAndGet(); + sendVideoBytes.addAndGet(evidence.getThumb().length); + zip.addFileToZip(name, evidence.getThumb()); + } + countVideoFiles.incrementAndGet(); + } else { + countImageThumbs.incrementAndGet(); + sendImageBytes.addAndGet(evidence.getThumb().length); + zip.addFileToZip(name, evidence.getThumb()); + } + + // Add a new 'name' to 'evidence' mapping to the queue + queue.put(name, evidence); + } + + private static String embeddingsToString(List embeddings) { + StringBuilder sb = new StringBuilder(); + for (int i = 0; i < embeddings.size(); i++) { + if (i > 0) { + sb.append(','); + } + sb.append(embeddings.get(i)); + } + return sb.toString(); + } +}