def setUp(self): """Put data in the local ndb table the tests to query from and set bandit selection environment variable.""" test_helpers.patch_environ(self) data = [] strategy1 = data_types.FuzzStrategyProbability() strategy1.strategy_name = 'corpus_mutations_ml_rnn,corpus_subset,' strategy1.probability = 0.33 strategy1.engine = 'afl' data.append(strategy1) strategy2 = data_types.FuzzStrategyProbability() strategy2.strategy_name = ('corpus_mutations_radamsa,corpus_subset,') strategy2.probability = 0.34 strategy2.engine = 'afl' data.append(strategy2) strategy3 = data_types.FuzzStrategyProbability() strategy3.strategy_name = ('corpus_subset,') strategy3.probability = 0.33 strategy3.engine = 'afl' data.append(strategy3) ndb.put_multi(data) distribution = fuzz_task.get_strategy_distribution_from_ndb() environment.set_value('USE_BANDIT_STRATEGY_SELECTION', True) environment.set_value('STRATEGY_SELECTION_DISTRIBUTION', distribution)
def _query_and_upload_strategy_probabilities(): """Uploads queried data into datastore. Calls query functions and uploads query results to datastore to use as new probabilities. Probabilities are based on new_edges feature.""" strategy_data = [] data = _query_multi_armed_bandit_probabilities() # TODO(mukundv): Update once we choose a temperature parameter for final # implementation. for row in data: curr_strategy = data_types.FuzzStrategyProbability() curr_strategy.strategy_name = str(row['strategy']) curr_strategy.probability_high_temperature = float( row['bandit_weight_high_temperature']) curr_strategy.probability_low_temperature = float( row['bandit_weight_low_temperature']) curr_strategy.probability_medium_temperature = float( row['bandit_weight_medium_temperature']) strategy_data.append(curr_strategy) ndb.delete_multi([ entity.key for entity in ndb_utils.get_all_from_model( data_types.FuzzStrategyProbability) ]) ndb.put_multi(strategy_data) _store_probabilities_in_bigquery(data)
def setUp(self): """Put data in the local ndb table the tests to query from and set bandit selection environment variable.""" test_helpers.patch_environ(self) data = [] strategy1 = data_types.FuzzStrategyProbability() strategy1.strategy_name = 'fork,corpus_subset,recommended_dict,' strategy1.probability_medium_temperature = 0.33 strategy1.probability_high_temperature = 0.33 strategy1.probability_low_temperature = 0.33 data.append(strategy1) strategy2 = data_types.FuzzStrategyProbability() strategy2.strategy_name = ('random_max_len,corpus_mutations_ml_rnn,' 'value_profile,recommended_dict,') strategy2.probability_medium_temperature = 0.34 strategy2.probability_high_temperature = 0.34 strategy2.probability_low_temperature = 0.34 data.append(strategy2) strategy3 = data_types.FuzzStrategyProbability() strategy3.strategy_name = ('corpus_mutations_radamsa,' 'random_max_len,corpus_subset,') strategy3.probability_medium_temperature = 0.33 strategy3.probability_high_temperature = 0.33 strategy3.probability_low_temperature = 0.33 data.append(strategy3) ndb.put_multi(data) distribution = fuzz_task.get_strategy_distribution_from_ndb() environment.set_value('USE_BANDIT_STRATEGY_SELECTION', True) environment.set_value('STRATEGY_SELECTION_DISTRIBUTION', distribution)
def execute(args): """Build keywords.""" count_diff = 0 query = data_types.Testcase.query().order(-data_types.Testcase.timestamp) for testcases in batcher.iterate(query, BATCH_SIZE): for testcase in testcases: before_testcase = to_dict(testcase) attribute_builder.populate(testcase) after_testcase = to_dict(testcase) diff = get_diff(before_testcase, after_testcase) if (count_diff % 10) == 0 and diff: print('Migrate (dry=%s) id:%s\n%s' % (not args.non_dry_run, testcase.key.id(), diff)) if diff: count_diff += 1 if args.non_dry_run: try: ndb.put_multi(testcases) except Exception: for testcase in testcases: try: testcase.put() except Exception: print('Error: %s %s' % (testcase.key.id(), sys.exc_info())) print('Done (count_diff=%d)' % count_diff)
def update_group_bug(group_id): """Update group bug information for a group.""" if not group_id: # No associated group, no work to do. Bail out. return testcases = get_testcases_in_group(group_id) if not testcases: # No group members found. Bail out. return group_bug_information = 0 for testcase in testcases: if not testcase.bug_information: continue issue_id = int(testcase.bug_information) if not group_bug_information: group_bug_information = issue_id else: group_bug_information = min(group_bug_information, issue_id) for testcase in testcases: testcase.group_bug_information = group_bug_information ndb.put_multi(testcases)
def update_mappings_for_fuzzer(fuzzer, mappings=None): """Clear existing mappings for a fuzzer, and replace them.""" if mappings is None: mappings = fuzzer.jobs query = data_types.FuzzerJob.query() query = query.filter(data_types.FuzzerJob.fuzzer == fuzzer.name) entities = ndb_utils.get_all_from_query(query) old_mappings = {} for entity in entities: old_mappings[(entity.job, entity.platform)] = entity new_mappings = [] for job_name in mappings: job = data_types.Job.query(data_types.Job.name == job_name).get() if not job: logs.log_error('An unknown job %s was selected for fuzzer %s.' % (job_name, fuzzer.name)) continue mapping = old_mappings.pop((job_name, job.platform), None) if mapping: continue mapping = data_types.FuzzerJob() mapping.fuzzer = fuzzer.name mapping.job = job_name mapping.platform = job.platform new_mappings.append(mapping) ndb.put_multi(new_mappings) ndb.delete_multi([m.key for m in list(old_mappings.values())])
def setUp(self): """Put data in the local ndb table the tests to query from and set bandit selection environment variable.""" test_helpers.patch_environ(self) os.environ['USE_BANDIT_STRATEGY_SELECTION'] = 'True' data = [] strategy1 = data_types.FuzzStrategyProbability() strategy1.strategy_name = 'fork,corpus_subset,recommended_dict,' strategy1.probability = 0.33 data.append(strategy1) strategy2 = data_types.FuzzStrategyProbability() strategy2.strategy_name = ('random_max_len,corpus_mutations_ml_rnn,' 'value_profile,recommended_dict,') strategy2.probability = .34 data.append(strategy2) strategy3 = data_types.FuzzStrategyProbability() strategy3.strategy_name = ('corpus_mutations_radamsa,' 'random_max_len,corpus_subset,') strategy3.probability = .33 data.append(strategy3) ndb.put_multi(data)
def _query_and_upload_strategy_probabilities(engine): """Uploads queried data into datastore. Calls query functions and uploads query results to datastore to use as new probabilities. Probabilities are based on new_edges feature.""" strategy_data = [] data = _query_multi_armed_bandit_probabilities(engine) logs.log('Queried distribution for {}.'.format(engine.name)) # TODO(mukundv): Update once we choose a temperature parameter for final # implementation. for row in data: curr_strategy = data_types.FuzzStrategyProbability() curr_strategy.strategy_name = str(row['strategy']) curr_strategy.probability = float(row['bandit_weight']) curr_strategy.engine = engine.name strategy_data.append(curr_strategy) query = data_types.FuzzStrategyProbability.query( data_types.FuzzStrategyProbability.engine == engine.name) ndb.delete_multi( [entity.key for entity in ndb_utils.get_all_from_query(query)]) ndb.put_multi(strategy_data) logs.log('Uploaded queried distribution to ndb for {}'.format(engine.name)) _store_probabilities_in_bigquery(engine, data) logs.log('Uploaded queried distribution to BigQuery for {}'.format( engine.name))
def update_platform_for_job(job_name, new_platform): """Update platform for all mappings for a particular job.""" query = data_types.FuzzerJob.query() query = query.filter(data_types.FuzzerJob.job == job_name) mappings = ndb_utils.get_all_from_query(query) new_mappings = [] for mapping in mappings: mapping.platform = new_platform new_mappings.append(mapping) ndb.put_multi(new_mappings)
def assign_hosts_to_workers(self): """Assign host instances to workers.""" all_host_names = set() for assignment in self.gce_project.host_worker_assignments: host_cluster = self.gce_project.get_cluster(assignment.host) worker_cluster = self.gce_project.get_cluster(assignment.worker) if host_cluster.gce_zone != worker_cluster.gce_zone: logging.error('Mismatching zones for %s and %s.', assignment.host, assignment.worker) continue if (host_cluster.instance_count * assignment.workers_per_host != worker_cluster.instance_count): logging.error( 'Invalid host/worker cluster size for %s and %s.', assignment.host, assignment.worker) continue if host_cluster.high_end != worker_cluster.high_end: logging.error('Mismatching high end setting for %s and %s', assignment.host, assignment.worker) continue manager = bot_manager.BotManager(self.gce_project.project_id, host_cluster.gce_zone) host_instance_group = manager.instance_group(host_cluster.name) if not host_instance_group.exists(): logging.error('Host instance group %s does not exist.', host_cluster.name) continue host_names = [ _instance_name_from_url(instance['instance']) for instance in host_instance_group.list_managed_instances() ] all_host_names.update(host_names) worker_instances = self.get_all_workers_in_cluster( manager, worker_cluster.name) if len(worker_instances) != worker_cluster.instance_count: logging.error( 'Actual number of worker instances for %s did not match. ' 'Expected %d, got %d.', worker_cluster.name, worker_cluster.instance_count, len(worker_instances)) continue new_assignments = self.do_assign_hosts_to_workers( host_names, worker_instances, assignment.workers_per_host) ndb.put_multi(new_assignments) self.cleanup_old_assignments(all_host_names)
def update_job_weight(job_name, multiplier): """Update a job weight.""" tool_name = environment.get_memory_tool_name(job_name) multiplier *= SANITIZER_WEIGHTS.get(tool_name, DEFAULT_SANITIZER_WEIGHT) query = data_types.FuzzerJob.query(data_types.FuzzerJob.job == job_name) changed_weights = [] for fuzzer_job in query: if fuzzer_job.multiplier != multiplier: fuzzer_job.multiplier = multiplier changed_weights.append(fuzzer_job) if changed_weights: ndb.put_multi(changed_weights)
def setUp(self): """Put data in the local ndb table the tests to query from.""" test_helpers.patch_environ(self) test_helpers.patch( self, ['bot.fuzzers.engine_common.decide_with_probability']) self.mock.decide_with_probability.return_value = True os.environ['USE_BANDIT_STRATEGY_SELECTION'] = 'True' data = [] strategy1 = data_types.FuzzStrategyProbability() strategy1.strategy_name = ('random_max_len,corpus_mutations_ml_rnn,' 'value_profile,recommended_dict,') strategy1.probability = 1 data.append(strategy1) ndb.put_multi(data)
def setUp(self): """Put data in the local ndb table the tests to query from.""" test_helpers.patch_environ(self) test_helpers.patch( self, ['bot.fuzzers.engine_common.decide_with_probability']) self.mock.decide_with_probability.return_value = True data = [] strategy1 = data_types.FuzzStrategyProbability() strategy1.strategy_name = 'corpus_mutations_ml_rnn,corpus_subset,' strategy1.probability = 1 strategy1.engine = 'afl' data.append(strategy1) ndb.put_multi(data) distribution = fuzz_task.get_strategy_distribution_from_ndb() environment.set_value('USE_BANDIT_STRATEGY_SELECTION', True) environment.set_value('STRATEGY_SELECTION_DISTRIBUTION', distribution)
def _query_and_upload_strategy_probabilities(): """Uploads queried data into datastore. Calls query functions and uploads query results to datastore to use as new probabilities. Probabilities are based on new_edges feature.""" strategy_data = [] data = _query_multi_armed_bandit_probabilities() for row in data: curr_strategy = data_types.FuzzStrategyProbability() curr_strategy.strategy_name = str(row['strategy']) curr_strategy.probability = float(row['bandit_weight']) strategy_data.append(curr_strategy) ndb.delete_multi([ entity.key for entity in ndb_utils.get_all_from_model( data_types.FuzzStrategyProbability) ]) ndb.put_multi(strategy_data) _store_probabilities_in_bigquery(data)
def setUp(self): """Put data in the local ndb table the tests to query from.""" test_helpers.patch_environ(self) test_helpers.patch( self, ['bot.fuzzers.engine_common.decide_with_probability']) self.mock.decide_with_probability.return_value = True data = [] strategy1 = data_types.FuzzStrategyProbability() strategy1.strategy_name = ('random_max_len,corpus_mutations_ml_rnn,' 'value_profile,recommended_dict,') strategy1.probability_medium_temperature = 1 strategy1.probability_high_temperature = 1 strategy1.probability_low_temperature = 1 data.append(strategy1) ndb.put_multi(data) distribution = fuzz_task.get_strategy_distribution_from_ndb() environment.set_value('USE_BANDIT_STRATEGY_SELECTION', True) environment.set_value('STRATEGY_SELECTION_DISTRIBUTION', distribution)
def update_admins(new_admins): """Update list of admins.""" existing_admins = ndb_utils.get_all_from_model(data_types.Admin) to_remove = [] existing_admin_emails = set() for admin in existing_admins: if admin.email not in new_admins: logs.log('Removing admin ' + admin.email) to_remove.append(admin.key) existing_admin_emails.add(admin.email) ndb.delete_multi(to_remove) to_add = [] for admin in new_admins: if admin not in existing_admin_emails: to_add.append(data_types.Admin(id=admin, email=admin)) logs.log('Adding admin ' + admin) ndb.put_multi(to_add)
def _process_project(project, bucket): """Collects coverage information for all fuzz targets in the given project and the total stats for the project.""" project_name = _basename(project) logs.log('Processing coverage for %s project.' % project_name) report_path = storage.get_cloud_storage_file_path(bucket, project) report_info = _read_json(report_path) # Iterate through report_info['fuzzer_stats_dir'] and prepare # CoverageInformation entities for invididual fuzz targets. entities = [] for fuzzer in storage.list_blobs( report_info['fuzzer_stats_dir'], recursive=False): entities.append( _process_fuzzer_stats(fuzzer, report_info, project_name, bucket)) logs.log('Processed coverage for %d targets in %s project.' % (len(entities), project_name)) # Prepare CoverageInformation entity for the total project stats. entities.append(_process_project_stats(report_info, project_name)) ndb.put_multi(entities)
def update_project_cpus(self): """Update CPU allocations for each project.""" all_projects = list(data_types.OssFuzzProject.query().order( data_types.OssFuzzProject.name)) self.cleanup_old_projects([project.name for project in all_projects]) projects = [ project for project in all_projects if not project.high_end ] high_end_projects = [ project for project in all_projects if project.high_end ] project_infos = [ self.get_or_create_project_info(project.name) for project in projects ] high_end_project_infos = [ self.get_or_create_project_info(project.name) for project in high_end_projects ] for project, project_info in itertools.chain( list(zip(projects, project_infos)), list(zip(high_end_projects, high_end_project_infos))): self.cleanup_clusters(project, project_info) self.start_thread_pool() # Calculate CPUs in each cluster. for cluster in self.gce_project.clusters: if not cluster.distribute: self.pending_updates.append( self.thread_pool.submit(self.update_cluster, cluster, cluster.name, cluster.instance_count)) continue if cluster.high_end: current_projects = high_end_projects current_project_infos = high_end_project_infos else: current_projects = projects current_project_infos = project_infos cpu_counts = self.distribute_cpus(current_projects, cluster.instance_count) # Resize projects starting with ones that reduce number of CPUs. This is # so that we always have quota when we're resizing a project cluster. # pylint: disable=cell-var-from-loop def _cpu_diff_key(index): cluster_info = current_project_infos[index].get_cluster_info( cluster.name) if cluster_info and cluster_info.cpu_count is not None: old_cpu_count = cluster_info.cpu_count else: old_cpu_count = 0 return cpu_counts[index] - old_cpu_count resize_order = sorted(list(range(len(cpu_counts))), key=_cpu_diff_key) for i in resize_order: project = current_projects[i] project_info = current_project_infos[i] self.update_project_cluster(project, project_info, cluster, cpu_counts[i], disk_size_gb=project.disk_size_gb) self.finish_updates() ndb.put_multi(project_infos) ndb.put_multi(high_end_project_infos)