| /** |
| * 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.management; |
| |
| import com.google.common.collect.ImmutableMap; |
| import org.adrianwalker.multilinestring.Multiline; |
| import org.apache.metron.common.configuration.enrichment.SensorEnrichmentConfig; |
| import org.apache.metron.common.configuration.enrichment.threatintel.RiskLevelRule; |
| import org.apache.metron.stellar.common.StellarProcessor; |
| import org.apache.metron.stellar.common.shell.VariableResult; |
| import org.apache.metron.stellar.dsl.*; |
| import org.apache.metron.threatintel.triage.ThreatTriageProcessor; |
| import org.junit.jupiter.api.BeforeEach; |
| import org.junit.jupiter.api.Test; |
| |
| import java.util.Collections; |
| import java.util.HashMap; |
| import java.util.List; |
| import java.util.Map; |
| |
| import static org.apache.metron.common.configuration.ConfigurationType.ENRICHMENT; |
| import static org.apache.metron.management.EnrichmentConfigFunctionsTest.emptyTransformationsConfig; |
| import static org.apache.metron.management.EnrichmentConfigFunctionsTest.toMap; |
| import static org.junit.jupiter.api.Assertions.*; |
| |
| public class ThreatTriageFunctionsTest { |
| |
| String configStr = emptyTransformationsConfig(); |
| Map<String, VariableResult> variables; |
| Context context = null; |
| |
| @BeforeEach |
| public void setup() { |
| variables = ImmutableMap.of( |
| "less", VariableResult.withExpression(true, "1 < 2"), |
| "greater", VariableResult.withExpression(false, "1 > 2") |
| ); |
| |
| context = new Context.Builder() |
| .with(Context.Capabilities.SHELL_VARIABLES, () -> variables) |
| .build(); |
| } |
| |
| public static List<RiskLevelRule> getTriageRules(String config) { |
| SensorEnrichmentConfig sensorConfig = (SensorEnrichmentConfig) ENRICHMENT.deserialize(config); |
| return sensorConfig.getThreatIntel().getTriageConfig().getRiskLevelRules(); |
| } |
| |
| private Object run(String rule, Map<String, Object> variables) { |
| StellarProcessor processor = new StellarProcessor(); |
| return processor.parse(rule, new DefaultVariableResolver(x -> variables.get(x),x -> variables.containsKey(x)), StellarFunctions.FUNCTION_RESOLVER(), context); |
| } |
| |
| private Object run(String rule) { |
| StellarProcessor processor = new StellarProcessor(); |
| return processor.parse(rule, new MapVariableResolver(Collections.emptyMap()), StellarFunctions.FUNCTION_RESOLVER(), context); |
| } |
| |
| /** |
| * Sequentially runs a set of expressions and returns the result of the last expression. |
| * @param expressions The set of expressions to execute. |
| * @return The result of running the last expression. |
| */ |
| private Object run(String... expressions) { |
| Object result = null; |
| |
| for(String expression: expressions) { |
| result = run(expression); |
| } |
| |
| return result; |
| } |
| |
| @Test |
| public void testSetAggregation() { |
| String newConfig = (String) run("THREAT_TRIAGE_SET_AGGREGATOR(config, 'MIN' )", toMap("config", configStr)); |
| SensorEnrichmentConfig sensorConfig = (SensorEnrichmentConfig) ENRICHMENT.deserialize(newConfig); |
| assertEquals("MIN", sensorConfig.getThreatIntel().getTriageConfig().getAggregator().toString()); |
| } |
| |
| |
| @Test |
| public void testSetAggregationWithEngine() { |
| // init the engine |
| ThreatTriageProcessor engine = (ThreatTriageProcessor) run("THREAT_TRIAGE_INIT()"); |
| Map<String, Object> vars = new HashMap<>(); |
| vars.put("engine", engine); |
| |
| // set the aggregator |
| String newConfig = (String) run("THREAT_TRIAGE_SET_AGGREGATOR(engine, 'MIN')", vars); |
| |
| // validate the return configuration |
| SensorEnrichmentConfig sensorConfig = (SensorEnrichmentConfig) ENRICHMENT.deserialize(newConfig); |
| assertEquals("MIN", sensorConfig.getThreatIntel().getTriageConfig().getAggregator().toString()); |
| |
| // validate that the engine was updated |
| assertEquals("MIN", engine.getSensorConfig().getThreatIntel().getTriageConfig().getAggregator().toString()); |
| } |
| |
| @Test |
| public void testAddEmpty() { |
| |
| String newConfig = (String) run( |
| "THREAT_TRIAGE_ADD(config, { 'rule' : SHELL_GET_EXPRESSION('less'), 'score' : 10 } )" |
| , toMap("config", configStr |
| ) |
| ); |
| |
| List<RiskLevelRule> triageRules = getTriageRules(newConfig); |
| assertEquals(1, triageRules.size()); |
| RiskLevelRule rule = triageRules.get(0); |
| assertEquals(variables.get("less").getExpression().get(), rule.getRule() ); |
| assertEquals("10", rule.getScoreExpression()); |
| } |
| |
| @Test |
| public void testAddEmptyWithEngine() { |
| // init the engine |
| ThreatTriageProcessor engine = (ThreatTriageProcessor) run("THREAT_TRIAGE_INIT()"); |
| Map<String, Object> vars = new HashMap<>(); |
| vars.put("engine", engine); |
| String newConfig = (String) run("THREAT_TRIAGE_ADD(engine, {'rule' : SHELL_GET_EXPRESSION('less'), 'score' : 10 } )", vars); |
| |
| // validate the returned configuration |
| List<RiskLevelRule> triageRules = getTriageRules(newConfig); |
| assertEquals(1, triageRules.size()); |
| |
| // validate that the engine was updated |
| assertEquals(1, engine.getSensorConfig().getThreatIntel().getTriageConfig().getRiskLevelRules().size()); |
| } |
| |
| @Test |
| public void testAddHasExisting() { |
| |
| String newConfig = (String) run( |
| "THREAT_TRIAGE_ADD(config, { 'rule' : SHELL_GET_EXPRESSION('less'), 'score' : 10, 'reason' : '2 + 2' } )" |
| , toMap("config", configStr |
| ) |
| ); |
| |
| newConfig = (String) run( |
| "THREAT_TRIAGE_ADD(config, { 'rule' : SHELL_GET_EXPRESSION('greater'), 'score' : 20 } )" |
| , toMap("config", newConfig |
| ) |
| ); |
| |
| List<RiskLevelRule> triageRules = getTriageRules(newConfig); |
| assertEquals(2, triageRules.size()); |
| RiskLevelRule less = triageRules.get(0); |
| assertEquals(variables.get("less").getExpression().get(), less.getRule()); |
| assertEquals("10", less.getScoreExpression()); |
| |
| RiskLevelRule greater = triageRules.get(1); |
| assertEquals(variables.get("greater").getExpression().get(), greater.getRule()); |
| assertEquals("20", greater.getScoreExpression()); |
| } |
| |
| @Test |
| public void testAddMalformed() { |
| assertThrows( |
| ParseException.class, |
| () -> |
| run( |
| "THREAT_TRIAGE_ADD(config, { 'rule': SHELL_GET_EXPRESSION('foo'), 'score' : 10 } )", |
| toMap("config", configStr))); |
| } |
| |
| @Test |
| public void testAddDuplicate() { |
| String newConfig = (String) run( |
| "THREAT_TRIAGE_ADD(config, { 'rule' : SHELL_GET_EXPRESSION('less'), 'score' : 10 } )" |
| , toMap("config", configStr |
| ) |
| ); |
| |
| newConfig = (String) run( |
| "THREAT_TRIAGE_ADD(config, { 'rule' : SHELL_GET_EXPRESSION('less'), 'score' : 10 } )" |
| , toMap("config",newConfig |
| ) |
| ); |
| |
| List<RiskLevelRule> triageRules = getTriageRules(newConfig); |
| assertEquals(1, triageRules.size()); |
| RiskLevelRule rule = triageRules.get(0); |
| assertEquals(variables.get("less").getExpression().get(), rule.getRule() ); |
| assertEquals("10", rule.getScoreExpression()); |
| } |
| |
| @Test |
| public void testAddMultiple() { |
| |
| // add a new rule |
| String newConfig = (String) run( |
| "THREAT_TRIAGE_ADD(config, { 'name':'rule1', 'rule':'value < 2', 'score':10 } )", |
| toMap("config", configStr)); |
| |
| // add another rule |
| newConfig = (String) run( |
| "THREAT_TRIAGE_ADD(config, { 'name':'rule2', 'rule':'value < 4', 'score':10 } )", |
| toMap("config", newConfig)); |
| |
| List<RiskLevelRule> triageRules = getTriageRules(newConfig); |
| assertEquals(2, triageRules.size()); |
| } |
| |
| @Test |
| public void testAddMultipleWithEngine() { |
| |
| // init the engine |
| ThreatTriageProcessor engine = (ThreatTriageProcessor) run("THREAT_TRIAGE_INIT()"); |
| Map<String, Object> vars = new HashMap<>(); |
| vars.put("engine", engine); |
| |
| // add a new rule |
| run("THREAT_TRIAGE_ADD(engine, { 'name':'rule1', 'rule':'value < 2', 'score':10 } )", vars); |
| |
| // add another rule |
| run("THREAT_TRIAGE_ADD(engine, { 'name':'rule2', 'rule':'value < 4', 'score':10 } )", vars); |
| |
| List<RiskLevelRule> triageRules = engine.getRiskLevelRules(); |
| assertEquals(2, triageRules.size()); |
| } |
| |
| @Test |
| public void testRemove() { |
| String newConfig = (String) run( |
| "THREAT_TRIAGE_ADD(config, [ { 'rule' : SHELL_GET_EXPRESSION('less'), 'score' : 10 }, { 'rule' : SHELL_GET_EXPRESSION('greater'), 'score' : 20 } ] )" |
| , toMap("config", configStr |
| ) |
| ); |
| |
| newConfig = (String) run( |
| "THREAT_TRIAGE_REMOVE(config, [ SHELL_GET_EXPRESSION('greater')] )" |
| , toMap("config",newConfig |
| ) |
| ); |
| |
| List<RiskLevelRule> triageRules = getTriageRules(newConfig); |
| assertEquals(1, triageRules.size()); |
| RiskLevelRule rule = triageRules.get(0); |
| assertEquals(variables.get("less").getExpression().get(), rule.getRule() ); |
| assertEquals("10", rule.getScoreExpression()); |
| } |
| |
| @Test |
| public void testRemoveWithEngine() { |
| // init the engine |
| ThreatTriageProcessor engine = (ThreatTriageProcessor) run("THREAT_TRIAGE_INIT()"); |
| |
| // set the aggregator |
| Map<String, Object> vars = new HashMap<>(); |
| vars.put("engine", engine); |
| |
| // add 2 rules |
| String newConfig = (String) run("THREAT_TRIAGE_ADD(engine, [" + |
| "{ 'rule' : SHELL_GET_EXPRESSION('less'), 'score' : 10 }, " + |
| "{ 'rule' : SHELL_GET_EXPRESSION('greater'), 'score' : 20 } ] )", vars); |
| |
| // remove 1 rule |
| newConfig = (String) run("THREAT_TRIAGE_REMOVE(engine, [ " + |
| "SHELL_GET_EXPRESSION('greater')] )", vars); |
| |
| List<RiskLevelRule> triageRules = engine.getRiskLevelRules(); |
| assertEquals(1, triageRules.size()); |
| RiskLevelRule rule = triageRules.get(0); |
| assertEquals(variables.get("less").getExpression().get(), rule.getRule() ); |
| assertEquals("10", rule.getScoreExpression()); |
| } |
| |
| @Test |
| public void testRemoveMultiple() { |
| String newConfig = (String) run( |
| "THREAT_TRIAGE_ADD(config, [ { 'rule' : SHELL_GET_EXPRESSION('less'), 'score' : 10 }, { 'rule' : SHELL_GET_EXPRESSION('greater'), 'score' : 20 } ] )" |
| , toMap("config", configStr ) |
| ); |
| |
| newConfig = (String) run( |
| "THREAT_TRIAGE_REMOVE(config, [ SHELL_GET_EXPRESSION('less'), SHELL_GET_EXPRESSION('greater')] )" |
| , toMap("config",newConfig |
| ) |
| ); |
| |
| List<RiskLevelRule> triageRules = getTriageRules(newConfig); |
| assertEquals(0, triageRules.size()); |
| } |
| |
| @Test |
| public void testRemoveMissing() { |
| |
| String newConfig = (String) run( |
| "THREAT_TRIAGE_ADD(config, [ { 'rule' : SHELL_GET_EXPRESSION('less'), 'score' : 10 }, { 'rule' : SHELL_GET_EXPRESSION('greater'), 'score' : 20 } ] )" |
| , toMap("config", configStr |
| ) |
| ); |
| |
| newConfig = (String) run( |
| "THREAT_TRIAGE_REMOVE(config, [ SHELL_GET_EXPRESSION('foo'), SHELL_GET_EXPRESSION('bar')] )" |
| , toMap("config",newConfig |
| ) |
| ); |
| |
| List<RiskLevelRule> triageRules = getTriageRules(newConfig); |
| assertEquals(2, triageRules.size()); |
| RiskLevelRule less = triageRules.get(0); |
| assertEquals(variables.get("less").getExpression().get(), less.getRule() ); |
| assertEquals("10", less.getScoreExpression()); |
| |
| RiskLevelRule greater = triageRules.get(1); |
| assertEquals(variables.get("greater").getExpression().get(), greater.getRule() ); |
| assertEquals("20", greater.getScoreExpression()); |
| } |
| |
| /** |
| ╔══════╤═════════╤═════════════╤═══════╤════════╗ |
| ║ Name │ Comment │ Triage Rule │ Score │ Reason ║ |
| ╠══════╪═════════╪═════════════╪═══════╪════════╣ |
| ║ │ │ 1 < 2 │ 10 │ 2 + 2 ║ |
| ╟──────┼─────────┼─────────────┼───────┼────────╢ |
| ║ │ │ 1 > 2 │ 20 │ ║ |
| ╚══════╧═════════╧═════════════╧═══════╧════════╝ |
| Aggregation: MAX*/ |
| @Multiline |
| static String testPrintExpected; |
| |
| @Test |
| public void testPrint() { |
| String newConfig = (String) run( |
| "THREAT_TRIAGE_ADD(config, [ " + |
| "{ 'rule' : SHELL_GET_EXPRESSION('less'), 'score' : 10, 'reason' : '2 + 2' }, " + |
| "{ 'rule' : SHELL_GET_EXPRESSION('greater'), 'score' : 20 } ] )" |
| , toMap("config", configStr |
| ) |
| ); |
| |
| String out = (String) run( |
| "THREAT_TRIAGE_PRINT(config)" |
| , toMap("config",newConfig |
| ) |
| ); |
| assertEquals(testPrintExpected, out); |
| } |
| |
| @Test |
| public void testPrintWithEngine() { |
| |
| // init the engine |
| ThreatTriageProcessor engine = (ThreatTriageProcessor) run("THREAT_TRIAGE_INIT()"); |
| Map<String, Object> vars = new HashMap<>(); |
| vars.put("engine", engine); |
| |
| // add 2 rules |
| run("THREAT_TRIAGE_ADD(engine, [ " + |
| "{ 'rule' : SHELL_GET_EXPRESSION('less'), 'score' : 10, 'reason' : '2 + 2' }, " + |
| "{ 'rule' : SHELL_GET_EXPRESSION('greater'), 'score' : 20 } ] )" |
| , vars); |
| |
| // print |
| String out = (String) run("THREAT_TRIAGE_PRINT(engine)", vars); |
| assertEquals(testPrintExpected, out); |
| } |
| |
| |
| /** |
| ╔══════╤═════════╤═════════════╤═══════╤════════╗ |
| ║ Name │ Comment │ Triage Rule │ Score │ Reason ║ |
| ╠══════╧═════════╧═════════════╧═══════╧════════╣ |
| ║ (empty) ║ |
| ╚═══════════════════════════════════════════════╝ |
| */ |
| @Multiline |
| static String testPrintEmptyExpected; |
| |
| @Test |
| public void testPrintEmpty() { |
| String out = (String) run( |
| "THREAT_TRIAGE_PRINT(config)" |
| , toMap("config",configStr |
| ) |
| ); |
| assertEquals(testPrintEmptyExpected, out); |
| } |
| |
| @Test |
| public void testPrintNull() { |
| Map<String,Object> variables = new HashMap<String,Object>(){{ |
| put("config", null); |
| }}; |
| assertThrows(ParseException.class, () -> run("THREAT_TRIAGE_PRINT(config)", variables)); |
| } |
| |
| /** |
| * { |
| * "timestamp": 1504548045948, |
| * "source.type": "example", |
| * "value": 22 |
| * } |
| */ |
| @Multiline |
| private String message; |
| |
| @Test |
| public void testTriageInitNoArg() { |
| Object result = run("THREAT_TRIAGE_INIT()"); |
| assertNotNull(result); |
| assertTrue(result instanceof ThreatTriageProcessor); |
| |
| // there should be no triage rules defined |
| ThreatTriageProcessor engine = (ThreatTriageProcessor) result; |
| assertEquals(0, engine.getRiskLevelRules().size()); |
| } |
| |
| @Test |
| public void testTriageInitWithArg() { |
| |
| // add a triage rule |
| String confWithRule = (String) run("THREAT_TRIAGE_ADD(conf, [{ 'rule': 'value > 0', 'score' : 10 } ])", |
| toMap("conf", configStr)); |
| |
| // initialize the engine |
| Object result = run("THREAT_TRIAGE_INIT(confWithRule)", |
| toMap("confWithRule", confWithRule)); |
| |
| assertNotNull(result); |
| assertTrue(result instanceof ThreatTriageProcessor); |
| |
| // validate that there is 1 triage rule |
| ThreatTriageProcessor engine = (ThreatTriageProcessor) result; |
| assertEquals(1, engine.getRiskLevelRules().size()); |
| } |
| |
| @Test |
| public void testTriageInitWithBadArg() { |
| assertThrows(ParseException.class, () -> run("THREAT_TRIAGE_INIT(missing)")); |
| } |
| |
| @Test |
| public void testTriageScoreWithNoRules() { |
| |
| // init the engine |
| Object engine = run("THREAT_TRIAGE_INIT()"); |
| |
| Map<String, Object> vars = new HashMap<>(); |
| vars.put("engine", engine); |
| vars.put("msg", message); |
| |
| // score the message |
| Object result = run("THREAT_TRIAGE_SCORE(msg, engine)", vars); |
| assertNotNull(result); |
| assertTrue(result instanceof Map); |
| |
| // validate the rules that were scored |
| Map<String, Object> score = (Map) result; |
| assertEquals(0, ((List) score.get(ThreatTriageFunctions.RULES_KEY)).size()); |
| |
| // validate the total score |
| Object totalScore = score.get(ThreatTriageFunctions.SCORE_KEY); |
| assertTrue(totalScore instanceof Double); |
| assertEquals(0.0, (Double) totalScore, 0.001); |
| } |
| |
| @Test |
| public void testTriageScoreWithRules() { |
| |
| // add a triage rule |
| String confWithRule = (String) run("THREAT_TRIAGE_ADD(conf, [{ 'rule': 'value > 0', 'score' : 10 }])", |
| toMap("conf", configStr)); |
| |
| // initialize the engine |
| Object engine = run("THREAT_TRIAGE_INIT(confWithRule)", |
| toMap("confWithRule", confWithRule)); |
| |
| Map<String, Object> vars = new HashMap<>(); |
| vars.put("engine", engine); |
| vars.put("msg", message); |
| |
| // score the message |
| Object result = run("THREAT_TRIAGE_SCORE(msg, engine)", vars); |
| assertNotNull(result); |
| assertTrue(result instanceof Map); |
| |
| // validate the rules that were scored |
| Map<String, Object> score = (Map) result; |
| assertEquals(1, ((List) score.get(ThreatTriageFunctions.RULES_KEY)).size()); |
| |
| // validate the total score |
| Object totalScore = score.get(ThreatTriageFunctions.SCORE_KEY); |
| assertTrue(totalScore instanceof Double); |
| assertEquals(10.0, (Double) totalScore, 0.001); |
| |
| // validate the aggregator |
| assertEquals("MAX", score.get(ThreatTriageFunctions.AGG_KEY)); |
| } |
| |
| @Test |
| public void testTriageWithScoreExpression() { |
| |
| // add a triage rule that uses an expression for the score |
| String confWithRule = (String) run("THREAT_TRIAGE_ADD(conf, [{ 'rule': 'value > 0', 'score' : 'value * 10' }])", |
| toMap("conf", configStr)); |
| |
| // initialize the engine |
| Object engine = run("THREAT_TRIAGE_INIT(confWithRule)", |
| toMap("confWithRule", confWithRule)); |
| |
| Map<String, Object> vars = new HashMap<>(); |
| vars.put("engine", engine); |
| vars.put("msg", message); |
| |
| // score the message |
| Object result = run("THREAT_TRIAGE_SCORE(msg, engine)", vars); |
| assertNotNull(result); |
| assertTrue(result instanceof Map); |
| |
| // validate the rules that were scored |
| Map<String, Object> score = (Map) result; |
| assertEquals(1, ((List) score.get(ThreatTriageFunctions.RULES_KEY)).size()); |
| |
| // validate the total score |
| Object totalScore = score.get(ThreatTriageFunctions.SCORE_KEY); |
| assertTrue(totalScore instanceof Double); |
| assertEquals(220.0, (Double) totalScore, 0.001); |
| |
| // validate the aggregator |
| assertEquals("MAX", score.get(ThreatTriageFunctions.AGG_KEY)); |
| } |
| |
| @Test |
| public void testTriageScoreWithNoMessage() { |
| |
| // add a triage rule |
| String confWithRule = (String) run("THREAT_TRIAGE_ADD(conf, [{ 'rule': 'value > 0', 'score' : 10 }])", |
| toMap("conf", configStr)); |
| |
| // initialize the engine |
| Object engine = run("THREAT_TRIAGE_INIT(confWithRule)", |
| toMap("confWithRule", confWithRule)); |
| |
| Map<String, Object> vars = new HashMap<>(); |
| vars.put("engine", engine); |
| |
| // score the message |
| assertThrows(Exception.class, () -> run("THREAT_TRIAGE_SCORE(11, engine)", vars)); |
| } |
| |
| @Test |
| public void testTriageConfig() { |
| |
| // init the engine |
| Object engine = run("THREAT_TRIAGE_INIT()"); |
| |
| Map<String, Object> vars = new HashMap<>(); |
| vars.put("engine", engine); |
| |
| // score the message |
| Object result = run("THREAT_TRIAGE_CONFIG(engine)", vars); |
| assertNotNull(result); |
| assertTrue(result instanceof String); |
| |
| // validate the configuration |
| String json = (String) result; |
| assertEquals(emptyTransformationsConfig(), json); |
| } |
| } |