Examples with CacheKey org.apache.metron.enrichment.cache.CacheKey used on opensource projects

Example 6 with CacheKey

use of org.apache.metron.enrichment.cache.CacheKey in project metron by apache.

the class SimpleHBaseAdapterTest method testMultiColumnFamiliesWrongCF.

@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));
    assertEquals(actualMessage, new JSONObject());
    actualMessage = sha.enrich(new CacheKey("ip_dst_addr", "10.0.2.4", broSc));
    assertNotNull(actualMessage);
    assertEquals(new JSONObject(new HashMap<String, Object>()), actualMessage);
}

Example 7 with CacheKey

use of org.apache.metron.enrichment.cache.CacheKey in project metron by apache.

the class ParallelEnricher method apply.

/**
 * Fully enriches a message.  Each enrichment is done in parallel via a threadpool.
 * Each enrichment is fronted with a LRU cache.
 *
 * @param message the message to enrich
 * @param strategy The enrichment strategy to use (e.g. enrichment or threat intel)
 * @param config The sensor enrichment config
 * @param perfLog The performance logger.  We log the performance for this call, the split portion and the enrichment portion.
 * @return the enrichment result
 */
public EnrichmentResult apply(JSONObject message, EnrichmentStrategies strategy, SensorEnrichmentConfig config, PerformanceLogger perfLog) throws ExecutionException, InterruptedException {
    if (message == null) {
        return null;
    }
    if (perfLog != null) {
        perfLog.mark("execute");
        if (perfLog.isDebugEnabled() && !cacheStats.isEmpty()) {
            CacheStats before = cacheStats.get(strategy);
            CacheStats after = concurrencyContext.getCache().stats();
            if (before != null && after != null) {
                CacheStats delta = after.minus(before);
                perfLog.log("cache", delta.toString());
            }
            cacheStats.put(strategy, after);
        }
    }
    String sensorType = MessageUtils.getSensorType(message);
    message.put(getClass().getSimpleName().toLowerCase() + ".splitter.begin.ts", "" + System.currentTimeMillis());
    // Split the message into individual tasks.
    // 
    // A task will either correspond to an enrichment adapter or,
    // in the case of Stellar, a stellar subgroup.  The tasks will be grouped by enrichment type (the key of the
    // tasks map).  Each JSONObject will correspond to a unit of work.
    Map<String, List<JSONObject>> tasks = splitMessage(message, strategy, config);
    message.put(getClass().getSimpleName().toLowerCase() + ".splitter.end.ts", "" + System.currentTimeMillis());
    message.put(getClass().getSimpleName().toLowerCase() + ".enrich.begin.ts", "" + System.currentTimeMillis());
    if (perfLog != null) {
        perfLog.mark("enrich");
    }
    List<CompletableFuture<JSONObject>> taskList = new ArrayList<>();
    List<Map.Entry<Object, Throwable>> errors = Collections.synchronizedList(new ArrayList<>());
    for (Map.Entry<String, List<JSONObject>> task : tasks.entrySet()) {
        // task is the list of enrichment tasks for the task.getKey() adapter
        EnrichmentAdapter<CacheKey> adapter = enrichmentsByType.get(task.getKey());
        if (adapter == null) {
            throw new IllegalStateException("Unable to find an adapter for " + task.getKey() + ", possible adapters are: " + Joiner.on(",").join(enrichmentsByType.keySet()));
        }
        message.put("adapter." + adapter.getClass().getSimpleName().toLowerCase() + ".begin.ts", "" + System.currentTimeMillis());
        for (JSONObject m : task.getValue()) {
            /* now for each unit of work (each of these only has one element in them)
         * the key is the field name and the value is value associated with that field.
         *
         * In the case of stellar enrichment, the field name is the subgroup name or empty string.
         * The value is the subset of the message needed for the enrichment.
         *
         * In the case of another enrichment (e.g. hbase), the field name is the field name being enriched.
         * The value is the corresponding value.
         */
            for (Object o : m.keySet()) {
                String field = (String) o;
                Object value = m.get(o);
                if (value == null) {
                    message.put("adapter." + adapter.getClass().getSimpleName().toLowerCase() + ".end.ts", "" + System.currentTimeMillis());
                    continue;
                }
                CacheKey cacheKey = new CacheKey(field, value, config);
                String prefix = adapter.getOutputPrefix(cacheKey);
                Supplier<JSONObject> supplier = () -> {
                    try {
                        JSONObject ret = concurrencyContext.getCache().get(cacheKey, new EnrichmentCallable(cacheKey, adapter));
                        if (ret == null) {
                            ret = new JSONObject();
                        }
                        // each enrichment has their own unique prefix to use to adjust the keys for the enriched fields.
                        JSONObject adjustedKeys = EnrichmentUtils.adjustKeys(new JSONObject(), ret, cacheKey.getField(), prefix);
                        adjustedKeys.put("adapter." + adapter.getClass().getSimpleName().toLowerCase() + ".end.ts", "" + System.currentTimeMillis());
                        return adjustedKeys;
                    } catch (Throwable e) {
                        JSONObject errorMessage = new JSONObject();
                        errorMessage.putAll(m);
                        errorMessage.put(Constants.SENSOR_TYPE, sensorType);
                        errors.add(new AbstractMap.SimpleEntry<>(errorMessage, new IllegalStateException(strategy + " error with " + task.getKey() + " failed: " + e.getMessage(), e)));
                        return new JSONObject();
                    }
                };
                // add the Future to the task list
                taskList.add(CompletableFuture.supplyAsync(supplier, ConcurrencyContext.getExecutor()));
            }
        }
    }
    if (taskList.isEmpty()) {
        message.put(getClass().getSimpleName().toLowerCase() + ".enrich.end.ts", "" + System.currentTimeMillis());
        return new EnrichmentResult(message, errors);
    }
    EnrichmentResult ret = new EnrichmentResult(all(taskList, message, (left, right) -> join(left, right)).get(), errors);
    ret.getResult().put(getClass().getSimpleName().toLowerCase() + ".enrich.end.ts", "" + System.currentTimeMillis());
    if (perfLog != null) {
        String key = message.get(Constants.GUID) + "";
        perfLog.log("enrich", "key={}, elapsed time to enrich", key);
        perfLog.log("execute", "key={}, elapsed time to run execute", key);
    }
    return ret;
}

Example 8 with CacheKey

use of org.apache.metron.enrichment.cache.CacheKey in project metron by apache.

the class ThreatIntelAdapterTest method testEnrichNonString.

@Test
public void testEnrichNonString() throws Exception {
    ThreatIntelAdapter tia = new ThreatIntelAdapter();
    tia.lookup = lookup;
    SensorEnrichmentConfig broSc = JSONUtils.INSTANCE.load(sourceConfigStr, SensorEnrichmentConfig.class);
    JSONObject actualMessage = tia.enrich(new CacheKey("ip_dst_addr", "10.0.2.3", broSc));
    assertNotNull(actualMessage);
    assertEquals(expectedMessage, actualMessage);
    actualMessage = tia.enrich(new CacheKey("ip_dst_addr", 10L, broSc));
    assertEquals(actualMessage, new JSONObject());
}

Example 9 with CacheKey

use of org.apache.metron.enrichment.cache.CacheKey in project metron by apache.

the class ParallelEnricherTest method setup.

@BeforeAll
public static void setup() {
    ConcurrencyContext infrastructure = new ConcurrencyContext();
    infrastructure.initialize(5, 100, 10, null, null, false);
    stellarContext = new Context.Builder().build();
    StellarFunctions.initialize(stellarContext);
    StellarAdapter adapter = new AccessLoggingStellarAdapter().ofType("ENRICHMENT");
    adapter.initializeAdapter(new HashMap<>());
    EnrichmentAdapter<CacheKey> dummy = new DummyEnrichmentAdapter();
    enrichmentsByType = ImmutableMap.of("stellar", adapter, "dummy", dummy);
    enricher = new ParallelEnricher(enrichmentsByType, infrastructure, false);
}

Example 10 with CacheKey

use of org.apache.metron.enrichment.cache.CacheKey in project metron by apache.

the class GenericEnrichmentBoltTest method test.

@Test
public void test() throws IOException {
    when(tuple.getSourceComponent()).thenReturn("unit test component");
    when(tuple.getSourceStreamId()).thenReturn("unit test stream");
    String key = "someKey";
    String enrichmentType = "enrichmentType";
    Enrichment<EnrichmentAdapter<CacheKey>> testEnrichment = new Enrichment<>();
    testEnrichment.setType(enrichmentType);
    testEnrichment.setAdapter(enrichmentAdapter);
    GenericEnrichmentBolt genericEnrichmentBolt = new GenericEnrichmentBolt("zookeeperUrl") {

        @Override
        protected void initializeStellar() {
        // do not initialize stellar here.
        }
    };
    genericEnrichmentBolt.setCuratorFramework(client);
    genericEnrichmentBolt.setZKCache(cache);
    genericEnrichmentBolt.getConfigurations().updateSensorEnrichmentConfig(sensorType, new FileInputStream(enrichmentConfigPath));
    HashMap<String, Object> globalConfig = new HashMap<>();
    String baseDir = UnitTestHelper.findDir(new File("../metron-enrichment-common"), "GeoLite");
    File geoHdfsFile = new File(new File(baseDir), "GeoLite2-City.mmdb.gz");
    globalConfig.put(GeoLiteCityDatabase.GEO_HDFS_FILE, geoHdfsFile.getAbsolutePath());
    genericEnrichmentBolt.getConfigurations().updateGlobalConfig(globalConfig);
    assertThrows(IllegalStateException.class, () -> genericEnrichmentBolt.prepare(new HashMap(), topologyContext, outputCollector), "Should fail if a maxCacheSize property is not set");
    genericEnrichmentBolt.withMaxCacheSize(100);
    assertThrows(IllegalStateException.class, () -> genericEnrichmentBolt.prepare(new HashMap(), topologyContext, outputCollector), "Should fail if a maxTimeRetain property is not set");
    genericEnrichmentBolt.withMaxTimeRetain(10000);
    assertThrows(IllegalStateException.class, () -> genericEnrichmentBolt.prepare(new HashMap(), topologyContext, outputCollector), "Should fail if an adapter is not set");
    genericEnrichmentBolt.withEnrichment(testEnrichment);
    when(enrichmentAdapter.initializeAdapter(globalConfig)).thenReturn(true);
    genericEnrichmentBolt.prepare(new HashMap(), topologyContext, outputCollector);
    verify(enrichmentAdapter, times(1)).initializeAdapter(globalConfig);
    when(enrichmentAdapter.initializeAdapter(globalConfig)).thenReturn(false);
    UnitTestHelper.setLog4jLevel(GenericEnrichmentBolt.class, Level.FATAL);
    assertThrows(IllegalStateException.class, () -> genericEnrichmentBolt.prepare(new HashMap(), topologyContext, outputCollector), "An exception should be thrown if enrichment adapter initialization fails");
    UnitTestHelper.setLog4jLevel(GenericEnrichmentBolt.class, Level.ERROR);
    genericEnrichmentBolt.declareOutputFields(declarer);
    verify(declarer, times(1)).declareStream(eq(enrichmentType), argThat(new FieldsMatcher("key", "message", "subgroup")));
    verify(declarer, times(1)).declareStream(eq("error"), argThat(new FieldsMatcher("message")));
    when(tuple.getStringByField("key")).thenReturn(null);
    UnitTestHelper.setLog4jLevel(GenericEnrichmentBolt.class, Level.FATAL);
    genericEnrichmentBolt.execute(tuple);
    UnitTestHelper.setLog4jLevel(GenericEnrichmentBolt.class, Level.ERROR);
    MetronError error = new MetronError().withErrorType(Constants.ErrorType.ENRICHMENT_ERROR).withThrowable(new Exception("Could not parse binary stream to JSON"));
    verify(outputCollector, times(1)).emit(eq(Constants.ERROR_STREAM), argThat(new MetronErrorJSONMatcher(error.getJSONObject())));
    when(tuple.getStringByField("key")).thenReturn(key);
    when(tuple.getValueByField("message")).thenReturn(originalMessage);
    when(enrichmentAdapter.enrich(any())).thenReturn(new JSONObject());
    genericEnrichmentBolt.execute(tuple);
    verify(outputCollector, times(1)).emit(eq(enrichmentType), argThat(new EnrichedMessageMatcher(key, new JSONObject(ImmutableMap.of("source.type", "test")))));
    reset(enrichmentAdapter);
    SensorEnrichmentConfig sensorEnrichmentConfig = SensorEnrichmentConfig.fromBytes(ConfigurationsUtils.readSensorEnrichmentConfigsFromFile(sampleConfigPath).get(sensorType));
    sensorEnrichmentConfig.getConfiguration().put(STELLAR_CONTEXT_CONF, genericEnrichmentBolt.getStellarContext());
    CacheKey cacheKey1 = new CacheKey("field1", "value1", sensorEnrichmentConfig);
    CacheKey cacheKey2 = new CacheKey("field2", "value2", sensorEnrichmentConfig);
    genericEnrichmentBolt.cache.invalidateAll();
    when(enrichmentAdapter.getOutputPrefix(cacheKey1)).thenReturn("field1");
    when(enrichmentAdapter.getOutputPrefix(cacheKey2)).thenReturn("field2");
    when(enrichmentAdapter.enrich(cacheKey1)).thenReturn(enrichedField1);
    when(enrichmentAdapter.enrich(cacheKey2)).thenReturn(enrichedField2);
    genericEnrichmentBolt.execute(tuple);
    verify(enrichmentAdapter, times(1)).logAccess(cacheKey1);
    verify(enrichmentAdapter, times(1)).logAccess(cacheKey2);
    verify(outputCollector, times(1)).emit(eq(enrichmentType), argThat(new EnrichedMessageMatcher(key, enrichedMessage)));
    reset(outputCollector);
    genericEnrichmentBolt.cache.invalidateAll();
    when(enrichmentAdapter.enrich(cacheKey1)).thenReturn(null);
    genericEnrichmentBolt.execute(tuple);
    error = new MetronError().withErrorType(Constants.ErrorType.ENRICHMENT_ERROR).withErrorFields(new HashSet<String>() {

        {
            add("field1");
        }
    }).addRawMessage(new JSONObject() {

        {
            put("field1", "value1");
            put("field2", "value2");
            put("source.type", "test");
        }
    }).withThrowable(new Exception("[Metron] Could not enrich string: value1"));
    verify(outputCollector, times(1)).emit(eq(Constants.ERROR_STREAM), argThat(new MetronErrorJSONMatcher(error.getJSONObject())));
}