Skip to content

Commit 1179030

Browse files
patrykkopycinskipatrykkopycinski
committed
refactor(detection-emulation): factory for per-family command tools
Creates createRunFamilyCommandTool factory that builds the schema, confirmation, and handler from a FamilyToolConfig object. All four per-family tools (process, file, network, execution) are now config-only modules (~50 lines each) delegating to the factory. Eliminates ~400 lines of duplicated handler/schema/confirmation logic. Adding a new family (e.g. registry) is now a one-file, config-only addition.
1 parent f6110d7 commit 1179030

5 files changed

Lines changed: 207 additions & 525 deletions

File tree

‎x-pack/solutions/security/plugins/security_solution/server/agent_builder/skills/detection_emulation/create_run_family_command_tool.ts‎

Lines changed: 127 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,127 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
import type { Logger } from '@kbn/core/server';
9+
import { z } from '@kbn/zod/v4';
10+
import { ToolType } from '@kbn/agent-builder-common';
11+
import type { BuiltinSkillBoundedTool } from '@kbn/agent-builder-server/skills';
12+
import { RunEmulationCommandInputSchema } from '../../../../common/detection_emulation/schemas/run_emulation_command_input';
13+
import { MAX_ENDPOINT_FANOUT } from '../../../../common/detection_emulation/schemas/constants';
14+
import type { ConfigType } from '../../../config';
15+
import type { EndpointAppContextService } from '../../../endpoint/endpoint_app_context_services';
16+
import type { SecuritySolutionPluginCoreSetupDependencies } from '../../../plugin_contract';
17+
import type { DetectionEmulationGuardrails } from '../../../lib/detection_emulation/execution/shared_guardrails';
18+
import { buildAgentBuilderActor } from '../../../lib/detection_emulation/execution/audit_context';
19+
import { withCommandGates } from './with_command_gates';
20+
import { buildEmulationConfirmation } from './build_emulation_confirmation';
21+
import { toolError } from './emulation_tool_errors';
22+
23+
export interface FamilyToolConfig {
24+
family: 'process' | 'file' | 'network' | 'execution';
25+
id: string;
26+
commands: readonly [string, ...string[]];
27+
description: string;
28+
commandFieldDescription: string;
29+
}
30+
31+
export interface RunFamilyCommandToolDeps {
32+
core: SecuritySolutionPluginCoreSetupDependencies;
33+
endpointService: EndpointAppContextService;
34+
config: ConfigType;
35+
logger: Logger;
36+
guardrails: DetectionEmulationGuardrails;
37+
}
38+
39+
export function createRunFamilyCommandTool(
40+
familyConfig: FamilyToolConfig,
41+
deps: RunFamilyCommandToolDeps
42+
): BuiltinSkillBoundedTool<any> {
43+
const { family, id, commands, description, commandFieldDescription } = familyConfig;
44+
const { core, endpointService, config, logger, guardrails } = deps;
45+
const { allowlist, rateLimiter, idempotencyCache } = guardrails;
46+
47+
const schema = z.object({
48+
emulationId: z.string().min(1).describe('Unique identifier for the emulation run.'),
49+
agentType: z
50+
.enum(['endpoint'])
51+
.default('endpoint')
52+
.describe(
53+
'EDR agent type. Currently only `endpoint` (Elastic Defend) is wired. Omit; defaults to `endpoint`.'
54+
),
55+
endpointIds: z
56+
.array(z.string().min(1))
57+
.min(1)
58+
.max(MAX_ENDPOINT_FANOUT, {
59+
message: `endpointIds must contain at most ${MAX_ENDPOINT_FANOUT} entries (MAX_ENDPOINT_FANOUT)`,
60+
})
61+
.describe(
62+
`Endpoint agent IDs to dispatch the action against (1–${MAX_ENDPOINT_FANOUT}). The fanout cap exists so a single call cannot N-multiply the per-host EDR rate budget by accident; if a user asks to dispatch against more than ${MAX_ENDPOINT_FANOUT} endpoints, split the request into sequential calls.`
63+
),
64+
command: z.enum(commands as any).describe(commandFieldDescription),
65+
parameters: z
66+
.record(z.string(), z.unknown())
67+
.optional()
68+
.describe(
69+
'Command-specific parameters (strictly validated server-side per command). See `command` description for the required shape. Every command additionally accepts an optional `{ comment: string }` attached to the response-action audit trail.'
70+
),
71+
});
72+
73+
return {
74+
id,
75+
type: ToolType.builtin,
76+
description,
77+
schema,
78+
confirmation: {
79+
askUser: 'once' as const,
80+
getConfirmation: ({ toolParams }: { toolParams: z.infer<typeof schema> }) =>
81+
buildEmulationConfirmation({
82+
family,
83+
emulationId: toolParams.emulationId,
84+
command: toolParams.command,
85+
endpointIds: toolParams.endpointIds,
86+
parameters: toolParams.parameters,
87+
}),
88+
},
89+
handler: async (
90+
rawParams: z.infer<typeof schema>,
91+
{ esClient, spaceId, request, runContext, callContext }: any
92+
) => {
93+
const { emulationId, agentType, command } = rawParams;
94+
95+
const strictParseResult = RunEmulationCommandInputSchema.safeParse(rawParams);
96+
if (!strictParseResult.success) {
97+
logger.warn(
98+
`Emulation command [${command}] for emulation [${emulationId}] rejected: invalid parameters for command (${strictParseResult.error.message})`
99+
);
100+
return toolError.invalidParameters({
101+
emulation_id: emulationId,
102+
agent_type: agentType,
103+
command,
104+
});
105+
}
106+
107+
const actorContext = buildAgentBuilderActor(runContext, callContext.toolCallId);
108+
109+
return withCommandGates(
110+
{
111+
core,
112+
endpointService,
113+
config,
114+
logger,
115+
allowlist,
116+
rateLimiter,
117+
idempotencyCache,
118+
request,
119+
esClient,
120+
spaceId,
121+
actorContext,
122+
},
123+
strictParseResult.data
124+
);
125+
},
126+
};
127+
}

‎x-pack/solutions/security/plugins/security_solution/server/agent_builder/skills/detection_emulation/run_execution_command_tool.ts‎

Lines changed: 20 additions & 125 deletions
Original file line numberDiff line numberDiff line change
@@ -5,86 +5,17 @@
55
* 2.0.
66
*/
77

8-
import type { Logger } from '@kbn/core/server';
9-
import { z } from '@kbn/zod/v4';
10-
import { ToolType } from '@kbn/agent-builder-common';
11-
import type { BuiltinSkillBoundedTool } from '@kbn/agent-builder-server/skills';
12-
import { RunEmulationCommandInputSchema } from '../../../../common/detection_emulation/schemas/run_emulation_command_input';
13-
import { MAX_ENDPOINT_FANOUT } from '../../../../common/detection_emulation/schemas/constants';
14-
import type { ConfigType } from '../../../config';
15-
import type { EndpointAppContextService } from '../../../endpoint/endpoint_app_context_services';
16-
import type { SecuritySolutionPluginCoreSetupDependencies } from '../../../plugin_contract';
17-
import type { DetectionEmulationGuardrails } from '../../../lib/detection_emulation/execution/shared_guardrails';
18-
import { buildAgentBuilderActor } from '../../../lib/detection_emulation/execution/audit_context';
19-
import { withCommandGates } from './with_command_gates';
20-
import { buildEmulationConfirmation } from './build_emulation_confirmation';
21-
import { toolError } from './emulation_tool_errors';
22-
23-
const EXECUTION_FAMILY_COMMANDS = ['execute', 'runscript', 'cancel'] as const;
24-
25-
/**
26-
* Tool boundary schema for the execution-family commands. See the
27-
* process-family tool docstring for why the boundary keeps
28-
* `parameters` opaque and the handler re-parses with the strict
29-
* discriminated union.
30-
*/
31-
const runExecutionCommandSchema = z.object({
32-
emulationId: z.string().min(1).describe('Unique identifier for the emulation run.'),
33-
agentType: z
34-
.enum(['endpoint'])
35-
.default('endpoint')
36-
.describe(
37-
'EDR agent type. Currently only `endpoint` (Elastic Defend) is wired. Omit; defaults to `endpoint`.'
38-
),
39-
endpointIds: z
40-
.array(z.string().min(1))
41-
.min(1)
42-
.max(MAX_ENDPOINT_FANOUT, {
43-
message: `endpointIds must contain at most ${MAX_ENDPOINT_FANOUT} entries (MAX_ENDPOINT_FANOUT)`,
44-
})
45-
.describe(
46-
`Endpoint agent IDs to dispatch the action against (1–${MAX_ENDPOINT_FANOUT}). The fanout cap exists so a single call cannot N-multiply the per-host EDR rate budget by accident; if a user asks to dispatch against more than ${MAX_ENDPOINT_FANOUT} endpoints, split the request into sequential calls.`
47-
),
48-
command: z.enum(EXECUTION_FAMILY_COMMANDS).describe(
49-
`Execution-family command (HIGHEST IMPACT — runs arbitrary code on the endpoint):
50-
- \`execute\` — \`{ command: string, timeout?: number }\` — run a shell command/executable
51-
- \`runscript\` — \`{ scriptId: string, scriptInput?: string, timeout?: number }\` — run a script-library entry
52-
- \`cancel\` — \`{ id: string }\` — cancel a previously-dispatched response action by id
53-
54-
Every command in this family ALSO accepts an optional \`comment: string\` in \`parameters\` — recorded against the response-action audit trail. Strongly recommended for \`execute\` and \`runscript\` so an auditor can see *why* the code ran (e.g. \`{ command: 'whoami', comment: 'verify hostname for rule X validation' }\`).`
55-
),
56-
parameters: z
57-
.record(z.string(), z.unknown())
58-
.optional()
59-
.describe(
60-
'Command-specific parameters (strictly validated server-side per command). See `command` description for the required shape. Every command additionally accepts an optional `{ comment: string }` attached to the response-action audit trail.'
61-
),
62-
});
63-
64-
export interface RunExecutionCommandToolDeps {
65-
core: SecuritySolutionPluginCoreSetupDependencies;
66-
endpointService: EndpointAppContextService;
67-
config: ConfigType;
68-
logger: Logger;
69-
/** See `RunProcessCommandToolDeps.guardrails`. */
70-
guardrails: DetectionEmulationGuardrails;
71-
}
72-
73-
/**
74-
* Execution-family runEmulationCommand tool. Covers `execute`,
75-
* `runscript`, and `cancel`. Shares the gate sequence with the other
76-
* three per-family tools via {@link withCommandGates}.
77-
*/
78-
export const createRunExecutionCommandTool = (
79-
deps: RunExecutionCommandToolDeps
80-
): BuiltinSkillBoundedTool<typeof runExecutionCommandSchema> => {
81-
const { core, endpointService, config, logger, guardrails } = deps;
82-
const { allowlist, rateLimiter, idempotencyCache } = guardrails;
83-
84-
return {
85-
id: 'security.detection-emulation.run-execution-command',
86-
type: ToolType.builtin,
87-
description: `Run an *execution-family* response action against one or more endpoints.
8+
import {
9+
createRunFamilyCommandTool,
10+
type FamilyToolConfig,
11+
type RunFamilyCommandToolDeps,
12+
} from './create_run_family_command_tool';
13+
14+
const EXECUTION_FAMILY_CONFIG: FamilyToolConfig = {
15+
family: 'execution',
16+
id: 'security.detection-emulation.run-execution-command',
17+
commands: ['execute', 'runscript', 'cancel'],
18+
description: `Run an *execution-family* response action against one or more endpoints.
8819
8920
Covers: \`execute\`, \`runscript\`, \`cancel\`.
9021
@@ -107,51 +38,15 @@ conversation before the first invocation. \`execute\` and \`runscript\` render
10738
a destructive (red) confirm button; \`cancel\` is treated as recoverable. If
10839
the user declines, do NOT retry; surface the cancellation and continue with
10940
unrelated work.`,
110-
schema: runExecutionCommandSchema,
111-
confirmation: {
112-
askUser: 'once',
113-
getConfirmation: ({ toolParams }) =>
114-
buildEmulationConfirmation({
115-
family: 'execution',
116-
emulationId: toolParams.emulationId,
117-
command: toolParams.command,
118-
endpointIds: toolParams.endpointIds,
119-
parameters: toolParams.parameters,
120-
}),
121-
},
122-
handler: async (rawParams, { esClient, spaceId, request, runContext, callContext }) => {
123-
const { emulationId, agentType, command } = rawParams;
41+
commandFieldDescription: `Execution-family command (HIGHEST IMPACT — runs arbitrary code on the endpoint):
42+
- \`execute\` — \`{ command: string, timeout?: number }\` — run a shell command/executable
43+
- \`runscript\` — \`{ scriptId: string, scriptInput?: string, timeout?: number }\` — run a script-library entry
44+
- \`cancel\` — \`{ id: string }\` — cancel a previously-dispatched response action by id
12445
125-
const strictParseResult = RunEmulationCommandInputSchema.safeParse(rawParams);
126-
if (!strictParseResult.success) {
127-
logger.warn(
128-
`Emulation command [${command}] for emulation [${emulationId}] rejected: invalid parameters for command (${strictParseResult.error.message})`
129-
);
130-
return toolError.invalidParameters({
131-
emulation_id: emulationId,
132-
agent_type: agentType,
133-
command,
134-
});
135-
}
46+
Every command in this family ALSO accepts an optional \`comment: string\` in \`parameters\` — recorded against the response-action audit trail. Strongly recommended for \`execute\` and \`runscript\` so an auditor can see *why* the code ran (e.g. \`{ command: 'whoami', comment: 'verify hostname for rule X validation' }\`).`,
47+
};
13648

137-
const actorContext = buildAgentBuilderActor(runContext, callContext.toolCallId);
49+
export type RunExecutionCommandToolDeps = RunFamilyCommandToolDeps;
13850

139-
return withCommandGates(
140-
{
141-
core,
142-
endpointService,
143-
config,
144-
logger,
145-
allowlist,
146-
rateLimiter,
147-
idempotencyCache,
148-
request,
149-
esClient,
150-
spaceId,
151-
actorContext,
152-
},
153-
strictParseResult.data
154-
);
155-
},
156-
};
157-
};
51+
export const createRunExecutionCommandTool = (deps: RunExecutionCommandToolDeps) =>
52+
createRunFamilyCommandTool(EXECUTION_FAMILY_CONFIG, deps);

0 commit comments

Comments
 (0)