From e8e865790d331d0bf7887519854bb4c4f3e5cf95 Mon Sep 17 00:00:00 2001 From: Nick Rolfe Date: Fri, 9 Oct 2026 15:36:17 +0100 Subject: [PATCH 1/2] MaD: update generator scripts to be format-agnostic In preparation for adding support for JSON data extensions, this commit updates the scripts so that they store and build models using an in-memory representation that is independent of the YAML format currently used. This means tuples are represented as Python lists, where previously they were created as YAML-formatted strings in the merging process. Scalar values inside those tuples are represented as Python strings and booleans. Only the DataExtensionFile.write_yaml function now has to deal with YAML formatting concerns. --- .../models-as-data/convert_extensions.py | 87 +++++-------- .../models-as-data/data_extension_file.py | 40 ++++++ misc/scripts/models-as-data/generate_mad.py | 115 +++++++----------- misc/scripts/models-as-data/helpers.py | 50 +++++++- 4 files changed, 155 insertions(+), 137 deletions(-) create mode 100644 misc/scripts/models-as-data/data_extension_file.py diff --git a/misc/scripts/models-as-data/convert_extensions.py b/misc/scripts/models-as-data/convert_extensions.py index 01a10dae5ffd..3578fa270826 100644 --- a/misc/scripts/models-as-data/convert_extensions.py +++ b/misc/scripts/models-as-data/convert_extensions.py @@ -8,25 +8,6 @@ import tempfile -def quote_if_needed(v): - # string columns - if type(v) is str: - return '"' + v + '"' - # bool column - return str(v) - - -def parseData(data): - rows = [{}, {}] - for row in data: - d = map(quote_if_needed, row) - provenance = row[-1] - targetRows = rows[1] if provenance.endswith("generated") else rows[0] - helpers.insert_update(targetRows, row[0], " - [" + ", ".join(d) + "]\n") - - return rows - - class Converter: def __init__(self, language, dbDir): self.language = language @@ -63,48 +44,38 @@ def runQuery(self, query): ) return helpers.readData(self.workDir, resultBqrs) - def asAddsTo(self, rows, predicate): - extensions = [{}, {}] - for i in range(2): - for key in rows[i]: - extensions[i][key] = helpers.addsToTemplate.format( - f"codeql/{self.language}-all", predicate, rows[i][key] - ) - - return extensions - - def getAddsTo(self, query, predicate): + def merge_query_results(self, query, predicate, mergers): data = self.runQuery(query) - rows = parseData(data) - return self.asAddsTo(rows, predicate) - - def makeContent(self): - summaries = self.getAddsTo("ExtractSummaries.ql", helpers.summaryModelPredicate) - sources = self.getAddsTo("ExtractSources.ql", helpers.sourceModelPredicate) - sinks = self.getAddsTo("ExtractSinks.ql", helpers.sinkModelPredicate) - neutrals = self.getAddsTo("ExtractNeutrals.ql", helpers.neutralModelPredicate) - return [ - helpers.merge(sources[0], sinks[0], summaries[0], neutrals[0]), - helpers.merge(sources[1], sinks[1], summaries[1], neutrals[1]), + for row in data: + provenance = row[-1] + namespace = row[0] + target_merger = ( + mergers[1] if provenance.endswith("generated") else mergers[0] + ) + target_merger.add_row(namespace, predicate, row) + + def make_extensions(self): + mergers = [ + helpers.ExtensionMerger(f"codeql/{self.language}-all"), + helpers.ExtensionMerger(f"codeql/{self.language}-all"), ] + self.merge_query_results( + "ExtractSummaries.ql", helpers.summaryModelPredicate, mergers + ) + self.merge_query_results( + "ExtractSources.ql", helpers.sourceModelPredicate, mergers + ) + self.merge_query_results("ExtractSinks.ql", helpers.sinkModelPredicate, mergers) + self.merge_query_results( + "ExtractNeutrals.ql", helpers.neutralModelPredicate, mergers + ) + return mergers + + def run(self): + mergers = self.make_extensions() - def save(self, extensions): # Create directory if it doesn't exist os.makedirs(self.extDir, exist_ok=True) - # Create a file for each namespace and save models. - extensionTemplate = """extensions: -{0}""" - for entry in extensions[0]: - with open(self.extDir + "/" + entry + self.modelFileExtension, "w") as f: - f.write(extensionTemplate.format(extensions[0][entry])) - - for entry in extensions[1]: - with open( - self.extDir + "/generated/" + entry + self.modelFileExtension, "w" - ) as f: - f.write(extensionTemplate.format(extensions[1][entry])) - - def run(self): - extensions = self.makeContent() - self.save(extensions) + mergers[0].save(self.extDir, self.modelFileExtension) + mergers[1].save(self.extDir + "/generated", self.modelFileExtension) diff --git a/misc/scripts/models-as-data/data_extension_file.py b/misc/scripts/models-as-data/data_extension_file.py new file mode 100644 index 000000000000..c4073df78462 --- /dev/null +++ b/misc/scripts/models-as-data/data_extension_file.py @@ -0,0 +1,40 @@ +import json + + +# Simple class for representing a data extension file and writing it to disk in a pretty-printed YAML format. +class DataExtensionFile: + def __init__(self, pack): + self.pack = pack + self.predicates = {} + + def add_rows(self, predicate, rows): + if predicate not in self.predicates: + self.predicates[predicate] = [] + self.predicates[predicate].extend(rows) + + def yaml_for_value(self, value): + # if is boolean: + if isinstance(value, bool): + # json.dumps would produce lower-case "true" or "false". Even though that's valid YAML, + # we have historically used the upper-case variants. + return "True" if value else "False" + return json.dumps(value) + + def yaml_for_row(self, row): + return f"[{', '.join(self.yaml_for_value(value) for value in row)}]" + + def yaml_for_predicate(self, predicate): + return f""" - addsTo: + pack: {self.pack} + extensible: {predicate} + data: + - {'\n - '.join(self.yaml_for_row(row) for row in self.predicates[predicate])} +""" + + def write_yaml(self, f): + # It would be preferable to use a YAML serialization library, but this lets us control the + # formatting e.g. to produce one line per tuple. + f.write("# THIS FILE IS AN AUTO-GENERATED MODELS AS DATA FILE. DO NOT EDIT.\n") + f.write("extensions:\n") + for predicate in sorted(self.predicates.keys()): + f.write(self.yaml_for_predicate(predicate)) diff --git a/misc/scripts/models-as-data/generate_mad.py b/misc/scripts/models-as-data/generate_mad.py index d78d5b10cc5d..17d8e4bdeecd 100755 --- a/misc/scripts/models-as-data/generate_mad.py +++ b/misc/scripts/models-as-data/generate_mad.py @@ -6,29 +6,8 @@ import subprocess import sys import tempfile -import re import argparse - -def quote_if_needed(row): - if row != "true" and row != "false": - return '"' + row + '"' - # subtypes column - return row[0].upper() + row[1:] - - -def parseData(data): - rows = {} - - for row in data: - d = row[0].split(";") - namespace = d[0] - d = map(quote_if_needed, d) - helpers.insert_update(rows, namespace, " - [" + ", ".join(d) + "]\n") - - return rows - - description = """\ This generates summary, source, sink and neutral models for the code in the database. The files will be placed in `LANGUAGE/ql/lib/ext/generated/modelgenerator/DIR`""" @@ -182,76 +161,62 @@ def runQuery(self, query): return helpers.readData(self.workDir, resultBqrs) - def asAddsTo(self, rows, predicate): - extensions = {} - for key in rows: - extensions[key] = helpers.addsToTemplate.format( - f"codeql/{self.language}-all", predicate, rows[key] - ) - return extensions - - def getAddsTo(self, query, predicate): + def convert_if_needed(self, value): + # Convert a string value to a boolean if it is "true" or "false". Otherwise, return the original + # value as a string. Since no extensible predicate currently uses integers, don't attempt to + # convert numeric strings to integers. + if value == "true": + return True + if value == "false": + return False + return value + + def merge_query_results(self, query, predicate, merger): data = self.runQuery(query) - rows = parseData(data) - if self.single_file and rows: - rows = {self.single_file: "".join(rows.values())} - return self.asAddsTo(rows, predicate) + for row in data: + d = row[0].split(";") + namespace = d[0] + d = list(map(self.convert_if_needed, d)) + merger.add_row(namespace, predicate, d) + + def make_extensions(self): + merger = helpers.ExtensionMerger(f"codeql/{self.language}-all") - def makeContent(self): - summaryAddsTo = {} if self.with_summaries: - summaryAddsTo = self.getAddsTo( - "CaptureSummaryModels.ql", helpers.summaryModelPredicate + self.merge_query_results( + "CaptureSummaryModels.ql", helpers.summaryModelPredicate, merger ) - sinkAddsTo = {} if self.with_sinks: - sinkAddsTo = self.getAddsTo( - "CaptureSinkModels.ql", helpers.sinkModelPredicate + self.merge_query_results( + "CaptureSinkModels.ql", helpers.sinkModelPredicate, merger ) - sourceAddsTo = {} if self.with_sources: - sourceAddsTo = self.getAddsTo( - "CaptureSourceModels.ql", helpers.sourceModelPredicate + self.merge_query_results( + "CaptureSourceModels.ql", helpers.sourceModelPredicate, merger ) - neutralAddsTo = {} if self.with_neutrals: - neutralAddsTo = self.getAddsTo( - "CaptureNeutralModels.ql", helpers.neutralModelPredicate + self.merge_query_results( + "CaptureNeutralModels.ql", helpers.neutralModelPredicate, merger ) - return helpers.merge(summaryAddsTo, sinkAddsTo, sourceAddsTo, neutralAddsTo) + return merger - def makeTypeBasedContent(self): + def make_type_based_extensions(self): + merger = helpers.ExtensionMerger(f"codeql/{self.language}-all") if self.with_typebased_summaries: - typeBasedSummaryAddsTo = self.getAddsTo( - "CaptureTypeBasedSummaryModels.ql", helpers.summaryModelPredicate + self.merge_query_results( + "CaptureTypeBasedSummaryModels.ql", + helpers.summaryModelPredicate, + merger, ) - else: - typeBasedSummaryAddsTo = {} - - return typeBasedSummaryAddsTo - - def save(self, extensions, extension): - # Create a file for each namespace and save models. - extensionTemplate = """# THIS FILE IS AN AUTO-GENERATED MODELS AS DATA FILE. DO NOT EDIT. -extensions: -{0}""" - for entry in extensions: - # Replace problematic characters with dashes, and collapse multiple dashes. - sanitizedEntry = re.sub( - r"-+", "-", entry.replace("/", "-").replace(":", "-") - ) - target = os.path.join(self.generated_frameworks, sanitizedEntry + extension) - with open(target, "w") as f: - f.write(extensionTemplate.format(extensions[entry])) - print("Models as data extensions written to " + target) + return merger def run(self): - content = self.makeContent() - typeBasedContent = self.makeTypeBasedContent() + merger = self.make_extensions() + type_based_extensions = self.make_type_based_extensions() if self.dry_run: print("Models as data extensions generated, but not written to file.") @@ -263,10 +228,12 @@ def run(self): or self.with_summaries or self.with_neutrals ): - self.save(content, ".model.yml") + merger.save(self.generated_frameworks, ".model.yml") if self.with_typebased_summaries: - self.save(typeBasedContent, ".typebased.model.yml") + type_based_extensions.save( + self.generated_frameworks, ".typebased.model.yml" + ) if __name__ == "__main__": diff --git a/misc/scripts/models-as-data/helpers.py b/misc/scripts/models-as-data/helpers.py index 580523f4bec6..372f66ac8438 100644 --- a/misc/scripts/models-as-data/helpers.py +++ b/misc/scripts/models-as-data/helpers.py @@ -2,17 +2,57 @@ import os import shutil import subprocess +import re +import data_extension_file # Shared strings. summaryModelPredicate = "summaryModel" sinkModelPredicate = "sinkModel" sourceModelPredicate = "sourceModel" neutralModelPredicate = "neutralModel" -addsToTemplate = """ - addsTo: - pack: {0} - extensible: {1} - data: -{2}""" + + +# Helper class for accumulating tuples grouped by namespace and predicate, and generating data +# extensions for them. +class ExtensionMerger: + def __init__(self, pack): + self.pack = pack + self.namespaces = {} + + def add_row(self, namespace, predicate, row): + if namespace not in self.namespaces: + self.namespaces[namespace] = {} + if predicate not in self.namespaces[namespace]: + self.namespaces[namespace][predicate] = [] + self.namespaces[namespace][predicate].append(row) + + # Helper function to yield unique elements from a sorted list. + def uniq(self, sorted_list): + last = None + for element in sorted_list: + if element == last: + continue + yield element + last = element + + def save(self, dir, file_extension): + # Create a file for each namespace and save models. + for namespace in self.namespaces: + # Sort and deduplicate rows for each predicate within this namespace. + for predicate in self.namespaces[namespace]: + l = self.namespaces[namespace][predicate] + self.namespaces[namespace][predicate] = list(self.uniq(sorted(l))) + extension = data_extension_file.DataExtensionFile(self.pack) + for predicate in self.namespaces[namespace]: + extension.add_rows(predicate, self.namespaces[namespace][predicate]) + # Replace problematic characters with dashes, and collapse multiple dashes. + sanitized_namespace = re.sub( + r"-+", "-", namespace.replace("/", "-").replace(":", "-") + ) + target = os.path.join(dir, f"{sanitized_namespace}{file_extension}") + with open(target, "w") as f: + extension.write_yaml(f) + print("Models as data extensions written to " + target) def remove_dir(dirName): From 95c94ef8b73dd975afc39299bfa79d13cceb8608 Mon Sep 17 00:00:00 2001 From: Nick Rolfe Date: Fri, 9 Oct 2026 16:02:04 +0100 Subject: [PATCH 2/2] MaD: update generator scripts to support writing JSON data extensions The `generate_mad.py` script accepts a flag that lets you choose between YAML or JSON, but YAML is still the default. We can change the default later, when CLI support for JSON is widely available. `convert_extensions.py` also still hard-codes `.yml` as the output extension/format. --- .../models-as-data/bulk_generate_mad.py | 9 +++--- .../models-as-data/data_extension_file.py | 30 ++++++++++++++++++- misc/scripts/models-as-data/generate_mad.py | 12 ++++++-- misc/scripts/models-as-data/helpers.py | 5 +++- 4 files changed, 48 insertions(+), 8 deletions(-) diff --git a/misc/scripts/models-as-data/bulk_generate_mad.py b/misc/scripts/models-as-data/bulk_generate_mad.py index e2de5038206e..905064706eb2 100755 --- a/misc/scripts/models-as-data/bulk_generate_mad.py +++ b/misc/scripts/models-as-data/bulk_generate_mad.py @@ -430,10 +430,11 @@ def download_and_decompress(analyzed_database: dict) -> str: def clean_up_mad_destination_for_project(config, name: str): target = pathlib.Path(config["destination"], name) if config.get("single-file", False): - target = target.with_suffix(".model.yml") - if target.exists(): - print(f"Deleting existing MaD file at {target}") - target.unlink() + for extension in (".model.json", ".model.yml"): + model_target = target.with_suffix(extension) + if model_target.exists(): + print(f"Deleting existing MaD file at {model_target}") + model_target.unlink() elif target.exists(): print(f"Deleting existing MaD directory at {target}") shutil.rmtree(target, ignore_errors=True) diff --git a/misc/scripts/models-as-data/data_extension_file.py b/misc/scripts/models-as-data/data_extension_file.py index c4073df78462..157e8a80008d 100644 --- a/misc/scripts/models-as-data/data_extension_file.py +++ b/misc/scripts/models-as-data/data_extension_file.py @@ -1,7 +1,7 @@ import json -# Simple class for representing a data extension file and writing it to disk in a pretty-printed YAML format. +# Simple class for representing a data extension file and writing it to disk in a pretty-printed YAML or JSON format. class DataExtensionFile: def __init__(self, pack): self.pack = pack @@ -38,3 +38,31 @@ def write_yaml(self, f): f.write("extensions:\n") for predicate in sorted(self.predicates.keys()): f.write(self.yaml_for_predicate(predicate)) + + def json_for_row(self, row): + return "[" + ", ".join(json.dumps(value) for value in row) + "]" + + def json_for_predicate(self, predicate): + return f""" {{ + "addsTo": {{ + "pack": "{self.pack}", + "extensible": "{predicate}" + }}, + "data": [ + {',\n '.join(self.json_for_row(row) for row in self.predicates[predicate])} + ] + }}""" + + def write_json(self, f): + # It would be preferable to use a serializer like Python's json.dumps, but it can't + # pretty-print with one line per tuple. + f.write("// THIS FILE IS AN AUTO-GENERATED MODELS AS DATA FILE. DO NOT EDIT.\n") + f.write("{\n") + f.write(' "extensions": [\n') + for i, predicate in enumerate(self.predicates): + if i > 0: + f.write(",\n") + f.write(self.json_for_predicate(predicate)) + f.write("\n") + f.write(" ]\n") + f.write("}") diff --git a/misc/scripts/models-as-data/generate_mad.py b/misc/scripts/models-as-data/generate_mad.py index 17d8e4bdeecd..99119009811b 100755 --- a/misc/scripts/models-as-data/generate_mad.py +++ b/misc/scripts/models-as-data/generate_mad.py @@ -119,6 +119,12 @@ def make(): "--single-file", help="Generate a single file with all models instead of separate files for each namespace, using provided argument as the base filename.", ) + p.add_argument( + "--extension-format", + choices=["json", "yaml"], + default="yaml", + help="Format for the generated data extension files (default %(default)s)", + ) generator = p.parse_args(namespace=Generator()) if ( @@ -222,17 +228,19 @@ def run(self): print("Models as data extensions generated, but not written to file.") sys.exit(0) + file_extension = ".json" if self.extension_format == "json" else ".yml" + if ( self.with_sinks or self.with_sources or self.with_summaries or self.with_neutrals ): - merger.save(self.generated_frameworks, ".model.yml") + merger.save(self.generated_frameworks, f".model{file_extension}") if self.with_typebased_summaries: type_based_extensions.save( - self.generated_frameworks, ".typebased.model.yml" + self.generated_frameworks, f".typebased.model{file_extension}" ) diff --git a/misc/scripts/models-as-data/helpers.py b/misc/scripts/models-as-data/helpers.py index 372f66ac8438..6a18f2827f76 100644 --- a/misc/scripts/models-as-data/helpers.py +++ b/misc/scripts/models-as-data/helpers.py @@ -51,7 +51,10 @@ def save(self, dir, file_extension): ) target = os.path.join(dir, f"{sanitized_namespace}{file_extension}") with open(target, "w") as f: - extension.write_yaml(f) + if file_extension.endswith(".json"): + extension.write_json(f) + else: + extension.write_yaml(f) print("Models as data extensions written to " + target)