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 get_flag_by_field(
name: str,
value: Union[str, int],
current_round: int,
) -> Optional[models.Flag]:
"""
Get flag by generic field.
:param name: field name to ask cache for
:param value: value of the field "field_name" to filter on
:param current_round: current round
:returns: Flag model instance with flag.field_name == field_value or None
"""
cached_key = CacheKeys.flags_cached()
with utils.redis_pipeline(transaction=True) as pipe:
cached, = pipe.exists(cached_key).execute()
if not cached:
cache_helper(
pipeline=pipe,
cache_key=cached_key,
cache_func=caching.cache_last_flags,
cache_args=(current_round, pipe),
)
flag_json, = pipe.get(CacheKeys.flag_by_field(name, value)).execute()
if not flag_json:
return None
flag = models.Flag.from_json(flag_json)
return flag
def get_current_game_config() -> models.GameConfig:
"""Get game config from cache is cached, cache it otherwise."""
with utils.redis_pipeline(transaction=True) as pipe:
cache_helper(
pipeline=pipe,
cache_key=CacheKeys.game_config(),
cache_func=caching.cache_game_config,
cache_args=(pipe,),
)
result, = pipe.get(CacheKeys.game_config()).execute()
game_config = models.GameConfig.from_json(result)
return game_config
def update_attack_data(current_round: int) -> None:
tasks = storage.tasks.get_tasks()
tasks = list(filter(lambda x: x.checker_provides_public_flag_data, tasks))
attack_data = storage.flags.get_attack_data(current_round, tasks)
with utils.redis_pipeline(transaction=False) as pipe:
pipe.set(CacheKeys.attack_data(), kjson.dumps(attack_data))
pipe.execute()
def get_cached_game_state() -> Optional[models.GameState]:
with storage.utils.redis_pipeline(transaction=False) as pipe:
state, = pipe.get(CacheKeys.game_state()).execute()
if not state:
return None
return models.GameState.from_json(state)
def try_add_stolen_flag(flag: models.Flag, attacker: int, current_round: int) -> bool:
"""
Flag validation function.
Checks that flag is valid for current round, adds it to cache for team,
atomically checking if it's already submitted by the attacker.
:param flag: Flag model instance
:param attacker: attacker team id
:param current_round: current round
"""
stolen_key = CacheKeys.team_stolen_flags(attacker)
with utils.redis_pipeline(transaction=True) as pipe:
# optimization of redis request count
cached_stolen = pipe.exists(stolen_key).execute()
if not cached_stolen:
cache_helper(
pipeline=pipe,
cache_key=stolen_key,
cache_func=caching.cache_last_stolen,
cache_args=(attacker, current_round, pipe),
)
is_new, = pipe.sadd(stolen_key, flag.id).execute()
return bool(is_new)
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 update_round(finished_round: int) -> None:
new_round = finished_round + 1
set_round_start(r=new_round)
update_real_round_in_db(new_round=new_round)
with utils.redis_pipeline(transaction=False) as pipe:
pipe.set(CacheKeys.current_round(), new_round)
pipe.execute()
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 get_real_round() -> int:
"""
Get real round of system (for flag submitting).
:returns: -1 if round not in cache, else round
"""
with utils.redis_pipeline(transaction=False) as pipe:
r, = pipe.get(CacheKeys.current_round()).execute()
return int(r or -1)
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 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 get_tasks() -> List[models.Task]:
"""Get list of tasks registered in database."""
key = CacheKeys.tasks()
with storage.utils.redis_pipeline(transaction=True) as pipe:
cache_helper(
pipeline=pipe,
cache_key=key,
cache_func=caching.cache_tasks,
cache_args=(pipe, ),
)
tasks, = pipe.smembers(key).execute()
tasks = list(models.Task.from_json(task) for task in tasks)
return tasks
def get_teams() -> List[models.Team]:
"""Get list of active teams."""
key = CacheKeys.teams()
with storage.utils.redis_pipeline(transaction=True) as pipe:
cache_helper(
pipeline=pipe,
cache_key=key,
cache_func=storage.caching.cache_teams,
cache_args=(pipe, ),
)
teams, = pipe.smembers(key).execute()
teams = list(models.Team.from_json(team) for team in teams)
return teams
def get_team_id_by_token(token: str) -> Optional[int]:
"""
Get team by token.
:param token: token string
:return: team id
"""
with storage.utils.redis_pipeline(transaction=False) as pipe:
team_id, = pipe.get(CacheKeys.team_by_token(token)).execute()
try:
team_id = int(team_id)
except (ValueError, TypeError):
return None
else:
return team_id
def update_task_status(
task_id: int,
team_id: int,
current_round: int,
checker_verdict: models.CheckerVerdict,
) -> None:
"""
Update task status in database.
:param task_id:
:param team_id:
:param current_round:
:param checker_verdict: instance of CheckerActionResult
"""
add = 0
public = checker_verdict.public_message
if checker_verdict.status == TaskStatus.UP:
add = 1
if checker_verdict.action == Action.PUT:
public = 'OK'
params = {
'round': current_round,
'task_id': task_id,
'team_id': team_id,
'status': checker_verdict.status.value,
'public_message': public,
'private_message': checker_verdict.private_message,
'command': checker_verdict.command,
'passed': add,
}
with storage.utils.db_cursor(dict_cursor=True) as (conn, curs):
curs.execute(_INSERT_TEAMTASKS_TO_LOG_QUERY, params)
curs.execute(_UPDATE_TEAMTASKS_QUERY, params)
data = curs.fetchone()
conn.commit()
data['round'] = current_round
with storage.utils.redis_pipeline(transaction=True) as pipe:
pipe.xadd(
CacheKeys.teamtasks(team_id, task_id),
dict(data),
maxlen=50,
approximate=False,
).execute()
def get_teamtasks_for_team(team_id: int) -> List[dict]:
"""Fetch teamtasks for team for all tasks."""
tasks = storage.tasks.get_tasks()
with storage.utils.redis_pipeline(transaction=False) as pipe:
for task in tasks:
pipe.xrevrange(CacheKeys.teamtasks(team_id, task.id))
data = pipe.execute()
data = sum(data, [])
results = []
for timestamp, record in data:
record['timestamp'] = timestamp
results.append(record)
return results
def get_last_teamtasks() -> List[dict]:
"""Fetch team tasks, last for each team for each task."""
teams = storage.teams.get_teams()
tasks = storage.tasks.get_tasks()
with storage.utils.redis_pipeline(transaction=True) as pipe:
for team in teams:
for task in tasks:
pipe.xrevrange(CacheKeys.teamtasks(team.id, task.id), count=1)
data = pipe.execute()
data = sum(data, [])
results = []
for timestamp, record in data:
record['timestamp'] = timestamp
results.append(record)
results = process_teamtasks(results)
return results
def flush_tasks_cache():
with utils.redis_pipeline(transaction=False) as pipe:
pipe.delete(CacheKeys.tasks()).execute()
def get_attack_data() -> str:
"""Get public flag ids for tasks that provide them, as json string."""
with utils.redis_pipeline(transaction=False) as pipe:
attack_data, = pipe.get(CacheKeys.attack_data()).execute()
return attack_data or 'null'
def cache_game_config(pipe: Pipeline) -> None:
"""Put game config to cache (without round or game_running)."""
game_config = game.get_db_game_config()
data = game_config.to_json()
pipe.set(CacheKeys.game_config(), data)
def set_round_start(r: int) -> None:
"""Set start time for round as str."""
cur_time = int(time.time())
with utils.redis_pipeline(transaction=False) as pipe:
pipe.set(CacheKeys.round_start(r), cur_time).execute()
def get_round_start(r: int) -> int:
"""Get start time for round as unix timestamp."""
with utils.redis_pipeline(transaction=False) as pipe:
start_time, = pipe.get(CacheKeys.round_start(r)).execute()
return int(start_time or 0)
