Shared: Update the model generator script to allow execution of the mixed model generator queries.

Author: michaelnebel

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

@@ -33,7 +33,8 @@ def __init__ (self, language):
         self.generateSources = False
         self.generateSummaries = False
         self.generateNeutrals = False
-        self.generateContentSummaries = False
+        self.generateMixedSummaries = False
+        self.generateMixedNeutrals = False
         self.generateTypeBasedSummaries = False
         self.dryRun = False
         self.dirname = "modelgenerator"
@@ -51,7 +52,8 @@ def printHelp(self):
     --with-sources
     --with-summaries
     --with-neutrals
-    --with-content-summaries (Experimental). May not be used in conjunction with --with-summaries
+    --with-mixed-summaries (Experimental). May not be used in conjunction with --with-summaries.
+    --with-mixed-neutrals (Experimental). Should only be used in conjunction with --with-mixed-summaries.
     --with-typebased-summaries (Experimental)
 If none of these flags are specified, all models are generated except for the type based models.
 
@@ -83,7 +85,7 @@ def make(language):
             generator.printHelp()
             sys.exit(0)
 
-        if "--with-summaries" in sys.argv and "--with-content-summaries" in sys.argv:
+        if "--with-summaries" in sys.argv and "--with-mixed-summaries" in sys.argv:
             generator.printHelp()
             sys.exit(0)
 
@@ -103,9 +105,13 @@ def make(language):
             sys.argv.remove("--with-neutrals")
             generator.generateNeutrals = True
 
-        if "--with-content-summaries" in sys.argv:
-            sys.argv.remove("--with-content-summaries")
-            generator.generateContentSummaries = True
+        if "--with-mixed-summaries" in sys.argv:
+            sys.argv.remove("--with-mixed-summaries")
+            generator.generateMixedSummaries = True
+
+        if "--with-mixed-neutrals" in sys.argv:
+            sys.argv.remove("--with-mixed-neutrals")
+            generator.generateMixedNeutrals = True
 
         if "--with-typebased-summaries" in sys.argv:
             sys.argv.remove("--with-typebased-summaries")
@@ -115,7 +121,13 @@ def make(language):
             sys.argv.remove("--dry-run")
             generator.dryRun = True
 
-        if not generator.generateSinks and not generator.generateSources and not generator.generateSummaries and not generator.generateNeutrals and not generator.generateTypeBasedSummaries and not generator.generateContentSummaries:
+        if (not generator.generateSinks and
+           not generator.generateSources and
+           not generator.generateSummaries and
+           not generator.generateNeutrals and
+           not generator.generateTypeBasedSummaries and
+           not generator.generateMixedSummaries and
+           not generator.generateMixedNeutrals):
             generator.generateSinks = generator.generateSources = generator.generateSummaries = generator.generateNeutrals = True
 
         n = len(sys.argv)
@@ -173,12 +185,17 @@ def makeContent(self):
         else:
             neutralAddsTo = { }
 
-        if self.generateContentSummaries:
-            contentSummaryAddsTo = self.getAddsTo("CaptureContentSummaryModels.ql", helpers.summaryModelPredicate)
+        if self.generateMixedSummaries:
+            mixedSummaryAddsTo = self.getAddsTo("CaptureMixedSummaryModels.ql", helpers.summaryModelPredicate)
+        else:
+            mixedSummaryAddsTo = { }
+
+        if self.generateMixedNeutrals:
+            mixedNeutralAddsTo = self.getAddsTo("CaptureMixedNeutralModels.ql", helpers.neutralModelPredicate)
         else:
-            contentSummaryAddsTo = { }
+            mixedNeutralAddsTo = { }
 
-        return helpers.merge(summaryAddsTo, contentSummaryAddsTo, sinkAddsTo, sourceAddsTo, neutralAddsTo)
+        return helpers.merge(summaryAddsTo, mixedSummaryAddsTo, sinkAddsTo, sourceAddsTo, neutralAddsTo, mixedNeutralAddsTo)
 
     def makeTypeBasedContent(self):
         if self.generateTypeBasedSummaries:
@@ -208,7 +225,12 @@ def run(self):
             print("Models as data extensions generated, but not written to file.")
             sys.exit(0)
 
-        if self.generateSinks or self.generateSources or self.generateSummaries or self.generateNeutrals:
+        if (self.generateSinks or
+           self.generateSources or
+           self.generateSummaries or
+           self.generateNeutrals or
+           self.generateMixedSummaries or
+           self.generatedMixedNeutrals):
             self.save(content, ".model.yml")
 
         if self.generateTypeBasedSummaries: