def populate_sample_data(session): team = Team(name="Blue Team 1", color="Blue") session.add(team) service = Service(name="Example Service 1", team=team, check_name="ICMP IPv4 Check", host='127.0.0.1') session.add(service) round_1 = Round(number=1) session.add(round_1) check_1 = Check(service=service, result=True, output='Good output', round=round_1) session.add(check_1) round_2 = Round(number=2) session.add(round_2) check_2 = Check(service=service, result=False, output='Bad output', round=round_2) session.add(check_2) session.commit() return team
def test_get_last_round_num_rounds_exist(self): service = generate_sample_model_tree('Service', self.db) round_5 = Round(number=5) self.db.save(round_5) round_6 = Round(number=6) self.db.save(round_6) assert Round.get_last_round_num() == 6
def test_get_last_round_num_rounds_exist(self): generate_sample_model_tree('Service', self.session) round_5 = Round(number=5) self.session.add(round_5) round_6 = Round(number=6) self.session.add(round_6) self.session.commit() assert Round.get_last_round_num() == 6
def test_checks_reversed(self): service = generate_sample_model_tree('Service', self.db) round_obj_1 = Round(number=1) round_obj_2 = Round(number=2) round_obj_3 = Round(number=3) self.db.save(round_obj_1) self.db.save(round_obj_2) self.db.save(round_obj_3) check_1 = Check(round=round_obj_1, service=service) self.db.save(check_1) check_2 = Check(round=round_obj_2, service=service) self.db.save(check_2) check_3 = Check(round=round_obj_3, service=service) self.db.save(check_3) assert service.checks_reversed == [check_3, check_2, check_1]
def update_team_score(team_id: int, first_round: int, last_round: int, add: bool = False) -> None: """Update the scores of a specific team. Note that it doesn't really make sense to perform an update that doesn't continue all the way to the most recent round of scoring, since that will be used as the basis for the team's score for the next round. :param team_id: The ID of the team whose scores should be updated. :type team_id: int :param first_round: The first round that scores will be recalculated for. :type first_round: int :param last_round: The last round (inclusive) that scores will be recalculated for. :type last_round: int :param add: Whether to add new score objects if they don't exist. Defaults to False. :tpe add: bool """ # Validate the first and last round if first_round < 1: first_round = 1 max_round = Round.get_last_round_num( ) # save value so we only it database once if last_round > max_round: last_round = max_round # Get score of previous round score = 0 if first_round > 1: prev_round = session.query(Round).filter_by(number=(first_round - 1)).first() prev_score = (session.query(Score).filter_by( team_id=team_id, round_id=prev_round.id).first()) score = prev_score.value # Get all services for the team services = session.query(Service).filter_by(team_id=team_id).all() # Iterate through each round for round_num in range(first_round, last_round + 1): round_obj = session.query(Round).filter_by(number=round_num).first() # Determine the check result for each service for service_obj in services: # Get the service's check check_obj = (session.query(Check).filter_by( service_id=service_obj.id, round_id=round_obj.id).first()) if check_obj.result: score += service_obj.points # Update the round's score score_obj = (session.query(Score).filter_by( team_id=team_id, round_id=round_obj.id).first()) if (score_obj is None) and add: score_obj = Score(value=score, team_id=team_id, round=round_obj) session.add(score_obj) elif score_obj is not None: score_obj.value = score session.commit()
def test_basic_round(self): round_obj = Round(number=5) self.session.add(round_obj) self.session.commit() assert round_obj.id is not None assert round_obj.number == 5 assert type(round_obj.local_round_start) is str
def __init__(self, total_rounds=0, round_time_sleep=180, worker_wait_time=30): self.checks = [] self.total_rounds = total_rounds self.round_time_sleep = round_time_sleep self.worker_wait_time = worker_wait_time self.config = config self.checks_location = self.config.checks_location self.checks_class_list = self.config.checks_class_list self.last_round = False self.rounds_run = 0 signal.signal(signal.SIGINT, partial(engine_sigint_handler, obj=self)) signal.signal(signal.SIGTERM, partial(engine_sigint_handler, obj=self)) self.db = DB() self.db.connect() self.current_round = Round.get_last_round_num() self.load_checks() self.round_running = False
def generate_sample_model_tree(model, session): # Team team = Team(name="Team 1", color="Blue") session.add(team) session.commit() if model == 'Team': return team # Users user = User(username="******" + str(random.randrange(10000)), password="******", team=team) session.add(user) session.commit() if model == 'User': return user # Services service = Service(name="ICMP IPv4", team=team, check_name="ICMP IPv4 Check", host='127.0.0.1') session.add(service) session.commit() if model == 'Service': return service # Environments environment = Environment(service=service, matching_regex='*') session.add(environment) session.commit() if model == 'Environment': return environment # Properties property_obj = Property(name="testproperty", value="testvalue", environment=environment) session.add(property_obj) session.commit() if model == 'Property': return property_obj # Rounds round_obj = Round(number=1) session.add(round_obj) session.commit() if model == 'Round': return round_obj # Checks check = Check(round=round_obj, service=service) session.add(check) session.commit() if model == 'Check': return check
def test_checks(self): service = generate_sample_model_tree('Service', self.db) round_obj = Round(number=5) self.db.save(round_obj) check_1 = Check(round=round_obj, service=service) self.db.save(check_1) check_2 = Check(round=round_obj, service=service) self.db.save(check_2) check_3 = Check(round=round_obj, service=service) self.db.save(check_3) assert round_obj.checks == [check_1, check_2, check_3]
def test_check_result_for_round_3_rounds(self): service = generate_sample_model_tree('Service', self.db) round_1 = Round(number=1) self.db.save(round_1) check_1 = Check(round=round_1, result=True, service=service) self.db.save(check_1) round_2 = Round(number=2) self.db.save(round_2) check_2 = Check(round=round_2, result=True, service=service) self.db.save(check_2) round_3 = Round(number=3) self.db.save(round_3) check_3 = Check(round=round_3, result=False, service=service) self.db.save(check_3) assert service.check_result_for_round(1) is True assert service.check_result_for_round(2) is True assert service.check_result_for_round(3) is False
def test_checks(self): service = generate_sample_model_tree('Service', self.session) round_obj = Round(number=5) self.session.add(round_obj) check_1 = Check(round=round_obj, service=service) self.session.add(check_1) check_2 = Check(round=round_obj, service=service) self.session.add(check_2) check_3 = Check(round=round_obj, service=service) self.session.add(check_3) self.session.commit() assert round_obj.checks == [check_1, check_2, check_3]
def admin_get_engine_stats(): if current_user.is_white_team: engine_stats = {} engine_stats['round_number'] = Round.get_last_round_num() engine_stats['num_passed_checks'] = Check.query.filter_by( result=True).count() engine_stats['num_failed_checks'] = Check.query.filter_by( result=False).count() engine_stats['total_checks'] = Check.query.count() return jsonify(engine_stats) else: return {'status': 'Unauthorized'}, 403
def admin_get_engine_stats(): if current_user.is_white_team: engine_stats = {} engine_stats["round_number"] = Round.get_last_round_num() engine_stats["num_passed_checks"] = (session.query(Check).filter_by( result=True).count()) engine_stats["num_failed_checks"] = (session.query(Check).filter_by( result=False).count()) engine_stats["total_checks"] = session.query(Check).count() return jsonify(engine_stats) else: return {"status": "Unauthorized"}, 403
def test_basic_check(self): service = generate_sample_model_tree('Service', self.db) round_obj = Round(number=1) self.db.save(round_obj) check = Check(round=round_obj, service=service, result=True, output="example_output") self.db.save(check) assert check.id is not None assert check.round == round_obj assert check.round_id == round_obj.id assert check.service == service assert check.service_id == service.id assert check.result is True assert check.output == 'example_output' assert check.completed is False
def overview_get_data(): # columns = get_table_columns() data = [] blue_teams = Team.get_all_blue_teams() last_round = Round.get_last_round_num() current_scores = ["Current Score"] current_places = ["Current Place"] service_ratios = ["Up/Down Ratio"] if len(blue_teams) > 0: for blue_team in blue_teams: num_up_services = (session.query( Service.team_id, ).join(Check).join(Round).filter( Check.result.is_(True)).filter( Service.team_id == blue_team.id).filter( Round.number == last_round).count()) num_down_services = (session.query( Service.team_id, ).join(Check).join(Round).filter( Check.result.is_(False)).filter( Service.team_id == blue_team.id).filter( Round.number == last_round).count()) current_scores.append(str(blue_team.current_score)) current_places.append(str(blue_team.place)) service_ratios.append("{0} / {1}".format(num_up_services, num_down_services)) data.append(current_scores) data.append(current_places) data.append(service_ratios) services = [] services.append([ service[0] for service in session.query(Service.name).distinct( Service.name).group_by(Service.name).all() ]) for blue_team in blue_teams: checks = (session.query(Check.result).join(Service).filter( Check.round_id == last_round).filter( Service.team_id == blue_team.id).order_by( Service.name).all()) services.append([check[0] for check in checks]) data += list( zip(*services)) # zip these to get the right datatables format return json.dumps({"data": data}) else: return "{}"
def test_finished(self): service = generate_sample_model_tree('Service', self.db) round_obj = Round(number=1) self.db.save(round_obj) check = Check(round=round_obj, service=service) self.db.save(check) assert check.result is None assert check.output == '' assert check.completed is False assert check.reason == '' check.finished(True, 'Successful Match', 'good output', 'example command') self.db.save(check) assert check.result is True assert check.output == 'good output' assert check.reason == 'Successful Match' assert check.command == 'example command' assert check.completed is True assert type(check.local_completed_timestamp) is str
def __init__(self, total_rounds=0): self.checks = [] self.total_rounds = total_rounds self.config = config self.checks_location = self.config.checks_location self.verify_settings() self.last_round = False self.rounds_run = 0 signal.signal(signal.SIGINT, partial(engine_sigint_handler, obj=self)) signal.signal(signal.SIGTERM, partial(engine_sigint_handler, obj=self)) self.session = session self.current_round = Round.get_last_round_num() self.load_checks() self.round_running = False
def get_round_scores(self, round_num: int) -> int: """Get the number of points a team earned during a specific round. :param round_num: The number of the round to check. :type round_num: int :return: How many points the team earned during the specified round. :rtype: int """ # TestTeam.test_get_round_scores requires we raise this exception if round_num > Round.get_last_round_num(): raise IndexError() # Get the current score and the previous score so we can subtract the two round_score = 0 prev_round_score = 0 for score in self.scores: if score.round.number == round_num: round_score = score.value elif score.round.number == round_num - 1: prev_round_score = score.value return round_score - prev_round_score
def test_current_score(self): team = generate_sample_model_tree('Team', self.session) round_obj = Round(number=1) self.session.add(round_obj) self.session.commit() service_1 = Service(name="Example Service 1", team=team, check_name="ICMP IPv4 Check", host='127.0.0.1') self.session.add(service_1) check_1 = Check(service=service_1, result=True, output='Good output', round=round_obj) self.session.add(check_1) service_2 = Service(name="Example Service 2", team=team, check_name="SSH IPv4 Check", host='127.0.0.2') self.session.add(service_2) check_2 = Check(service=service_2, result=True, output='Good output', round=round_obj) self.session.add(check_2) service_3 = Service(name="Example Service 3", team=team, check_name="SSH IPv4 Check", host='127.0.0.3') self.session.add(service_3) check_3 = Check(service=service_3, result=False, output='bad output', round=round_obj) self.session.add(check_3) self.session.commit() util.update_team_score(team.id, 1, 1, add=True) assert team.current_score == 200
def update_scores() -> List[int]: """Update the scores of any teams that are listed in the teams_to_update setting. :return: A list of the IDs for teams whose scores were updated. :rtype: List[int] """ # Get the list of team IDs teams_to_update = session.query(Setting).filter_by( name="teams_to_update").first() if teams_to_update is None: return [] # If the setting doesn't exist, skip it team_string = teams_to_update.value teams = [team_id for team_id in team_string.split(",")] for team in teams: # Update all rounds update_team_score(team, 1, Round.get_last_round_num()) # Clear the team queue teams_to_update = '' session.commit() return teams
def test_basic_round(self): round_obj = Round(number=5) self.db.save(round_obj) assert round_obj.id is not None assert round_obj.number == 5 assert type(round_obj.local_round_start) is str
def test_place(self): round_1 = Round(number=1) round_2 = Round(number=2) self.session.add(round_1) self.session.add(round_2) self.session.commit() team_1 = Team(name="Blue Team 1", color="Blue") self.session.add(team_1) service_1 = Service(name="Example Service 1", team=team_1, check_name="ICMP IPv4 Check", host='127.0.0.1') self.session.add(service_1) check_1 = Check(service=service_1, result=True, output='Good output', round=round_1) check_2 = Check(service=service_1, result=True, output='Good output', round=round_2) self.session.add(check_1) self.session.add(check_2) self.session.commit() team_2 = Team(name="Blue Team 2", color="Blue") self.session.add(team_2) service_1 = Service(name="Example Service 1", team=team_2, check_name="ICMP IPv4 Check", host='127.0.0.1') self.session.add(service_1) check_1 = Check(service=service_1, result=True, output='Good output', round=round_1) check_2 = Check(service=service_1, result=True, output='Good output', round=round_2) self.session.add(check_1) self.session.add(check_2) self.session.commit() team_3 = Team(name="Blue Team 3", color="Blue") self.session.add(team_3) service_1 = Service(name="Example Service 1", team=team_3, check_name="ICMP IPv4 Check", host='127.0.0.1') self.session.add(service_1) check_1 = Check(service=service_1, result=True, output='Good output', round=round_1) check_2 = Check(service=service_1, result=False, output='Good output', round=round_2) self.session.add(check_1) self.session.add(check_2) self.session.commit() util.update_team_score(team_1.id, 1, 2, add=True) util.update_team_score(team_2.id, 1, 2, add=True) util.update_team_score(team_3.id, 1, 2, add=True) assert team_1.place == 1 assert team_2.place == 1 assert team_3.place == 3
def run(self): while (not self.last_round) and (self.rounds_run < self.total_rounds or self.total_rounds == 0): self.current_round += 1 logger.info("Running round: " + str(self.current_round)) self.round_running = True self.rounds_run += 1 services = self.db.session.query(Service).all()[:] random.shuffle(services) task_ids = {} for service in services: check_class = self.check_name_to_obj(service.check_name) if check_class is None: raise LookupError( "Unable to map service to check code for " + str(service.check_name)) logger.info("Adding " + service.team.name + ' - ' + service.name + " check to queue") environment = random.choice(service.environments) check_obj = check_class(environment) command_str = check_obj.command() job = Job(environment_id=environment.id, command=command_str) task = execute_command.delay(job) team_name = environment.service.team.name if team_name not in task_ids: task_ids[team_name] = [] task_ids[team_name].append(task.id) # We store the list of tasks in the db, so that the web app # can consume them and can dynamically update a progress bar task_ids_str = json.dumps(task_ids) latest_kb = KB(name='task_ids', value=task_ids_str, round_num=self.current_round) self.db.save(latest_kb) pending_tasks = self.all_pending_tasks(task_ids) while pending_tasks: waiting_info = "Waiting for all jobs to finish (sleeping " + str( self.worker_wait_time) + " seconds)" waiting_info += " " + str( len(pending_tasks)) + " left in queue." logger.info(waiting_info) self.sleep(self.worker_wait_time) pending_tasks = self.all_pending_tasks(task_ids) logger.info("All jobs have finished for this round") logger.info("Determining check results and saving to db") round_obj = Round(number=self.current_round) self.db.save(round_obj) # We keep track of the number of passed and failed checks per round # so we can report a little bit at the end of each round teams = {} for team_name, task_ids in task_ids.items(): for task_id in task_ids: task = execute_command.AsyncResult(task_id) environment = self.db.session.query(Environment).get( task.result['environment_id']) if task.result['errored_out']: result = False reason = 'Task Timed Out' else: if re.search(environment.matching_regex, task.result['output']): result = True reason = "Successful Content Match" else: result = False reason = 'Unsuccessful Content Match' if environment.service.team.name not in teams: teams[environment.service.team.name] = { "Success": [], "Failed": [], } if result: teams[environment.service.team.name]['Success'].append( environment.service.name) else: teams[environment.service.team.name]['Failed'].append( environment.service.name) check = Check(service=environment.service, round=round_obj) check.finished(result=result, reason=reason, output=task.result['output'], command=task.result['command']) self.db.save(check) logger.info("Finished Round " + str(self.current_round)) logger.info("Round Stats:") for team_name in sorted(teams): stat_string = " " + team_name stat_string += " Success: " + str( len(teams[team_name]['Success'])) stat_string += ", Failed: " + str( len(teams[team_name]['Failed'])) if len(teams[team_name]['Failed']) > 0: stat_string += ' ' + str(teams[team_name]['Failed']) logger.info(stat_string) self.round_running = False if not self.last_round: logger.info("Sleeping in between rounds (" + str(self.round_time_sleep) + " seconds)") self.sleep(self.round_time_sleep)
def test_round_num(self): assert Round.get_last_round_num() == NUM_OVERALL_ROUNDS, \ "Expecting only {0} of rounds to have run...".format(NUM_OVERALL_ROUNDS)
def test_overview_data(self): # create 1 white team, 1 red team, 5 blue teams, 3 services per team, 5 checks per service # White Team self.db.save(Team(name="whiteteam", color="White")) # Red Team self.db.save(Team(name="redteam", color="Red")) teams = [] last_check_results = { "team1": { "FTPUpload": True, "DNS": True, "ICMP": True, }, "team2": { "FTPUpload": False, "DNS": True, "ICMP": True, }, "team3": { "FTPUpload": False, "DNS": True, "ICMP": False, }, "team4": { "FTPUpload": True, "DNS": False, "ICMP": False, }, "team5": { "FTPUpload": False, "DNS": False, "ICMP": False, }, } for team_num in range(1, 6): team = Team(name="team" + str(team_num), color="Blue") icmp_service = Service(name="ICMP", team=team, check_name="ICMP IPv4 Check", ip_address="127.0.0.1") dns_service = Service(name="DNS", team=team, check_name="DNSCheck", ip_address="8.8.8.8", port=53) ftp_upload_service = Service(name='FTPUpload', team=team, check_name='FTPUploadCheck', ip_address='1.2.3.4', port=21) self.db.save(icmp_service) self.db.save(dns_service) self.db.save(ftp_upload_service) # 5 rounds of checks round_1 = Round(number=1) icmp_check_1 = Check(round=round_1, service=icmp_service, result=True, output="example_output") dns_check_1 = Check(round=round_1, service=dns_service, result=False, output="example_output") ftp_upload_check_1 = Check(round=round_1, service=ftp_upload_service, result=True, output="example_output") self.db.save(round_1) self.db.save(icmp_check_1) self.db.save(dns_check_1) self.db.save(ftp_upload_check_1) round_2 = Round(number=2) icmp_check_2 = Check(round=round_2, service=icmp_service, result=True, output="example_output") dns_check_2 = Check(round=round_2, service=dns_service, result=True, output="example_output") ftp_upload_check_2 = Check(round=round_2, service=ftp_upload_service, result=True, output="example_output") self.db.save(round_2) self.db.save(icmp_check_2) self.db.save(dns_check_2) self.db.save(ftp_upload_check_2) round_3 = Round(number=3) icmp_check_3 = Check(round=round_3, service=icmp_service, result=True, output="example_output") dns_check_3 = Check(round=round_3, service=dns_service, result=False, output="example_output") ftp_upload_check_3 = Check(round=round_3, service=ftp_upload_service, result=True, output="example_output") self.db.save(round_3) self.db.save(icmp_check_3) self.db.save(dns_check_3) self.db.save(ftp_upload_check_3) round_4 = Round(number=4) icmp_check_4 = Check(round=round_4, service=icmp_service, result=False, output="example_output") dns_check_4 = Check(round=round_4, service=dns_service, result=False, output="example_output") ftp_upload_check_4 = Check(round=round_4, service=ftp_upload_service, result=False, output="example_output") self.db.save(round_4) self.db.save(icmp_check_4) self.db.save(dns_check_4) self.db.save(ftp_upload_check_4) round_5 = Round(number=5) icmp_check_5 = Check(round=round_5, service=icmp_service, result=last_check_results[team.name]['ICMP'], output="example_output") dns_check_5 = Check(round=round_5, service=dns_service, result=last_check_results[team.name]['DNS'], output="example_output") ftp_upload_check_5 = Check(round=round_5, service=ftp_upload_service, result=last_check_results[team.name]['FTPUpload'], output="example_output") self.db.save(round_5) self.db.save(icmp_check_5) self.db.save(dns_check_5) self.db.save(ftp_upload_check_5) self.db.save(team) teams.append(team) overview_data = self.client.get('/api/overview/data') overview_dict = json.loads(overview_data.data.decode('utf8')) expected_dict = { 'team1': { 'FTPUpload': {'passing': True, 'ip_address': '1.2.3.4', 'port': 21}, 'DNS': {'passing': True, 'ip_address': '8.8.8.8', 'port': 53}, 'ICMP': {'passing': True, 'ip_address': '127.0.0.1', 'port': 0}}, 'team2': { 'FTPUpload': {'passing': False, 'ip_address': '1.2.3.4', 'port': 21}, 'DNS': {'passing': True, 'ip_address': '8.8.8.8', 'port': 53}, 'ICMP': {'passing': True, 'ip_address': '127.0.0.1', 'port': 0}}, 'team3': { 'FTPUpload': {'passing': False, 'ip_address': '1.2.3.4', 'port': 21}, 'DNS': {'passing': True, 'ip_address': '8.8.8.8', 'port': 53}, 'ICMP': {'passing': False, 'ip_address': '127.0.0.1', 'port': 0}}, 'team4': { 'FTPUpload': {'passing': True, 'ip_address': '1.2.3.4', 'port': 21}, 'DNS': {'passing': False, 'ip_address': '8.8.8.8', 'port': 53}, 'ICMP': {'passing': False, 'ip_address': '127.0.0.1', 'port': 0}}, 'team5': { 'FTPUpload': {'passing': False, 'ip_address': '1.2.3.4', 'port': 21}, 'DNS': {'passing': False, 'ip_address': '8.8.8.8', 'port': 53}, 'ICMP': {'passing': False, 'ip_address': '127.0.0.1', 'port': 0} } } assert sorted(overview_dict.items()) == sorted(expected_dict.items())
def test_get_last_round_num_non_exist(self): assert Round.get_last_round_num() == 0
def run(self): if self.total_rounds == 0: logger.info("Running engine for unlimited rounds") else: logger.info("Running engine for {0} round(s)".format(self.total_rounds)) while not self.is_last_round(): self.current_round += 1 logger.info("Running round: " + str(self.current_round)) self.round_running = True self.rounds_run += 1 services = self.session.query(Service).all()[:] random.shuffle(services) task_ids = {} for service in services: check_class = self.check_name_to_obj(service.check_name) if check_class is None: raise LookupError("Unable to map service to check code for " + str(service.check_name)) logger.debug("Adding " + service.team.name + ' - ' + service.name + " check to queue") environment = random.choice(service.environments) check_obj = check_class(environment) command_str = check_obj.command() job = Job(environment_id=environment.id, command=command_str) task = execute_command.apply_async(args=[job], queue=service.worker_queue) team_name = environment.service.team.name if team_name not in task_ids: task_ids[team_name] = [] task_ids[team_name].append(task.id) # This array keeps track of all current round objects # incase we need to backout any changes to prevent # inconsistent check results cleanup_items = [] try: # We store the list of tasks in the db, so that the web app # can consume them and can dynamically update a progress bar task_ids_str = json.dumps(task_ids) latest_kb = KB(name='task_ids', value=task_ids_str, round_num=self.current_round) cleanup_items.append(latest_kb) self.session.add(latest_kb) self.session.commit() pending_tasks = self.all_pending_tasks(task_ids) while pending_tasks: worker_refresh_time = int(Setting.get_setting('worker_refresh_time').value) waiting_info = "Waiting for all jobs to finish (sleeping " + str(worker_refresh_time) + " seconds)" waiting_info += " " + str(len(pending_tasks)) + " left in queue." logger.info(waiting_info) self.sleep(worker_refresh_time) pending_tasks = self.all_pending_tasks(task_ids) logger.info("All jobs have finished for this round") logger.info("Determining check results and saving to db") round_obj = Round(number=self.current_round) cleanup_items.append(round_obj) self.session.add(round_obj) self.session.commit() # We keep track of the number of passed and failed checks per round # so we can report a little bit at the end of each round teams = {} # Used so we import the finished checks at the end of the round finished_checks = [] for team_name, task_ids in task_ids.items(): for task_id in task_ids: task = execute_command.AsyncResult(task_id) environment = self.session.query(Environment).get(task.result['environment_id']) if task.result['errored_out']: result = False reason = CHECK_TIMED_OUT_TEXT else: if re.search(environment.matching_content, task.result['output']): result = True reason = CHECK_SUCCESS_TEXT else: result = False reason = CHECK_FAILURE_TEXT if environment.service.team.name not in teams: teams[environment.service.team.name] = { "Success": [], "Failed": [], } if result: teams[environment.service.team.name]['Success'].append(environment.service.name) else: teams[environment.service.team.name]['Failed'].append(environment.service.name) check = Check(service=environment.service, round=round_obj) # Grab the first 35,000 characters of output so it'll fit into our TEXT column, # which maxes at 2^32 (65536) characters check.finished(result=result, reason=reason, output=task.result['output'][:35000], command=task.result['command']) finished_checks.append(check) for finished_check in finished_checks: cleanup_items.append(finished_check) self.session.add(finished_check) self.session.commit() except Exception as e: # We got an error while writing to db (could be normal docker stop command) # but we gotta clean up any trace of the current round so when we startup # again, we're at a consistent state logger.error('Error received while writing check results to db') logger.exception(e) logger.error('Ending round and cleaning up the db') for cleanup_item in cleanup_items: try: self.session.delete(cleanup_item) self.session.commit() except Exception: pass sys.exit(1) logger.info("Finished Round " + str(self.current_round)) logger.info("Round Stats:") for team_name in sorted(teams): stat_string = " " + team_name stat_string += " Success: " + str(len(teams[team_name]['Success'])) stat_string += ", Failed: " + str(len(teams[team_name]['Failed'])) if len(teams[team_name]['Failed']) > 0: stat_string += ' (' + ', '.join(teams[team_name]['Failed']) + ')' logger.info(stat_string) logger.info("Updating Caches") update_all_cache() self.round_running = False if not self.is_last_round(): round_time_sleep = int(Setting.get_setting('round_time_sleep').value) logger.info("Sleeping in between rounds (" + str(round_time_sleep) + " seconds)") self.sleep(round_time_sleep) logger.info("Engine finished running")
def test_init_round(self): round_obj = Round(number=5) assert round_obj.id is None assert round_obj.number == 5