Dlp

core/src/main/java/com/predic8/membrane/core/DLP.java

@@ -0,0 +1,109 @@
+package com.predic8.membrane.core;
+
+import com.fasterxml.jackson.core.JsonFactory;
+import com.fasterxml.jackson.core.JsonParser;
+import com.fasterxml.jackson.core.JsonToken;
+import com.predic8.membrane.core.http.Message;
+
+import java.io.InputStreamReader;
+import java.io.IOException;
+import java.util.*;
+import java.util.logging.Logger;
+
+public class DLP {
+
+    private static final Logger logger = Logger.getLogger(DLP.class.getName());
+
+    private final Map<String, String> riskDict;
+    private static final JsonFactory JSON_FACTORY = new JsonFactory();
+
+    public DLP(Map<String, String> riskDict) {
+        this.riskDict = riskDict;
+    }
+
+    public Map<String, Object> analyze(Message msg) {
+        return evaluateRisk(extractFieldNames(msg));
+    }
+
+    private Set<String> extractFieldNames(Message msg) {
+        Set<String> fieldNames = new HashSet<>();
+        try {
+            JsonParser parser = JSON_FACTORY.createParser(new InputStreamReader(msg.getBodyAsStreamDecoded(), msg.getCharset()));
+            Deque<String> contextStack = new ArrayDeque<>();
+            String currentFieldName = null;
+
+            while (parser.nextToken() != null) {
+                JsonToken token = parser.getCurrentToken();
+
+                switch (token) {
+                    case START_OBJECT:
+                    case START_ARRAY:
+                        if (currentFieldName != null) {
+                            contextStack.push(currentFieldName);
+                            currentFieldName = null;
+                        }
+                        break;
+
+                    case END_OBJECT:
+                    case END_ARRAY:
+                        if (!contextStack.isEmpty()) {
+                            contextStack.pop();
+                        }
+                        break;
+
+                    case FIELD_NAME:
+                        currentFieldName = parser.currentName();
+                        fieldNames.add(buildFullPath(contextStack, currentFieldName).toLowerCase());
+                        break;
+
+                    default:
+                        break;
+                }
+            }
+        } catch (IOException e) {
+            throw new RuntimeException("Error extracting JSON field names", e);
+        }
+        return fieldNames;
+    }
+
+    private String buildFullPath(Deque<String> contextStack, String currentFieldName) {
+        List<String> path = new ArrayList<>(contextStack);
+        path.add(currentFieldName);
+        return String.join(".", path);
+    }
+
+
+    private Map<String, Object> evaluateRisk(Set<String> fieldNames) {
+        Map<String, String> matchedFields = new HashMap<>();
+        int high = 0;
+        int medium = 0;
+        int low = 0;
+        int unclassified = 0;
+
+        for (String field : fieldNames) {
+            String risk = riskDict.getOrDefault(field, "unclassified");
+            matchedFields.put(field, risk);
+
+            switch (risk.toLowerCase()) {
+                case "high": high++; break;
+                case "medium": medium++; break;
+                case "low": low++; break;
+                default: unclassified++; break;
+            }
+        }
+
+        logger.info(String.format(
+                "High: %d, Medium: %d, Low: %d, Unclassified: %d",
+                high, medium, low, unclassified
+        ));
+
+        Map<String, Object> result = new LinkedHashMap<>();
+        result.put("high_risk", high);
+        result.put("medium_risk", medium);
+        result.put("low_risk", low);
+        result.put("unclassified", unclassified);
+        result.put("matched_fields", matchedFields);
+
+        return result;
+    }
+}

core/src/test/java/com/predic8/membrane/core/DLPTest.java

@@ -0,0 +1,64 @@
+package com.predic8.membrane.core;
+
+import com.predic8.membrane.core.http.Message;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.io.ByteArrayInputStream;
+import java.nio.charset.StandardCharsets;
+import java.util.HashMap;
+import java.util.Map;
+
+import static org.junit.jupiter.api.Assertions.*;
+import static org.mockito.Mockito.*;
+
+class DLPTest {
+
+    private Message mockMessage;
+    private Map<String, String> riskDict;
+    String json = """
+                   {
+                   "user": {
+                     "email": "[email protected]",
+                     "profile": {
+                       "firstName": "John",
+                       "lastName": "Doe"
+                     }
+                   },
+                   "active": true
+                 }
+                """;
+
+    @BeforeEach
+    void setUp() {
+        mockMessage = mock(Message.class);
+        when(mockMessage.getBodyAsStreamDecoded())
+                .thenReturn(new ByteArrayInputStream(json.getBytes(StandardCharsets.UTF_8)));
+        when(mockMessage.getCharset())
+                .thenReturn(String.valueOf(StandardCharsets.UTF_8));
+
+
+        riskDict = new HashMap<>();
+        riskDict.put("user.email", "High");
+        riskDict.put("user.profile", "High");
+        riskDict.put("active", "Low");
+        riskDict.put("nothing", "Low");
+    }
+
+    @Test
+    @SuppressWarnings("unchecked")
+    void analyze() {
+        DLP dlp = new DLP(riskDict);
+        Map<String, Object> result = dlp.analyze(mockMessage);
+
+        assertEquals(2, result.get("high_risk"));
+        assertEquals(0, result.get("medium_risk"));
+        assertEquals(1, result.get("low_risk"));
+        assertEquals(3, result.get("unclassified"));
+
+        Map<String, String> matchedFields = (Map<String, String>) result.get("matched_fields");
+        assertEquals("High", matchedFields.get("user.email"));
+        assertEquals("High", matchedFields.get("user.profile"));
+        assertEquals("Low", matchedFields.get("active"));
+    }
+}