Skip to content

Commit 523c55b

Browse files
committed
feat: support dict-style steps with per-step metadata in adversary profiles
Extends adversary atomic_ordering to accept dict-style steps alongside legacy string ability IDs. Each dict step can specify executor_facts per platform for fact injection into command templates. Updates the atomic planner to use step_idx for correct ordering of repeated abilities, and the planning service to generate links with metadata-injected facts.
1 parent a29906d commit 523c55b

5 files changed

Lines changed: 182 additions & 38 deletions

File tree

AdversarySCHEMA.yml

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
# Adversary Profile Schema - Extended Format
2+
#
3+
# atomic_ordering accepts both legacy string ability IDs and dict-style
4+
# steps with per-step metadata. The two formats can be mixed freely.
5+
#
6+
# Dict-style step fields:
7+
# ability_id (required) - UUID of the ability to execute
8+
# metadata (optional) - per-step configuration
9+
# executor_facts - map of platform -> list of {trait, value} dicts
10+
# used to inject facts into command templates
11+
12+
adversary_id: 95cad57a-8e4f-4ba0-87cf-03d8a7cad0a0
13+
name: Example Adversary
14+
description: Demonstrates the extended adversary profile schema
15+
objective: 495a9828-cab1-44dd-a0ca-66e58177d8cc
16+
atomic_ordering:
17+
# Dict-style step with per-platform executor facts
18+
- ability_id: 36eecb80-ede3-442b-8774-956e906aff02
19+
metadata:
20+
executor_facts:
21+
linux:
22+
- trait: time.sleep
23+
value: '10'
24+
# Same ability can appear multiple times with different metadata
25+
- ability_id: 36eecb80-ede3-442b-8774-956e906aff02
26+
metadata:
27+
executor_facts:
28+
linux:
29+
- trait: time.sleep
30+
value: '3'
31+
# Legacy string format still works
32+
# - <ability-uuid>

app/objects/c_adversary.py

Lines changed: 18 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,9 @@ class Meta:
1717
adversary_id = ma.fields.String()
1818
name = ma.fields.String()
1919
description = ma.fields.String()
20-
atomic_ordering = ma.fields.List(ma.fields.String())
20+
atomic_ordering = ma.fields.List(
21+
ma.fields.Raw(), # Accepts either str (ability_id) or dict (step with metadata)
22+
)
2123
objective = ma.fields.String()
2224
tags = ma.fields.List(ma.fields.String(), allow_none=True)
2325
has_repeatable_abilities = ma.fields.Boolean(dump_only=True)
@@ -59,16 +61,18 @@ class Adversary(FirstClassObjectInterface, BaseObject):
5961
def unique(self):
6062
return self.hash('%s' % self.adversary_id)
6163

62-
def __init__(self, name='', adversary_id='', description='', atomic_ordering=(), objective='', tags=None, plugin=''):
64+
def __init__(self, name='', adversary_id='', description='', atomic_ordering=(), objective='', tags=None,
65+
plugin='', metadata=None, **_):
6366
super().__init__()
6467
self.adversary_id = adversary_id if adversary_id else str(uuid.uuid4())
6568
self.name = name
6669
self.description = description
67-
self.atomic_ordering = atomic_ordering
70+
self.atomic_ordering = list(atomic_ordering or [])
6871
self.objective = objective or DEFAULT_OBJECTIVE_ID
6972
self.tags = set(tags) if tags else set()
7073
self.has_repeatable_abilities = False
7174
self.plugin = plugin
75+
self.metadata = metadata or {}
7276

7377
def store(self, ram):
7478
existing = self.retrieve(ram['adversaries'], self.unique)
@@ -85,9 +89,11 @@ def store(self, ram):
8589
return existing
8690

8791
def verify(self, log, abilities, objectives):
88-
for ability_id in self.atomic_ordering:
89-
if not next((ability for ability in abilities if ability.ability_id == ability_id), None):
90-
log.warning('Ability referenced in adversary %s but not found: %s', self.adversary_id, ability_id)
92+
for step in self.atomic_ordering:
93+
ability_id = step if isinstance(step, str) else step.get('ability_id')
94+
if not any(ability.ability_id == ability_id for ability in abilities):
95+
log.warning('Ability referenced in adversary %s but not found: %s',
96+
self.adversary_id, ability_id)
9197

9298
if not self.objective:
9399
self.objective = DEFAULT_OBJECTIVE_ID
@@ -108,4 +114,9 @@ async def which_plugin(self):
108114
return self.plugin
109115

110116
def check_repeatable_abilities(self, ability_list):
111-
return any(ab.repeatable for ab_id in self.atomic_ordering for ab in ability_list if ab.ability_id == ab_id)
117+
for step in self.atomic_ordering:
118+
ability_id = step if isinstance(step, str) else step.get('ability_id')
119+
for ab in ability_list:
120+
if ab.ability_id == ability_id and ab.repeatable:
121+
return True
122+
return False

app/objects/secondclass/c_link.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@ class Meta:
4949
output = ma.fields.String()
5050
deadman = ma.fields.Boolean()
5151
agent_reported_time = ma.fields.DateTime(format=BaseObject.TIME_FORMAT, load_default=None)
52+
metadata = ma.fields.Dict(keys=ma.fields.String(), values=ma.fields.Raw(), allow_none=True)
5253

5354
@ma.pre_load()
5455
def fix_ability(self, link, **_):
@@ -160,7 +161,7 @@ def is_global_variable(cls, variable):
160161
return variable in cls.RESERVED
161162

162163
def __init__(self, command='', plaintext_command='', paw='', ability=None, executor=None, status=-3, score=0, jitter=0, cleanup=0, id='',
163-
pin=0, host=None, deadman=False, used=None, relationships=None, agent_reported_time=None):
164+
pin=0, host=None, deadman=False, used=None, relationships=None, agent_reported_time=None, metadata=None):
164165
super().__init__()
165166
self.id = str(id)
166167
self.command = command
@@ -186,6 +187,7 @@ def __init__(self, command='', plaintext_command='', paw='', ability=None, execu
186187
self.output = False
187188
self.deadman = deadman
188189
self.agent_reported_time = agent_reported_time
190+
self.metadata = metadata or {}
189191

190192
def __eq__(self, other):
191193
if isinstance(other, Link):

app/planners/atomic.py

Lines changed: 15 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -32,11 +32,19 @@ async def _get_links(self, agent=None):
3232
return await self.planning_svc.get_links(operation=self.operation, agent=agent)
3333

3434
# Given list of links, returns the link that appears first in the adversary's atomic ordering.
35-
async def _get_next_atomic_link(self, links):
36-
abil_id_to_link = dict()
37-
for link in links:
38-
abil_id_to_link[link.ability.ability_id] = link
35+
async def _get_next_atomic_link(self, possible_links):
36+
# Try to match based on explicit step_idx (set by _generate_new_links for dict-style steps)
37+
link_lookup = {link.step_idx: link for link in possible_links if hasattr(link, 'step_idx')}
38+
for idx, _ in enumerate(self.operation.adversary.atomic_ordering):
39+
if idx in link_lookup:
40+
return link_lookup[idx]
41+
42+
# Fallback to ability_id matching (legacy string-style ordering)
43+
abil_id_to_link = {link.ability.ability_id: link for link in possible_links}
3944
candidate_ids = set(abil_id_to_link.keys())
40-
for ab_id in self.operation.adversary.atomic_ordering:
41-
if ab_id in candidate_ids:
42-
return abil_id_to_link[ab_id]
45+
for step in self.operation.adversary.atomic_ordering:
46+
ability_id = step if isinstance(step, str) else step.get('ability_id')
47+
if ability_id in candidate_ids:
48+
return abil_id_to_link[ability_id]
49+
50+
return None

app/service/planning_svc.py

Lines changed: 114 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
from app.objects.secondclass.c_fact import Fact
12
from app.objects.secondclass.c_link import Link
23
from app.service.interfaces.i_planning_svc import PlanningServiceInterface
34
from app.utility.base_planning_svc import BasePlanningService
@@ -142,7 +143,8 @@ async def get_links(self, operation, buckets=None, agent=None, trim=True):
142143
143144
For an operation and agent combination, create links (that can
144145
be executed). When no agent is supplied, links for all agents
145-
are returned.
146+
are returned. Supports both legacy string ability IDs and
147+
dict-style steps with per-step metadata in atomic_ordering.
146148
147149
:param operation: Operation to generate links for
148150
:type operation: Operation
@@ -158,25 +160,56 @@ async def get_links(self, operation, buckets=None, agent=None, trim=True):
158160
:type trim: bool, optional
159161
:return: a list of links sorted by score and atomic ordering
160162
"""
161-
ao = operation.adversary.atomic_ordering
162-
abilities = await self.get_service('data_svc') \
163-
.locate('abilities', match=dict(ability_id=tuple(ao)))
163+
raw_steps = operation.adversary.atomic_ordering
164+
step_entries = []
165+
unique_ids = set()
166+
167+
for idx, step in enumerate(raw_steps):
168+
if isinstance(step, str):
169+
ability_id = step
170+
metadata = {}
171+
else:
172+
ability_id = step.get('ability_id')
173+
metadata = step.get('metadata', {})
174+
175+
step_entries.append({
176+
'step_idx': idx,
177+
'ability_id': ability_id,
178+
'metadata': metadata,
179+
})
180+
unique_ids.add(ability_id)
181+
182+
all_abilities = await self.get_service('data_svc').locate(
183+
'abilities', match=dict(ability_id=tuple(unique_ids))
184+
)
185+
ability_map = {a.ability_id: a for a in all_abilities}
186+
187+
steps_with_abilities = []
188+
for step in step_entries:
189+
ability = ability_map.get(step['ability_id'])
190+
if not ability:
191+
continue
192+
steps_with_abilities.append({
193+
'step_idx': step['step_idx'],
194+
'ability': ability,
195+
'metadata': step['metadata'],
196+
})
197+
164198
if buckets:
165-
# buckets specified - get all links for given buckets,
166-
# (still in underlying atomic adversary order)
167-
t = []
168-
for bucket in buckets:
169-
t.extend([ab for ab in abilities for b in ab.buckets if b == bucket])
170-
abilities = t
199+
steps_with_abilities = [
200+
s for s in steps_with_abilities
201+
if any(bucket in s['ability'].buckets for bucket in buckets)
202+
]
203+
171204
links = []
172205
if agent:
173-
links.extend(await self.generate_and_trim_links(agent, operation, abilities, trim))
206+
links.extend(await self.generate_and_trim_links(agent, operation, steps_with_abilities, trim))
174207
else:
175208
agent_links = []
176209
for agent in operation.agents:
177-
agent_links.append(await self.generate_and_trim_links(agent, operation, abilities, trim))
210+
agent_links.append(await self.generate_and_trim_links(agent, operation, steps_with_abilities, trim))
178211
links = await self._remove_links_of_duplicate_singletons(agent_links)
179-
self.log.debug('Generated %s usable links' % (len(links)))
212+
self.log.debug('Generated %s usable links', len(links))
180213
return await self.sort_links(links)
181214

182215
async def get_cleanup_links(self, operation, agent=None):
@@ -332,33 +365,91 @@ async def _check_and_generate_cleanup_links(self, agent, operation):
332365
link_status=operation.link_status())
333366
return agent_cleanup_links
334367

335-
async def _generate_new_links(self, operation, agent, abilities, link_status):
336-
"""Generate links with given status
368+
async def _generate_new_links(self, operation, agent, steps_with_abilities, link_status):
369+
"""Generate links with given status, supporting per-step metadata.
370+
371+
Each entry in steps_with_abilities is a dict with keys:
372+
step_idx - ordinal position in the adversary profile
373+
ability - resolved Ability object
374+
metadata - dict of per-step metadata (may contain executor_facts)
337375
338376
:param operation: Operation to generate links on
339377
:type operation: Operation
340378
:param agent: Agent to generate links on
341379
:type agent: Agent
342-
:param agent: Abilities to generate links for
343-
:type agent: list(Ability)
380+
:param steps_with_abilities: Steps with resolved abilities and metadata
381+
:type steps_with_abilities: list(dict)
344382
:param link_status: Link status, referencing link state dict
345383
:type link_status: int
346384
:return: Links for agent
347385
:rtype: list(Link)
348386
"""
349387
links = []
350-
for ability in await agent.capabilities(abilities):
388+
for step in steps_with_abilities:
389+
idx = step['step_idx']
390+
ability = step['ability']
391+
step_meta = step.get('metadata', {})
392+
393+
supported_abilities = await agent.capabilities([ability])
394+
if ability.ability_id not in [a.ability_id for a in supported_abilities]:
395+
continue
396+
351397
executor = await agent.get_preferred_executor(ability)
352398
if not executor:
353399
continue
400+
354401
await self._call_ability_plugin_hooks(ability, executor)
355-
if executor.command:
356-
link = Link.load(dict(command=self.encode_string(executor.test), paw=agent.paw, score=0,
357-
ability=ability, executor=executor, status=link_status,
358-
jitter=self.jitter(operation.jitter)))
359-
links.append(link)
402+
403+
# Collect per-step executor_facts if present
404+
fact_dicts = []
405+
exec_facts = step_meta.get('executor_facts', {}).get(executor.platform, [])
406+
for f in exec_facts:
407+
try:
408+
fact_dicts.append({f['trait'].strip(): f['value'].strip('" ')})
409+
except (KeyError, AttributeError) as e:
410+
self.log.warning('Skipping malformed fact in step %s: %s (%s)', idx, f, e)
411+
412+
link = Link.load(dict(
413+
paw=agent.paw,
414+
score=0,
415+
ability=ability,
416+
executor=executor,
417+
status=link_status,
418+
jitter=self.jitter(operation.jitter),
419+
))
420+
link.step_idx = idx
421+
422+
if fact_dicts:
423+
for fd in fact_dicts:
424+
for trait, value in fd.items():
425+
link.used.append(Fact(trait=trait, value=value, source=operation.id))
426+
injected = self._inject_facts(executor.test, fact_dicts)
427+
else:
428+
injected, _, used = await self._build_single_test_variant(executor.test, [], executor.name)
429+
for f in used:
430+
link.used.append(f)
431+
432+
link.command = self.encode_string(injected)
433+
links.append(link)
434+
360435
return links
361436

437+
@staticmethod
438+
def _inject_facts(template, fact_dicts):
439+
"""Replace #{trait} placeholders in the template using fact dicts.
440+
441+
:param template: Command template with #{trait} placeholders
442+
:type template: str
443+
:param fact_dicts: List of {trait: value} dicts
444+
:type fact_dicts: list(dict)
445+
:return: Template with placeholders replaced
446+
:rtype: str
447+
"""
448+
for entry in fact_dicts:
449+
for trait, value in entry.items():
450+
template = template.replace(f'#{{{trait}}}', value)
451+
return template
452+
362453
async def _generate_cleanup_links(self, operation, agent, link_status):
363454
"""Generate cleanup links with given status
364455

0 commit comments

Comments
 (0)