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)