Skip to content

Commit 20414c0

Browse files
michaelnebelmichaelnebel
committed
Java: Move Generator creation into class definition.
1 parent 3f33cdf commit 20414c0

File tree

1 file changed

+116
-112
lines changed

1 file changed

+116
-112
lines changed

‎java/ql/src/utils/model-generator/generate_flow_model.py

100755100644
Lines changed: 116 additions & 112 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,10 @@
1313
class Generator:
1414
def __init__ (self, language):
1515
self.language = language
16+
self.generateSinks = False
17+
self.generateSources = False
18+
self.generateSummaries = False
19+
self.dryRun = False
1620

1721
def printHelp(self):
1822
print(f"""Usage:
@@ -38,99 +42,97 @@ def printHelp(self):
3842
Requirements: `codeql` should both appear on your path.
3943
""")
4044

41-
generator = Generator(language)
4245

43-
if any(s == "--help" for s in sys.argv):
44-
generator.printHelp()
45-
sys.exit(0)
46-
47-
generateSinks = False
48-
generateSources = False
49-
generateSummaries = False
50-
dryRun = False
51-
52-
if "--with-sinks" in sys.argv:
53-
sys.argv.remove("--with-sinks")
54-
generateSinks = True
55-
56-
if "--with-sources" in sys.argv:
57-
sys.argv.remove("--with-sources")
58-
generateSources = True
59-
60-
if "--with-summaries" in sys.argv:
61-
sys.argv.remove("--with-summaries")
62-
generateSummaries = True
63-
64-
if "--dry-run" in sys.argv:
65-
sys.argv.remove("--dry-run")
66-
dryRun = True
67-
68-
if not generateSinks and not generateSources and not generateSummaries:
69-
generateSinks = generateSources = generateSummaries = True
70-
71-
if len(sys.argv) != 3:
72-
generator.printHelp()
73-
sys.exit(1)
74-
75-
codeQlRoot = subprocess.check_output(
76-
["git", "rev-parse", "--show-toplevel"]).decode("utf-8").strip()
77-
targetQll = sys.argv[2]
78-
if not targetQll.endswith(".qll"):
79-
targetQll += ".qll"
80-
filename = os.path.basename(targetQll)
81-
shortname = filename[:-4]
82-
generatedFrameworks = os.path.join(
83-
codeQlRoot, f"{language}/ql/lib/semmle/code/{language}/frameworks/")
84-
frameworkTarget = os.path.join(generatedFrameworks, targetQll)
85-
86-
workDir = tempfile.mkdtemp()
87-
os.makedirs(generatedFrameworks, exist_ok=True)
88-
89-
90-
def runQuery(infoMessage, query):
91-
print("########## Querying " + infoMessage + "...")
92-
database = sys.argv[1]
93-
queryFile = os.path.join(os.path.dirname(
94-
__file__), query)
95-
resultBqrs = os.path.join(workDir, "out.bqrs")
96-
cmd = ['codeql', 'query', 'run', queryFile, '--database',
97-
database, '--output', resultBqrs, '--threads', '8']
98-
99-
ret = subprocess.call(cmd)
100-
if ret != 0:
101-
print("Failed to generate " + infoMessage +
102-
". Failed command was: " + shlex.join(cmd))
103-
sys.exit(1)
104-
return readRows(resultBqrs)
105-
106-
107-
def readRows(bqrsFile):
108-
generatedJson = os.path.join(workDir, "out.json")
109-
cmd = ['codeql', 'bqrs', 'decode', bqrsFile,
110-
'--format=json', '--output', generatedJson]
111-
ret = subprocess.call(cmd)
112-
if ret != 0:
113-
print("Failed to decode BQRS. Failed command was: " + shlex.join(cmd))
114-
sys.exit(1)
115-
116-
with open(generatedJson) as f:
117-
results = json.load(f)
118-
119-
try:
120-
results['#select']['tuples']
121-
except KeyError:
122-
print('Unexpected JSON output - no tuples found')
123-
exit(1)
124-
125-
rows = ""
126-
for (row) in results['#select']['tuples']:
127-
rows += " \"" + row[0] + "\",\n"
128-
129-
return rows[:-2]
130-
131-
132-
def asCsvModel(superclass, kind, rows):
133-
classTemplate = """
46+
def setenvironment(self, target, database):
47+
self.codeQlRoot = subprocess.check_output(["git", "rev-parse", "--show-toplevel"]).decode("utf-8").strip()
48+
if not target.endswith(".qll"):
49+
target += ".qll"
50+
self.filename = os.path.basename(target)
51+
self.shortname = self.filename[:-4]
52+
self.database = database
53+
self.generatedFrameworks = os.path.join(
54+
self.codeQlRoot, f"{self.language}/ql/lib/semmle/code/{self.language}/frameworks/")
55+
self.frameworkTarget = os.path.join(self.generatedFrameworks, target)
56+
57+
self.workDir = tempfile.mkdtemp()
58+
os.makedirs(self.generatedFrameworks, exist_ok=True)
59+
60+
@staticmethod
61+
def make(language):
62+
generator = Generator(language)
63+
if any(s == "--help" for s in sys.argv):
64+
generator.printHelp()
65+
sys.exit(0)
66+
67+
if "--with-sinks" in sys.argv:
68+
sys.argv.remove("--with-sinks")
69+
generator.generateSinks = True
70+
71+
if "--with-sources" in sys.argv:
72+
sys.argv.remove("--with-sources")
73+
generator.generateSources = True
74+
75+
if "--with-summaries" in sys.argv:
76+
sys.argv.remove("--with-summaries")
77+
generator.generateSummaries = True
78+
79+
if "--dry-run" in sys.argv:
80+
sys.argv.remove("--dry-run")
81+
generator.dryRun = True
82+
83+
if not generator.generateSinks and not generator.generateSources and not generator.generateSummaries:
84+
generator.generateSinks = generator.generateSources = generator.generateSummaries = True
85+
86+
if len(sys.argv) != 3:
87+
generator.printHelp()
88+
sys.exit(1)
89+
90+
generator.setenvironment(sys.argv[2], sys.argv[1])
91+
return generator
92+
93+
def runQuery(self, infoMessage, query):
94+
print("########## Querying " + infoMessage + "...")
95+
queryFile = os.path.join(os.path.dirname(
96+
__file__), query)
97+
resultBqrs = os.path.join(self.workDir, "out.bqrs")
98+
cmd = ['codeql', 'query', 'run', queryFile, '--database',
99+
self.database, '--output', resultBqrs, '--threads', '8']
100+
101+
ret = subprocess.call(cmd)
102+
if ret != 0:
103+
print("Failed to generate " + infoMessage +
104+
". Failed command was: " + shlex.join(cmd))
105+
sys.exit(1)
106+
return self.readRows(resultBqrs)
107+
108+
109+
def readRows(self, bqrsFile):
110+
generatedJson = os.path.join(self.workDir, "out.json")
111+
cmd = ['codeql', 'bqrs', 'decode', bqrsFile,
112+
'--format=json', '--output', generatedJson]
113+
ret = subprocess.call(cmd)
114+
if ret != 0:
115+
print("Failed to decode BQRS. Failed command was: " + shlex.join(cmd))
116+
sys.exit(1)
117+
118+
with open(generatedJson) as f:
119+
results = json.load(f)
120+
121+
try:
122+
results['#select']['tuples']
123+
except KeyError:
124+
print('Unexpected JSON output - no tuples found')
125+
exit(1)
126+
127+
rows = ""
128+
for (row) in results['#select']['tuples']:
129+
rows += " \"" + row[0] + "\",\n"
130+
131+
return rows[:-2]
132+
133+
134+
def asCsvModel(self, superclass, kind, rows):
135+
classTemplate = """
134136
private class {0}{1}Csv extends {2} {{
135137
override predicate row(string row) {{
136138
row =
@@ -139,54 +141,56 @@ def asCsvModel(superclass, kind, rows):
139141
]
140142
}}
141143
}}
142-
"""
143-
if rows.strip() == "":
144-
return ""
145-
return classTemplate.format(shortname[0].upper() + shortname[1:], kind.capitalize(), superclass, rows)
144+
"""
145+
if rows.strip() == "":
146+
return ""
147+
return classTemplate.format(self.shortname[0].upper() + self.shortname[1:], kind.capitalize(), superclass, rows)
148+
146149

150+
generator = Generator.make(language)
147151

148-
if generateSummaries:
149-
summaryRows = runQuery("summary models", "CaptureSummaryModels.ql")
150-
summaryCsv = asCsvModel("SummaryModelCsv", "summary", summaryRows)
152+
if generator.generateSummaries:
153+
summaryRows = generator.runQuery("summary models", "CaptureSummaryModels.ql")
154+
summaryCsv = generator.asCsvModel("SummaryModelCsv", "summary", summaryRows)
151155
else:
152156
summaryCsv = ""
153157

154-
if generateSinks:
155-
sinkRows = runQuery("sink models", "CaptureSinkModels.ql")
156-
sinkCsv = asCsvModel("SinkModelCsv", "sinks", sinkRows)
158+
if generator.generateSinks:
159+
sinkRows = generator.runQuery("sink models", "CaptureSinkModels.ql")
160+
sinkCsv = generator.asCsvModel("SinkModelCsv", "sinks", sinkRows)
157161
else:
158162
sinkCsv = ""
159163

160-
if generateSources:
161-
sourceRows = runQuery("source models", "CaptureSourceModels.ql")
162-
sourceCsv = asCsvModel("SourceModelCsv", "sources", sourceRows)
164+
if generator.generateSources:
165+
sourceRows = generator.runQuery("source models", "CaptureSourceModels.ql")
166+
sourceCsv = generator.asCsvModel("SourceModelCsv", "sources", sourceRows)
163167
else:
164168
sourceCsv = ""
165169

166170
qllContents = f"""
167-
/** Definitions of taint steps in the {shortname} framework */
171+
/** Definitions of taint steps in the {generator.shortname} framework */
168172
169-
import {language}
170-
private import semmle.code.{language}.dataflow.ExternalFlow
173+
import {generator.language}
174+
private import semmle.code.{generator.language}.dataflow.ExternalFlow
171175
172176
{sinkCsv}
173177
{sourceCsv}
174178
{summaryCsv}
175179
176180
"""
177181

178-
if dryRun:
182+
if generator.dryRun:
179183
print("CSV Models generated, but not written to file.")
180184
sys.exit(0)
181185

182-
with open(frameworkTarget, "w") as frameworkQll:
186+
with open(generator.frameworkTarget, "w") as frameworkQll:
183187
frameworkQll.write(qllContents)
184188

185-
cmd = ['codeql', 'query', 'format', '--in-place', frameworkTarget]
189+
cmd = ['codeql', 'query', 'format', '--in-place', generator.frameworkTarget]
186190
ret = subprocess.call(cmd)
187191
if ret != 0:
188192
print("Failed to format query. Failed command was: " + shlex.join(cmd))
189193
sys.exit(1)
190194

191195
print("")
192-
print("CSV model written to " + frameworkTarget)
196+
print("CSV model written to " + generator.frameworkTarget)

0 commit comments

Comments
 (0)