diff --git a/build_tech_assets.py b/build_tech_assets.py index 24a3772..7f50a5f 100644 --- a/build_tech_assets.py +++ b/build_tech_assets.py @@ -153,3 +153,43 @@ def build_storage_tm(name: str, asset_type: str) -> tuple: tag_list = storage_dict["tags"] return storage_asset_yaml, tag_list + + +def build_db_tm(name: str, asset_type: str) -> tuple: + db_dict = { + "name": name, + "type": asset_type.split("/")[0], + "description": "A Microsoft SQL Database.", + "size": "service", + "technology": "database", + "machine": "virtual", + "tags": [name, "azure", "azure-sql", "sql", "microsoft-sql", "database", asset_type], + } + with open("yaml-templates/technical_asset_template.yaml") as template_file: + template_str = template_file.read() + tech_asset_template = Template(template_str, autoescape=True) + db_asset_yaml = tech_asset_template.render(db_dict) + + tag_list = db_dict["tags"] + + return db_asset_yaml, tag_list + + +def build_vm_tm(name: str, asset_type: str) -> tuple: + vm_dict = { + "name": name, + "type": asset_type.split("/")[0], + "description": "An Azure virtual machine.", + "size": "system", + "technology": "web-server", + "machine": "virtual", + "tags": [name, "azure", "azure-virtual-machine", "virtual-machine", "vm", asset_type], + } + with open("yaml-templates/technical_asset_template.yaml") as template_file: + template_str = template_file.read() + tech_asset_template = Template(template_str, autoescape=True) + vm_asset_yaml = tech_asset_template.render(vm_dict) + + tag_list = vm_dict["tags"] + + return vm_asset_yaml, tag_list \ No newline at end of file diff --git a/dfe_threagile.py b/dfe_threagile.py index 4ba45ea..2a1afc2 100644 --- a/dfe_threagile.py +++ b/dfe_threagile.py @@ -12,6 +12,8 @@ build_cache_tm, build_app_service_tm, build_storage_tm, + build_db_tm, + build_vm_tm ) from build_data_assets import ( build_client_app_data_asset, @@ -101,6 +103,22 @@ def produce_assets() -> list: all_tech_tags.append(tag) print(storage_yaml) + case "microsoft.sql/servers/databases": + db_yaml, tag_list = build_db_tm(name, asset_type) + yaml_list.append(db_yaml) + + for tag in tag_list: + all_tech_tags.append(tag) + + print(db_yaml) + case "microsoft.compute/virtualmachines": + vm_yaml, tag_list = build_vm_tm(name, asset_type) + yaml_list.append(vm_yaml) + + for tag in tag_list: + all_tech_tags.append(tag) + + print(vm_yaml) return yaml_list, all_tech_tags @@ -228,7 +246,7 @@ def data_assets() -> list: def template_inject( - yaml_list: list, data_list: list, all_tags: list, risks: list = [] + yaml_list: list, data_list: list, all_tags: list, risks: list = [], autoescape: bool = True ) -> str: with open("yaml-templates/threagile-example-model-template.yaml") as template_file: template_str = template_file.read() @@ -369,24 +387,13 @@ def produce_asset_lists() -> tuple: risks = read_risks_json("/app/work/output/risks.json") - final_with_risks = template_inject(yaml_list, data_list, all_tags, risks) + final_with_risks = template_inject(yaml_list, data_list, all_tags, risks, autoescape=False) try: with open( "/app/work/yaml-templates/dfe-threagile-final.yaml", "x" ) as yaml_file: yaml_file.write(final_with_risks) - - with open( - "/app/work/yaml-templates/dfe-threagile-final.yaml", "r" - ) as yaml_file: - yaml_contents = yaml_file.read() - pattern = re.compile(re.escape("rating: >")) - updated_contents = pattern.sub("rating: >", yaml_contents) - with open( - "/app/work/yaml-templates/dfe-threagile-final.yaml", "w" - ) as yaml_file: - yaml_file.write(updated_contents) except FileExistsError: print("File exists, overwriting...") with open( @@ -394,17 +401,6 @@ def produce_asset_lists() -> tuple: ) as yaml_file: yaml_file.write(final_with_risks) - with open( - "/app/work/yaml-templates/dfe-threagile-final.yaml", "r" - ) as yaml_file: - yaml_contents = yaml_file.read() - pattern = re.compile(re.escape("rating: >")) - updated_contents = pattern.sub("rating: >", yaml_contents) - with open( - "/app/work/yaml-templates/dfe-threagile-final.yaml", "w" - ) as yaml_file: - yaml_file.write(updated_contents) - os.system( "threagile -verbose -model /app/work/yaml-templates/dfe-threagile-final.yaml -output /app/work/output" )