diff --git a/README.md b/README.md index 6079366..3f1d726 100644 --- a/README.md +++ b/README.md @@ -57,7 +57,7 @@ Examles of config: ### Alerts `yarn start help` - To see the commands list -`yarn start generate-alerts ` - Generate *n* alerts +`yarn start generate-alerts -n -h -u `yarn start delete-alerts` - Delete all alerts @@ -76,7 +76,7 @@ To modify alert document, you can change `createAlert.ts` file. Example list of command for testing Risk Score API woth 10.000 alerts. ``` yarn start delete-alerts -yarn start generate-alerts 10000 +yarn start generate-alerts -n 10000 -h 100 -u 100 yarn start test-risk-score ``` diff --git a/package.json b/package.json index d828047..f86da8d 100644 --- a/package.json +++ b/package.json @@ -16,6 +16,7 @@ "@faker-js/faker": "^8.0.0", "@types/lodash-es": "^4.17.12", "chalk": "^5.2.0", + "cli-progress": "^3.12.0", "commander": "^10.0.0", "conf": "^11.0.1", "fp-ts": "^2.16.5", @@ -24,17 +25,19 @@ "lodash-es": "^4.17.21", "moment": "^2.29.4", "node-fetch": "^3.3.1", + "p-map": "^7.0.2", "tsx": "^4.7.1", "url-join": "^5.0.0", "uuid": "^9.0.1" }, "devDependencies": { - "eslint": "<9.0.0", + "@types/cli-progress": "^3.11.5", "@types/inquirer": "^9.0.7", "@types/uuid": "^9.0.8", "@typescript-eslint/eslint-plugin": "^7.7.1", "@typescript-eslint/parser": "^7.7.1", "esbuild": "^0.20.2", + "eslint": "<9.0.0", "eslint-config-airbnb": "^19.0.4", "eslint-config-airbnb-typescript": "^18.0.0", "typescript": "^5.3.3" diff --git a/src/commands/documents.ts b/src/commands/documents.ts index 818c562..23f8432 100644 --- a/src/commands/documents.ts +++ b/src/commands/documents.ts @@ -1,17 +1,19 @@ import createAlerts from '../createAlerts'; import createEvents from '../createEvents'; -import alertMappings from '../mappings/alertMappings.json' assert { type: 'json' }; import eventMappings from '../mappings/eventMappings.json' assert { type: 'json' }; import { getEsClient, indexCheck } from './utils/index'; import { getConfig } from '../get_config'; import { MappingTypeMapping, BulkOperationContainer } from '@elastic/elasticsearch/lib/api/types'; +import pMap from 'p-map'; +import _ from 'lodash'; +import cliProgress from 'cli-progress'; +import { faker } from '@faker-js/faker'; const config = getConfig(); const client = getEsClient(); const ALERT_INDEX = '.alerts-security.alerts-default'; -const EVENT_INDEX = config.eventIndex; const generateDocs = async ({ createDocs, amount, index }: {createDocs: DocumentCreator; amount: number; index: string}) => { if (!client) { @@ -28,21 +30,39 @@ const generateDocs = async ({ createDocs, amount, index }: {createDocs: Document index ); try { - const result = await client.bulk({ body: docs, refresh: true }); + const result = await bulkUpsert(docs); generated += result.items.length / 2; - console.log( - `${result.items.length} documents created, ${amount - generated} left` - ); } catch (err) { console.log('Error: ', err); + process.exit(1); } } }; +const bulkUpsert = async (docs: unknown[]) => { + if (!client) { + throw new Error('failed to create ES client'); + } + try { + return client.bulk({ body: docs, refresh: true }); + } catch (err) { + console.log('Error: ', err); + process.exit(1); + } +}; + interface DocumentCreator { (descriptor: { id_field: string, id_value: string }): object; } +const alertToBatchOps = (alert: object, index: string): unknown[] => { + return [ + { index: { _index: index } }, + { ...alert }, + ]; + +} + const createDocuments = (n: number, generated: number, createDoc: DocumentCreator, index: string): unknown[] => { return Array(n) .fill(null) @@ -64,29 +84,71 @@ const createDocuments = (n: number, generated: number, createDoc: DocumentCreato }; -export const generateAlerts = async (n: number) => { - await indexCheck(ALERT_INDEX, alertMappings as MappingTypeMapping); +export const generateAlerts = async (alertCount: number, hostCount: number, userCount: number) => { + + if (userCount > alertCount) { + console.log('User count should be less than alert count'); + process.exit(1); + } - console.log('Generating alerts...'); + if (hostCount > alertCount) { + console.log('Host count should be less than alert count'); + process.exit(1); + } - await generateDocs({ - createDocs: createAlerts, - amount: n, - index: ALERT_INDEX, - }); + console.log(`Generating ${alertCount} alerts containing ${hostCount} hosts and ${userCount} users.`); + const concurrency = 10; // how many batches to send in parallel + const batchSize = 2500; // number of alerts in a batch + const no_overrides = {}; + + const batchOpForIndex = ({ userName, hostName } : { userName: string, hostName: string }) => alertToBatchOps(createAlerts(no_overrides, { userName, hostName }), ALERT_INDEX); + + + console.log('Generating entity names...'); + const userNames = Array.from({ length: userCount }, () => faker.internet.userName()); + const hostNames = Array.from({ length: hostCount }, () => faker.internet.domainName()); + + console.log('Assigning entity names...') + const alertEntityNames = Array.from({ length: alertCount }, (_, i) => ({ + userName: userNames[i % userCount], + hostName: hostNames[i % hostCount], + })); + + console.log('Entity names assigned. Batching...'); + const operationBatches = _.chunk(alertEntityNames, batchSize).map((batch) => + batch.flatMap(batchOpForIndex) + ); + + console.log('Batching complete. Sending to ES...'); + + console.log(`Sending in ${operationBatches.length} batches of ${batchSize} alerts, with up to ${concurrency} batches in parallel\n\n`); + const progress = new cliProgress.SingleBar({}, cliProgress.Presets.shades_classic); + + progress.start(operationBatches.length, 0); + + await pMap( + operationBatches, + async (operations) => { + await bulkUpsert(operations); + progress.increment(); + }, + { concurrency } + ); - console.log('Finished gerating alerts'); + progress.stop(); }; +// this creates asset criticality not events? export const generateEvents = async (n: number) => { - await indexCheck(EVENT_INDEX, eventMappings as MappingTypeMapping); + if(!config.eventIndex) { throw new Error('eventIndex not defined in config'); } + await indexCheck(config.eventIndex, eventMappings as MappingTypeMapping); console.log('Generating events...'); await generateDocs({ createDocs: createEvents, amount: n, - index: EVENT_INDEX, + index: config.eventIndex, }); console.log('Finished generating events'); @@ -112,10 +174,10 @@ export const generateGraph = async ({ users = 100, maxHosts = 3 }) => { for (let j = 0; j < maxHosts; j++) { const alert = createAlerts({ host: { - name: `Host ${i}${j}`, + name: 'Host mark', }, user: { - name: `User ${i}`, + name: 'User pablo', }, }); userCluster.push(alert); @@ -179,11 +241,12 @@ export const deleteAllAlerts = async () => { export const deleteAllEvents = async () => { console.log('Deleting all events...'); + if (!config.eventIndex) { throw new Error('eventIndex not defined in config'); } try { console.log('Deleted all events'); if (!client) throw new Error; await client.deleteByQuery({ - index: EVENT_INDEX, + index: config.eventIndex, refresh: true, body: { query: { diff --git a/src/createAlerts.ts b/src/createAlerts.ts index 6fadfc7..6ba147b 100644 --- a/src/createAlerts.ts +++ b/src/createAlerts.ts @@ -1,7 +1,16 @@ import { faker } from '@faker-js/faker'; -function baseCreateAlerts() { +function baseCreateAlerts({ + userName = 'user-1', + hostName = 'host-1', +} : { + userName?: string, + hostName?: string, +} = { +}) { return { + 'host.name': hostName, + 'user.name': userName, 'kibana.alert.start': '2023-04-11T20:18:15.816Z', 'kibana.alert.last_detected': '2023-04-11T20:18:15.816Z', 'kibana.version': '8.7.0', @@ -59,7 +68,7 @@ function baseCreateAlerts() { 'kibana.alert.status': 'active', 'kibana.alert.workflow_status': 'open', 'kibana.alert.depth': 1, - 'kibana.alert.reason': 'event on Host 4 created low alert 1.', + 'kibana.alert.reason': 'event on ' + hostName + 'created low alert 1.', 'kibana.alert.severity': 'low', 'kibana.alert.risk_score': 21, 'kibana.alert.rule.actions': [], @@ -96,6 +105,13 @@ function baseCreateAlerts() { } } -export default function createAlerts(override: O): O & ReturnType { - return { ...baseCreateAlerts(), ...override }; +export default function createAlerts(override: O, { + userName, + hostName, +} : { + userName?: string, + hostName?: string, +} = { +}): O & ReturnType { + return { ...baseCreateAlerts({ userName, hostName}), ...override }; } diff --git a/src/get_config.ts b/src/get_config.ts index ca230de..512a07a 100644 --- a/src/get_config.ts +++ b/src/get_config.ts @@ -18,8 +18,8 @@ const Node = t.union([NodeWithCredentials, NodeWithAPIKey]); const Config = t.type({ elastic: Node, kibana: Node, - eventIndex: t.string, - eventDateOffsetHours: t.number, + eventIndex: t.union([t.string, t.undefined]), + eventDateOffsetHours: t.union([t.number, t.undefined]), }); type ConfigType = t.TypeOf; diff --git a/src/index.ts b/src/index.ts index 744f2d2..77c8f1a 100644 --- a/src/index.ts +++ b/src/index.ts @@ -17,9 +17,16 @@ import { ENTITY_STORE_OPTIONS, generateNewSeed } from './constants'; program .command('generate-alerts') - .argument('', 'integer argument', parseInt) + .option('-n ', 'number of alerts') + .option('-h ', 'number of hosts') + .option('-u ', 'number of users') .description('Generate fake alerts') - .action(generateAlerts); + .action((options) => { + const alertsCount = parseInt(options.n || 1); + const hostCount = parseInt(options.h || 1); + const userCount = parseInt(options.u || 1); + generateAlerts(alertsCount, userCount, hostCount); + }); program .command('generate-events') diff --git a/yarn.lock b/yarn.lock index 3b4f935..c368103 100644 --- a/yarn.lock +++ b/yarn.lock @@ -346,6 +346,13 @@ "@nodelib/fs.scandir" "2.1.5" fastq "^1.6.0" +"@types/cli-progress@^3.11.5": + version "3.11.5" + resolved "https://registry.yarnpkg.com/@types/cli-progress/-/cli-progress-3.11.5.tgz#9518c745e78557efda057e3f96a5990c717268c3" + integrity sha512-D4PbNRbviKyppS5ivBGyFO29POlySLmA2HyUFE4p5QGazAMM3CwkKWcvTl8gvElSuxRh6FPKL8XmidX873ou4g== + dependencies: + "@types/node" "*" + "@types/inquirer@^9.0.7": version "9.0.7" resolved "https://registry.yarnpkg.com/@types/inquirer/-/inquirer-9.0.7.tgz#61bb8d0e42f038b9a1738b08fba7fa98ad9b4b24" @@ -650,6 +657,13 @@ cli-cursor@^3.1.0: dependencies: restore-cursor "^3.1.0" +cli-progress@^3.12.0: + version "3.12.0" + resolved "https://registry.yarnpkg.com/cli-progress/-/cli-progress-3.12.0.tgz#807ee14b66bcc086258e444ad0f19e7d42577942" + integrity sha512-tRkV3HJ1ASwm19THiiLIXLO7Im7wlTuKnvkYaTkyoAPefqjNg7W7DHKUlGRxy9vxDvbyCYQkQozvptuMkGCg8A== + dependencies: + string-width "^4.2.3" + cli-spinners@^2.5.0: version "2.9.2" resolved "https://registry.yarnpkg.com/cli-spinners/-/cli-spinners-2.9.2.tgz#1773a8f4b9c4d6ac31563df53b3fc1d79462fe41" @@ -1570,6 +1584,11 @@ p-locate@^5.0.0: dependencies: p-limit "^3.0.2" +p-map@^7.0.2: + version "7.0.2" + resolved "https://registry.yarnpkg.com/p-map/-/p-map-7.0.2.tgz#7c5119fada4755660f70199a66aa3fe2f85a1fe8" + integrity sha512-z4cYYMMdKHzw4O5UkWJImbZynVIo0lSGTXc7bzB1e/rrDqkgGUNysK/o4bTr+0+xKvvLoTyGqYC4Fgljy9qe1Q== + parent-module@^1.0.0: version "1.0.1" resolved "https://registry.yarnpkg.com/parent-module/-/parent-module-1.0.1.tgz#691d2709e78c79fae3a156622452d00762caaaa2"