Coverage for graphqler / utils / stats.py: 96%
222 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-03-20 10:09 -0400
1import json
2import cloudpickle as pickle
3import pprint
4import time
5from pathlib import Path
6from typing import Self
8from graphqler import config
9from graphqler.fuzzer.engine.types import Result
10from graphqler.graph import Node
12from .file_utils import initialize_file, intialize_file_if_not_exists, recreate_path, get_or_create_file
13from .singleton import singleton
14import os
15import re
18@singleton
19class Stats :
20 ### PUT THE STATS YOU WANT HERE
21 file_path = "/tmp/stats.txt" # This gets overriden by the set_file_path function
22 endpoint_results_dir = "/tmp/endpoint_results"
23 unique_responses_file_path = "/tmp/unique_responses.txt"
24 start_time: float = time.time()
25 http_status_codes: dict[str, dict[str, int]] = {}
26 successful_nodes: dict[str, int] = {}
27 failed_nodes: dict[str, int] = {}
28 results: dict[str, set[Result]] = {} # Mapping of query/muation to results for that node
29 unique_responses: dict[str, list[str]] = {} # Mapping of response to endpoints (query/mutation)
30 number_of_queries: int = 0
31 number_of_mutations: int = 0
32 number_of_objects: int = 0
33 number_of_successes: int = 0
34 number_of_failures: int = 0
35 vulnerabilities = {} # Mapping of vulnerability to node name, and if it's a potentiall or confirmed vulnerability
36 node_timings: dict[str, list[float]] = {} # Mapping of node name to list of elapsed times in seconds
38 # Detection stats
39 is_introspection_available: bool = False
41 def __init__(self):
42 self.http_status_codes = {}
43 self.node_timings = {}
44 self.pickle_save_path = Path(config.OUTPUT_DIRECTORY) / config.SERIALIZED_DIR_NAME / config.STATS_PICKLE_FILE_NAME
46 def load(self) -> Self:
47 """Loads the stats from the pickle file"""
48 self.__load_pickle()
49 return self
51 def add_successful_node(self, node: Node):
52 """Adds a new successful node to the succesful stats
54 Args:
55 node (Node): A graphqler node
56 """
57 key_name = f"{node.graphql_type}|{node.name}"
58 self.number_of_successes += 1
59 if key_name in self.successful_nodes:
60 self.successful_nodes[key_name] += 1
61 else:
62 self.successful_nodes[key_name] = 1
63 self.save()
65 def add_failed_node(self, node: Node):
66 """Adds a new failed node to the internal failed stats
68 Args:
69 node (Node): A graphqler node
70 """
71 key_name = f"{node.graphql_type}|{node.name}"
72 self.number_of_failures += 1
73 if key_name in self.failed_nodes:
74 self.failed_nodes[key_name] += 1
75 else:
76 self.failed_nodes[key_name] = 1
77 self.save()
79 def add_http_status_code(self, payload_name: str, status_code: int):
80 """Adds the http status code to stats
82 Args:
83 payload_name (str): The name of the query or mutation
84 status_code (int): The status code
85 """
86 status_code_str = str(status_code)
87 if status_code_str in self.http_status_codes.keys():
88 if payload_name in self.http_status_codes[status_code_str]:
89 self.http_status_codes[status_code_str][payload_name] += 1
90 else:
91 self.http_status_codes[status_code_str][payload_name] = 1
92 else:
93 self.http_status_codes[status_code_str] = {payload_name: 1}
94 self.save()
96 def set_file_paths(self, working_dir: str):
97 """
99 Args:
100 working_dir (str): _description_
101 """
102 # Do the stats file first
103 initialize_file(Path(working_dir) / config.STATS_FILE_NAME)
104 self.file_path = Path(working_dir) / config.STATS_FILE_NAME
106 # JSON report path (machine-readable)
107 json_file_name = config.STATS_FILE_NAME.replace(".txt", ".json") if config.STATS_FILE_NAME.endswith(".txt") else config.STATS_FILE_NAME + ".json"
108 initialize_file(Path(working_dir) / json_file_name)
109 self.json_file_path = Path(working_dir) / json_file_name
111 # Do the endpoint results directory
112 self.endpoint_results_dir = Path(working_dir) / config.ENDPOINT_RESULTS_DIR_NAME
113 recreate_path(self.endpoint_results_dir)
115 # Do the unique responses file
116 self.unique_responses_file_path = Path(working_dir) / config.UNIQUE_RESPONSES_FILE_NAME
117 initialize_file(self.unique_responses_file_path)
119 def print_running_stats(self):
120 """Function to print stats during runtime (not saved to file)"""
121 print(f"Number of success: {self.number_of_successes}", end="")
122 print("|", end="")
123 print(f"Number of failures: {self.number_of_failures}", end="")
124 print("\r", end="", flush=True)
126 def add_vulnerability(self, vulnerability_name: str, node_name: str, is_vulnerable: bool, potentially_vulnerable: bool = False):
127 """Whether a detection was detected or not -- if already detected, it will stay detected
129 Args:
130 detection_name (str): name of the detection
131 detected (bool): whether the detection was detected or not
132 """
133 if vulnerability_name not in self.vulnerabilities:
134 self.vulnerabilities[vulnerability_name] = {}
136 if node_name in self.vulnerabilities[vulnerability_name]:
137 self.vulnerabilities[vulnerability_name][node_name]["potentially_vulnerable"] = (
138 potentially_vulnerable | self.vulnerabilities[vulnerability_name][node_name]["potentially_vulnerable"]
139 )
140 self.vulnerabilities[vulnerability_name][node_name]["is_vulnerable"] = is_vulnerable | self.vulnerabilities[vulnerability_name][node_name]["is_vulnerable"]
141 else:
142 self.vulnerabilities[vulnerability_name][node_name] = {}
143 self.vulnerabilities[vulnerability_name][node_name]["potentially_vulnerable"] = potentially_vulnerable
144 self.vulnerabilities[vulnerability_name][node_name]["is_vulnerable"] = is_vulnerable
146 def get_formatted_vulnerabilites(self) -> str:
147 """Returns the formatted vulnerabilities
149 Returns:
150 str: The formatted vulnerabilities
151 """
152 formatted_vulnerabilities = ""
153 for vulnerability_name, nodes in self.vulnerabilities.items():
154 vulnerable_nodes = ""
155 for node_name, vulnerability in nodes.items():
156 if vulnerability["is_vulnerable"] or vulnerability["potentially_vulnerable"]:
157 if vulnerability["is_vulnerable"]:
158 vulnerable_nodes += f" ❗'{node_name}' - Is vulnerable\n"
159 else:
160 vulnerable_nodes += f" 🔍'{node_name}' - Is potentially vulnerable \n"
161 if vulnerable_nodes != "":
162 formatted_vulnerabilities += f"\n{vulnerability_name}:\n"
163 formatted_vulnerabilities += vulnerable_nodes
164 return formatted_vulnerabilities
166 def record_node_timing(self, node: Node, elapsed_seconds: float):
167 """Records the elapsed time for a node execution
169 Args:
170 node (Node): The node that was executed
171 elapsed_seconds (float): Time taken in seconds
172 """
173 key_name = f"{node.graphql_type}|{node.name}"
174 if key_name not in self.node_timings:
175 self.node_timings[key_name] = []
176 self.node_timings[key_name].append(elapsed_seconds)
178 def update_stats_from_result(self, node, result: Result) -> None:
179 """Parses the result and adds it to the stats
181 Args:
182 result (Result): the result
183 """
184 result_status = result.success
186 # Update success / fail stats first
187 if result_status:
188 self.add_successful_node(node)
189 else:
190 self.add_failed_node(node)
192 # Update results
193 if node.name in self.results:
194 self.results[node.name].add(result)
195 else:
196 self.results[node.name] = {result}
198 # Update unique responses
199 if str(result.graphql_response) in self.unique_responses:
200 self.unique_responses[str(result.graphql_response)].append(node.name)
201 else:
202 self.unique_responses[str(result.graphql_response)] = [node.name]
204 def get_number_of_successful_mutations_and_queries(self) -> tuple[int, int]:
205 """Returns the number of successful mutations and queries"""
206 number_success_of_mutations_and_queries = 0
207 num_mutations_and_queries = self.number_of_mutations + self.number_of_queries
208 for action, num_success in self.successful_nodes.items():
209 action_name = action.split("|")[0]
210 if action_name == "Mutation" or action_name == "Query":
211 if num_success > 0:
212 number_success_of_mutations_and_queries += 1
213 return number_success_of_mutations_and_queries, num_mutations_and_queries
215 def get_number_of_failed_mutations_and_queries(self) -> tuple[int, int]:
216 """Returns the number of failed EXTERNAL mutations and queries"""
217 number_failed_of_mutations_and_queries = 0
218 num_mutations_and_queries = self.number_of_mutations + self.number_of_queries
219 for action, num_failed in self.failed_nodes.items():
220 action_name = action.split("|")[0]
221 if action_name == "Mutation" or action_name == "Query":
222 if num_failed > 0:
223 number_failed_of_mutations_and_queries += 1
224 return number_failed_of_mutations_and_queries, num_mutations_and_queries
226 def print_results(self):
227 print("\n----------------------RESULTS-------------------------")
228 print("Unique success nodes:")
229 pprint.pprint(self.successful_nodes)
230 print("Unique failed nodes:")
231 pprint.pprint(self.failed_nodes)
232 number_success_of_mutations_and_queries, num_mutations_and_queries = self.get_number_of_successful_mutations_and_queries()
233 number_failed_of_mutations_and_queries, num_mutations_and_queries = self.get_number_of_failed_mutations_and_queries()
234 print(f"(RESULTS): Time taken: {time.time() - self.start_time} seconds")
235 print(f"(RESULTS): Number of queries: {self.number_of_queries}")
236 print(f"(RESULTS): Number of mutations: {self.number_of_mutations}")
237 print(f"(RESULTS): Number of objects: {self.number_of_objects}")
238 print(f"(RESULTS): Number of unique query/mutation successes: {number_success_of_mutations_and_queries}/{num_mutations_and_queries}")
239 print(f"(RESULTS): Number of unique external query/mutation failures: {number_failed_of_mutations_and_queries}/{num_mutations_and_queries}")
240 print(f"(RESULTS): Please check {self.file_path} for more information regarding the run")
241 if len(self.vulnerabilities) > 0:
242 print("----------------------DETECTED VULNS-------------------------")
243 print(self.get_formatted_vulnerabilites())
244 print("---------------------------------------------------------")
246 def save(self):
247 """Saves the stats into the stats text file
248 """
249 number_success_of_mutations_and_queries, num_mutations_and_queries = self.get_number_of_successful_mutations_and_queries()
250 number_failed_of_mutations_and_queries, num_mutations_and_queries = self.get_number_of_failed_mutations_and_queries()
251 with open(self.file_path, "w") as f:
252 f.write("\n===================HTTP Status Codes===================\n")
253 f.write(json.dumps(self.http_status_codes, indent=4))
254 f.write("\n===================Successful Nodes===================\n")
255 f.write(json.dumps(self.successful_nodes, indent=4))
256 f.write("\n===================Failed Nodes===================\n")
257 f.write(json.dumps(self.failed_nodes, indent=4))
258 f.write("\n===================General stats ===================\n")
259 f.write(f"\nTime taken: {str(time.time() - self.start_time)} seconds")
260 f.write(f"\nNumber of unique query/mutation successes: {number_success_of_mutations_and_queries}/{num_mutations_and_queries}")
261 f.write(f"\nNumber of unique external query/mutation failures: {number_failed_of_mutations_and_queries}/{num_mutations_and_queries}")
262 f.write(f"\nNumber of queries: {self.number_of_queries}")
263 f.write(f"\nNumber of mutations: {self.number_of_mutations}")
264 f.write(f"\nNumber of objects: {self.number_of_objects}")
265 f.write(f"\nNumber of successes: {self.number_of_successes}")
266 f.write(f"\nNumber of failures: {self.number_of_failures}")
267 if len(self.vulnerabilities) > 0:
268 f.write("\n===================Detected Vulnerabilities===================\n")
269 f.write(json.dumps(self.vulnerabilities, indent=4))
270 self.save_endpoint_results()
271 self.save_unique_response()
272 self.save_json()
274 # Saves the pickle file as well
275 self.__save_pickle()
277 def save_json(self):
278 """Saves a machine-readable JSON report alongside the text stats file"""
279 json_path = getattr(self, "json_file_path", None)
280 if json_path is None:
281 return
282 number_success_of_mutations_and_queries, num_mutations_and_queries = self.get_number_of_successful_mutations_and_queries()
283 number_failed_of_mutations_and_queries, _ = self.get_number_of_failed_mutations_and_queries()
284 report = {
285 "time_taken_seconds": time.time() - self.start_time,
286 "number_of_queries": self.number_of_queries,
287 "number_of_mutations": self.number_of_mutations,
288 "number_of_objects": self.number_of_objects,
289 "number_of_successes": self.number_of_successes,
290 "number_of_failures": self.number_of_failures,
291 "unique_successful_nodes": number_success_of_mutations_and_queries,
292 "unique_failed_nodes": number_failed_of_mutations_and_queries,
293 "total_nodes": num_mutations_and_queries,
294 "http_status_codes": self.http_status_codes,
295 "successful_nodes": self.successful_nodes,
296 "failed_nodes": self.failed_nodes,
297 "vulnerabilities": self.vulnerabilities,
298 "node_timings": self.node_timings,
299 }
300 with open(json_path, "w") as f:
301 json.dump(report, f, indent=4)
303 def save_endpoint_results(self):
304 """Reads the results, for each node in the node name -> results, create a directory for the
305 result type, then a file for the response code, and append the payload and the response to the file.
306 """
307 unique_results = {}
308 # Filter out for only unique results
309 for node_name, results in self.results.items():
310 # If the node name has slashes, replace them with underscores
311 node_name = node_name.replace("/", "_")
313 if os.name == "nt":
314 # Replace characters that are invalid in Windows filenames
315 node_name = re.sub(r'[\\/:*?"<>|]', "_", node_name)
317 for result in results:
318 result_type = "success" if result.success else "failure"
319 result_file_path = Path(self.endpoint_results_dir) / node_name / result_type / f"{result.status_code}"
321 payload_string = str(result.payload)
322 if result_file_path not in unique_results:
323 unique_results[result_file_path] = {payload_string: result.graphql_response}
324 else:
325 if payload_string not in unique_results[result_file_path]:
326 unique_results[result_file_path][payload_string] = result.graphql_response
328 # Write the unique results to the file
329 for result_file_path, payloads in unique_results.items():
330 intialize_file_if_not_exists(result_file_path)
331 for payload, response in payloads.items():
332 with open(result_file_path, "a") as f:
333 f.write("------------------Payload:-------------------\n")
334 f.write(f"{payload}\n")
335 f.write("------------------Response:-------------------\n")
336 f.write(f"{response}\n")
338 def save_unique_response(self):
339 """Saves the unique responses to a file"""
340 with open(Path(self.unique_responses_file_path), "w") as f:
341 for response, endpoints in self.unique_responses.items():
342 f.write(f"Response: {response}\n")
343 f.write(f"Endpoints: {endpoints}\n")
345 def __save_pickle(self):
346 """Saves the stats to a pickle file"""
347 self.pickle_save_path = get_or_create_file(self.pickle_save_path)
348 with open(self.pickle_save_path, "wb") as file:
349 pickle.dump(self, file)
351 def __load_pickle(self):
352 """Loads the stats from a pickle file"""
353 if self.pickle_save_path.exists():
354 with open(self.pickle_save_path, "rb") as file:
355 loaded_stats = pickle.load(file)
356 self.__dict__.update(loaded_stats.__dict__)