mal-toolbox 0.1.11__py3-none-any.whl → 0.2.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -4,12 +4,12 @@ MAL-Toolbox Language Graph Module
4
4
 
5
5
  from __future__ import annotations
6
6
 
7
- import copy
8
7
  import logging
9
8
  import json
10
9
  import zipfile
11
10
 
12
11
  from dataclasses import dataclass, field
12
+ from functools import cached_property
13
13
  from typing import Any, Optional
14
14
 
15
15
  from maltoolbox.file_utils import (
@@ -26,44 +26,62 @@ from ..exceptions import (
26
26
 
27
27
  logger = logging.getLogger(__name__)
28
28
 
29
+
30
+ def disaggregate_attack_step_full_name(
31
+ attack_step_full_name: str) -> list[str]:
32
+ return attack_step_full_name.split(':')
33
+
34
+
29
35
  @dataclass
30
36
  class LanguageGraphAsset:
31
- name: Optional[str] = None
32
- associations: list[LanguageGraphAssociation] = field(default_factory=lambda: [])
33
- attack_steps: list[LanguageGraphAttackStep] = field(default_factory=lambda: [])
34
- description: dict = field(default_factory=lambda: {})
35
- # MAL languages currently do not support multiple inheritance, but this is
36
- # futureproofing at its most hopeful.
37
- super_assets: list = field(default_factory=lambda: [])
38
- sub_assets: list = field(default_factory=lambda: [])
37
+ name: str
38
+ own_associations: dict[str, LanguageGraphAssociation] = \
39
+ field(default_factory = dict)
40
+ attack_steps: dict[str, LanguageGraphAttackStep] = \
41
+ field(default_factory = dict)
42
+ info: dict = field(default_factory = dict)
43
+ own_super_asset: Optional[LanguageGraphAsset] = None
44
+ own_sub_assets: set[LanguageGraphAsset] = field(default_factory = set)
45
+ own_variables: dict = field(default_factory = dict)
39
46
  is_abstract: Optional[bool] = None
40
47
 
48
+
41
49
  def to_dict(self) -> dict:
42
50
  """Convert LanguageGraphAsset to dictionary"""
43
51
  node_dict: dict[str, Any] = {
44
52
  'name': self.name,
45
- 'associations': [],
46
- 'attack_steps': [],
47
- 'description': self.description,
48
- 'super_assets': [],
49
- 'sub_assets': []
53
+ 'associations': {},
54
+ 'attack_steps': {},
55
+ 'info': self.info,
56
+ 'super_asset': self.own_super_asset.name \
57
+ if self.own_super_asset else "",
58
+ 'sub_assets': [asset.name for asset in self.own_sub_assets],
59
+ 'variables': {},
60
+ 'is_abstract': self.is_abstract
50
61
  }
51
62
 
52
- for assoc in self.associations:
53
- node_dict['associations'].append((assoc.name,
54
- assoc.left_field.fieldname,
55
- assoc.right_field.fieldname))
56
- for attack_step in self.attack_steps:
57
- node_dict['attack_steps'].append(attack_step.name)
58
- for super_asset in self.super_assets:
59
- node_dict['super_assets'].append(super_asset.name)
60
- for sub_asset in self.sub_assets:
61
- node_dict['sub_assets'].append(sub_asset.name)
63
+ for assoc in self.own_associations.values():
64
+ node_dict['associations'][assoc.full_name] = assoc.to_dict()
65
+ for attack_step in self.attack_steps.values():
66
+ node_dict['attack_steps'][attack_step.name] = \
67
+ attack_step.to_dict()
68
+ for variable_name, (var_target_asset, var_expr_chain) in \
69
+ self.own_variables.items():
70
+ node_dict['variables'][variable_name] = (
71
+ var_target_asset.name,
72
+ var_expr_chain.to_dict()
73
+ )
62
74
  return node_dict
63
75
 
76
+
64
77
  def __repr__(self) -> str:
65
78
  return str(self.to_dict())
66
79
 
80
+
81
+ def __hash__(self):
82
+ return hash(self.name)
83
+
84
+
67
85
  def is_subasset_of(self, target_asset: LanguageGraphAsset) -> bool:
68
86
  """
69
87
  Check if an asset extends the target asset through inheritance.
@@ -76,15 +94,15 @@ class LanguageGraphAsset:
76
94
  True if this asset extends the target_asset via inheritance.
77
95
  False otherwise.
78
96
  """
79
- current_assets = [self]
80
- while (current_assets):
81
- current_asset = current_assets.pop()
97
+ current_asset: Optional[LanguageGraphAsset] = self
98
+ while (current_asset):
82
99
  if current_asset == target_asset:
83
100
  return True
84
- current_assets.extend(current_asset.super_assets)
101
+ current_asset = current_asset.own_super_asset
85
102
  return False
86
103
 
87
- def get_all_subassets(self) -> list[LanguageGraphAsset]:
104
+ @cached_property
105
+ def sub_assets(self) -> set[LanguageGraphAsset]:
88
106
  """
89
107
  Return a list of all of the assets that directly or indirectly extend
90
108
  this asset.
@@ -92,15 +110,18 @@ class LanguageGraphAsset:
92
110
  Return:
93
111
  A list of all of the assets that extend this asset plus itself.
94
112
  """
95
- current_assets = [self]
96
- subassets = [self]
97
- while (current_assets):
98
- current_asset = current_assets.pop()
99
- current_assets.extend(current_asset.sub_assets)
100
- subassets.extend(current_asset.sub_assets)
101
- return subassets
102
-
103
- def get_all_superassets(self) -> list[LanguageGraphAsset]:
113
+ subassets: list[LanguageGraphAsset] = []
114
+ for subasset in self.own_sub_assets:
115
+ subassets.extend(subasset.sub_assets)
116
+
117
+ subassets.extend(self.own_sub_assets)
118
+ subassets.append(self)
119
+
120
+ return set(subassets)
121
+
122
+
123
+ @cached_property
124
+ def super_assets(self) -> list[LanguageGraphAsset]:
104
125
  """
105
126
  Return a list of all of the assets that this asset directly or
106
127
  indirectly extends.
@@ -108,24 +129,79 @@ class LanguageGraphAsset:
108
129
  Return:
109
130
  A list of all of the assets that this asset extends plus itself.
110
131
  """
111
- current_assets = [self]
112
- superassets = [self]
113
- while (current_assets):
114
- current_asset = current_assets.pop()
115
- current_assets.extend(current_asset.super_assets)
116
- superassets.extend(current_asset.super_assets)
132
+ current_asset: Optional[LanguageGraphAsset] = self
133
+ superassets = []
134
+ while (current_asset):
135
+ superassets.append(current_asset)
136
+ current_asset = current_asset.own_super_asset
117
137
  return superassets
118
138
 
139
+
140
+ @cached_property
141
+ def associations(self) -> dict[str, LanguageGraphAssociation]:
142
+ """
143
+ Return a list of all of the associations that belong to this asset
144
+ directly or indirectly via inheritance.
145
+
146
+ Return:
147
+ A list of all of the associations that apply to this asset, either
148
+ directly or via inheritance.
149
+ """
150
+
151
+ associations = dict(self.own_associations)
152
+ if self.own_super_asset:
153
+ associations |= self.own_super_asset.associations
154
+ return associations
155
+
156
+
157
+ @property
158
+ def variables(self) -> dict[str, ExpressionsChain]:
159
+ """
160
+ Return a list of all of the variables that belong to this asset
161
+ directly or indirectly via inheritance.
162
+
163
+ Return:
164
+ A list of all of the variables that apply to this asset, either
165
+ directly or via inheritance.
166
+ """
167
+
168
+ all_vars = dict(self.own_variables)
169
+ if self.own_super_asset:
170
+ all_vars |= self.own_super_asset.variables
171
+ return all_vars
172
+
173
+
174
+ def get_variable(
175
+ self,
176
+ var_name: str,
177
+ ) -> Optional[tuple]:
178
+ """
179
+ Return a variable matching the given name if the asset or any of its
180
+ super assets has its definition.
181
+
182
+ Return:
183
+ A tuple containing the target asset and expressions chain to it if the
184
+ variable was defined.
185
+ None otherwise.
186
+ """
187
+ current_asset: Optional[LanguageGraphAsset] = self
188
+ while (current_asset):
189
+ if var_name in current_asset.own_variables:
190
+ return current_asset.own_variables[var_name]
191
+ current_asset = current_asset.own_super_asset
192
+ return None
193
+
194
+
119
195
  def get_all_common_superassets(
120
196
  self, other: LanguageGraphAsset
121
- ) -> set[Optional[str]]:
197
+ ) -> set[str]:
122
198
  """Return a set of all common ancestors between this asset
123
199
  and the other asset given as parameter"""
124
200
  self_superassets = set(
125
- asset.name for asset in self.get_all_superassets()
201
+ asset.name for asset in self.super_assets
126
202
  )
127
203
  other_superassets = set(
128
- asset.name for asset in other.get_all_superassets()
204
+ asset.name for asset in other.super_assets
129
205
  )
130
206
  return self_superassets.intersection(other_superassets)
131
207
 
@@ -143,12 +219,13 @@ class LanguageGraphAssociation:
143
219
  name: str
144
220
  left_field: LanguageGraphAssociationField
145
221
  right_field: LanguageGraphAssociationField
146
- description: Optional[dict] = None
222
+ info: dict = field(default_factory = dict)
147
223
 
148
224
  def to_dict(self) -> dict:
149
225
  """Convert LanguageGraphAssociation to dictionary"""
150
- node_dict = {
226
+ assoc_dict = {
151
227
  'name': self.name,
228
+ 'info': self.info,
152
229
  'left': {
153
230
  'asset': self.left_field.asset.name,
154
231
  'fieldname': self.left_field.fieldname,
@@ -160,15 +237,31 @@ class LanguageGraphAssociation:
160
237
  'fieldname': self.right_field.fieldname,
161
238
  'min': self.right_field.minimum,
162
239
  'max': self.right_field.maximum
163
- },
164
- 'description': self.description
240
+ }
165
241
  }
166
242
 
167
- return node_dict
243
+ return assoc_dict
244
+
168
245
 
169
246
  def __repr__(self) -> str:
170
247
  return str(self.to_dict())
171
248
 
249
+
250
+ @property
251
+ def full_name(self) -> str:
252
+ """
253
+ Return the full name of the association. This is a combination of the
254
+ association name, left field name, left asset type, right field name,
255
+ and right asset type.
256
+ """
257
+ full_name = '%s_%s_%s' % (
258
+ self.name,\
259
+ self.left_field.fieldname,\
260
+ self.right_field.fieldname
261
+ )
262
+ return full_name
263
+
264
+
172
265
  def contains_fieldname(self, fieldname: str) -> bool:
173
266
  """
174
267
  Check if the association contains the field name given as a parameter.
@@ -184,6 +277,7 @@ class LanguageGraphAssociation:
184
277
  return True
185
278
  return False
186
279
 
280
+
187
281
  def contains_asset(self, asset: Any) -> bool:
188
282
  """
189
283
  Check if the association matches the asset given as a parameter. A
@@ -201,6 +295,7 @@ class LanguageGraphAssociation:
201
295
  return True
202
296
  return False
203
297
 
298
+
204
299
  def get_opposite_fieldname(self, fieldname: str) -> str:
205
300
  """
206
301
  Return the opposite field name if the association contains the field
@@ -221,6 +316,7 @@ class LanguageGraphAssociation:
221
316
  logger.error(msg, fieldname, self.name)
222
317
  raise LanguageGraphAssociationError(msg % (fieldname, self.name))
223
318
 
319
+
224
320
  def get_opposite_asset(
225
321
  self, asset: LanguageGraphAsset
226
322
  ) -> Optional[LanguageGraphAsset]:
@@ -255,15 +351,29 @@ class LanguageGraphAttackStep:
255
351
  name: str
256
352
  type: str
257
353
  asset: LanguageGraphAsset
258
- ttc: dict = field(default_factory = lambda: {})
259
- children: dict = field(default_factory = lambda: {})
260
- parents: dict = field(default_factory = lambda: {})
261
- description: dict = field(default_factory = lambda: {})
262
- attributes: Optional[dict] = None
354
+ ttc: dict = field(default_factory = dict)
355
+ overrides: bool = False
356
+ children: dict = field(default_factory = dict)
357
+ parents: dict = field(default_factory = dict)
358
+ info: dict = field(default_factory = dict)
359
+ inherits: Optional[LanguageGraphAttackStep] = None
360
+ tags: set = field(default_factory = set)
361
+ _attributes: Optional[dict] = None
362
+
363
+
364
+ def __hash__(self):
365
+ return hash(self.full_name)
366
+
263
367
 
264
368
  @property
265
- def qualified_name(self) -> str:
266
- return f"{self.asset.name}:{self.name}"
369
+ def full_name(self) -> str:
370
+ """
371
+ Return the full name of the attack step. This is a combination of the
372
+ asset type name to which the attack step belongs and attack step name
373
+ itself.
374
+ """
375
+ full_name = self.asset.name + ':' + self.name
376
+ return full_name
267
377
 
268
378
  def to_dict(self) -> dict:
269
379
  node_dict: dict[Any, Any] = {
@@ -273,95 +383,138 @@ class LanguageGraphAttackStep:
273
383
  'ttc': self.ttc,
274
384
  'children': {},
275
385
  'parents': {},
276
- 'description': self.description
386
+ 'info': self.info,
387
+ 'overrides': self.overrides,
388
+ 'inherits': self.inherits.full_name if self.inherits else None,
389
+ 'tags': list(self.tags)
277
390
  }
278
391
 
279
392
  for child in self.children:
280
393
  node_dict['children'][child] = []
281
- for (_, dep_chain) in self.children[child]:
282
- if dep_chain:
394
+ for (_, expr_chain) in self.children[child]:
395
+ if expr_chain:
283
396
  node_dict['children'][child].append(
284
- dep_chain.to_dict())
397
+ expr_chain.to_dict())
285
398
  else:
286
399
  node_dict['children'][child].append(None)
287
400
 
288
401
  for parent in self.parents:
289
402
  node_dict['parents'][parent] = []
290
- for (_, dep_chain) in self.parents[parent]:
291
- if dep_chain:
403
+ for (_, expr_chain) in self.parents[parent]:
404
+ if expr_chain:
292
405
  node_dict['parents'][parent].append(
293
- dep_chain.to_dict())
406
+ expr_chain.to_dict())
294
407
  else:
295
408
  node_dict['parents'][parent].append(None)
296
409
 
410
+ if hasattr(self, 'own_requires'):
411
+ node_dict['requires'] = []
412
+ for requirement in self.own_requires:
413
+ node_dict['requires'].append(requirement.to_dict())
414
+
297
415
  return node_dict
298
416
 
417
+
418
+ @cached_property
419
+ def requires(self):
420
+ if not hasattr(self, 'own_requires'):
421
+ requirements = []
422
+ else:
423
+ requirements = self.own_requires
424
+
425
+ if self.inherits:
426
+ requirements.extend(self.inherits.requires)
427
+ return requirements
428
+
429
+
299
430
  def __repr__(self) -> str:
300
431
  return str(self.to_dict())
301
432
 
302
433
 
303
- class DependencyChain:
304
- def __init__(self, type: str, next_link: Optional[DependencyChain]):
434
+ class ExpressionsChain:
435
+ def __init__(self,
436
+ type: str,
437
+ left_link: Optional[ExpressionsChain] = None,
438
+ right_link: Optional[ExpressionsChain] = None,
439
+ sub_link: Optional[ExpressionsChain] = None,
440
+ fieldname: Optional[str] = None,
441
+ association = None,
442
+ subtype = None
443
+ ):
305
444
  self.type = type
306
- self.next_link: Optional[DependencyChain] = next_link
307
- self.fieldname: str = ""
308
- self.association: Optional[LanguageGraphAssociation] = None
309
- self.left_chain: Optional[DependencyChain] = None
310
- self.right_chain: Optional[DependencyChain] = None
311
- self.subtype: Optional[Any] = None
312
- self.current_link: Optional[DependencyChain] = None
313
-
314
- def __iter__(self):
315
- self.current_link = self
316
- return self
317
-
318
- def __next__(self):
319
- if self.current_link:
320
- dep_chain = self.current_link
321
- self.current_link = self.current_link.next_link
322
- return dep_chain
323
- raise StopIteration
445
+ self.left_link: Optional[ExpressionsChain] = left_link
446
+ self.right_link: Optional[ExpressionsChain] = right_link
447
+ self.sub_link: Optional[ExpressionsChain] = sub_link
448
+ self.fieldname: Optional[str] = fieldname
449
+ self.association: Optional[LanguageGraphAssociation] = association
450
+ self.subtype: Optional[Any] = subtype
451
+
324
452
 
325
453
  def to_dict(self) -> dict:
326
- """Convert DependencyChain to dictionary"""
454
+ """Convert ExpressionsChain to dictionary"""
327
455
  match (self.type):
328
- case 'union' | 'intersection' | 'difference':
329
- return {self.type: {
330
- 'left': self.left_chain.to_dict()
331
- if self.left_chain else {},
332
- 'right': self.right_chain.to_dict()
333
- if self.right_chain else {}
334
- }
456
+ case 'union' | 'intersection' | 'difference' | 'collect':
457
+ return {
458
+ self.type: {
459
+ 'left': self.left_link.to_dict()
460
+ if self.left_link else {},
461
+ 'right': self.right_link.to_dict()
462
+ if self.right_link else {}
463
+ },
464
+ 'type': self.type
335
465
  }
336
466
 
337
467
  case 'field':
338
468
  if not self.association:
339
- raise LanguageGraphAssociationError("Missing association for dep chain")
340
- return {self.association.name:
341
- {'fieldname': self.fieldname,
342
- 'next_association':
343
- self.next_link.to_dict()
344
- if self.next_link else {}
345
- }
346
- }
469
+ raise LanguageGraphAssociationError(
470
+ "Missing association for expressions chain"
471
+ )
472
+ if self.fieldname == self.association.left_field.fieldname:
473
+ asset_type = self.association.left_field.asset.name
474
+ elif self.fieldname == self.association.right_field.fieldname:
475
+ asset_type = self.association.right_field.asset.name
476
+ else:
477
+ raise LanguageGraphException(
478
+ 'Failed to find fieldname "%s" in association:\n%s' %
479
+ (
480
+ self.fieldname,
481
+ json.dumps(self.association.to_dict(),
482
+ indent = 2)
483
+ )
484
+ )
485
+
486
+ return {
487
+ self.association.full_name:
488
+ {
489
+ 'fieldname': self.fieldname,
490
+ 'asset type': asset_type
491
+ },
492
+ 'type': self.type
493
+ }
347
494
 
348
495
  case 'transitive':
349
- return {'transitive':
350
- self.next_link.to_dict()
351
- if self.next_link else {}
496
+ if not self.sub_link:
497
+ raise LanguageGraphException(
498
+ "No sub link for transitive expressions chain"
499
+ )
500
+ return {
501
+ 'transitive': self.sub_link.to_dict(),
502
+ 'type': self.type
352
503
  }
353
504
 
354
505
  case 'subType':
355
506
  if not self.subtype:
356
507
  raise LanguageGraphException(
357
- "No subtype for dependency chain"
508
+ "No subtype for expressions chain"
358
509
  )
359
- if not self.next_link:
510
+ if not self.sub_link:
360
511
  raise LanguageGraphException(
361
- "No next link for subtype dependency chain"
512
+ "No sub link for subtype expressions chain"
362
513
  )
363
- return {'subType': self.subtype.name,
364
- 'expression': self.next_link.to_dict()
514
+ return {
515
+ 'subType': self.subtype.name,
516
+ 'expression': self.sub_link.to_dict(),
517
+ 'type': self.type
365
518
  }
366
519
 
367
520
  case _:
@@ -369,22 +522,112 @@ class DependencyChain:
369
522
  logger.error(msg, self.type)
370
523
  raise LanguageGraphAssociationError(msg % self.type)
371
524
 
525
+ @classmethod
526
+ def _from_dict(cls,
527
+ serialized_expr_chain: dict,
528
+ lang_graph: LanguageGraph,
529
+ ) -> Optional[ExpressionsChain]:
530
+ """Create LanguageGraph from dict
531
+ Args:
532
+ serialized_expr_chain - expressions chain in dict format
533
+ lang_graph - the LanguageGraph that contains the assets,
534
+ associations, and attack steps relevant for
535
+ the expressions chain
536
+ """
537
+
538
+ if serialized_expr_chain is None or not serialized_expr_chain:
539
+ return None
540
+
541
+ if 'type' not in serialized_expr_chain:
542
+ logger.debug(json.dumps(serialized_expr_chain, indent = 2))
543
+ msg = 'Missing expressions chain type!'
544
+ logger.error(msg)
545
+ raise LanguageGraphAssociationError(msg)
546
+ return None
547
+
548
+ expr_chain_type = serialized_expr_chain['type']
549
+ match (expr_chain_type):
550
+ case 'union' | 'intersection' | 'difference' | 'collect':
551
+ left_link = cls._from_dict(
552
+ serialized_expr_chain[expr_chain_type]['left'],
553
+ lang_graph
554
+ )
555
+ right_link = cls._from_dict(
556
+ serialized_expr_chain[expr_chain_type]['right'],
557
+ lang_graph
558
+ )
559
+ new_expr_chain = ExpressionsChain(
560
+ type = expr_chain_type,
561
+ left_link = left_link,
562
+ right_link = right_link
563
+ )
564
+ return new_expr_chain
565
+
566
+ case 'field':
567
+ assoc_name = list(serialized_expr_chain.keys())[0]
568
+ target_asset = lang_graph.assets[\
569
+ serialized_expr_chain[assoc_name]['asset type']]
570
+ new_expr_chain = ExpressionsChain(
571
+ type = 'field',
572
+ association = target_asset.associations[assoc_name],
573
+ fieldname = serialized_expr_chain[assoc_name]['fieldname']
574
+ )
575
+ return new_expr_chain
576
+
577
+ case 'transitive':
578
+ sub_link = cls._from_dict(
579
+ serialized_expr_chain['transitive'],
580
+ lang_graph
581
+ )
582
+ new_expr_chain = ExpressionsChain(
583
+ type = 'transitive',
584
+ sub_link = sub_link
585
+ )
586
+ return new_expr_chain
587
+
588
+ case 'subType':
589
+ sub_link = cls._from_dict(
590
+ serialized_expr_chain['expression'],
591
+ lang_graph
592
+ )
593
+ subtype_name = serialized_expr_chain['subType']
594
+ if subtype_name in lang_graph.assets:
595
+ subtype_asset = lang_graph.assets[subtype_name]
596
+ else:
597
+ msg = 'Failed to find subtype %s'
598
+ logger.error(msg % subtype_name)
599
+ raise LanguageGraphException(msg % subtype_name)
600
+
601
+ new_expr_chain = ExpressionsChain(
602
+ type = 'subType',
603
+ sub_link = sub_link,
604
+ subtype = subtype_asset
605
+ )
606
+ return new_expr_chain
607
+
608
+ case _:
609
+ msg = 'Unknown expressions chain type %s!'
610
+ logger.error(msg, serialized_expr_chain['type'])
611
+ raise LanguageGraphAssociationError(msg %
612
+ serialized_expr_chain['type'])
613
+
614
+
372
615
  def __repr__(self) -> str:
373
616
  return str(self.to_dict())
374
617
 
375
618
 
376
619
  class LanguageGraph():
377
620
  """Graph representation of a MAL language"""
378
- def __init__(self, lang: dict):
379
- self.assets: list = []
380
- self.associations: list = []
381
- self.attack_steps: list = []
382
- self._lang_spec: dict = lang
383
- self.metadata = {
384
- "version": lang["defines"]["version"],
385
- "id": lang["defines"]["id"],
386
- }
387
- self._generate_graph()
621
+ def __init__(self, lang: Optional[dict] = None):
622
+ self.assets: dict = {}
623
+ if lang is not None:
624
+ self._lang_spec: dict = lang
625
+ self.metadata = {
626
+ "version": lang["defines"]["version"],
627
+ "id": lang["defines"]["id"],
628
+ }
629
+ self._generate_graph()
630
+
388
631
 
389
632
  @classmethod
390
633
  def from_mal_spec(cls, mal_spec_file: str) -> LanguageGraph:
@@ -397,6 +640,7 @@ class LanguageGraph():
397
640
  logger.info("Loading mal spec %s", mal_spec_file)
398
641
  return LanguageGraph(MalCompiler().compile(mal_spec_file))
399
642
 
643
+
400
644
  @classmethod
401
645
  def from_mar_archive(cls, mar_archive: str) -> LanguageGraph:
402
646
  """
@@ -411,39 +655,236 @@ class LanguageGraph():
411
655
  langspec = archive.read('langspec.json')
412
656
  return LanguageGraph(json.loads(langspec))
413
657
 
658
+
414
659
  def _to_dict(self):
415
660
  """Converts LanguageGraph into a dict"""
416
- serialized_assets = []
417
- for asset in self.assets:
418
- serialized_assets.append(asset.to_dict())
419
- serialized_associations = []
420
- for associations in self.associations:
421
- serialized_associations.append(associations.to_dict())
422
- serialized_attack_steps = []
423
- for attack_step in self.attack_steps:
424
- serialized_attack_steps.append(attack_step.to_dict())
425
661
 
426
662
  logger.debug(
427
- 'Serializing %s assets, %s associations, and %s attack steps',
428
- len(serialized_assets), len(serialized_associations),
429
- len(serialized_attack_steps)
663
+ 'Serializing %s assets.', len(self.assets.items())
430
664
  )
431
665
 
432
- serialized_graph = {
433
- 'Assets': serialized_assets,
434
- 'Associations': serialized_associations,
435
- 'Attack Steps': serialized_attack_steps
436
- }
666
+ serialized_graph = {}
667
+ for asset in self.assets.values():
668
+ serialized_graph[asset.name] = asset.to_dict()
669
+
437
670
  return serialized_graph
438
671
 
672
+
439
673
  def save_to_file(self, filename: str) -> None:
440
674
  """Save to json/yml depending on extension"""
441
675
  return save_dict_to_file(filename, self._to_dict())
442
676
 
677
+
443
678
  @classmethod
444
- def _from_dict(cls, serialized_object: dict) -> None:
445
- raise NotImplementedError(
446
- "Converting from dict feature is not implemented yet")
679
+ def _from_dict(cls, serialized_graph: dict) -> LanguageGraph:
680
+ """Create LanguageGraph from dict
681
+ Args:
682
+ serialized_graph - LanguageGraph in dict format
683
+ """
684
+
685
+ logger.debug('Create language graph from dictionary.')
686
+ lang_graph = LanguageGraph()
687
+ # Recreate all of the assets
688
+ for asset_dict in serialized_graph.values():
689
+ logger.debug(
690
+ 'Create asset language graph nodes for asset %s',
691
+ asset_dict['name']
692
+ )
693
+ asset_node = LanguageGraphAsset(
694
+ name = asset_dict['name'],
695
+ own_associations = {},
696
+ attack_steps = {},
697
+ info = asset_dict['info'],
698
+ own_super_asset = None,
699
+ own_sub_assets = set(),
700
+ own_variables = {},
701
+ is_abstract = asset_dict['is_abstract']
702
+ )
703
+ lang_graph.assets[asset_dict['name']] = asset_node
704
+
705
+ # Relink assets based on inheritance
706
+ for asset_dict in serialized_graph.values():
707
+ asset = lang_graph.assets[asset_dict['name']]
708
+ super_asset_name = asset_dict['super_asset']
709
+ if not super_asset_name:
710
+ continue
711
+
712
+ super_asset = lang_graph.assets[super_asset_name]
713
+ if not super_asset:
714
+ msg = 'Failed to find super asset "%s" for asset "%s"!'
715
+ logger.error(
716
+ msg, asset_dict["super_asset"], asset_dict["name"])
717
+ raise LanguageGraphSuperAssetNotFoundError(
718
+ msg % (asset_dict["super_asset"], asset_dict["name"]))
719
+
720
+ super_asset.own_sub_assets.add(asset)
721
+ asset.own_super_asset = super_asset
722
+
723
+ # Generate all of the association nodes of the language graph.
724
+ for asset_dict in serialized_graph.values():
725
+ logger.debug(
726
+ 'Create association language graph nodes for asset %s',
727
+ asset_dict['name']
728
+ )
729
+
730
+ asset = lang_graph.assets[asset_dict['name']]
731
+ for association in asset_dict['associations'].values():
732
+ left_asset = lang_graph.assets[association['left']['asset']]
733
+ if not left_asset:
734
+ msg = 'Left asset "%s" for association "%s" not found!'
735
+ logger.error(
736
+ msg, association['left']['asset'],
737
+ association['name'])
738
+ raise LanguageGraphAssociationError(
739
+ msg % (association['left']['asset'],
740
+ association['name']))
741
+
742
+ right_asset = lang_graph.assets[association['right']['asset']]
743
+ if not right_asset:
744
+ msg = 'Right asset "%s" for association "%s" not found!'
745
+ logger.error(
746
+ msg, association['right']['asset'],
747
+ association['name'])
748
+ raise LanguageGraphAssociationError(
749
+ msg % (association['right']['asset'],
750
+ association['name'])
751
+ )
752
+
753
+ assoc_node = LanguageGraphAssociation(
754
+ name = association['name'],
755
+ left_field = LanguageGraphAssociationField(
756
+ left_asset,
757
+ association['left']['fieldname'],
758
+ association['left']['min'],
759
+ association['left']['max']),
760
+ right_field = LanguageGraphAssociationField(
761
+ right_asset,
762
+ association['right']['fieldname'],
763
+ association['right']['min'],
764
+ association['right']['max']),
765
+ info = association['info']
766
+ )
767
+
768
+ # Add the association to the left and right asset
769
+ associated_assets = [left_asset, right_asset]
770
+ while associated_assets != []:
771
+ asset = associated_assets.pop()
772
+ if assoc_node.full_name not in asset.own_associations:
773
+ asset.own_associations[assoc_node.full_name] = assoc_node
774
+
775
+ # Recreate the variables
776
+ for asset_dict in serialized_graph.values():
777
+ asset = lang_graph.assets[asset_dict['name']]
778
+ for variable_name, var_target in asset_dict['variables'].items():
779
+ (target_asset_name, expr_chain_dict) = var_target
780
+ target_asset = lang_graph.assets[target_asset_name]
781
+ expr_chain = ExpressionsChain._from_dict(
782
+ expr_chain_dict,
783
+ lang_graph
784
+ )
785
+ asset.own_variables[variable_name] = (target_asset, expr_chain)
786
+
787
+ # Recreate the attack steps
788
+ for asset_dict in serialized_graph.values():
789
+ asset = lang_graph.assets[asset_dict['name']]
790
+ logger.debug(
791
+ 'Create attack steps language graph nodes for asset %s',
792
+ asset_dict['name']
793
+ )
794
+ for attack_step_dict in asset_dict['attack_steps'].values():
795
+ attack_step_node = LanguageGraphAttackStep(
796
+ name = attack_step_dict['name'],
797
+ type = attack_step_dict['type'],
798
+ asset = asset,
799
+ ttc = attack_step_dict['ttc'],
800
+ overrides = attack_step_dict['overrides'],
801
+ children = {},
802
+ parents = {},
803
+ info = attack_step_dict['info'],
804
+ tags = set(attack_step_dict['tags'])
805
+ )
806
+ asset.attack_steps[attack_step_dict['name']] = \
807
+ attack_step_node
808
+
809
+ # Relink attack steps based on inheritence
810
+ for asset_dict in serialized_graph.values():
811
+ asset = lang_graph.assets[asset_dict['name']]
812
+ for attack_step_dict in asset_dict['attack_steps'].values():
813
+ if 'inherits' in attack_step_dict and \
814
+ attack_step_dict['inherits'] is not None:
815
+ attack_step = asset.attack_steps[
816
+ attack_step_dict['name']]
817
+ ancestor_asset_name, ancestor_attack_step_name = \
818
+ disaggregate_attack_step_full_name(
819
+ attack_step_dict['inherits'])
820
+ ancestor_asset = lang_graph.assets[ancestor_asset_name]
821
+ ancestor_attack_step = ancestor_asset.attack_steps[\
822
+ ancestor_attack_step_name]
823
+ attack_step.inherits = ancestor_attack_step
824
+
825
+
826
+ # Relink attack steps based on expressions chains
827
+ for asset_dict in serialized_graph.values():
828
+ asset = lang_graph.assets[asset_dict['name']]
829
+ for attack_step_dict in asset_dict['attack_steps'].values():
830
+ attack_step = asset.attack_steps[attack_step_dict['name']]
831
+ for child_target in attack_step_dict['children'].items():
832
+ target_full_attack_step_name = child_target[0]
833
+ expr_chains = child_target[1]
834
+ target_asset_name, target_attack_step_name = \
835
+ disaggregate_attack_step_full_name(
836
+ target_full_attack_step_name)
837
+ target_asset = lang_graph.assets[target_asset_name]
838
+ target_attack_step = target_asset.attack_steps[
839
+ target_attack_step_name]
840
+ for expr_chain_dict in expr_chains:
841
+ expr_chain = ExpressionsChain._from_dict(
842
+ expr_chain_dict,
843
+ lang_graph
844
+ )
845
+ if target_attack_step.full_name in attack_step.children:
846
+ attack_step.children[target_attack_step.full_name].\
847
+ append((target_attack_step, expr_chain))
848
+ else:
849
+ attack_step.children[target_attack_step.full_name] = \
850
+ [(target_attack_step, expr_chain)]
851
+
852
+ for parent_target in attack_step_dict['parents'].items():
853
+ target_full_attack_step_name = parent_target[0]
854
+ expr_chains = parent_target[1]
855
+ target_asset_name, target_attack_step_name = \
856
+ disaggregate_attack_step_full_name(
857
+ target_full_attack_step_name)
858
+ target_asset = lang_graph.assets[target_asset_name]
859
+ target_attack_step = target_asset.attack_steps[
860
+ target_attack_step_name]
861
+ for expr_chain_dict in expr_chains:
862
+ expr_chain = ExpressionsChain._from_dict(
863
+ expr_chain_dict,
864
+ lang_graph
865
+ )
866
+ if target_attack_step.full_name in attack_step.parents:
867
+ attack_step.parents[target_attack_step.full_name].\
868
+ append((target_attack_step, expr_chain))
869
+ else:
870
+ attack_step.parents[target_attack_step.full_name] = \
871
+ [(target_attack_step, expr_chain)]
872
+
873
+ # Recreate the requirements of exist and notExist attack steps
874
+ if attack_step.type == 'exist' or \
875
+ attack_step.type == 'notExist':
876
+ if 'requires' in attack_step_dict:
877
+ expr_chains = attack_step_dict['requires']
878
+ attack_step.own_requires = []
879
+ for expr_chain_dict in expr_chains:
880
+ expr_chain = ExpressionsChain._from_dict(
881
+ expr_chain_dict,
882
+ lang_graph
883
+ )
884
+ attack_step.own_requires.append(expr_chain)
885
+
886
+ return lang_graph
887
+
447
888
 
448
889
  @classmethod
449
890
  def load_from_file(cls, filename: str) -> LanguageGraph:
@@ -453,17 +894,23 @@ class LanguageGraph():
453
894
  lang_graph = cls.from_mal_spec(filename)
454
895
  elif filename.endswith('.mar'):
455
896
  lang_graph = cls.from_mar_archive(filename)
456
- elif filename.endswith(('yaml', 'yml')):
897
+ elif filename.endswith(('.yaml', '.yml')):
457
898
  lang_graph = cls._from_dict(load_dict_from_yaml_file(filename))
458
- elif filename.endswith(('json')):
899
+ elif filename.endswith(('.json')):
459
900
  lang_graph = cls._from_dict(load_dict_from_json_file(filename))
901
+ else:
902
+ raise TypeError(
903
+ "Unknown file extension, expected json/mal/mar/yml/yaml"
904
+ )
460
905
 
461
906
  if lang_graph:
462
907
  return lang_graph
908
+ else:
909
+ raise LanguageGraphException(
910
+ f'Failed to load language graph from file "{filename}".'
911
+ )
912
+
463
913
 
464
- raise TypeError(
465
- "Unknown file extension, expected json/mal/mar/yml/yaml"
466
- )
467
914
 
468
915
  def save_language_specification_to_json(self, filename: str) -> None:
469
916
  """
@@ -477,23 +924,21 @@ class LanguageGraph():
477
924
  with open(filename, 'w', encoding='utf-8') as file:
478
925
  json.dump(self._lang_spec, file, indent=4)
479
926
 
927
+
480
928
  def process_step_expression(self,
481
- lang: dict,
482
929
  target_asset,
483
- dep_chain,
930
+ expr_chain,
484
931
  step_expression: dict
485
932
  ) -> tuple:
486
933
  """
487
934
  Recursively process an attack step expression.
488
935
 
489
936
  Arguments:
490
- lang - A dictionary representing the MAL language
491
- specification.
492
937
  target_asset - The asset type that this step expression should
493
938
  apply to. Initially it will contain the asset
494
939
  type to which the attack step belongs.
495
- dep_chain - A dependency chain of linked of associations and
496
- set operations from the attack step to its
940
+ expr_chain - A expressions chain of linked of associations
941
+ and set operations from the attack step to its
497
942
  parent attack step.
498
943
  Note: This was done for the parent attack step
499
944
  because it was easier to construct recursively
@@ -518,19 +963,29 @@ class LanguageGraph():
518
963
  # The attack step expression just adds the name of the attack
519
964
  # step. All other step expressions only modify the target
520
965
  # asset and parent associations chain.
521
- return (target_asset,
522
- dep_chain,
523
- step_expression['name'])
966
+ return (
967
+ target_asset,
968
+ None,
969
+ step_expression['name']
970
+ )
524
971
 
525
972
  case 'union' | 'intersection' | 'difference':
526
973
  # The set operators are used to combine the left hand and right
527
974
  # hand targets accordingly.
528
- lh_target_asset, lh_dep_chain, _ = self.process_step_expression(
529
- lang, target_asset, dep_chain, step_expression['lhs'])
530
- rh_target_asset, rh_dep_chain, _ = self.process_step_expression(
531
- lang, target_asset, dep_chain, step_expression['rhs'])
975
+ lh_target_asset, lh_expr_chain, _ = self.process_step_expression(
976
+ target_asset,
977
+ expr_chain,
978
+ step_expression['lhs']
979
+ )
980
+ rh_target_asset, rh_expr_chain, _ = \
981
+ self.process_step_expression(
982
+ target_asset,
983
+ expr_chain,
984
+ step_expression['rhs']
985
+ )
532
986
 
533
- if not lh_target_asset.get_all_common_superassets(rh_target_asset):
987
+ if not lh_target_asset.get_all_common_superassets(
988
+ rh_target_asset):
534
989
  logger.error(
535
990
  "Set operation attempted between targets that"
536
991
  " do not share any common superassets: %s and %s!",
@@ -538,30 +993,32 @@ class LanguageGraph():
538
993
  )
539
994
  return (None, None, None)
540
995
 
541
- new_dep_chain = DependencyChain(
996
+ new_expr_chain = ExpressionsChain(
542
997
  type = step_expression['type'],
543
- next_link = None)
544
- new_dep_chain.left_chain = lh_dep_chain
545
- new_dep_chain.right_chain = rh_dep_chain
546
- return (lh_target_asset,
547
- new_dep_chain,
548
- None)
998
+ left_link = lh_expr_chain,
999
+ right_link = rh_expr_chain
1000
+ )
1001
+ return (
1002
+ lh_target_asset,
1003
+ new_expr_chain,
1004
+ None
1005
+ )
549
1006
 
550
1007
  case 'variable':
551
- # Fetch the step expression associated with the variable from
552
- # the language specification and resolve that.
553
- variable_step_expr = self._get_variable_for_asset_type_by_name(
554
- target_asset.name, step_expression['name'])
555
- if variable_step_expr:
556
- return self.process_step_expression(
557
- lang,
558
- target_asset,
559
- dep_chain,
560
- variable_step_expr)
561
-
1008
+ var_name = step_expression['name']
1009
+ var_target_asset, var_expr_chain = self._resolve_variable(
1010
+ target_asset, var_name)
1011
+ var_target_asset, var_expr_chain = \
1012
+ target_asset.get_variable(var_name)
1013
+ if var_expr_chain is not None:
1014
+ return (
1015
+ var_target_asset,
1016
+ var_expr_chain,
1017
+ None
1018
+ )
562
1019
  else:
563
1020
  logger.error(
564
- 'Failed to find variable %s for %s',
1021
+ 'Failed to find variable \"%s\" for %s',
565
1022
  step_expression["name"], target_asset.name
566
1023
  )
567
1024
  return (None, None, None)
@@ -578,7 +1035,7 @@ class LanguageGraph():
578
1035
  return (None, None, None)
579
1036
 
580
1037
  new_target_asset = None
581
- for association in target_asset.associations:
1038
+ for association in target_asset.associations.values():
582
1039
  if (association.left_field.fieldname == fieldname and \
583
1040
  target_asset.is_subasset_of(
584
1041
  association.right_field.asset)):
@@ -590,15 +1047,16 @@ class LanguageGraph():
590
1047
  new_target_asset = association.right_field.asset
591
1048
 
592
1049
  if new_target_asset:
593
- new_dep_chain = DependencyChain(
1050
+ new_expr_chain = ExpressionsChain(
594
1051
  type = 'field',
595
- next_link = dep_chain)
596
- new_dep_chain.fieldname = \
597
- association.get_opposite_fieldname(fieldname)
598
- new_dep_chain.association = association
599
- return (new_target_asset,
600
- new_dep_chain,
601
- None)
1052
+ fieldname = fieldname,
1053
+ association = association
1054
+ )
1055
+ return (
1056
+ new_target_asset,
1057
+ new_expr_chain,
1058
+ None
1059
+ )
602
1060
  logger.error(
603
1061
  'Failed to find field "%s" on asset "%s"!',
604
1062
  fieldname, target_asset.name
@@ -609,18 +1067,22 @@ class LanguageGraph():
609
1067
  # Create a transitive tuple entry that applies to the next
610
1068
  # component of the step expression.
611
1069
  result_target_asset, \
612
- result_dep_chain, \
1070
+ result_expr_chain, \
613
1071
  attack_step = \
614
- self.process_step_expression(lang,
1072
+ self.process_step_expression(
615
1073
  target_asset,
616
- dep_chain,
617
- step_expression['stepExpression'])
618
- new_dep_chain = DependencyChain(
1074
+ expr_chain,
1075
+ step_expression['stepExpression']
1076
+ )
1077
+ new_expr_chain = ExpressionsChain(
619
1078
  type = 'transitive',
620
- next_link = result_dep_chain)
621
- return (result_target_asset,
622
- new_dep_chain,
623
- attack_step)
1079
+ sub_link = result_expr_chain
1080
+ )
1081
+ return (
1082
+ result_target_asset,
1083
+ new_expr_chain,
1084
+ attack_step
1085
+ )
624
1086
 
625
1087
  case 'subType':
626
1088
  # Create a subType tuple entry that applies to the next
@@ -628,19 +1090,20 @@ class LanguageGraph():
628
1090
  # asset to the subasset.
629
1091
  subtype_name = step_expression['subType']
630
1092
  result_target_asset, \
631
- result_dep_chain, \
1093
+ result_expr_chain, \
632
1094
  attack_step = \
633
- self.process_step_expression(lang,
1095
+ self.process_step_expression(
634
1096
  target_asset,
635
- dep_chain,
636
- step_expression['stepExpression'])
637
-
638
- subtype_asset = next((asset for asset in self.assets if asset.name == subtype_name), None)
1097
+ expr_chain,
1098
+ step_expression['stepExpression']
1099
+ )
639
1100
 
640
- if not subtype_asset:
641
- msg = 'Failed to find subtype attackstep "{subtype_name}"'
642
- logger.error(msg)
643
- raise LanguageGraphException(msg)
1101
+ if subtype_name in self.assets:
1102
+ subtype_asset = self.assets[subtype_name]
1103
+ else:
1104
+ msg = 'Failed to find subtype %s'
1105
+ logger.error(msg % subtype_name)
1106
+ raise LanguageGraphException(msg % subtype_name)
644
1107
 
645
1108
  if not subtype_asset.is_subasset_of(result_target_asset):
646
1109
  logger.error(
@@ -650,31 +1113,48 @@ class LanguageGraph():
650
1113
  )
651
1114
  return (None, None, None)
652
1115
 
653
- new_dep_chain = DependencyChain(
1116
+ new_expr_chain = ExpressionsChain(
654
1117
  type = 'subType',
655
- next_link = result_dep_chain)
656
- new_dep_chain.subtype = subtype_asset
657
- return (subtype_asset,
658
- new_dep_chain,
659
- attack_step)
1118
+ sub_link = result_expr_chain,
1119
+ subtype = subtype_asset
1120
+ )
1121
+ return (
1122
+ subtype_asset,
1123
+ new_expr_chain,
1124
+ attack_step
1125
+ )
660
1126
 
661
1127
  case 'collect':
662
1128
  # Apply the right hand step expression to left hand step
663
1129
  # expression target asset and parent associations chain.
664
- (lh_target_asset, lh_dep_chain, _) = \
665
- self.process_step_expression(lang,
1130
+ (lh_target_asset, lh_expr_chain, _) = \
1131
+ self.process_step_expression(
666
1132
  target_asset,
667
- dep_chain,
668
- step_expression['lhs'])
1133
+ expr_chain,
1134
+ step_expression['lhs']
1135
+ )
669
1136
  (rh_target_asset,
670
- rh_dep_chain,
1137
+ rh_expr_chain,
671
1138
  rh_attack_step_name) = \
672
- self.process_step_expression(lang,
1139
+ self.process_step_expression(
673
1140
  lh_target_asset,
674
- lh_dep_chain,
675
- step_expression['rhs'])
676
- return (rh_target_asset, rh_dep_chain,
677
- rh_attack_step_name)
1141
+ None,
1142
+ step_expression['rhs']
1143
+ )
1144
+ if rh_expr_chain:
1145
+ new_expr_chain = ExpressionsChain(
1146
+ type = 'collect',
1147
+ left_link = lh_expr_chain,
1148
+ right_link = rh_expr_chain
1149
+ )
1150
+ else:
1151
+ new_expr_chain = lh_expr_chain
1152
+
1153
+ return (
1154
+ rh_target_asset,
1155
+ new_expr_chain,
1156
+ rh_attack_step_name
1157
+ )
678
1158
 
679
1159
  case _:
680
1160
  logger.error(
@@ -682,17 +1162,18 @@ class LanguageGraph():
682
1162
  )
683
1163
  return (None, None, None)
684
1164
 
685
- def reverse_dep_chain(
1165
+
1166
+ def reverse_expr_chain(
686
1167
  self,
687
- dep_chain: Optional[DependencyChain],
688
- reverse_chain: Optional[DependencyChain]
689
- ) -> Optional[DependencyChain]:
1168
+ expr_chain: Optional[ExpressionsChain],
1169
+ reverse_chain: Optional[ExpressionsChain]
1170
+ ) -> Optional[ExpressionsChain]:
690
1171
  """
691
1172
  Recursively reverse the associations chain. From parent to child or
692
1173
  vice versa.
693
1174
 
694
1175
  Arguments:
695
- dep_chain - A chain of nested tuples that specify the
1176
+ expr_chain - A chain of nested tuples that specify the
696
1177
  associations and set operations chain from an
697
1178
  attack step to its connected attack step.
698
1179
  reverse_chain - A chain of nested tuples that represents the
@@ -701,70 +1182,103 @@ class LanguageGraph():
701
1182
  Return:
702
1183
  The resulting reversed associations chain.
703
1184
  """
704
- if not dep_chain:
1185
+ if not expr_chain:
705
1186
  return reverse_chain
706
1187
  else:
707
- match (dep_chain.type):
708
- case 'union' | 'intersection' | 'difference':
1188
+ match (expr_chain.type):
1189
+ case 'union' | 'intersection' | 'difference' | 'collect':
709
1190
  left_reverse_chain = \
710
- self.reverse_dep_chain(dep_chain.left_chain,
1191
+ self.reverse_expr_chain(expr_chain.left_link,
711
1192
  reverse_chain)
712
1193
  right_reverse_chain = \
713
- self.reverse_dep_chain(dep_chain.right_chain,
1194
+ self.reverse_expr_chain(expr_chain.right_link,
714
1195
  reverse_chain)
715
- new_dep_chain = DependencyChain(
716
- type = dep_chain.type,
717
- next_link = None)
718
- new_dep_chain.left_chain = left_reverse_chain
719
- new_dep_chain.right_chain = right_reverse_chain
720
- return new_dep_chain
1196
+ if expr_chain.type == 'collect':
1197
+ new_expr_chain = ExpressionsChain(
1198
+ type = expr_chain.type,
1199
+ left_link = right_reverse_chain,
1200
+ right_link = left_reverse_chain
1201
+ )
1202
+ else:
1203
+ new_expr_chain = ExpressionsChain(
1204
+ type = expr_chain.type,
1205
+ left_link = left_reverse_chain,
1206
+ right_link = right_reverse_chain
1207
+ )
1208
+
1209
+ return new_expr_chain
721
1210
 
722
1211
  case 'transitive':
723
- result_reverse_chain = self.reverse_dep_chain(
724
- dep_chain.next_link, reverse_chain)
725
- new_dep_chain = DependencyChain(
1212
+ result_reverse_chain = self.reverse_expr_chain(
1213
+ expr_chain.sub_link, reverse_chain)
1214
+ new_expr_chain = ExpressionsChain(
726
1215
  type = 'transitive',
727
- next_link = result_reverse_chain)
728
- return new_dep_chain
1216
+ sub_link = result_reverse_chain
1217
+ )
1218
+ return new_expr_chain
729
1219
 
730
1220
  case 'field':
731
- association = dep_chain.association
1221
+ association = expr_chain.association
732
1222
 
733
1223
  if not association:
734
1224
  raise LanguageGraphException(
735
- "Missing association for dep chain"
1225
+ "Missing association for expressions chain"
1226
+ )
1227
+
1228
+ if not expr_chain.fieldname:
1229
+ raise LanguageGraphException(
1230
+ "Missing field name for expressions chain"
736
1231
  )
737
1232
 
738
1233
  opposite_fieldname = association.get_opposite_fieldname(
739
- dep_chain.fieldname)
740
- new_dep_chain = DependencyChain(
1234
+ expr_chain.fieldname)
1235
+ new_expr_chain = ExpressionsChain(
741
1236
  type = 'field',
742
- next_link = reverse_chain
1237
+ association = association,
1238
+ fieldname = opposite_fieldname
743
1239
  )
744
- new_dep_chain.fieldname = opposite_fieldname
745
- new_dep_chain.association = association
746
- return self.reverse_dep_chain(
747
- dep_chain.next_link,
748
- new_dep_chain
749
- )
1240
+ return new_expr_chain
750
1241
 
751
1242
  case 'subType':
752
- result_reverse_chain = self.reverse_dep_chain(
753
- dep_chain.next_link,
1243
+ result_reverse_chain = self.reverse_expr_chain(
1244
+ expr_chain.sub_link,
754
1245
  reverse_chain
755
1246
  )
756
- new_dep_chain = DependencyChain(
1247
+ new_expr_chain = ExpressionsChain(
757
1248
  type = 'subType',
758
- next_link = result_reverse_chain
1249
+ sub_link = result_reverse_chain,
1250
+ subtype = expr_chain.subtype
759
1251
  )
760
- new_dep_chain.subtype = dep_chain.subtype
761
- return new_dep_chain
762
- # return reverse_chain
1252
+ return new_expr_chain
763
1253
 
764
1254
  case _:
765
1255
  msg = 'Unknown assoc chain element "%s"'
766
- logger.error(msg, dep_chain.type)
767
- raise LanguageGraphAssociationError(msg % dep_chain.type)
1256
+ logger.error(msg, expr_chain.type)
1257
+ raise LanguageGraphAssociationError(msg % expr_chain.type)
1258
+
1259
+ def _resolve_variable(self, asset, var_name) -> tuple:
1260
+ """
1261
+ Resolve a variable for a specific asset by variable name.
1262
+
1263
+ Arguments:
1264
+ asset - a language graph asset to which the variable belongs
1265
+ var_name - a string representing the variable name
1266
+
1267
+ Return:
1268
+ A tuple containing the target asset and expressions chain required to
1269
+ reach it.
1270
+ """
1271
+ if var_name not in asset.variables:
1272
+ var_expr = self._get_var_expr_for_asset(asset.name, var_name)
1273
+ target_asset, expr_chain, _ = self.process_step_expression(
1274
+ asset,
1275
+ None,
1276
+ var_expr
1277
+ )
1278
+ asset.own_variables[var_name] = (target_asset, expr_chain)
1279
+ return (target_asset, expr_chain)
1280
+ return asset.variables[var_name]
1281
+
768
1282
 
769
1283
  def _generate_graph(self) -> None:
770
1284
  """
@@ -772,41 +1286,40 @@ class LanguageGraph():
772
1286
  given in the constructor.
773
1287
  """
774
1288
  # Generate all of the asset nodes of the language graph.
775
- for asset in self._lang_spec['assets']:
1289
+ for asset_dict in self._lang_spec['assets']:
776
1290
  logger.debug(
777
1291
  'Create asset language graph nodes for asset %s',
778
- asset["name"]
1292
+ asset_dict['name']
779
1293
  )
780
1294
  asset_node = LanguageGraphAsset(
781
- name = asset['name'],
782
- associations = [],
783
- attack_steps = [],
784
- description = asset['meta'],
785
- super_assets = [],
786
- sub_assets = [],
787
- is_abstract = asset['isAbstract']
1295
+ name = asset_dict['name'],
1296
+ own_associations = {},
1297
+ attack_steps = {},
1298
+ info = asset_dict['meta'],
1299
+ own_super_asset = None,
1300
+ own_sub_assets = set(),
1301
+ own_variables = {},
1302
+ is_abstract = asset_dict['isAbstract']
788
1303
  )
789
- self.assets.append(asset_node)
1304
+ self.assets[asset_dict['name']] = asset_node
790
1305
 
791
1306
  # Link assets based on inheritance
792
- for asset_info in self._lang_spec['assets']:
793
- asset = next((asset for asset in self.assets \
794
- if asset.name == asset_info['name']), None)
795
- if asset_info['superAsset']:
796
- super_asset = next((asset for asset in self.assets \
797
- if asset.name == asset_info['superAsset']), None)
1307
+ for asset_dict in self._lang_spec['assets']:
1308
+ asset = self.assets[asset_dict['name']]
1309
+ if asset_dict['superAsset']:
1310
+ super_asset = self.assets[asset_dict['superAsset']]
798
1311
  if not super_asset:
799
1312
  msg = 'Failed to find super asset "%s" for asset "%s"!'
800
1313
  logger.error(
801
- msg, asset_info["superAsset"], asset_info["name"])
1314
+ msg, asset_dict["superAsset"], asset_dict["name"])
802
1315
  raise LanguageGraphSuperAssetNotFoundError(
803
- msg % (asset_info["superAsset"], asset_info["name"]))
1316
+ msg % (asset_dict["superAsset"], asset_dict["name"]))
804
1317
 
805
- super_asset.sub_assets.append(asset)
806
- asset.super_assets.append(super_asset)
1318
+ super_asset.own_sub_assets.add(asset)
1319
+ asset.own_super_asset = super_asset
807
1320
 
808
1321
  # Generate all of the association nodes of the language graph.
809
- for asset in self.assets:
1322
+ for asset in self.assets.values():
810
1323
  logger.debug(
811
1324
  'Create association language graph nodes for asset %s',
812
1325
  asset.name
@@ -814,8 +1327,7 @@ class LanguageGraph():
814
1327
 
815
1328
  associations = self._get_associations_for_asset_type(asset.name)
816
1329
  for association in associations:
817
- left_asset = next((asset for asset in self.assets \
818
- if asset.name == association['leftAsset']), None)
1330
+ left_asset = self.assets[association['leftAsset']]
819
1331
  if not left_asset:
820
1332
  msg = 'Left asset "%s" for association "%s" not found!'
821
1333
  logger.error(
@@ -823,8 +1335,7 @@ class LanguageGraph():
823
1335
  raise LanguageGraphAssociationError(
824
1336
  msg % (association["leftAsset"], association["name"]))
825
1337
 
826
- right_asset = next((asset for asset in self.assets \
827
- if asset.name == association['rightAsset']), None)
1338
+ right_asset = self.assets[association['rightAsset']]
828
1339
  if not right_asset:
829
1340
  msg = 'Right asset "%s" for association "%s" not found!'
830
1341
  logger.error(
@@ -833,17 +1344,6 @@ class LanguageGraph():
833
1344
  msg % (association["rightAsset"], association["name"])
834
1345
  )
835
1346
 
836
- # Technically we should be more exhaustive and check the
837
- # flipped version too and all of the fieldnames as well.
838
- assoc_node = next((assoc for assoc in self.associations \
839
- if assoc.name == association['name'] and
840
- assoc.left_field.asset == left_asset and
841
- assoc.right_field.asset == right_asset),
842
- None)
843
- if assoc_node:
844
- # The association was already created, skip it
845
- continue
846
-
847
1347
  assoc_node = LanguageGraphAssociation(
848
1348
  name = association['name'],
849
1349
  left_field = LanguageGraphAssociationField(
@@ -856,124 +1356,214 @@ class LanguageGraph():
856
1356
  association['rightField'],
857
1357
  association['rightMultiplicity']['min'],
858
1358
  association['rightMultiplicity']['max']),
859
- description = association['meta']
1359
+ info = association['meta']
860
1360
  )
861
1361
 
862
- # Add the association to the left and right asset and all of
863
- # the assets that inherit them
1362
+ # Add the association to the left and right asset
864
1363
  associated_assets = [left_asset, right_asset]
865
1364
  while associated_assets != []:
866
1365
  asset = associated_assets.pop()
867
- associated_assets.extend(asset.sub_assets)
868
- if assoc_node not in asset.associations:
869
- asset.associations.append(assoc_node)
1366
+ if assoc_node.full_name not in asset.own_associations:
1367
+ asset.own_associations[assoc_node.full_name] = assoc_node
1368
+
1369
+ # Set the variables
1370
+ for asset in self.assets.values():
1371
+ for variable in self._get_variables_for_asset_type(asset.name):
1372
+ if logger.isEnabledFor(logging.DEBUG):
1373
+ # Avoid running json.dumps when not in debug
1374
+ logger.debug(
1375
+ 'Processing Variable Expression:\n%s',
1376
+ json.dumps(variable, indent = 2)
1377
+ )
1378
+ self._resolve_variable(asset, variable['name'])
870
1379
 
871
- self.associations.append(assoc_node)
872
1380
 
873
1381
  # Generate all of the attack step nodes of the language graph.
874
- for asset in self.assets:
1382
+ for asset in self.assets.values():
875
1383
  logger.debug(
876
1384
  'Create attack steps language graph nodes for asset %s',
877
1385
  asset.name
878
1386
  )
879
1387
  attack_steps = self._get_attacks_for_asset_type(asset.name)
880
- for attack_step_name, attack_step_attribs in attack_steps.items():
1388
+ for attack_step_attribs in attack_steps.values():
881
1389
  logger.debug(
882
1390
  'Create attack step language graph nodes for %s',
883
- attack_step_name
1391
+ attack_step_attribs['name']
884
1392
  )
885
1393
 
886
1394
  attack_step_node = LanguageGraphAttackStep(
887
- name = attack_step_name,
1395
+ name = attack_step_attribs['name'],
888
1396
  type = attack_step_attribs['type'],
889
1397
  asset = asset,
890
1398
  ttc = attack_step_attribs['ttc'],
1399
+ overrides = attack_step_attribs['reaches']['overrides'] \
1400
+ if attack_step_attribs['reaches'] else False,
891
1401
  children = {},
892
1402
  parents = {},
893
- description = attack_step_attribs['meta']
1403
+ info = attack_step_attribs['meta'],
1404
+ tags = set(attack_step_attribs['tags'])
1405
+ )
1406
+ attack_step_node._attributes = attack_step_attribs
1407
+ asset.attack_steps[attack_step_attribs['name']] = \
1408
+ attack_step_node
1409
+
1410
+ # Create the inherited attack steps
1411
+ assets = list(self.assets.values())
1412
+ while len(assets) > 0:
1413
+ asset = assets.pop(0)
1414
+ if asset.own_super_asset in assets:
1415
+ # The asset still has super assets that should be resolved
1416
+ # first, moved it to the back.
1417
+ assets.append(asset)
1418
+ else:
1419
+ if asset.own_super_asset:
1420
+ for attack_step in \
1421
+ asset.own_super_asset.attack_steps.values():
1422
+ if attack_step.name not in asset.attack_steps:
1423
+ attack_step_node = LanguageGraphAttackStep(
1424
+ name = attack_step.name,
1425
+ type = attack_step.type,
1426
+ asset = asset,
1427
+ ttc = attack_step.ttc,
1428
+ overrides = False,
1429
+ children = {},
1430
+ parents = {},
1431
+ info = attack_step.info,
1432
+ tags = set(attack_step.tags)
1433
+ )
1434
+ attack_step_node.inherits = attack_step
1435
+ asset.attack_steps[attack_step.name] = attack_step_node
1436
+ elif asset.attack_steps[attack_step.name].overrides:
1437
+ # The inherited attack step was already overridden.
1438
+ continue
1439
+ else:
1440
+ asset.attack_steps[attack_step.name].inherits = \
1441
+ attack_step
1442
+ asset.attack_steps[attack_step.name].tags |= \
1443
+ attack_step.tags
1444
+ asset.attack_steps[attack_step.name].info |= \
1445
+ attack_step.info
1446
+
1447
+ # Then, link all of the attack step nodes according to their
1448
+ # associations.
1449
+ for asset in self.assets.values():
1450
+ for attack_step in asset.attack_steps.values():
1451
+ logger.debug(
1452
+ 'Determining children for attack step %s',
1453
+ attack_step.name
894
1454
  )
895
- attack_step_node.attributes = attack_step_attribs
896
- asset.attack_steps.append(attack_step_node)
897
- self.attack_steps.append(attack_step_node)
898
1455
 
899
- # Then, link all of the attack step nodes according to their associations.
900
- for attack_step in self.attack_steps:
901
- logger.debug(
902
- 'Determining children for attack step %s',
903
- attack_step.name
904
- )
905
- step_expressions = \
906
- attack_step.attributes['reaches']['stepExpressions'] if \
907
- attack_step.attributes['reaches'] else []
908
-
909
- for step_expression in step_expressions:
910
- # Resolve each of the attack step expressions listed for this
911
- # attack step to determine children.
912
- (target_asset, dep_chain, attack_step_name) = \
913
- self.process_step_expression(self._lang_spec,
914
- attack_step.asset,
915
- None,
916
- step_expression)
917
- if not target_asset:
918
- msg = 'Failed to find target asset to link with for ' \
919
- 'step expression:\n%s'
920
- raise LanguageGraphStepExpressionError(
921
- msg % json.dumps(step_expression, indent = 2)
922
- )
1456
+ if attack_step._attributes is None:
1457
+ # This is simply an empty inherited attack step
1458
+ continue
923
1459
 
924
- target_attack_step = next((attack_step \
925
- for attack_step in target_asset.attack_steps \
926
- if attack_step.name == attack_step_name), None)
927
-
928
- if not target_attack_step:
929
- msg = 'Failed to find target attack step %s on %s to ' \
930
- 'link with for step expression:\n%s'
931
- raise LanguageGraphStepExpressionError(
932
- msg % (
933
- attack_step_name,
934
- target_asset.name,
935
- json.dumps(step_expression, indent = 2)
1460
+ step_expressions = \
1461
+ attack_step._attributes['reaches']['stepExpressions'] if \
1462
+ attack_step._attributes['reaches'] else []
1463
+
1464
+ for step_expression in step_expressions:
1465
+ # Resolve each of the attack step expressions listed for
1466
+ # this attack step to determine children.
1467
+ (target_asset, expr_chain, target_attack_step_name) = \
1468
+ self.process_step_expression(
1469
+ attack_step.asset,
1470
+ None,
1471
+ step_expression
1472
+ )
1473
+ if not target_asset:
1474
+ msg = 'Failed to find target asset to link with for ' \
1475
+ 'step expression:\n%s'
1476
+ raise LanguageGraphStepExpressionError(
1477
+ msg % json.dumps(step_expression, indent = 2)
936
1478
  )
937
- )
938
1479
 
939
- # It is easier to create the parent associations chain due to
940
- # the left-hand first progression.
941
- if attack_step.name in target_attack_step.parents:
942
- target_attack_step.parents[attack_step.name].append(
943
- (attack_step, dep_chain))
944
- else:
945
- target_attack_step.parents[attack_step.name] = \
946
- [(attack_step, dep_chain)]
947
- # Reverse the parent associations chain to get the child
948
- # associations chain.
949
- if target_attack_step.name in attack_step.children:
950
- attack_step.children[target_attack_step.name].append(
951
- (target_attack_step,
952
- self.reverse_dep_chain(dep_chain,
953
- None)))
954
- else:
955
- attack_step.children[target_attack_step.name] = \
956
- [(target_attack_step,
957
- self.reverse_dep_chain(dep_chain,
958
- None))]
1480
+ target_asset_attack_steps = target_asset.attack_steps
1481
+ if target_attack_step_name not in \
1482
+ target_asset_attack_steps:
1483
+ msg = 'Failed to find target attack step %s on %s to ' \
1484
+ 'link with for step expression:\n%s'
1485
+ raise LanguageGraphStepExpressionError(
1486
+ msg % (
1487
+ target_attack_step_name,
1488
+ target_asset.name,
1489
+ json.dumps(step_expression, indent = 2)
1490
+ )
1491
+ )
1492
+
1493
+ target_attack_step = target_asset_attack_steps[
1494
+ target_attack_step_name]
1495
+
1496
+ # Link to the children target attack steps
1497
+ if target_attack_step.full_name in attack_step.children:
1498
+ attack_step.children[target_attack_step.full_name].\
1499
+ append((target_attack_step, expr_chain))
1500
+ else:
1501
+ attack_step.children[target_attack_step.full_name] = \
1502
+ [(target_attack_step, expr_chain)]
1503
+ # Reverse the children associations chains to get the
1504
+ # parents associations chain.
1505
+ if attack_step.full_name in target_attack_step.parents:
1506
+ target_attack_step.parents[attack_step.full_name].\
1507
+ append((attack_step,
1508
+ self.reverse_expr_chain(expr_chain,
1509
+ None)))
1510
+ else:
1511
+ target_attack_step.parents[attack_step.full_name] = \
1512
+ [(attack_step,
1513
+ self.reverse_expr_chain(expr_chain,
1514
+ None))]
1515
+
1516
+ # Evaluate the requirements of exist and notExist attack steps
1517
+ if attack_step.type == 'exist' or \
1518
+ attack_step.type == 'notExist':
1519
+ step_expressions = \
1520
+ attack_step._attributes['requires']['stepExpressions'] \
1521
+ if attack_step._attributes['requires'] else []
1522
+ if not step_expressions:
1523
+ msg = 'Failed to find requirements for attack step' \
1524
+ ' "%s" of type "%s":\n%s'
1525
+ raise LanguageGraphStepExpressionError(
1526
+ msg % (
1527
+ attack_step.name,
1528
+ attack_step.type,
1529
+ json.dumps(attack_step._attributes, indent = 2)
1530
+ )
1531
+ )
1532
+
1533
+ attack_step.own_requires = []
1534
+ for step_expression in step_expressions:
1535
+ _, \
1536
+ result_expr_chain, \
1537
+ _ = \
1538
+ self.process_step_expression(
1539
+ attack_step.asset,
1540
+ None,
1541
+ step_expression
1542
+ )
1543
+ attack_step.own_requires.append(
1544
+ self.reverse_expr_chain(result_expr_chain, None))
959
1545
 
960
1546
  def _get_attacks_for_asset_type(self, asset_type: str) -> dict:
961
1547
  """
962
1548
  Get all Attack Steps for a specific Class
963
1549
 
964
1550
  Arguments:
965
- asset_type - a string representing the class for which we want to list
966
- the possible attack steps
1551
+ asset_type - a string representing the class for which we want to
1552
+ list the possible attack steps
967
1553
 
968
1554
  Return:
969
- A dictionary representing the set of possible attacks for the specified
970
- class. Each key in the dictionary is an attack name and is associated
971
- with a dictionary containing other characteristics of the attack such as
972
- type of attack, TTC distribution, child attack steps and other information
1555
+ A dictionary representing the set of possible attacks for the
1556
+ specified class. Each key in the dictionary is an attack name and is
1557
+ associated with a dictionary containing other characteristics of the
1558
+ attack such as type of attack, TTC distribution, child attack steps
1559
+ and other information
973
1560
  """
974
1561
  attack_steps: dict = {}
975
1562
  try:
976
- asset = next((asset for asset in self._lang_spec['assets'] if asset['name'] == asset_type))
1563
+ asset = next(
1564
+ asset for asset in self._lang_spec['assets'] \
1565
+ if asset['name'] == asset_type
1566
+ )
977
1567
  except StopIteration:
978
1568
  logger.error(
979
1569
  'Failed to find asset type %s when looking'
@@ -983,33 +1573,10 @@ class LanguageGraph():
983
1573
 
984
1574
  logger.debug(
985
1575
  'Get attack steps for %s asset from '
986
- 'language specification.', asset["name"]
1576
+ 'language specification.', asset['name']
987
1577
  )
988
- if asset['superAsset']:
989
- logger.debug('Asset extends another one, fetch the superclass '\
990
- 'attack steps for it.')
991
- attack_steps = self._get_attacks_for_asset_type(asset['superAsset'])
992
-
993
- for step in asset['attackSteps']:
994
- if step['name'] not in attack_steps:
995
- attack_steps[step['name']] = copy.deepcopy(step)
996
- elif not step['reaches']:
997
- # This attack step does not lead to any attack steps
998
- continue
999
- elif step['reaches']['overrides'] == True:
1000
- attack_steps[step['name']] = copy.deepcopy(step)
1001
- else:
1002
- if attack_steps[step['name']]['reaches'] is not None and \
1003
- 'stepExpressions' in \
1004
- attack_steps[step['name']]['reaches']:
1005
- attack_steps[step['name']]['reaches']['stepExpressions'].\
1006
- extend(step['reaches']['stepExpressions'])
1007
- else:
1008
- attack_steps[step['name']]['reaches'] = {
1009
- 'overrides': False,
1010
- 'stepExpressions': step['reaches']['stepExpressions']
1011
- }
1012
1578
 
1579
+ attack_steps = {step['name']: step for step in asset['attackSteps']}
1013
1580
 
1014
1581
  return attack_steps
1015
1582
 
@@ -1018,14 +1585,15 @@ class LanguageGraph():
1018
1585
  Get all Associations for a specific Class
1019
1586
 
1020
1587
  Arguments:
1021
- asset_type - a string representing the class for which we want to list
1022
- the associations
1588
+ asset_type - a string representing the class for which we want to
1589
+ list the associations
1023
1590
 
1024
1591
  Return:
1025
1592
  A dictionary representing the set of associations for the specified
1026
1593
  class. Each key in the dictionary is an attack name and is associated
1027
- with a dictionary containing other characteristics of the attack such as
1028
- type of attack, TTC distribution, child attack steps and other information
1594
+ with a dictionary containing other characteristics of the attack such
1595
+ as type of attack, TTC distribution, child attack steps and other
1596
+ information
1029
1597
  """
1030
1598
  logger.debug(
1031
1599
  'Get associations for %s asset from '
@@ -1033,8 +1601,8 @@ class LanguageGraph():
1033
1601
  )
1034
1602
  associations: list = []
1035
1603
 
1036
- asset = next((asset for asset in self._lang_spec['assets'] if asset['name'] == \
1037
- asset_type), None)
1604
+ asset = next((asset for asset in self._lang_spec['assets'] \
1605
+ if asset['name'] == asset_type), None)
1038
1606
  if not asset:
1039
1607
  logger.error(
1040
1608
  'Failed to find asset type %s when '
@@ -1042,10 +1610,6 @@ class LanguageGraph():
1042
1610
  )
1043
1611
  return associations
1044
1612
 
1045
- if asset['superAsset']:
1046
- logger.debug('Asset extends another one, fetch the superclass '\
1047
- 'associations for it.')
1048
- associations.extend(self._get_associations_for_asset_type(asset['superAsset']))
1049
1613
  assoc_iter = (assoc for assoc in self._lang_spec['associations'] \
1050
1614
  if assoc['leftAsset'] == asset_type or \
1051
1615
  assoc['rightAsset'] == asset_type)
@@ -1056,43 +1620,56 @@ class LanguageGraph():
1056
1620
 
1057
1621
  return associations
1058
1622
 
1059
- def _get_variable_for_asset_type_by_name(
1060
- self, asset_type: str, variable_name: str) -> dict:
1623
+ def _get_variables_for_asset_type(
1624
+ self, asset_type: str) -> dict:
1061
1625
  """
1062
1626
  Get a variables for a specific asset type by name.
1063
- NOTE: Variables are the ones specified in MAL through `let` statements
1627
+ Note: Variables are the ones specified in MAL through `let` statements
1064
1628
 
1065
1629
  Arguments:
1066
1630
  asset_type - a string representing the type of asset which
1067
- contains the variable
1068
- variable_name - the name of the variable to search for
1631
+ contains the variables
1069
1632
 
1070
1633
  Return:
1071
- A dictionary representing the step expressions for the specified variable.
1634
+ A dictionary representing the step expressions for the variables
1635
+ belonging to the asset.
1072
1636
  """
1073
1637
 
1074
- asset = next((asset for asset in self._lang_spec['assets'] if asset['name'] == \
1075
- asset_type), None)
1076
- if not asset:
1077
- msg = 'Failed to find asset type %s when looking for variable.'
1638
+ asset_dict = next((asset for asset in self._lang_spec['assets'] \
1639
+ if asset['name'] == asset_type), None)
1640
+ if not asset_dict:
1641
+ msg = 'Failed to find asset type %s in language specification '\
1642
+ 'when looking for variables.'
1078
1643
  logger.error(msg, asset_type)
1079
1644
  raise LanguageGraphException(msg % asset_type)
1080
1645
 
1081
- variable_dict = next((variable for variable in \
1082
- asset['variables'] if variable['name'] == variable_name), None)
1083
- if not variable_dict:
1084
- if asset['superAsset']:
1085
- variable_dict = self._get_variable_for_asset_type_by_name(asset['superAsset'],
1086
- variable_name)
1087
- if variable_dict:
1088
- return variable_dict
1089
- else:
1090
- msg = 'Failed to find variable %s in %s lang specification.'
1091
- logger.error(msg, variable_name, asset_type)
1092
- raise LanguageGraphException(
1093
- msg % (variable_name, asset_type))
1646
+ return asset_dict['variables']
1094
1647
 
1095
- return variable_dict['stepExpression']
1648
+ def _get_var_expr_for_asset(
1649
+ self, asset_type: str, var_name) -> dict:
1650
+ """
1651
+ Get a variable for a specific asset type by variable name.
1652
+
1653
+ Arguments:
1654
+ asset_type - a string representing the type of asset which
1655
+ contains the variable
1656
+ var_name - a string representing the variable name
1657
+
1658
+ Return:
1659
+ A dictionary representing the step expression for the variable.
1660
+ """
1661
+
1662
+ vars_dict = self._get_variables_for_asset_type(asset_type)
1663
+
1664
+ var_expr = next((var_entry['stepExpression'] for var_entry \
1665
+ in vars_dict if var_entry['name'] == var_name), None)
1666
+
1667
+ if not var_expr:
1668
+ msg = 'Failed to find variable name "%s" in language '\
1669
+ 'specification when looking for variables for "%s" asset.'
1670
+ logger.error(msg, var_name, asset_type)
1671
+ raise LanguageGraphException(msg % (var_name, asset_type))
1672
+ return var_expr
1096
1673
 
1097
1674
  def regenerate_graph(self) -> None:
1098
1675
  """
@@ -1100,30 +1677,9 @@ class LanguageGraph():
1100
1677
  given in the constructor.
1101
1678
  """
1102
1679
 
1103
- self.assets = []
1104
- self.associations = []
1105
- self.attack_steps = []
1680
+ self.assets = {}
1106
1681
  self._generate_graph()
1107
1682
 
1108
- def get_asset_by_name(
1109
- self,
1110
- asset_name
1111
- ) -> Optional[LanguageGraphAsset]:
1112
- """
1113
- Get an asset based on its name
1114
-
1115
- Arguments:
1116
- asset_name - a string containing the asset name
1117
-
1118
- Return:
1119
- The asset matching the name.
1120
- None if there is no match.
1121
- """
1122
- for asset in self.assets:
1123
- if asset.name == asset_name:
1124
- return asset
1125
-
1126
- return None
1127
1683
 
1128
1684
  def get_association_by_fields_and_assets(
1129
1685
  self,
@@ -1145,30 +1701,21 @@ class LanguageGraph():
1145
1701
  The association matching the fieldnames and asset types.
1146
1702
  None if there is no match.
1147
1703
  """
1148
- first_asset = self.get_asset_by_name(first_asset_name)
1149
- if first_asset is None:
1150
- raise LookupError(
1151
- f'Failed to find asset with name \"{first_asset_name}\" in '
1152
- 'the language graph.'
1153
- )
1154
-
1155
- second_asset = self.get_asset_by_name(second_asset_name)
1156
- if second_asset is None:
1157
- raise LookupError(
1158
- f'Failed to find asset with name \"{second_asset_name}\" in '
1159
- 'the language graph.'
1160
- )
1704
+ first_asset = self.assets[first_asset_name]
1705
+ second_asset = self.assets[second_asset_name]
1161
1706
 
1162
- for assoc in self.associations:
1707
+ for assoc in first_asset.associations.values():
1163
1708
  logger.debug(
1164
- 'Compare ("%s", "%s", "%s", "%s") to ("%s", "%s", "%s", "%s").',
1709
+ 'Compare ("%s", "%s", "%s", "%s") to '
1710
+ '("%s", "%s", "%s", "%s").',
1165
1711
  first_asset_name, first_field,
1166
1712
  second_asset_name, second_field,
1167
1713
  assoc.left_field.asset.name, assoc.left_field.fieldname,
1168
1714
  assoc.right_field.asset.name, assoc.right_field.fieldname
1169
1715
  )
1170
1716
 
1171
- # If the asset and fields match either way we accept it as a match.
1717
+ # If the asset and fields match either way we accept it as a
1718
+ # match.
1172
1719
  if assoc.left_field.fieldname == first_field and \
1173
1720
  assoc.right_field.fieldname == second_field and \
1174
1721
  first_asset.is_subasset_of(assoc.left_field.asset) and \