Skip to content

Add LookupDict for faster and simpler lookups #105

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
andrewbwm wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Add LookupDict for faster and simpler lookups #105

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
42af1b6
Add and use LookupDict
Jan 16, 2025
6593af6
Merge file_utils into utils
nkakouros Feb 17, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 2 additions & 3 deletions maltoolbox/attackgraph/analyzers/apriori.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -178,13 +178,12 @@ def prune_unviable_and_unnecessary_nodes(graph: AttackGraph) -> None:
graph - the attack graph for which we wish to remove the
the nodes which are not viable or necessary.
"""
logger.debug(
'Prune unviable and unnecessary nodes from the attack graph.')
logger.debug('Prune unviable and unnecessary nodes from the attack graph.')

nodes_to_remove = set()
for node in graph.nodes.values():
if node.type in ('or', 'and') and \
(not node.is_viable or not node.is_necessary):
not node.is_viable or not node.is_necessary:
nodes_to_remove.add(node)

# Do the removal separatly so we don't remove
Expand Down
79 changes: 23 additions & 56 deletions maltoolbox/attackgraph/attackgraph.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,12 +2,12 @@
MAL-Toolbox Attack Graph Module
"""
from __future__ import annotations
import copy
import logging
import json
import sys
import zipfile

from copy import deepcopy
from itertools import chain
from typing import TYPE_CHECKING

Expand All @@ -20,10 +20,11 @@
from ..model import Model
from ..language import (LanguageGraph, ExpressionsChain,
LanguageGraphAttackStep, disaggregate_attack_step_full_name)
from ..file_utils import (
from ..utils import (
load_dict_from_json_file,
load_dict_from_yaml_file,
save_dict_to_file
save_dict_to_file,
LookupDict,
)


Expand Down Expand Up @@ -82,11 +83,10 @@ def create_attack_graph(
class AttackGraph():
"""Graph representation of attack steps"""
def __init__(self, lang_graph, model: Optional[Model] = None):
self.nodes: dict[int, AttackGraphNode] = {}
self.nodes: LookupDict[int, AttackGraphNode] = LookupDict()
self.attackers: dict[int, Attacker] = {}
# Dictionaries used in optimization to get nodes and attackers by id
# or full name faster
self._full_name_to_node: dict[str, AttackGraphNode] = {}

self.model = model
self.lang_graph = lang_graph
Expand Down Expand Up @@ -123,38 +123,24 @@ def __deepcopy__(self, memo):
return memo[id(self)]

copied_attackgraph = AttackGraph(self.lang_graph)
copied_attackgraph.model = self.model

copied_attackgraph.nodes = {}
memo[id(self)] = copied_attackgraph

# Deep copy nodes
for node_id, node in self.nodes.items():
copied_node = copy.deepcopy(node, memo)
copied_attackgraph.nodes[node_id] = copied_node
copied_attackgraph.model = self.model
copied_attackgraph.nodes = deepcopy(self.nodes, memo)
copied_attackgraph.attackers = deepcopy(self.attackers, memo)

# Re-link node references
for node in self.nodes.values():
nid = id(node)
if node.parents:
memo[id(node)].parents = copy.deepcopy(node.parents, memo)
memo[nid].parents = deepcopy(node.parents, memo)
if node.children:
memo[id(node)].children = copy.deepcopy(node.children, memo)

# Deep copy attackers
for attacker_id, attacker in self.attackers.items():
copied_attacker = copy.deepcopy(attacker, memo)
copied_attackgraph.attackers[attacker_id] = copied_attacker

# Re-link attacker references
for node in self.nodes.values():
memo[nid].children = deepcopy(node.children, memo)
if node.compromised_by:
Comment thread
nkakouros marked this conversation as resolved.
memo[id(node)].compromised_by = copy.deepcopy(
node.compromised_by, memo)
memo[nid].compromised_by = deepcopy(node.compromised_by, memo)

# Copy lookup dicts
copied_attackgraph._full_name_to_node = \
copy.deepcopy(self._full_name_to_node, memo)
copied_attackgraph.attackers = deepcopy(self.attackers, memo)

# Copy counters
copied_attackgraph.next_node_id = self.next_node_id
copied_attackgraph.next_attacker_id = self.next_attacker_id

Expand Down Expand Up @@ -189,8 +175,9 @@ def _from_dict(
# Recreate asset links if model is available.
node_asset = None
if model and 'asset' in node_dict:
node_asset = model.get_asset_by_name(node_dict['asset'])
if node_asset is None:
try:
node_asset = model.assets.fetch("name", node_dict['asset'])
except KeyError:
msg = ('Failed to find asset with name "%s"'
' when loading from attack graph dict')
logger.error(msg, node_dict["asset"])
Expand Down Expand Up @@ -293,21 +280,6 @@ def load_from_file(
return cls._from_dict(serialized_attack_graph,
lang_graph, model = model)

def get_node_by_full_name(self, full_name: str) -> Optional[AttackGraphNode]:

@mrkickling mrkickling Feb 13, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I personally liked the methods get_node_by_id and get_node_by_full_name. They made it clear what attributes you could fetch a node by. Now the user would do attack_graph.nodes.fetch(attr, value) without knowing which attributes are supported. Or are all attributes of nodes supported? If so, could we still have the get_node_by_x and use the LookupDict in those? Or at least have AttackGraph.get_node_by_attribute_value(attr: str, value: Any) so the user does not have to go through attack_graph.nodes.fetch() ?

@nkakouros nkakouros Feb 13, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

All valid attributes are supported. Removing those methods is what this PR does. Having them back and then going through the lookup dict undoes this.

I don't like the get methods as:

  • they are too java-ish
  • they take up lines of code for no reason
  • whenever you need to get an item by another attribute that is not yet implemented you have to add new methods. With the lookup dict you can use whatever attribute you like whenever you like.
  • they don't support comparisons other than equals
  • each one needs its own unit test
  • they make a bug more probable
  • they make refactorings and optimizations harder as changes need to be made all over the code
  • they don't support plural as in fetch all that match
  • they require to bloat the class with as many private attributes.

As long as the user understands what fetch() or lookup() do, they won't have difficulty in reading the code. They are also cognitively more light as you don't have to read and read long lines, long method names etc. You read nodes.fetch(key, value) and you understand the same thing as "get node by key" but via a single consistent API that works across any object and that you read and learn once as compared to reading 10 different methods per class.

@mrkickling mrkickling Feb 19, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am ok with removing the get_by_id and get_by_full_name, but we should still have one method in attack graph which uses the LookupDict so users don't have to interact with it directly.

As long as the user understands what fetch() or lookup() do, they won't have difficulty in reading the code

Perhaps, but I think having a method that wraps around nodes.fetch() called something like attack_graph.get_nodes_by_attribute_value(attribute: str, value: str) would be easier to find/understand. And from the AttackGraphs perspective, that is what the LookupDict.fetch does.

Just to illustrate my thinking:

From the users perspective if we have no method in AttackGraph:

  • They have an AttackGraph and want to fetch a node with a certain full_name
  • They might look for a method in the attack graph that does this
  • They will not find it, so they will think that they need to implement something themselves
  • They look at attackgraph.nodes and realize it is a LookupDict
  • They try to understand how the LookupDict works and find the .fetch()
  • They use the .fetch('full_name', x)

If we instead had a method attack_graph.get_nodes_by_attribute_value(attribute: str, value: str) in the attack_graph (which uses the LookupDict.fetch), they would find it more easily.

I am skeptical of making the API less explicit by having common methods in instance variables (.nodes).

@nkakouros nkakouros Feb 19, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is where documentation helps, with examples how you do things. And IDE completion is also useful here, it will show a fetch or lookup method with full docstrings and type hints.

I do not favor adding such a method. This code:

get_node_by_attribute_value(attribute, value)

I think does not add much compared to:

nodes.fetch(attribute, value)

especially once you have seen, read and used it once.

nodes is LookupDict and it supports an interface. If it were a dict it would support a different one (a subset). Do you think of adding wrappers around the dict.get() method or the dict.remove() method in that case? If not, why do it now? It is a custom API that not everyone is familiar with but so is the rest of the codebase and people have to study it first. It's not some hard piece of code, just that you use nodes.fetch to get by an attribute.

Also, adding a get_ method would provide users with two ways to do the same thing which is not ideal.

A further motivation for doing this PR was that this LookupDict will be used elsewhere too, not just for nodes. Nodes is just the first application of this. For example, in model we have get_asset_by_id and get_asset_by_id. These can be removed both just by using a LookupDict. The interface will be one the users already now.

@mrkickling mrkickling Feb 19, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, I would add wrappers around the dicts if they were dicts. That is how it worked before with get_node_by_id and get_node_by_full_name works like this I believe (the dicts were private variables of the AttackGraph).

I agree the LookupDict could be used in more places. But I mean that in the AttackGraph it has only one use, which is to fetch nodes by attribute values / key values, which is the context the method (name) would add.

But I don't seem to be able to convince you, so maybe I will just back off.

@nkakouros nkakouros Feb 19, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We have a different programming style. I believe the changes will increase the quality of the python code, but that may not be the case at all. As a maintainer, you have a greater say, but I would hate to write back that kind of code. Maybe you want to add these changes or maybe @andrewbwm can help break the deadlock? But many of the changes I propose are of this kind, removing things that I find verbose and not pythonic and accepting some but not others will lead to an inconsistent style in the code. My mistake is that I didn't sync with you first before embarking on this. I was driven by a desire to make the code run faster and to do that it helps to have it simplified and more "linear".

@mrkickling mrkickling Feb 20, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, that is it. I don't think it is easy to say who is right or wrong, we just prioritize different things. Yes, maybe @andrewbwm can chime in to give some third perspective!

@nkakouros nkakouros Feb 20, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@andrewbwm 🔔 or 🔕 ?

"""
Return the attack node that matches the full name provided.

Arguments:
full_name - the full name of the attack graph node we are looking
for

Return:
The attack step node that matches the given full name.
"""

logger.debug(f'Looking up node with full name "%s"', full_name)
return self._full_name_to_node.get(full_name)

def attach_attackers(self) -> None:
"""
Create attackers and their entry point nodes and attach them to the
Expand Down Expand Up @@ -340,7 +312,7 @@ def attach_attackers(self) -> None:
for (asset, attack_steps) in attacker_info.entry_points:
for attack_step in attack_steps:
full_name = asset.name + ':' + attack_step
ag_node = self.get_node_by_full_name(full_name)
ag_node = self.nodes.fetch("full_name", full_name)
if not ag_node:
logger.warning(
'Failed to find attacker entry point '
Expand Down Expand Up @@ -620,8 +592,8 @@ def _generate_graph(self) -> None:
if target_asset is not None:
target_node_full_name = target_asset.name + \
':' + target_attack_step.name
target_node = self.get_node_by_full_name(
target_node_full_name)
target_node = self.nodes.fetch("full_name",
target_node_full_name)
if target_node is None:
msg = ('Failed to find target node '
'"%s" to link with for attack '
Expand Down Expand Up @@ -665,7 +637,7 @@ def regenerate_graph(self) -> None:
the MAL language specification provided at initialization.
"""

self.nodes = {}
self.nodes = LookupDict()
self.attackers = {}
self._generate_graph()

Expand Down Expand Up @@ -713,7 +685,6 @@ def add_node(
node_id
))


node = AttackGraphNode(
node_id = node_id,
lg_attack_step = lg_attack_step,
Expand All @@ -722,8 +693,7 @@ def add_node(
existence_status = existence_status
)

self.nodes[node_id] = node
self._full_name_to_node[node.full_name] = node
self.nodes[node.id] = node

return node

Expand All @@ -740,10 +710,7 @@ def remove_node(self, node: AttackGraphNode) -> None:
for parent in node.parents:
parent.children.remove(node)

if not isinstance(node.id, int):
raise ValueError(f'Invalid node id.')
del self.nodes[node.id]
del self._full_name_to_node[node.full_name]
self.nodes.pop(node.id, None)

def add_attacker(
self,
Expand Down
2 changes: 2 additions & 0 deletions maltoolbox/attackgraph/node.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -118,6 +118,8 @@ def __deepcopy__(self, memo) -> AttackGraphNode:

return copied_node

def __hash__(self):
return hash(self.full_name)

def is_compromised(self) -> bool:
"""
Expand Down
62 changes: 0 additions & 62 deletions maltoolbox/file_utils.py
View file
Open in desktop

This file was deleted.

8 changes: 4 additions & 4 deletions maltoolbox/language/languagegraph.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,17 +12,17 @@
from functools import cached_property
from typing import Any, Optional

from maltoolbox.file_utils import (
load_dict_from_yaml_file, load_dict_from_json_file,
save_dict_to_file
)
from .compiler import MalCompiler
from ..exceptions import (
LanguageGraphAssociationError,
LanguageGraphStepExpressionError,
LanguageGraphException,
LanguageGraphSuperAssetNotFoundError
)
from ..utils import (
load_dict_from_yaml_file, load_dict_from_json_file,
save_dict_to_file
)

logger = logging.getLogger(__name__)

Expand Down
55 changes: 8 additions & 47 deletions maltoolbox/model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,12 +9,14 @@
from typing import TYPE_CHECKING
import math

from .file_utils import (
from .utils import (
load_dict_from_json_file,
load_dict_from_yaml_file,
save_dict_to_file
save_dict_to_file,
LookupDict,
)


from . import __version__
from .exceptions import ModelException

Expand Down Expand Up @@ -146,8 +148,7 @@ def __init__(
):

self.name = name
self.assets: dict[int, ModelAsset] = {}
self._name_to_asset:dict[str, ModelAsset] = {} # optimization
self.assets: LookupDict[int, ModelAsset] = LookupDict()
self.attackers: list[AttackerAttachment] = []
self.lang_graph = lang_graph
self.maltoolbox_version: str = mt_version
Expand Down Expand Up @@ -224,7 +225,6 @@ def add_asset(
'Add "%s"(%d) to model "%s".', name, asset_id, self.name
)
self.assets[asset_id] = asset
self._name_to_asset[name] = asset

return asset

Expand Down Expand Up @@ -262,7 +262,6 @@ def remove_asset(self, asset: ModelAsset) -> None:
attacker.entry_points.remove(entry_point_tuple)

del self.assets[asset.id]
del self._name_to_asset[asset.name]


def add_attacker(
Expand All @@ -288,44 +287,6 @@ def add_attacker(
self.attackers.append(attacker)


def get_asset_by_id(
self, asset_id: int
) -> Optional[ModelAsset]:
"""
Find an asset in the model based on its id.

Arguments:
asset_id - the id of the asset we are looking for

Return:
An asset matching the id if it exists in the model.
"""
logger.debug(
'Get asset with id %d from model "%s".',
asset_id, self.name
)
return self.assets.get(asset_id, None)


def get_asset_by_name(
self, asset_name: str
) -> Optional[ModelAsset]:
"""
Find an asset in the model based on its name.

Arguments:
asset_name - the name of the asset we are looking for

Return:
An asset matching the name if it exists in the model.
"""
logger.debug(
'Get asset with name "%s" from model "%s".',
asset_name, self.name
)
return self._name_to_asset.get(asset_name, None)


def get_attacker_by_id(
self, attacker_id: int
) -> Optional[AttackerAttachment]:
Expand Down Expand Up @@ -471,9 +432,9 @@ def _from_dict(
attacker.entry_points = []
for asset_name, entry_points_dict in \
attackers_info[attacker_id]['entry_points'].items():
target_asset = model.get_asset_by_id(
entry_points_dict['asset_id'])
if target_asset is None:
try:
target_asset = model.assets[entry_points_dict['asset_id']]
except KeyError:
raise LookupError(
'Asset "%s"(%d) is not part of model "%s".' % (
asset_name,
Expand Down
Loading