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
