def get_attack_data(
current_round: int,
tasks: List[models.Task],
) -> Dict[str, DefaultDict[int, List[str]]]:
"""
Get unexpired flags for round.
:returns: flags in format {task.name: {team.ip: [flag.public_data]}}
"""
task_ids = tuple(task.id for task in tasks)
task_names = {task.id: task.name for task in tasks}
config = game.get_current_game_config()
need_round = current_round - config.flag_lifetime
if task_ids:
with utils.db_cursor() as (_, curs):
curs.execute(_GET_UNEXPIRED_FLAGS_QUERY, (need_round, task_ids))
flags = curs.fetchall()
else:
flags = []
data: Dict[str, DefaultDict[int, List[str]]] = {
task_names[task_id]: defaultdict(list) for task_id in task_ids
}
for flag in flags:
ip, task_id, flag_data = flag
data[task_names[task_id]][ip].append(flag_data)
return data
def get_game_running() -> bool:
"""Get current game_running value from db."""
with utils.db_cursor() as (_, curs):
curs.execute(_GET_GAME_RUNNING_QUERY)
game_running, = curs.fetchone()
return game_running
def get_random_round_flag(
team_id: int,
task_id: int,
from_round: int,
current_round: int) -> Optional[models.Flag]:
"""
Get random flag for team generated for specified round and task.
:param team_id: team id
:param task_id: task id
:param from_round: round to fetch flag for
:param current_round: current round
:returns: Flag mode instance or None if no flag from rounds exist
"""
with utils.db_cursor() as (_, curs):
curs.execute(
_GET_RANDOM_ROUND_FLAG_QUERY,
{
'round': from_round,
'team_id': team_id,
'task_id': task_id,
}
)
result = curs.fetchone()
if not result:
return None
return get_flag_by_id(result[0], current_round)
def get_db_game_config() -> models.GameConfig:
"""Get game config from database."""
with utils.db_cursor(dict_cursor=True) as (_, curs):
curs.execute(_GET_GAME_CONFIG_QUERY)
result = curs.fetchone()
return models.GameConfig.from_dict(result)
def get_real_round_from_db() -> int:
"""Get real round from database. Fully persistent to use with game management."""
with utils.db_cursor() as (_, curs):
curs.execute(_CURRENT_REAL_ROUND_QUERY)
r, = curs.fetchone()
return r
def cache_last_flags(current_round: int, pipe: Pipeline) -> None:
"""
Cache all generated flags from last "flag_lifetime" rounds.
Just adds commands to pipeline stack, don't forget to execute afterwards.
:param current_round: current round
:param pipe: redis connection to add command to
"""
game_config = game.get_current_game_config()
expires = game_config.flag_lifetime * game_config.round_time * 2
with utils.db_cursor(dict_cursor=True) as (_, curs):
curs.execute(
_SELECT_ALL_LAST_FLAGS_QUERY,
{'round': current_round - game_config.flag_lifetime},
)
flags = curs.fetchall()
flag_models = list(models.Flag.from_dict(data) for data in flags)
pipe.set(CacheKeys.flags_cached(), 1)
for flag in flag_models:
pipe.set(CacheKeys.flag_by_id(flag.id), flag.to_json(), ex=expires)
pipe.set(CacheKeys.flag_by_str(flag.flag), flag.to_json(), ex=expires)
def cache_last_stolen(team_id: int, current_round: int,
pipe: Pipeline) -> None:
"""
Caches stolen flags from "flag_lifetime" rounds.
Just adds commands to pipeline stack, don't forget to execute afterwards.
:param team_id: attacker team id
:param current_round: current round
:param pipe: redis connection to add command to
"""
game_config = game.get_current_game_config()
with utils.db_cursor() as (_, curs):
curs.execute(
_SELECT_LAST_STOLEN_TEAM_FLAGS_QUERY,
{
'round': current_round - game_config.flag_lifetime,
'attacker_id': team_id,
},
)
flags = curs.fetchall()
key = CacheKeys.team_stolen_flags(team_id)
pipe.delete(key)
if flags:
pipe.sadd(key, *(flag[0] for flag in flags))
def get_last_run(schedule_id: str) -> Optional[datetime]:
with utils.db_cursor() as (_, curs):
curs.execute(
SELECT_LAST_RUN,
{'id': schedule_id},
)
result = curs.fetchone()
return result[0] if result else None
def set_last_run(schedule_id: str, last_run: datetime) -> None:
with utils.db_cursor() as (conn, curs):
curs.execute(
UPDATE_LAST_RUN,
{
'id': schedule_id,
'last_run': last_run
},
)
conn.commit()
def cache_tasks(pipe: Pipeline) -> None:
"""
Put active tasks table data from database to cache.
Just adds commands to pipeline stack don't forget to execute afterwards.
"""
with utils.db_cursor(dict_cursor=True) as (_, curs):
curs.execute(models.Task.get_select_active_query())
tasks = curs.fetchall()
tasks = list(models.Task.from_dict(task) for task in tasks)
key = CacheKeys.tasks()
pipe.delete(key)
if tasks:
pipe.sadd(key, *(task.to_json() for task in tasks))
def cache_teams(pipe: Pipeline) -> None:
"""
Put "teams" table data from database to cache.
Just adds commands to pipeline stack, don't forget to execute afterwards.
"""
with utils.db_cursor(dict_cursor=True) as (_, curs):
curs.execute(models.Team.get_select_active_query())
teams = curs.fetchall()
teams = list(models.Team.from_dict(team) for team in teams)
key = CacheKeys.teams()
pipe.delete(key)
if teams:
pipe.sadd(key, *[team.to_json() for team in teams])
for team in teams:
pipe.set(CacheKeys.team_by_token(team.token), team.id)
def add_flag(flag: models.Flag) -> models.Flag:
"""
Inserts a newly generated flag into the database and cache.
:param flag: Flag model instance to be inserted
:returns: flag with set "id" field
"""
with utils.db_cursor() as (conn, curs):
flag.insert(curs)
conn.commit()
game_config = game.get_current_game_config()
expires = game_config.flag_lifetime * game_config.round_time * 2
with utils.redis_pipeline(transaction=True) as pipe:
pipe.set(CacheKeys.flag_by_id(flag.id), flag.to_json(), ex=expires)
pipe.set(CacheKeys.flag_by_str(flag.flag), flag.to_json(), ex=expires)
pipe.execute()
return flag
def handle_attack(attacker_id: int, flag_str: str,
current_round: int) -> models.AttackResult:
"""
Main routine for attack validation & state change.
Checks flag, locks team for update, calls rating recalculation,
then publishes redis message about stolen flag.
:param attacker_id: id of the attacking team
:param flag_str: flag to be checked
:param current_round: round of the attack
:return: attacker rating change
"""
result = models.AttackResult(attacker_id=attacker_id)
try:
if current_round == -1:
raise FlagExceptionEnum.GAME_NOT_AVAILABLE
flag = storage.flags.get_flag_by_str(
flag_str=flag_str,
current_round=current_round,
)
if flag is None:
raise FlagExceptionEnum.FLAG_INVALID
if flag.team_id == attacker_id:
raise FlagExceptionEnum.FLAG_YOUR_OWN
game_config = game.get_current_game_config()
if current_round - flag.round > game_config.flag_lifetime:
raise FlagExceptionEnum.FLAG_TOO_OLD
result.victim_id = flag.team_id
result.task_id = flag.task_id
success = storage.flags.try_add_stolen_flag(
flag=flag,
attacker=attacker_id,
current_round=current_round,
)
if not success:
raise FlagExceptionEnum.FLAG_YOUR_OWN
except exceptions.FlagSubmitException as e:
result.submit_ok = False
result.message = str(e)
else:
result.submit_ok = True
with utils.db_cursor() as (conn, curs):
curs.callproc(
"recalculate_rating",
(
attacker_id,
flag.team_id,
flag.task_id,
flag.id,
),
)
attacker_delta, victim_delta = curs.fetchone()
conn.commit()
result.attacker_delta = attacker_delta
result.victim_delta = victim_delta
result.message = f'Flag accepted! Earned {attacker_delta} flag points!'
return result
def set_game_running(new_value: bool) -> None:
"""Update game_running value in db."""
with utils.db_cursor() as (conn, curs):
curs.execute(_SET_GAME_RUNNING_QUERY, {'value': new_value})
conn.commit()
def update_real_round_in_db(new_round: int) -> None:
"""Update real_round of game config stored in DB."""
with utils.db_cursor() as (conn, curs):
curs.execute(_UPDATE_REAL_ROUND_QUERY, {'round': new_round})
conn.commit()