diff --git a/metron-platform/metron-common/README.md b/metron-platform/metron-common/README.md index 60bd28715b..678c6a2483 100644 --- a/metron-platform/metron-common/README.md +++ b/metron-platform/metron-common/README.md @@ -68,18 +68,28 @@ The configuration is a complex JSON object with the following top level fields: ###The `enrichment` Configuration -| Field | Description | Example | -|------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|---------------------------------------------------------------------------| +| Field | Description | Example | +|------------------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|------------------------------------------------------------------| | `fieldToTypeMap` | In the case of a simple HBase enrichment (i.e. a key/value lookup), the mapping between fields and the enrichment types associated with those fields must be known. This enrichment type is used as part of the HBase key. | `"fieldToTypeMap" : { "ip_src_addr" : [ "asset_enrichment" ] }` | -| `fieldMap` | The map of enrichment bolts names to fields in the JSON messages.,Each field is sent to the enrichment referenced in the key. | `"fieldMap": {"hbaseEnrichment": ["ip_src_addr","ip_dst_addr"]}` | +| `fieldMap` | The map of enrichment bolts names to fields in the JSON messages.,Each field is sent to the enrichment referenced in the key. | `"fieldMap": {"hbaseEnrichment": ["ip_src_addr","ip_dst_addr"]}` | +| `config` | The general configuration for the enrichment | `"config": {"typeToColumnFamily": { "asset_enrichment","cf" } }` | + +The `config` map is intended to house enrichment specific configuration. +For instance, for the `hbaseEnrichment`, the mappings between the +enrichment types to the column families is specified. ###The `threatIntel` Configuration -| Field | Description | Example | -|------------------|--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------| +| Field | Description | Example | +|------------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------| | `fieldToTypeMap` | In the case of a simple HBase threat intel enrichment (i.e. a key/value lookup), the mapping between fields and the enrichment types associated with those fields must be known. This enrichment type is used as part of the HBase key. | `"fieldToTypeMap" : { "ip_src_addr" : [ "malicious_ips" ] }` | | `fieldMap` | The map of threat intel enrichment bolts names to fields in the JSON messages. Each field is sent to the threat intel enrichment bolt referenced in the key. | `"fieldMap": {"hbaseThreatIntel": ["ip_src_addr","ip_dst_addr"]}` | -| `triageConfig` | The configuration of the threat triage scorer. In the situation where a threat is detected, a score is assigned to the message and embedded in the indexed message. | `"riskLevelRules" : { "IN_SUBNET(ip_dst_addr, '192.168.0.0/24')" : 10 }` | +| `triageConfig` | The configuration of the threat triage scorer. In the situation where a threat is detected, a score is assigned to the message and embedded in the indexed message. | `"riskLevelRules" : { "IN_SUBNET(ip_dst_addr, '192.168.0.0/24')" : 10 }` | +| `config` | The general configuration for the Threat Intel | `"config": {"typeToColumnFamily": { "malicious_ips","cf" } }` | + +The `config` map is intended to house threat intel specific configuration. +For instance, for the `hbaseThreatIntel` threat intel adapter, the mappings between the +enrichment types to the column families is specified. The `triageConfig` field is also a complex field and it bears some description: diff --git a/metron-platform/metron-common/src/main/java/org/apache/metron/common/configuration/enrichment/EnrichmentConfig.java b/metron-platform/metron-common/src/main/java/org/apache/metron/common/configuration/enrichment/EnrichmentConfig.java index af6c148eb6..932b0704f4 100644 --- a/metron-platform/metron-common/src/main/java/org/apache/metron/common/configuration/enrichment/EnrichmentConfig.java +++ b/metron-platform/metron-common/src/main/java/org/apache/metron/common/configuration/enrichment/EnrichmentConfig.java @@ -25,6 +25,15 @@ public class EnrichmentConfig { private Map> fieldMap = new HashMap<>(); private Map> fieldToTypeMap = new HashMap<>(); + private Map config = new HashMap<>(); + + public Map getConfig() { + return config; + } + + public void setConfig(Map config) { + this.config = config; + } public Map> getFieldMap() { return fieldMap; @@ -42,6 +51,15 @@ public void setFieldToTypeMap(Map> fieldToTypeMap) { this.fieldToTypeMap = fieldToTypeMap; } + @Override + public String toString() { + return "EnrichmentConfig{" + + "fieldMap=" + fieldMap + + ", fieldToTypeMap=" + fieldToTypeMap + + ", config=" + config + + '}'; + } + @Override public boolean equals(Object o) { if (this == o) return true; @@ -50,7 +68,9 @@ public boolean equals(Object o) { EnrichmentConfig that = (EnrichmentConfig) o; if (getFieldMap() != null ? !getFieldMap().equals(that.getFieldMap()) : that.getFieldMap() != null) return false; - return getFieldToTypeMap() != null ? getFieldToTypeMap().equals(that.getFieldToTypeMap()) : that.getFieldToTypeMap() == null; + if (getFieldToTypeMap() != null ? !getFieldToTypeMap().equals(that.getFieldToTypeMap()) : that.getFieldToTypeMap() != null) + return false; + return getConfig() != null ? getConfig().equals(that.getConfig()) : that.getConfig() == null; } @@ -58,14 +78,7 @@ public boolean equals(Object o) { public int hashCode() { int result = getFieldMap() != null ? getFieldMap().hashCode() : 0; result = 31 * result + (getFieldToTypeMap() != null ? getFieldToTypeMap().hashCode() : 0); + result = 31 * result + (getConfig() != null ? getConfig().hashCode() : 0); return result; } - - @Override - public String toString() { - return "EnrichmentConfig{" + - "fieldMap=" + fieldMap + - ", fieldToTypeMap=" + fieldToTypeMap + - '}'; - } } diff --git a/metron-platform/metron-data-management/src/test/java/org/apache/metron/dataloads/hbase/mr/LeastRecentlyUsedPrunerIntegrationTest.java b/metron-platform/metron-data-management/src/test/java/org/apache/metron/dataloads/hbase/mr/LeastRecentlyUsedPrunerIntegrationTest.java index 42ec575cdb..93c216cffd 100644 --- a/metron-platform/metron-data-management/src/test/java/org/apache/metron/dataloads/hbase/mr/LeastRecentlyUsedPrunerIntegrationTest.java +++ b/metron-platform/metron-data-management/src/test/java/org/apache/metron/dataloads/hbase/mr/LeastRecentlyUsedPrunerIntegrationTest.java @@ -37,6 +37,7 @@ import org.apache.metron.enrichment.lookup.accesstracker.PersistentAccessTracker; import org.apache.metron.dataloads.bulk.ThreatIntelBulkLoader; import org.apache.metron.dataloads.nonbulk.taxii.TaxiiLoader; +import org.apache.metron.enrichment.lookup.handler.KeyWithContext; import org.junit.After; import org.junit.Assert; import org.junit.Before; @@ -118,7 +119,7 @@ public void test() throws Exception { ) ) ); - Assert.assertTrue(lookup.exists((EnrichmentKey)k, testTable, true)); + Assert.assertTrue(lookup.exists((EnrichmentKey)k, new EnrichmentLookup.HBaseContext(testTable, cf), true)); } pat.persist(true); for(LookupKey k : goodKeysOtherHalf) { @@ -129,7 +130,7 @@ public void test() throws Exception { ) ) ); - Assert.assertTrue(lookup.exists((EnrichmentKey)k, testTable, true)); + Assert.assertTrue(lookup.exists((EnrichmentKey)k, new EnrichmentLookup.HBaseContext(testTable, cf), true)); } testUtil.flush(); Assert.assertFalse(lookup.getAccessTracker().hasSeen(goodKeysHalf.get(0))); @@ -153,10 +154,10 @@ public void test() throws Exception { Job job = LeastRecentlyUsedPruner.createJob(config, tableName, cf, atTableName, atCF, ts); Assert.assertTrue(job.waitForCompletion(true)); for(LookupKey k : goodKeys) { - Assert.assertTrue(lookup.exists((EnrichmentKey)k, testTable, true)); + Assert.assertTrue(lookup.exists((EnrichmentKey)k, new EnrichmentLookup.HBaseContext(testTable, cf), true)); } for(LookupKey k : badKey) { - Assert.assertFalse(lookup.exists((EnrichmentKey)k, testTable, true)); + Assert.assertFalse(lookup.exists((EnrichmentKey)k, new EnrichmentLookup.HBaseContext(testTable, cf), true)); } } diff --git a/metron-platform/metron-enrichment/src/main/java/org/apache/metron/enrichment/adapters/simplehbase/SimpleHBaseAdapter.java b/metron-platform/metron-enrichment/src/main/java/org/apache/metron/enrichment/adapters/simplehbase/SimpleHBaseAdapter.java index 22629a492b..a67015041d 100644 --- a/metron-platform/metron-enrichment/src/main/java/org/apache/metron/enrichment/adapters/simplehbase/SimpleHBaseAdapter.java +++ b/metron-platform/metron-enrichment/src/main/java/org/apache/metron/enrichment/adapters/simplehbase/SimpleHBaseAdapter.java @@ -70,9 +70,11 @@ public JSONObject enrich(CacheKey value) { try { for (LookupKV kv : lookup.get(Iterables.transform(enrichmentTypes - , new EnrichmentUtils.TypeToKey(value.getValue()) + , new EnrichmentUtils.TypeToKey( value.getValue() + , lookup.getTable() + , value.getConfig().getEnrichment() + ) ) - , lookup.getTable() , false ) ) diff --git a/metron-platform/metron-enrichment/src/main/java/org/apache/metron/enrichment/adapters/threatintel/ThreatIntelAdapter.java b/metron-platform/metron-enrichment/src/main/java/org/apache/metron/enrichment/adapters/threatintel/ThreatIntelAdapter.java index ee5636bc17..5183272cb1 100644 --- a/metron-platform/metron-enrichment/src/main/java/org/apache/metron/enrichment/adapters/threatintel/ThreatIntelAdapter.java +++ b/metron-platform/metron-enrichment/src/main/java/org/apache/metron/enrichment/adapters/threatintel/ThreatIntelAdapter.java @@ -74,9 +74,11 @@ public JSONObject enrich(CacheKey value) { try { for (Boolean isThreat : lookup.exists(Iterables.transform(enrichmentTypes - , new EnrichmentUtils.TypeToKey(value.getValue()) + , new EnrichmentUtils.TypeToKey(value.getValue() + , lookup.getTable() + , value.getConfig().getThreatIntel() + ) ) - , lookup.getTable() , false ) ) diff --git a/metron-platform/metron-enrichment/src/main/java/org/apache/metron/enrichment/lookup/EnrichmentLookup.java b/metron-platform/metron-enrichment/src/main/java/org/apache/metron/enrichment/lookup/EnrichmentLookup.java index f43f85409a..3ee9754fca 100644 --- a/metron-platform/metron-enrichment/src/main/java/org/apache/metron/enrichment/lookup/EnrichmentLookup.java +++ b/metron-platform/metron-enrichment/src/main/java/org/apache/metron/enrichment/lookup/EnrichmentLookup.java @@ -17,6 +17,7 @@ */ package org.apache.metron.enrichment.lookup; +import com.google.common.collect.Iterables; import org.apache.hadoop.hbase.client.Get; import org.apache.hadoop.hbase.client.HTableInterface; import org.apache.hadoop.hbase.client.Result; @@ -25,41 +26,65 @@ import org.apache.metron.enrichment.converter.EnrichmentKey; import org.apache.metron.enrichment.converter.EnrichmentValue; import org.apache.metron.enrichment.lookup.accesstracker.AccessTracker; +import org.apache.metron.enrichment.lookup.handler.KeyWithContext; import java.io.IOException; import java.util.ArrayList; +import java.util.Collections; +import java.util.Iterator; import java.util.List; -public class EnrichmentLookup extends Lookup> implements AutoCloseable { +public class EnrichmentLookup extends Lookup> implements AutoCloseable { - public static class Handler implements org.apache.metron.enrichment.lookup.handler.Handler> { + public static class HBaseContext { + private HTableInterface table; + private String columnFamily; + public HBaseContext(HTableInterface table, String columnFamily) { + this.table = table; + this.columnFamily = columnFamily; + } + + public HTableInterface getTable() { return table; } + public String getColumnFamily() { return columnFamily; } + } + + public static class Handler implements org.apache.metron.enrichment.lookup.handler.Handler> { String columnFamily; HbaseConverter converter = new EnrichmentConverter(); public Handler(String columnFamily) { this.columnFamily = columnFamily; } + + private String getColumnFamily(HBaseContext context) { + return context.getColumnFamily() == null?columnFamily:context.getColumnFamily(); + } + @Override - public boolean exists(EnrichmentKey key, HTableInterface table, boolean logAccess) throws IOException { - return table.exists(converter.toGet(columnFamily, key)); + public boolean exists(EnrichmentKey key, HBaseContext context, boolean logAccess) throws IOException { + return context.getTable().exists(converter.toGet(getColumnFamily(context), key)); } @Override - public LookupKV get(EnrichmentKey key, HTableInterface table, boolean logAccess) throws IOException { - return converter.fromResult(table.get(converter.toGet(columnFamily, key)), columnFamily); + public LookupKV get(EnrichmentKey key, HBaseContext context, boolean logAccess) throws IOException { + return converter.fromResult(context.getTable().get(converter.toGet(getColumnFamily(context), key)), getColumnFamily(context)); } - private List keysToGets(Iterable keys) { + private List keysToGets(Iterable> keys) { List ret = new ArrayList<>(); - for(EnrichmentKey key : keys) { - ret.add(converter.toGet(columnFamily, key)); + for(KeyWithContext key : keys) { + ret.add(converter.toGet(getColumnFamily(key.getContext()), key.getKey())); } return ret; } @Override - public Iterable exists(Iterable key, HTableInterface table, boolean logAccess) throws IOException { + public Iterable exists(Iterable> key, boolean logAccess) throws IOException { List ret = new ArrayList<>(); + if(Iterables.isEmpty(key)) { + return Collections.emptyList(); + } + HTableInterface table = Iterables.getFirst(key, null).getContext().getTable(); for(boolean b : table.existsAll(keysToGets(key))) { ret.add(b); } @@ -67,14 +92,19 @@ public Iterable exists(Iterable key, HTableInterface tab } @Override - public Iterable> get( Iterable keys - , HTableInterface table + public Iterable> get( Iterable> keys , boolean logAccess ) throws IOException { + if(Iterables.isEmpty(keys)) { + return Collections.emptyList(); + } + HTableInterface table = Iterables.getFirst(keys, null).getContext().getTable(); List> ret = new ArrayList<>(); + Iterator> keyWithContextIterator = keys.iterator(); for(Result result : table.get(keysToGets(keys))) { - ret.add(converter.fromResult(result, columnFamily)); + HBaseContext context = keyWithContextIterator.next().getContext(); + ret.add(converter.fromResult(result, getColumnFamily(context))); } return ret; } diff --git a/metron-platform/metron-enrichment/src/main/java/org/apache/metron/enrichment/lookup/Lookup.java b/metron-platform/metron-enrichment/src/main/java/org/apache/metron/enrichment/lookup/Lookup.java index 265fccde30..1ec8c9572d 100644 --- a/metron-platform/metron-enrichment/src/main/java/org/apache/metron/enrichment/lookup/Lookup.java +++ b/metron-platform/metron-enrichment/src/main/java/org/apache/metron/enrichment/lookup/Lookup.java @@ -19,6 +19,7 @@ import org.apache.metron.enrichment.lookup.accesstracker.AccessTracker; import org.apache.metron.enrichment.lookup.handler.Handler; +import org.apache.metron.enrichment.lookup.handler.KeyWithContext; import java.io.IOException; @@ -68,24 +69,24 @@ public RESULT_T get(KEY_T key, CONTEXT_T context, boolean logAccess) throws IOEx } @Override - public Iterable exists(Iterable key, CONTEXT_T context, boolean logAccess) throws IOException { + public Iterable exists(Iterable> key, boolean logAccess) throws IOException { if(logAccess) { - for (KEY_T k : key) { - accessTracker.logAccess(k); + for (KeyWithContext k : key) { + accessTracker.logAccess(k.getKey()); } } - return lookupHandler.exists(key, context, logAccess); + return lookupHandler.exists(key, logAccess); } @Override - public Iterable get(Iterable key, CONTEXT_T context, boolean logAccess) throws IOException { + public Iterable get(Iterable> key, boolean logAccess) throws IOException { if(logAccess) { - for (KEY_T k : key) { - accessTracker.logAccess(k); + for (KeyWithContext k : key) { + accessTracker.logAccess(k.getKey()); } } - return lookupHandler.get(key, context, logAccess); + return lookupHandler.get(key, logAccess); } @Override diff --git a/metron-platform/metron-enrichment/src/main/java/org/apache/metron/enrichment/lookup/handler/Handler.java b/metron-platform/metron-enrichment/src/main/java/org/apache/metron/enrichment/lookup/handler/Handler.java index 4ba4de464a..0477a749a7 100644 --- a/metron-platform/metron-enrichment/src/main/java/org/apache/metron/enrichment/lookup/handler/Handler.java +++ b/metron-platform/metron-enrichment/src/main/java/org/apache/metron/enrichment/lookup/handler/Handler.java @@ -24,6 +24,6 @@ public interface Handler extends AutoCloseable{ boolean exists(KEY_T key, CONTEXT_T context, boolean logAccess) throws IOException; RESULT_T get(KEY_T key, CONTEXT_T context, boolean logAccess) throws IOException; - Iterable exists(Iterable key, CONTEXT_T context, boolean logAccess) throws IOException; - Iterable get(Iterable key, CONTEXT_T context, boolean logAccess) throws IOException; + Iterable exists(Iterable> key, boolean logAccess) throws IOException; + Iterable get(Iterable> key, boolean logAccess) throws IOException; } diff --git a/metron-platform/metron-enrichment/src/main/java/org/apache/metron/enrichment/lookup/handler/KeyWithContext.java b/metron-platform/metron-enrichment/src/main/java/org/apache/metron/enrichment/lookup/handler/KeyWithContext.java new file mode 100644 index 0000000000..0150711e13 --- /dev/null +++ b/metron-platform/metron-enrichment/src/main/java/org/apache/metron/enrichment/lookup/handler/KeyWithContext.java @@ -0,0 +1,30 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.metron.enrichment.lookup.handler; + +public class KeyWithContext { + private KEY_T key; + private CONTEXT_T context; + public KeyWithContext(KEY_T key, CONTEXT_T context) { + this.key = key; + this.context = context; + } + public KEY_T getKey() { return key; } + public CONTEXT_T getContext() { return context; } +} diff --git a/metron-platform/metron-enrichment/src/main/java/org/apache/metron/enrichment/utils/EnrichmentUtils.java b/metron-platform/metron-enrichment/src/main/java/org/apache/metron/enrichment/utils/EnrichmentUtils.java index 655188c3de..ab3d462d3b 100644 --- a/metron-platform/metron-enrichment/src/main/java/org/apache/metron/enrichment/utils/EnrichmentUtils.java +++ b/metron-platform/metron-enrichment/src/main/java/org/apache/metron/enrichment/utils/EnrichmentUtils.java @@ -21,11 +21,19 @@ import com.google.common.base.Joiner; import com.google.common.base.Splitter; import com.google.common.collect.Iterables; +import org.apache.hadoop.hbase.client.HTableInterface; +import org.apache.metron.common.configuration.enrichment.EnrichmentConfig; +import org.apache.metron.common.configuration.enrichment.SensorEnrichmentConfig; +import org.apache.metron.enrichment.lookup.EnrichmentLookup; +import org.apache.metron.enrichment.lookup.handler.KeyWithContext; import org.apache.metron.hbase.TableProvider; import org.apache.metron.enrichment.converter.EnrichmentKey; +import sun.management.Sensor; import javax.annotation.Nullable; import java.lang.reflect.InvocationTargetException; +import java.util.HashMap; +import java.util.Map; public class EnrichmentUtils { @@ -35,19 +43,52 @@ public static String getEnrichmentKey(String enrichmentName, String field) { return Joiner.on(".").join(new String[]{KEY_PREFIX, enrichmentName, field}); } - public static class TypeToKey implements Function { + public static class TypeToKey implements Function> { private final String indicator; - - public TypeToKey(String indicator) { + private final EnrichmentConfig config; + private final HTableInterface table; + public TypeToKey(String indicator, HTableInterface table, EnrichmentConfig config) { this.indicator = indicator; - + this.config = config; + this.table = table; } @Nullable @Override - public EnrichmentKey apply(@Nullable String enrichmentType) { - return new EnrichmentKey(enrichmentType, indicator); + public KeyWithContext apply(@Nullable String enrichmentType) { + EnrichmentKey key = new EnrichmentKey(enrichmentType, indicator); + EnrichmentLookup.HBaseContext context = new EnrichmentLookup.HBaseContext(table, getColumnFamily(enrichmentType, config)); + return new KeyWithContext<>(key, context); + } + } + private static ThreadLocal>> typeToCFs = new ThreadLocal>>() { + @Override + protected Map>initialValue() { + return new HashMap<>(); + } + }; + + public static final String TYPE_TO_COLUMN_FAMILY_CONF = "typeToColumnFamily"; + public static String getColumnFamily(String enrichmentType, EnrichmentConfig config) { + Object o = config.getConfig().get(TYPE_TO_COLUMN_FAMILY_CONF); + if(o == null) { + return null; + } + else { + Map cfMap = typeToCFs.get().get(o); + if(cfMap == null) { + cfMap = new HashMap<>(); + if(o instanceof Map) { + Map map = (Map) o; + for(Object key : map.keySet()) { + cfMap.put(key.toString(), map.get(key).toString()); + } + } + typeToCFs.get().put(o, cfMap); + } + return cfMap.get(enrichmentType); } } + public static String toTopLevelField(String field) { if(field == null) { return null; diff --git a/metron-platform/metron-enrichment/src/test/java/org/apache/metron/enrichment/adapters/simplehbase/SimpleHBaseAdapterTest.java b/metron-platform/metron-enrichment/src/test/java/org/apache/metron/enrichment/adapters/simplehbase/SimpleHBaseAdapterTest.java index e2c630986e..1a958df8b7 100644 --- a/metron-platform/metron-enrichment/src/test/java/org/apache/metron/enrichment/adapters/simplehbase/SimpleHBaseAdapterTest.java +++ b/metron-platform/metron-enrichment/src/test/java/org/apache/metron/enrichment/adapters/simplehbase/SimpleHBaseAdapterTest.java @@ -18,6 +18,7 @@ package org.apache.metron.enrichment.adapters.simplehbase; +import com.google.common.collect.ImmutableMap; import org.adrianwalker.multilinestring.Multiline; import org.apache.metron.common.configuration.enrichment.SensorEnrichmentConfig; import org.apache.metron.enrichment.bolt.CacheKey; @@ -43,17 +44,22 @@ public class SimpleHBaseAdapterTest { private String cf = "cf"; + private String cf1 = "cf1"; private String atTableName = "tracker"; private final String hbaseTableName = "enrichments"; private EnrichmentLookup lookup; private static final String PLAYFUL_CLASSIFICATION_TYPE = "playful_classification"; + private static final String CF1_CLASSIFICATION_TYPE = "cf1"; + private static final Map CF1_ENRICHMENT = new HashMap() {{ + put("key", "value"); + }}; private static final Map PLAYFUL_ENRICHMENT = new HashMap() {{ put("orientation", "north"); }}; /** { - "10.0.2.3.orientation":"north" + "playful_classification.orientation":"north" } */ @Multiline @@ -65,19 +71,37 @@ public class SimpleHBaseAdapterTest { "batchSize": 5, "enrichment": { "fieldMap": { - "geo": ["ip_dst_addr", "ip_src_addr"], - "host": ["host"] + "hbaseEnrichment" : [ "ip_dst_addr" ] }, "fieldToTypeMap": { - "ip_dst_addr" : [ "10.0.2.3" ], - "ip_src_addr" : [ "10.3.30.120" ] + "ip_dst_addr" : [ "playful_classification", "cf1" ] } } } */ @Multiline private String sourceConfigStr; - + /** + { + "index": "bro", + "batchSize": 5, + "enrichment": { + "fieldMap": { + "hbaseEnrichment" : [ "ip_dst_addr" ] + }, + "fieldToTypeMap": { + "ip_dst_addr" : [ "playful_classification", "cf1" ] + }, + "config" : { + "typeToColumnFamily" : { + "cf1" : "cf1" + } + } + } + } + */ + @Multiline + private String sourceConfigWithCFStr; private JSONObject expectedMessage; @Before @@ -85,11 +109,17 @@ public void setup() throws Exception { final MockHTable trackerTable = (MockHTable) MockHTable.Provider.addToCache(atTableName, cf); final MockHTable hbaseTable = (MockHTable) MockHTable.Provider.addToCache(hbaseTableName, cf); EnrichmentHelper.INSTANCE.load(hbaseTable, cf, new ArrayList>() {{ - add(new LookupKV<>(new EnrichmentKey("10.0.2.3", "10.0.2.3") + add(new LookupKV<>(new EnrichmentKey(PLAYFUL_CLASSIFICATION_TYPE, "10.0.2.3") , new EnrichmentValue(PLAYFUL_ENRICHMENT) ) ); }}); + EnrichmentHelper.INSTANCE.load(hbaseTable, cf1, new ArrayList>() {{ + add(new LookupKV<>(new EnrichmentKey(CF1_CLASSIFICATION_TYPE, "10.0.2.4") + , new EnrichmentValue(CF1_ENRICHMENT) + ) + ); + }}); BloomAccessTracker bat = new BloomAccessTracker(hbaseTableName, 100, 0.03); PersistentAccessTracker pat = new PersistentAccessTracker(hbaseTableName, "0", trackerTable, cf, bat, 0L); lookup = new EnrichmentLookup(hbaseTable, cf, pat); @@ -109,6 +139,29 @@ public void testEnrich() throws Exception { Assert.assertEquals(expectedMessage, actualMessage); } + @Test + public void testMultiColumnFamilies() throws Exception { + SimpleHBaseAdapter sha = new SimpleHBaseAdapter(); + sha.lookup = lookup; + SensorEnrichmentConfig broSc = JSONUtils.INSTANCE.load(sourceConfigWithCFStr, SensorEnrichmentConfig.class); + JSONObject actualMessage = sha.enrich(new CacheKey("test", "test", broSc)); + Assert.assertEquals(actualMessage, new JSONObject()); + actualMessage = sha.enrich(new CacheKey("ip_dst_addr", "10.0.2.4", broSc)); + Assert.assertNotNull(actualMessage); + Assert.assertEquals(new JSONObject(ImmutableMap.of("cf1.key", "value")), actualMessage); + } + + @Test + public void testMultiColumnFamiliesWrongCF() throws Exception { + SimpleHBaseAdapter sha = new SimpleHBaseAdapter(); + sha.lookup = lookup; + SensorEnrichmentConfig broSc = JSONUtils.INSTANCE.load(sourceConfigStr, SensorEnrichmentConfig.class); + JSONObject actualMessage = sha.enrich(new CacheKey("test", "test", broSc)); + Assert.assertEquals(actualMessage, new JSONObject()); + actualMessage = sha.enrich(new CacheKey("ip_dst_addr", "10.0.2.4", broSc)); + Assert.assertNotNull(actualMessage); + Assert.assertEquals(new JSONObject(new HashMap()), actualMessage); + } @Test(expected = Exception.class) public void testInitializeAdapter() { SimpleHBaseConfig config = new SimpleHBaseConfig();