[Entity Analytics] Move scripted metric painless scripts to static file & remove category based weighting

x-pack/plugins/security_solution/server/lib/entity_analytics/risk_score/calculate_risk_scores.ts

@@ -13,10 +13,7 @@ import type {
 import type { ElasticsearchClient, Logger } from '@kbn/core/server';
 import {
   ALERT_RISK_SCORE,
-  ALERT_RULE_NAME,
-  ALERT_UUID,
   ALERT_WORKFLOW_STATUS,
-  EVENT_KIND,
 } from '@kbn/rule-registry-plugin/common/technical_rule_data_field_names';
 import {
   type AfterKeys,
@@ -25,6 +22,7 @@ import {
   type RiskScore,
   getRiskLevel,
   RiskCategories,
+  RiskWeightTypes,
 } from '../../../../common/entity_analytics/risk_engine';
 import { withSecuritySpan } from '../../../utils/with_security_span';
 import type { AssetCriticalityRecord } from '../../../../common/api/entity_analytics';
@@ -35,25 +33,18 @@ import {
   normalize,
 } from '../asset_criticality/helpers';
 import { getAfterKeyForIdentifierType, getFieldForIdentifier } from './helpers';
-import {
-  buildCategoryCountDeclarations,
-  buildCategoryAssignment,
-  buildCategoryScoreDeclarations,
-  buildWeightingOfScoreByCategory,
-  getGlobalWeightForIdentifierType,
-} from './risk_weights';
 import type {
   CalculateRiskScoreAggregations,
   CalculateScoresParams,
   CalculateScoresResponse,
   RiskScoreBucket,
 } from '../types';
 import {
-  MAX_INPUTS_COUNT,
   RISK_SCORING_INPUTS_COUNT_MAX,
   RISK_SCORING_SUM_MAX,
   RISK_SCORING_SUM_VALUE,
 } from './constants';
+import { getPainlessScripts, type PainlessScripts } from './painless';
 
 const formatForResponse = ({
   bucket,
@@ -116,67 +107,22 @@ const filterFromRange = (range: CalculateScoresParams['range']): QueryDslQueryCo
   range: { '@timestamp': { lt: range.end, gte: range.start } },
 });
 
-const buildReduceScript = ({
-  globalIdentifierTypeWeight,
-}: {
-  globalIdentifierTypeWeight?: number;
-}): string => {
-  return `
-    Map results = new HashMap();
-    List inputs = [];
-    for (state in states) {
-      inputs.addAll(state.inputs)
-    }
-    Collections.sort(inputs, (a, b) -> b.get('weighted_score').compareTo(a.get('weighted_score')));
-
-    double num_inputs_to_score = Math.min(inputs.length, params.max_risk_inputs_per_identity);
-    results['notes'] = [];
-    if (num_inputs_to_score == params.max_risk_inputs_per_identity) {
-      results['notes'].add('Number of risk inputs (' + inputs.length + ') exceeded the maximum allowed (' + params.max_risk_inputs_per_identity + ').');
-    }
-
-    ${buildCategoryScoreDeclarations()}
-    ${buildCategoryCountDeclarations()}
-
-    double total_score = 0;
-    double current_score = 0;
-    List risk_inputs = [];
-    for (int i = 0; i < num_inputs_to_score; i++) {
-      current_score = inputs[i].weighted_score / Math.pow(i + 1, params.p);
-
-      if (i < ${MAX_INPUTS_COUNT}) {
-        inputs[i]["contribution"] = 100 * current_score / params.risk_cap;
-        risk_inputs.add(inputs[i]);
-      }
-
-      ${buildCategoryAssignment()}
-      total_score += current_score;
-    }
-
-    ${globalIdentifierTypeWeight != null ? `total_score *= ${globalIdentifierTypeWeight};` : ''}
-    double score_norm = 100 * total_score / params.risk_cap;
-    results['score'] = total_score;
-    results['normalized_score'] = score_norm;
-    results['risk_inputs'] = risk_inputs;
-
-    return results;
-  `;
-};
-
 const buildIdentifierTypeAggregation = ({
   afterKeys,
   identifierType,
   pageSize,
   weights,
   alertSampleSizePerShard,
+  scriptedMetricPainless,
 }: {
   afterKeys: AfterKeys;
   identifierType: IdentifierType;
   pageSize: number;
   weights?: RiskWeights;
   alertSampleSizePerShard: number;
+  scriptedMetricPainless: PainlessScripts;
 }): AggregationsAggregationContainer => {
-  const globalIdentifierTypeWeight = getGlobalWeightForIdentifierType({ identifierType, weights });
+  const globalIdentifierTypeWeight = getGlobalWeightForIdentifierType(identifierType, weights);
   const identifierField = getFieldForIdentifier(identifierType);
 
   return {
@@ -202,33 +148,16 @@ const buildIdentifierTypeAggregation = ({
         aggs: {
           risk_details: {
             scripted_metric: {
-              init_script: 'state.inputs = []',
-              map_script: `
-                Map fields = new HashMap();
-                String category = doc['${EVENT_KIND}'].value;
-                double score = doc['${ALERT_RISK_SCORE}'].value;
-                double weighted_score = 0.0;
-
-                fields.put('time', doc['@timestamp'].value);
-                fields.put('rule_name', doc['${ALERT_RULE_NAME}'].value);
-
-                fields.put('category', category);
-                fields.put('index', doc['_index'].value);
-                fields.put('id', doc['${ALERT_UUID}'].value);
-                fields.put('score', score);
-
-                ${buildWeightingOfScoreByCategory({ userWeights: weights, identifierType })}
-                fields.put('weighted_score', weighted_score);
-
-                state.inputs.add(fields);
-              `,
-              combine_script: 'return state;',
+              init_script: scriptedMetricPainless.init,
+              map_script: scriptedMetricPainless.map,
+              combine_script: scriptedMetricPainless.combine,
               params: {
                 max_risk_inputs_per_identity: RISK_SCORING_INPUTS_COUNT_MAX,
                 p: RISK_SCORING_SUM_VALUE,
                 risk_cap: RISK_SCORING_SUM_MAX,
+                global_identifier_type_weight: globalIdentifierTypeWeight,
               },
-              reduce_script: buildReduceScript({ globalIdentifierTypeWeight }),
+              reduce_script: scriptedMetricPainless.reduce,
             },
           },
         },
@@ -284,6 +213,12 @@ const processScores = async ({
   });
 };
 
+export const getGlobalWeightForIdentifierType = (
+  identifierType: IdentifierType,
+  weights?: RiskWeights
+): number | undefined =>
+  weights?.find((weight) => weight.type === RiskWeightTypes.global)?.[identifierType];
+
 export const calculateRiskScores = async ({
   afterKeys: userAfterKeys,
   assetCriticalityService,
@@ -305,6 +240,7 @@ export const calculateRiskScores = async ({
 } & CalculateScoresParams): Promise<CalculateScoresResponse> =>
   withSecuritySpan('calculateRiskScores', async () => {
     const now = new Date().toISOString();
+    const scriptedMetricPainless = await getPainlessScripts();
     const filter = [
       filterFromRange(range),
       { bool: { must_not: { term: { [ALERT_WORKFLOW_STATUS]: 'closed' } } } },
@@ -343,6 +279,7 @@ export const calculateRiskScores = async ({
           pageSize,
           weights,
           alertSampleSizePerShard,
+          scriptedMetricPainless,
         });
         return aggs;
       }, {} as Record<string, AggregationsAggregationContainer>),

x-pack/plugins/security_solution/server/lib/entity_analytics/risk_score/painless/index.test.ts

@@ -0,0 +1,24 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { getPainlessScripts } from '.';
+
+describe('getPainlessScripts', () => {
+  // to update snapshot run `yarn test:jest x-pack/plugins/security_solution/server/lib/entity_analytics/risk_score/painless/index.test.ts -u`
+  test('Scripts should not have changed. If this change is intentional, ensure that Serverless scripted metric allowlists are updated', async () => {
+    const scripts = await getPainlessScripts();
+
+    expect(scripts).toMatchInlineSnapshot(`
+      Object {
+        "combine": "return state;",
+        "init": "state.inputs = []",
+        "map": "Map fields = new HashMap();String category = doc['event.kind'].value;double score = doc['kibana.alert.risk_score'].value;fields.put('time', doc['@timestamp'].value);fields.put('rule_name', doc['kibana.alert.rule.name'].value);fields.put('category', category);fields.put('index', doc['_index'].value);fields.put('id', doc['kibana.alert.uuid'].value);fields.put('score', score);state.inputs.add(fields); ",
+        "reduce": "Map results = new HashMap();List inputs = [];for (state in states) { inputs.addAll(state.inputs)}Collections.sort(inputs, (a, b) -> b.get('score').compareTo(a.get('score')));double num_inputs_to_score = Math.min(inputs.length, params.max_risk_inputs_per_identity);results['notes'] = [];if (num_inputs_to_score == params.max_risk_inputs_per_identity) { results['notes'].add('Number of risk inputs (' + inputs.length + ') exceeded the maximum allowed (' + params.max_risk_inputs_per_identity + ').');}results['category_1_score'] = 0.0;results['category_1_count'] = 0;double total_score = 0;double current_score = 0;List risk_inputs = [];for (int i = 0; i < num_inputs_to_score; i++) { current_score = inputs[i].score / Math.pow(i + 1, params.p); if (i < 10) { inputs[i][\\"contribution\\"] = 100 * current_score / params.risk_cap; risk_inputs.add(inputs[i]); } if (inputs[i].category == 'signal') { results['category_1_score'] += current_score; results['category_1_count'] += 1; } total_score += current_score;}if (params.containsKey('global_identifier_type_weight') && params.global_identifier_type_weight != null) { total_score *= params.global_identifier_type_weight;}double score_norm = 100 * total_score / params.risk_cap;results['score'] = total_score;results['normalized_score'] = score_norm;results['risk_inputs'] = risk_inputs;return results;",
+      }
+    `);
+  });
+});

x-pack/plugins/security_solution/server/lib/entity_analytics/risk_score/painless/index.ts

@@ -0,0 +1,37 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import fs from 'fs';
+
+const PHASES = ['init', 'map', 'combine', 'reduce'] as const;
+
+type Phase = typeof PHASES[number];
+export type PainlessScripts = Record<Phase, string>;
+
+const removeNewlines = (content: string) => content.replace(/\n/g, '');
+const condenseMultipleSpaces = (content: string) => content.replace(/\s+/g, ' ');
+const minifyContent = (content: string) => condenseMultipleSpaces(removeNewlines(content));
+const readScript = async (phase: Phase) => {
+  const content = await fs.promises.readFile(`${__dirname}/risk_scoring_${phase}.painless`, 'utf8');
+  return minifyContent(content);
+};
+
+let cache: PainlessScripts | undefined;
+
+export const getPainlessScripts = async (): Promise<PainlessScripts> => {
+  if (cache) {
+    return cache;
+  }
+
+  const [init, map, combine, reduce] = await Promise.all(PHASES.map(readScript));
+
+  // The cache will only ever have one value, so we can safely update it
+  // un-atomicly without worrying about lost updates.
+  // eslint-disable-next-line require-atomic-updates
+  cache = { init, map, combine, reduce };
+  return cache;
+};

x-pack/plugins/security_solution/server/lib/entity_analytics/risk_score/painless/risk_scoring_combine.painless

@@ -0,0 +1 @@
+return state;

x-pack/plugins/security_solution/server/lib/entity_analytics/risk_score/painless/risk_scoring_init.painless

@@ -0,0 +1 @@
+state.inputs = []

x-pack/plugins/security_solution/server/lib/entity_analytics/risk_score/painless/risk_scoring_map.painless

@@ -0,0 +1,10 @@
+Map fields = new HashMap();
+String category = doc['event.kind'].value;
+double score = doc['kibana.alert.risk_score'].value;
+fields.put('time', doc['@timestamp'].value);
+fields.put('rule_name', doc['kibana.alert.rule.name'].value);
+fields.put('category', category);
+fields.put('index', doc['_index'].value);
+fields.put('id', doc['kibana.alert.uuid'].value);
+fields.put('score', score);
+state.inputs.add(fields);      

x-pack/plugins/security_solution/server/lib/entity_analytics/risk_score/painless/risk_scoring_reduce.painless

@@ -0,0 +1,43 @@
+Map results = new HashMap();
+List inputs = [];
+for (state in states) {
+  inputs.addAll(state.inputs)
+}
+double num_inputs_to_score = Math.min(inputs.length, params.max_risk_inputs_per_identity);
+results['notes'] = [];
+if (num_inputs_to_score == params.max_risk_inputs_per_identity) {
+  results['notes'].add('Number of risk inputs (' + inputs.length + ') exceeded the maximum allowed (' + params.max_risk_inputs_per_identity + ').');
+}
+
+results['category_1_score'] = 0.0;
+results['category_1_count'] = 0;
+
+double total_score = 0;
+double current_score = 0;
+List risk_inputs = [];
+for (int i = 0; i < num_inputs_to_score; i++) {
+  current_score = inputs[i].score / Math.pow(i + 1, params.p);
+
+  if (i < 10) {
+    inputs[i]["contribution"] = 100 * current_score / params.risk_cap;
+    risk_inputs.add(inputs[i]);
+  }
+
+  if (inputs[i].category == 'signal') {
+    results['category_1_score'] += current_score; results['category_1_count'] += 1;
+  }
+
+  total_score += current_score;
+}
+
+if (params.containsKey('global_identifier_type_weight') && params.global_identifier_type_weight != null) {
+  total_score *= params.global_identifier_type_weight;
+}
+
+double score_norm = 100 * total_score / params.risk_cap;
+results['score'] = total_score;
+results['normalized_score'] = score_norm;
+results['risk_inputs'] = risk_inputs;
+
+return results;
+