1#!/usr/bin/env python3 2# 3# Copyright 2024 The Chromium Authors 4# Use of this source code is governed by a BSD-style license that can be 5# found in the LICENSE file. 6 7import argparse 8import collections 9import copy 10import os 11import sys 12import re 13 14 15def _GetDirAbove(dirname: str): 16 """Returns the directory "above" this file containing |dirname| (which must 17 also be "above" this file).""" 18 path = os.path.abspath(__file__) 19 while True: 20 path, tail = os.path.split(path) 21 if not tail: 22 return None 23 if tail == dirname: 24 return path 25 26 27SOURCE_DIR = _GetDirAbove('testing') 28 29sys.path.insert(1, os.path.join(SOURCE_DIR, 'third_party')) 30sys.path.append(os.path.join(SOURCE_DIR, 'build')) 31 32import action_helpers 33import jinja2 34 35_C_STR_TRANS = str.maketrans({ 36 '\n': '\\n', 37 '\r': '\\r', 38 '\t': '\\t', 39 '\"': '\\\"', 40 '\\': '\\\\' 41}) 42 43 44def c_escape(v: str) -> str: 45 return v.translate(_C_STR_TRANS) 46 47 48def main(): 49 parser = argparse.ArgumentParser( 50 description= 51 'Generate the necessary files for DomatoLPM to function properly.') 52 parser.add_argument('-p', 53 '--path', 54 required=True, 55 help='The path to the template file.') 56 parser.add_argument('-f', 57 '--file-format', 58 required=True, 59 help='The path (file format) where the generated files' 60 ' should be written to.') 61 parser.add_argument('-n', 62 '--name', 63 required=True, 64 help='The name of the fuzzer.') 65 66 parser.add_argument('-g', '--grammar', action='append') 67 68 parser.add_argument('-d', 69 '--generated-dir', 70 required=True, 71 help='The path to the target gen directory.') 72 73 args = parser.parse_args() 74 template_str = '' 75 with open(args.path, 'r') as f: 76 template_str = f.read() 77 78 grammars = [{ 79 'proto_type': repr.split(':')[0], 80 'proto_name': repr.split(':')[1] 81 } for repr in args.grammar] 82 grammar_types = [f'<{grammar["proto_type"]}>' for grammar in grammars] 83 84 # This splits the template into fuzzing tags, so that we know where we need 85 # to insert grammar results. 86 splitted_template = re.split('|'.join([f'({g})' for g in grammar_types]), 87 template_str) 88 splitted_template = [a for a in splitted_template if a is not None] 89 90 grammar_elements = [] 91 counter = collections.defaultdict(int) 92 for elt in splitted_template: 93 if elt in grammar_types: 94 g = next(g for g in grammars if f'<{g["proto_type"]}>' == elt) 95 g = copy.deepcopy(g) 96 g['is_str'] = False 97 counter[elt] += 1 98 c = counter[elt] 99 g['proto_field_name'] = f'{g["proto_name"]}{c}' 100 grammar_elements.append(g) 101 else: 102 grammar_elements.append({'is_str': True, 'content': c_escape(elt)}) 103 104 template_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 105 'templates') 106 environment = jinja2.Environment(loader=jinja2.FileSystemLoader(template_dir)) 107 rendering_context = { 108 'template_path': args.generated_dir, 109 'template_name': args.name, 110 'grammars': grammars, 111 'grammar_elements': grammar_elements, 112 } 113 template = environment.get_template('domatolpm_fuzzer.proto.tmpl') 114 with action_helpers.atomic_output(f'{args.file_format}.proto', mode='w') as f: 115 f.write(template.render(rendering_context)) 116 template = environment.get_template('domatolpm_fuzzer.h.tmpl') 117 with action_helpers.atomic_output(f'{args.file_format}.h', mode='w') as f: 118 f.write(template.render(rendering_context)) 119 template = environment.get_template('domatolpm_fuzzer.cc.tmpl') 120 with action_helpers.atomic_output(f'{args.file_format}.cc', mode='w') as f: 121 f.write(template.render(rendering_context)) 122 123 124if __name__ == '__main__': 125 main() 126