diff --git a/dffml/source/dataset/threat_modeling.py b/dffml/source/dataset/threat_modeling.py new file mode 100644 index 0000000000..d97f240d25 --- /dev/null +++ b/dffml/source/dataset/threat_modeling.py @@ -0,0 +1,43 @@ +import json +import pathlib + +from ...record import Record +from ..memory import MemorySource +from .base import dataset_source + + +@dataset_source("owasp.threat-dragon") +async def threat_dragon( + filepath: pathlib.Path, feature_name: str = "threat_model", +): + r""" + Examples + -------- + + .. code-block:: console + :test: + + $ dffml list records -sources threat_model=owasp.threat-dragon \ + -source-threat_model-filepath /home/pdxjohnny/Documents/python/living-threat-models/models/good.json + + >>> from dffml.noasync import load + >>> from dffml import iris_training + >>> + >>> records = list(load(iris_training.source())) + >>> print(len(records)) + 120 + >>> records[0].export() + {'key': '0', 'features': {'SepalLength': 6.4, 'SepalWidth': 2.8, 'PetalLength': 5.6, 'PetalWidth': 2.2, 'classification': 2}, 'extra': {}} + """ + contents = filepath.read_text() + threat_model_dict = json.loads(contents) + # TODO(security) Validate JSON schema + title = threat_model_dict["summary"]["title"] + yield MemorySource( + records=[ + Record( + key=title, + data={"features": {feature_name: threat_model_dict,},}, + ) + ], + ) diff --git a/setup.py b/setup.py index e9fba62a06..47c5955473 100644 --- a/setup.py +++ b/setup.py @@ -94,6 +94,7 @@ class InstallException(Exception): "dir = dffml.source.dir:DirectorySource", "dataframe = dffml.source.dataframe:DataFrameSource", "iris.training = dffml.source.dataset.iris:iris_training.source", + "owasp.threat-dragon = dffml.source.dataset.threat_modeling:threat_dragon.source", ], "dffml.port": ["json = dffml.port.json:JSON"], "dffml.service.cli": ["dev = dffml.service.dev:Develop"],