Coverage for graphqler / graph / graph_generator.py: 91%

1"""GraphGenerator: Creates a networkx graph and stores it in a pickle file for use later on during fuzzing 

2The linker does the following: 

3- Serialize all the objects (Objects, Queries, Mutations, InputObjects, Enums) 

4- Generate a graph of object dependencies 

5- Attach queries to the object node 

6- Attach mutations related to the object node 

7 

8!Note!: We decide to not link object-objects together here as it is not relevant for graph traversal 

9""" 

10 

11from pathlib import Path 

12from graphqler.utils.file_utils import read_yaml_to_dict 

13from graphqler import config 

14from .node import Node 

15from .utils import draw_graph 

16 

17import networkx 

18 

19 

20class GraphGenerator: 

21 def __init__(self, save_path: str): 

22 self.save_path = save_path 

23 self.compiled_queries_save_path = Path(save_path) / config.COMPILED_QUERIES_FILE_NAME 

24 self.compiled_objects_save_path = Path(save_path) / config.COMPILED_OBJECTS_FILE_NAME 

25 self.compiled_mutations_save_path = Path(save_path) / config.COMPILED_MUTATIONS_FILE_NAME 

26 self.dependency_graph_visualization_save_path = Path(save_path) / config.GRAPH_VISUALIZATION_OUTPUT 

27 

28 self.compiled_queries = read_yaml_to_dict(self.compiled_queries_save_path) 

29 self.compiled_objects = read_yaml_to_dict(self.compiled_objects_save_path) 

30 self.compiled_mutations = read_yaml_to_dict(self.compiled_mutations_save_path) 

31 

32 self.dependency_graph = networkx.DiGraph() 

33 

34 def get_dependency_graph(self) -> networkx.DiGraph: 

35 """Runs the graph generator and returns the graph 

36 

37 Returns: 

38 networkx.DiGraph: The directed graph 

39 """ 

40 self.run() 

41 return self.dependency_graph 

42 

43 def draw_dependency_graph(self): 

44 """Draws the dependency graph based on the GRAPH_VISUALIZATION_OUTPUT constant""" 

45 draw_graph(self.dependency_graph, self.dependency_graph_visualization_save_path) 

46 

47 def run(self): 

48 """Generates the graph, creating nodes and creating edges between nodes. 

49 3 types of nodes (Objects, Queries, Mutations) 

50 """ 

51 

52 """1. Create query nodes""" 

53 query_nodes = {} 

54 for query_name, query_body in self.compiled_queries.items(): 

55 query_nodes[query_name] = Node("Query", query_name, query_body) 

56 

57 """2. Create mutation nodes""" 

58 mutation_nodes = {} 

59 for mutation_name, mutation_body in self.compiled_mutations.items(): 

60 mutation_node = Node("Mutation", mutation_name, mutation_body) 

61 mutation_node.set_mutation_type(mutation_body["mutationType"]) 

62 mutation_nodes[mutation_name] = mutation_node 

63 

64 """3. Create object nodes""" 

65 object_nodes = {} 

66 for object_name, object_body in self.compiled_objects.items(): 

67 object_nodes[object_name] = Node("Object", object_name, object_body) 

68 

69 """4. Add all nodes to the graph""" 

70 self.dependency_graph.add_nodes_from(query_nodes.values()) 

71 self.dependency_graph.add_nodes_from(mutation_nodes.values()) 

72 self.dependency_graph.add_nodes_from(object_nodes.values()) 

73 

74 """5. Link objects and mutations together""" 

75 self.create_object_mutation_edges(object_nodes, mutation_nodes) 

76 

77 """6. Link objects and queries together""" 

78 self.create_object_query_edges(object_nodes, query_nodes) 

79 

80 def create_object_mutation_edges(self, object_nodes: dict, mutation_nodes: dict): 

81 """Updates the dependency graph with edges between objects and mutations. 3 cases: 

82 Case 1: M -> O | When object(O) depends on mutation(M), means O has M in its "associatedMutations", weight 100 

83 Case 2: O -> M | When mutation(M) depends on object(O), means M has O in its "hardDependsOn", weight 100 

84 Case 3: O -> M | When mutation(M) depends on object(O), means M has O in its "softDependsOn", weight 1 

85 

86 Args: 

87 object_nodes (dict): Mapping of object_name -> object node 

88 mutation_nodes (dict): Mapping of mutation name -> mutation node 

89 """ 

90 # Case 1 

91 for object_name, object_node in object_nodes.items(): 

92 object_information = self.compiled_objects[object_name] 

93 if not object_information["associatedMutatations"]: 

94 continue # skip if this object doesn't have any associated mutations 

95 

96 for associated_mutation_name in object_information["associatedMutatations"]: 

97 mutation_node = mutation_nodes[associated_mutation_name] 

98 self.dependency_graph.add_edge(mutation_node, object_node, weight=100) 

99 

100 # Case 2 

101 for mutation_name, mutation_node in mutation_nodes.items(): 

102 mutation_information = self.compiled_mutations[mutation_name] 

103 if not mutation_information["hardDependsOn"]: 

104 continue # skip if this mutation doesn't have any hardDependsOn 

105 

106 if mutation_information["hardDependsOn"]: 

107 for input_name, object_name in mutation_information["hardDependsOn"].items(): 

108 if object_name != "UNKNOWN": 

109 object_node = object_nodes[object_name] 

110 self.dependency_graph.add_edge(object_node, mutation_node, weight=100) 

111 

112 # Case 3 

113 for mutation_name, mutation_node in mutation_nodes.items(): 

114 mutation_information = self.compiled_mutations[mutation_name] 

115 if not mutation_information["softDependsOn"]: 

116 continue # skip if this mutation doesn't have any hardDependsOn 

117 

118 if mutation_information["softDependsOn"]: 

119 for input_name, object_name in mutation_information["softDependsOn"].items(): 

120 if object_name != "UNKNOWN": 

121 object_node = object_nodes[object_name] 

122 self.dependency_graph.add_edge(object_node, mutation_node, weight=1) 

123 

124 def create_object_query_edges(self, object_nodes: dict, query_nodes: dict): 

125 """Updates the dependency graph with edges in between objects and queries. 3 cases: 

126 Case 1: M -> O | When object(O) is produced by query(Q), means O has Q in its "associatedQueries", weight 100 

127 Case 2: O -> Q | When query(Q) depends on object(O), means Q has O in its "hardDependsOn", weight 100 

128 Case 3: O -> Q | When query(Q) depends on object(O), means Q has O in its "softDependsOn", weight 1 

129 

130 Args: 

131 object_nodes (dict): Mapping of object_name -> object node 

132 query_nodes (dict): Mapping of query_name -> query node 

133 """ 

134 # Case 1 

135 for object_name, object_node in object_nodes.items(): 

136 object_information = self.compiled_objects[object_name] 

137 if not object_information["associatedQueries"]: 

138 continue # skip if this object doesn't have any associatedQueries 

139 

140 for associated_query_name in object_information["associatedQueries"]: 

141 query_node = query_nodes[associated_query_name] 

142 self.dependency_graph.add_edge(query_node, object_node, weight=100) 

143 

144 # Case 2 

145 for query_name, query_node in query_nodes.items(): 

146 query_information = self.compiled_queries[query_name] 

147 if not query_information["hardDependsOn"]: 

148 continue # skip if this querry doesn't have any hardDependsOn 

149 

150 for input_name, object_name in query_information["hardDependsOn"].items(): 

151 if object_name != "UNKNOWN": 

152 object_node = object_nodes[object_name] 

153 self.dependency_graph.add_edge(object_node, query_node, weight=100) 

154 

155 # Case 3 

156 for query_name, query_node in query_nodes.items(): 

157 query_information = self.compiled_queries[query_name] 

158 if not query_information["softDependsOn"]: 

159 continue # skip if this querry doesn't have any hardDependsOn 

160 

161 for input_name, object_name in query_information["softDependsOn"].items(): 

162 if object_name != "UNKNOWN": 

163 object_node = object_nodes[object_name] 

164 self.dependency_graph.add_edge(object_node, query_node, weight=1)