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.

Author: nickrolfe

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

@@ -7,26 +7,6 @@
 import sys
 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 +43,30 @@ 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,39 @@
+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))

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,50 @@ 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(":", "-")
-            )
-            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)
+            self.merge_query_results("CaptureTypeBasedSummaryModels.ql", helpers.summaryModelPredicate, merger)
+        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 +216,10 @@ 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__":

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

@@ -2,17 +2,56 @@
 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):