MaD: update generator scripts to support writing JSON data extensions

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)

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)

misc/scripts/models-as-data/data_extension_file.py

@@ -0,0 +1,68 @@
+import json
+
+
+# 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
+        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))
+
+    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("}")

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`"""
@@ -140,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 (
@@ -182,91 +167,81 @@ 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
-            )
-        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(":", "-")
+            self.merge_query_results(
+                "CaptureTypeBasedSummaryModels.ql",
+                helpers.summaryModelPredicate,
+                merger,
             )
-            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.")
             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
         ):
-            self.save(content, ".model.yml")
+            merger.save(self.generated_frameworks, f".model{file_extension}")
 
         if self.with_typebased_summaries:
-            self.save(typeBasedContent, ".typebased.model.yml")
+            type_based_extensions.save(
+                self.generated_frameworks, f".typebased.model{file_extension}"
+            )
 
 
 if __name__ == "__main__":