Coverage for graphqler / __main__.py: 36%

1""" 

2Graphler - main start 

3""" 

4 

5import sys 

6import argparse 

7import pprint 

8import importlib.metadata 

9import cloudpickle as pickle 

10 

11from graphqler.compiler.compiler import Compiler 

12from graphqler.fuzzer import Fuzzer, IDORFuzzer 

13from graphqler.graph import GraphGenerator 

14from graphqler.utils.stats import Stats 

15from graphqler.utils.cli_utils import set_auth_token_constant, is_compiled 

16from graphqler.utils.config_handler import parse_config, set_config, generate_new_config, does_config_file_exist_in_path 

17from graphqler.utils.file_utils import get_or_create_directory 

18from graphqler import config 

19 

20 

21def run_compile_mode(path: str, url: str): 

22 """Runs the full compilation pipeline by delegating to compile-graph then compile-chains. 

23 

24 Args: 

25 path (str): Directory for all compilation outputs to be saved to 

26 url (str): URL of the target 

27 """ 

28 print("(C) In compile mode!") 

29 run_compile_graph_mode(path, url) 

30 run_compile_chains_mode(path) 

31 print("(C) Complete compilation phase") 

32 

33 

34def run_compile_graph_mode(path: str, url: str): 

35 """Runs only the introspection / parsing / resolving steps and generates the dependency graph. 

36 

37 Use this when you want to regenerate the dependency graph without re-running 

38 chain generation, or when you plan to run ``compile-chains`` separately. 

39 

40 Args: 

41 path (str): Directory for all compilation outputs to be saved to 

42 url (str): URL of the target 

43 """ 

44 print("(C) In compile-graph mode!") 

45 Compiler(path, url).run() 

46 

47 print("(C) Finished compiling, starting graph generator") 

48 graph_generator = GraphGenerator(path) 

49 graph = graph_generator.get_dependency_graph() 

50 graph_generator.draw_dependency_graph() 

51 

52 print("(C) Found", len(graph.nodes), "nodes and", len(graph.edges), "edges") 

53 print("(C) Complete graph compilation phase (chains not generated)") 

54 

55 

56def run_compile_chains_mode(path: str): 

57 """Generates (or re-generates) fuzzing chains from an already-compiled graph. 

58 

59 Requires that ``compile`` or ``compile-graph`` has been run first so that the 

60 compiled YAML files and dependency graph are present on disk. 

61 

62 Args: 

63 path (str): Directory used during the original compilation. 

64 """ 

65 from graphqler.chains import ChainGenerator 

66 

67 print("(C) In compile-chains mode!") 

68 dependency_graph = GraphGenerator(path).get_dependency_graph() 

69 in_degrees = dict(dependency_graph.in_degree()) 

70 if not in_degrees: 

71 print("(C) Dependency graph is empty — no chains generated") 

72 return 

73 

74 min_degree = min(in_degrees.values()) 

75 starter_nodes = [node for node, degree in in_degrees.items() if degree == min_degree] 

76 

77 chain_generator = ChainGenerator() 

78 chain_generator.generate(dependency_graph, starter_nodes) 

79 chain_generator.save_to_yaml(path) 

80 print(f"(C) Generated {len(chain_generator.chains)} chains") 

81 print("(C) Chain generation complete") 

82 

83 

84def run_fuzz_mode(path: str, url: str): 

85 """Runs the program in fuzz mode 

86 

87 Args: 

88 path (str): Directory for all compilation outputs to be saved to 

89 url (str): URL of the target 

90 """ 

91 print("(F) Initializing stats file") 

92 stats = Stats() 

93 stats.set_file_paths(path) 

94 

95 print("(F) Starting fuzzer") 

96 if not config.USE_OBJECTS_BUCKET: 

97 print("(F) Not using Objects Bucket") 

98 

99 if config.USE_DEPENDENCY_GRAPH: 

100 print("(F) Running in dependency graph mode") 

101 Fuzzer(path, url).run() 

102 else: 

103 print("(F) Not using dependency graph") 

104 Fuzzer(path, url).run_no_dfs() 

105 

106 print("(F) Complete fuzzing phase") 

107 

108 

109def run_idor_mode(path: str, url: str): 

110 print("(F) Running IDOR fuzzer") 

111 try: 

112 with open(f"{path}/{config.OBJECTS_BUCKET_PICKLE_FILE_NAME}", "rb") as f: 

113 objects_bucket = pickle.load(f) 

114 possible_idor_nodes = IDORFuzzer(path, url, objects_bucket).run() 

115 print("Possible IDOR nodes:") 

116 pprint.pprint(possible_idor_nodes) 

117 except FileNotFoundError: 

118 print("(F) Error: objects_bucket.pkl not found") 

119 return 

120 

121 

122def run_single_mode(path: str, url: str, name: str): 

123 print("(F) Running single mode") 

124 Fuzzer(path, url).run_single(name) 

125 

126 

127def main(args: dict): 

128 # Run either compilation or fuzzing mode 

129 if 'mode' not in args or not args['mode']: 

130 print("Please provide a mode to run the program in") 

131 sys.exit(1) 

132 

133 # compile-chains works from disk — URL not needed; all other modes require it 

134 if args['mode'] != "compile-chains" and not args.get('url'): 

135 print(f"--url is required for mode '{args['mode']}'") 

136 sys.exit(1) 

137 

138 # If not compile mode, check if compiled directory exists 

139 if args['mode'] not in ["compile", "compile-graph", "compile-chains", "run", "single", "idor"] and (not is_compiled(args['path']) or not is_compiled(config.OUTPUT_DIRECTORY)): 

140 print("(!) Compiled directory does not exist, please run in compile mode first") 

141 sys.exit(1) 

142 

143 # Set the path if provided and create the directory if it doesn't exist 

144 if 'path' in args and args['path']: 

145 config.OUTPUT_DIRECTORY = args['path'] 

146 get_or_create_directory(config.OUTPUT_DIRECTORY) 

147 

148 # Set proxy if provided 

149 if 'proxy' in args and args['proxy']: 

150 config.PROXY = args['proxy'] 

151 

152 # Set auth token if provided 

153 if 'auth' in args and args['auth']: 

154 set_auth_token_constant(args['auth']) 

155 

156 # Parse config if provided 

157 if 'config' in args and args['config']: 

158 print("(P) Using provided config file") 

159 new_config = parse_config(args['config']) 

160 set_config(new_config) 

161 elif does_config_file_exist_in_path(args['path']): 

162 print("(P) Using config file in path") 

163 new_config = parse_config(f"{args['path']}/{config.CONFIG_FILE_NAME}") 

164 set_config(new_config) 

165 else: 

166 print("(P) Generating new config") 

167 generate_new_config(f"{args['path']}/{config.CONFIG_FILE_NAME}") 

168 

169 # Parse plugins if defined 

170 if 'plugins_path' in args and args['plugins_path']: 

171 config.PLUGINS_PATH = args['plugins_path'] 

172 print(f"(P) Using plugins from {config.PLUGINS_PATH}") 

173 

174 # Apply LLM CLI overrides — these take precedence over config file values 

175 if args.get('use_llm'): 

176 config.USE_LLM = True 

177 print("(P) LLM mode enabled via CLI flag") 

178 if args.get('llm_model'): 

179 config.LLM_MODEL = args['llm_model'] 

180 if args.get('llm_api_key'): 

181 config.LLM_API_KEY = args['llm_api_key'] 

182 if args.get('llm_base_url'): 

183 config.LLM_BASE_URL = args['llm_base_url'] 

184 if args.get('llm_max_retries') is not None: 

185 config.LLM_MAX_RETRIES = args['llm_max_retries'] 

186 

187 # Apply mutation CLI override 

188 if args.get('disable_mutations'): 

189 config.DISABLE_MUTATIONS = True 

190 print("(P) Mutation fuzzing disabled — only Query chains will be generated") 

191 

192 # Start the program 

193 if args['mode'] == "compile": 

194 run_compile_mode(config.OUTPUT_DIRECTORY, args['url']) 

195 elif args['mode'] == "compile-graph": 

196 run_compile_graph_mode(config.OUTPUT_DIRECTORY, args['url']) 

197 elif args['mode'] == "compile-chains": 

198 run_compile_chains_mode(config.OUTPUT_DIRECTORY) 

199 elif args['mode'] == "fuzz": 

200 run_fuzz_mode(config.OUTPUT_DIRECTORY, args['url']) 

201 elif args['mode'] == "run": 

202 run_compile_mode(config.OUTPUT_DIRECTORY, args['url']) 

203 run_fuzz_mode(config.OUTPUT_DIRECTORY, args['url']) 

204 elif args['mode'] == "idor": 

205 run_idor_mode(config.OUTPUT_DIRECTORY, args['url']) 

206 elif args['mode'] == "single": 

207 if 'node' not in args or not args['node']: 

208 print("Please provide a node to run in single mode") 

209 sys.exit(1) 

210 run_single_mode(args['path'], args['url'], args['node']) 

211 

212 

213# If running as a CLI 

214if __name__ == "__main__": 

215 # If version, display version and exit 

216 if "--version" in sys.argv: 

217 version = importlib.metadata.version("GraphQLer") 

218 print(version) 

219 sys.exit(0) 

220 

221 # Parse arguments 

222 parser = argparse.ArgumentParser() 

223 parser.add_argument("--url", help="remote host URL (required for all modes except compile-chains)", required=False) 

224 parser.add_argument("--path", help=f"directory location for files to be saved-to/used-from. Defaults to {config.OUTPUT_DIRECTORY}", required=False) 

225 parser.add_argument("--config", help="TOML configuration file for the program", required=False) 

226 parser.add_argument("--mode", help="mode to run the program in", choices=["compile", "compile-graph", "compile-chains", "fuzz", "idor", "run", "single"], required=True) 

227 parser.add_argument("--auth", help="authentication token Example: 'Bearer arandompat-abcdefgh'", required=False) 

228 parser.add_argument("--proxy", help="proxy to use for requests (ie. http://127.0.0.1:8080)", required=False) 

229 parser.add_argument("--node", help="node to run (only used in single mode)", required=False) 

230 parser.add_argument("--plugins-path", help="path to plugins directory", required=False) 

231 parser.add_argument("--use-llm", help="enable LLM-based dependency graph inference (requires LLM_MODEL and credentials)", action="store_true", default=False) 

232 parser.add_argument("--llm-model", help="litellm model string, e.g. 'gpt-4o-mini', 'ollama/llama3', 'anthropic/claude-3-5-haiku-20241022'", required=False) 

233 parser.add_argument("--llm-api-key", help="API key for the LLM provider (or set OPENAI_API_KEY / ANTHROPIC_API_KEY env var)", required=False) 

234 parser.add_argument("--llm-base-url", help="custom base URL for LLM endpoint (required for Ollama and LiteLLM proxies)", required=False) 

235 parser.add_argument("--llm-max-retries", help="number of retries when LLM returns non-JSON (default: 2)", type=int, required=False) 

236 parser.add_argument("--disable-mutations", help="only generate and run Query chains — all Mutation nodes are excluded from fuzzing", action="store_true", default=False) 

237 parser.add_argument("--version", help="display version", action="store_true") 

238 

239 args = parser.parse_args() 

240 args_as_dict = vars(args) 

241 

242 # Some massaging 

243 if args_as_dict['path'] is None: 

244 args_as_dict['path'] = config.OUTPUT_DIRECTORY 

245 main(args_as_dict)