def cache_last_stolen(team_id: int, round: int, pipeline):
"""Put stolen flags for attacker team from last "flag_lifetime" rounds to cache
:param team_id: attacker team id
:param round: current round
:param pipeline: redis connection to add command to
Just adds commands to pipeline stack, don't forget to execute afterwards
"""
game_config = config.get_game_config()
conn = storage.get_db_pool().getconn()
curs = conn.cursor()
curs.execute(_SELECT_LAST_STOLEN_TEAM_FLAGS_QUERY,
(round - game_config['flag_lifetime'], team_id))
flags = curs.fetchall()
curs.close()
storage.get_db_pool().putconn(conn)
pipeline.delete(f'team:{team_id}:cached:stolen',
f'team:{team_id}:stolen_flags')
if flags:
pipeline.sadd(f'team:{team_id}:stolen_flags',
*[flag_id for flag_id, in flags])
pipeline.set(f'team:{team_id}:cached:stolen', 1)
def cache_last_flags(round: int, pipeline):
"""Put all generated flags from last "flag_lifetime" rounds to cache
:param round: current round
:param pipeline: redis connection to add command to
Just adds commands to pipeline stack, don't forget to execute afterwards
"""
game_config = config.get_game_config()
conn = storage.get_db_pool().getconn()
curs = conn.cursor(cursor_factory=extras.DictCursor)
curs.execute(_SELECT_ALL_LAST_FLAGS_QUERY,
(round - game_config['flag_lifetime'], ))
flags = curs.fetchall()
curs.close()
storage.get_db_pool().putconn(conn)
pipeline.delete('flags:cached')
flag_models = []
for flag_dict in flags:
flag = helplib.models.Flag.from_dict(flag_dict)
flag_models.append(flag)
if flag_models:
pipeline.delete(*[
f'team:{flag.team_id}:task:{flag.task_id}:round_flags:{flag.round}'
for flag in flag_models
])
for flag in flag_models:
pipeline.set(f'flag:id:{flag.id}', flag.to_json())
pipeline.set(f'flag:str:{flag.flag}', flag.to_json())
pipeline.sadd(
f'team:{flag.team_id}:task:{flag.task_id}:round_flags:{flag.round}',
flag.id)
pipeline.set('flags:cached', 1)
def startup(**_kwargs):
"""Task to run on start of celery, schedules game start"""
game_config = config.get_game_config()
logger.info(f'Received game config: {game_config}')
start_game.apply_async(eta=game_config['start_time'], )
round = storage.game.get_real_round()
if not round or round == -1:
storage.caching.cache_teamtasks(round=0)
game_state = storage.game.get_game_state(round=0)
if not game_state:
logger.warning('Initial game_state missing')
else:
with storage.get_redis_storage().pipeline(
transaction=True) as pipeline:
logger.info(
f"Initializing game_state with {game_state.to_dict()}")
pipeline.set('game_state', game_state.to_json())
pipeline.publish('scoreboard', game_state.to_json())
pipeline.execute()
def get_action(put_ok, team_json, task_json, round):
"""Run "get" checker action
:param put_ok: boolean passed by "put" action in chain to indicate successful check and put
:param team_json: json-dumped team
:param task_json: json-dumped task
:param round: current round
If "check" or "put" actions fail, get is not run.
It runs `task.gets` times, each time a flag is chosen randomly from last "flag_lifetime" rounds
"""
if not put_ok:
return False
team = models.Team.from_json(team_json)
task = models.Task.from_json(task_json)
flag_lifetime = config.get_game_config()['flag_lifetime']
rounds_to_check = list(
set(max(1, round - x) for x in range(0, flag_lifetime)))
random.shuffle(rounds_to_check)
rounds_to_check = rounds_to_check[:task.gets]
logger.info(
f'Running GET on rounds {rounds_to_check} for team {team.id} task {task.id}'
)
checker_verdict = models.CheckerVerdict(
status=TaskStatus.UP,
public_message='',
private_message='',
command=[],
)
for get_round in rounds_to_check:
flag = storage.flags.get_random_round_flag(
team_id=team.id,
task_id=task.id,
round=get_round,
current_round=round,
)
if not flag:
checker_verdict.status = TaskStatus.CORRUPT
checker_verdict.private_message = f'No flags from round {get_round}'
checker_verdict.public_message = f'Could not get flag'
else:
checker_verdict = checkers.run_get_command(
checker_path=task.checker,
env_path=task.env_path,
host=team.ip,
flag=flag,
team_name=team.name,
timeout=task.checker_timeout,
logger=logger,
)
if checker_verdict.status != TaskStatus.UP:
break
storage.tasks.update_task_status(
task_id=task.id,
team_id=team.id,
checker_verdict=checker_verdict,
round=round,
)
return checker_verdict.status == TaskStatus.UP
def run():
conn = storage.get_db_pool().getconn()
curs = conn.cursor()
create_query_path = os.path.join(SCRIPTS_DIR, 'create_query.sql')
create_query = open(create_query_path).read()
curs.execute(create_query)
curs.execute(_CONFIG_INITIALIZATION_QUERY, (0, 0))
teams_config = config.get_teams_config()
teams = []
for team_conf in teams_config:
team_token = secrets.token_hex(8)
team = models.Team(id=None, **team_conf, token=team_token)
curs.execute(_TEAM_INSERT_QUERY, (team.name, team.ip, team_token))
team.id, = curs.fetchone()
teams.append(team)
tasks_config = config.get_tasks_config()
tasks = []
game_config = config.get_game_config()
global_env_path = game_config['env_path']
checkers_path = game_config['checkers_path']
global_default_score = game_config['default_score']
for task_conf in tasks_config:
if 'env_path' not in task_conf:
task_conf['env_path'] = global_env_path
if 'default_score' not in task_conf:
task_conf['default_score'] = global_default_score
task_conf['checker'] = os.path.join(checkers_path,
task_conf['checker'])
task = models.Task(id=None, **task_conf)
curs.execute(_TASK_INSERT_QUERY, (
task.name,
task.checker,
task.gets,
task.puts,
task.places,
task.checker_timeout,
task.env_path,
int(task.checker_returns_flag_id),
))
task.id, = curs.fetchone()
tasks.append(task)
for team in teams:
for task in tasks:
curs.execute(_TEAMTASK_INSERT_QUERY,
(task.id, team.id, 0, task.default_score, -1))
conn.commit()
curs.close()
storage.get_db_pool().putconn(conn)
storage.caching.cache_teamtasks(round=0)
game_state = storage.game.get_game_state(round=0)
with storage.get_redis_storage().pipeline(transaction=True) as pipeline:
pipeline.set('game_state', game_state.to_json())
pipeline.publish('scoreboard', game_state.to_json())
pipeline.execute()
from celery import Celery
import config
celery_config = config.get_celery_config()
app = Celery(
__name__,
include=[
'celery_tasks.tasks',
],
)
game_config = config.get_game_config()
app.conf.beat_schedule = {
'process_round': {
'task': 'celery_tasks.tasks.process_round',
'schedule': game_config['round_time'],
},
}
app.conf.update(celery_config)