diff --git a/butler.py b/butler.py index 2d2c886be52..644603b127e 100644 --- a/butler.py +++ b/butler.py @@ -132,6 +132,25 @@ def _add_weights_fuzzer_subparser(weights_subparsers): aggregate_parser.add_argument( '-j', '--jobs', help='Which jobs to aggregate.', nargs='+') + set_parser = subparsers.add_parser( + 'set', help='Set the weight of a FuzzerJob entry.') + set_parser.add_argument( + '-f', + '--fuzzer', + help='The fuzzer field of the entry to modify.', + required=True) + set_parser.add_argument( + '-j', + '--job', + help='The job field of the entry to modify.', + required=True) + set_parser.add_argument( + '-w', + '--weight', + help='The new weight to set.', + type=float, + required=True) + def _add_weights_batches_subparser(weights_subparsers): """Adds a parser for the `weights fuzzer-batch` command.""" diff --git a/src/local/butler/weights.py b/src/local/butler/weights.py index f80d6ddb468..5b42a8340ca 100644 --- a/src/local/butler/weights.py +++ b/src/local/butler/weights.py @@ -39,8 +39,8 @@ def _iter_weights( fuzzer_jobs: Sequence[data_types.FuzzerJob]) -> Sequence[float]: - for fj in fuzzer_jobs: - yield fj.actual_weight + for fuzzer_job in fuzzer_jobs: + yield fuzzer_job.actual_weight def _sum_weights(fuzzer_jobs: Sequence[data_types.FuzzerJob]) -> float: @@ -123,7 +123,8 @@ def _display_fuzzer_jobs(fuzzer_jobs: Sequence[data_types.FuzzerJob], printer = _print_with_prefix(prefix) fuzzer_jobs = list(fuzzer_jobs) - fuzzer_jobs.sort(key=lambda fj: fj.actual_weight, reverse=True) + fuzzer_jobs.sort( + key=lambda fuzzer_job: fuzzer_job.actual_weight, reverse=True) total_weight = _sum_weights(fuzzer_jobs) @@ -316,6 +317,45 @@ def _aggregate_fuzzer_jobs( _print_stats(others, total_weight) +def _set_fuzzer_job_weight( + fuzzer: str, + job: str, + weight: float, +) -> None: + """Sets the matching FuzzerJob's weight to the given value.""" + fuzzer_jobs = list( + data_types.FuzzerJob.query(data_types.FuzzerJob.fuzzer == fuzzer, + data_types.FuzzerJob.job == job)) + + if not fuzzer_jobs: + print('No matching FuzzerJob entries found for ' + + f'fuzzer {fuzzer} and job {job}') + return + + if len(fuzzer_jobs) > 1: + print('Bailing out! Multiple FuzzerJob entries found for ' + + f'fuzzer {fuzzer} and job {job}: {fuzzer_jobs}') + return + + fuzzer_job = fuzzer_jobs[0] + + print(f'Fuzzer: {fuzzer_job.fuzzer}') + print(f'Job: {fuzzer_job.job}') + print(f'Platform: {fuzzer_job.platform}') + print(f'Multiplier: {fuzzer_job.multiplier}') + print(f'Old weight: {fuzzer_job.weight}') + print(f'-> New weight: {weight}') + + answer = input('Do you want to apply this mutation? [y,n] ') + if answer.lower() != 'y': + print('Not applying mutation.') + return + + fuzzer_job.weight = weight + fuzzer_job.put() + print('Mutation applied.') + + def _set_fuzz_target_job_weight( fuzz_target_name: str, job: str, @@ -362,6 +402,8 @@ def _execute_fuzzer_command(args) -> None: raise TypeError(f'--format {repr(args.format)} unrecognized') elif cmd == 'aggregate': _aggregate_fuzzer_jobs(args.platform, fuzzers=args.fuzzers, jobs=args.jobs) + elif cmd == 'set': + _set_fuzzer_job_weight(args.fuzzer, args.job, args.weight) else: raise TypeError(f'weights fuzzer command {repr(cmd)} unrecognized')