def start_game(): """Starts game Sets `game_running` in DB """ logger.info('Starting game') with storage.get_redis_storage().pipeline(transaction=True) as pipeline: with locking.acquire_redis_lock(pipeline, 'game_starting_lock'): already_started = storage.game.get_game_running() if already_started: logger.info('Game already started') return storage.game.set_game_running(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 update_task_status(task_id: int, team_id: int, round: int, checker_verdict: models.CheckerVerdict): """ Update task status in database :param task_id: :param team_id: :param 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' with storage.db_cursor(dict_cursor=True) as (conn, curs): curs.callproc( 'update_teamtasks_status', ( round, team_id, task_id, checker_verdict.status.value, add, public, checker_verdict.private_message, checker_verdict.command, ) ) data = curs.fetchone() conn.commit() data['round'] = round with storage.get_redis_storage().pipeline(transaction=True) as pipeline: pipeline.xadd(f'teamtasks:{team_id}:{task_id}', dict(data), maxlen=50, approximate=False).execute()
def get_random_round_flag(team_id: int, task_id: int, round: int, current_round: int) -> Optional[helplib.models.Flag]: """Get random flag for team generated for specified round and task :param team_id: team id :param task_id: task id :param round: round to fetch flag for :param current_round: current round :return: Flag mode instance or None if no flag from rounds exist """ with storage.get_redis_storage().pipeline(transaction=True) as pipeline: cache_helper( pipeline=pipeline, cache_key='flags:cached', cache_func=caching.cache_last_flags, cache_args=(current_round, pipeline), ) flags, = pipeline.smembers( f'team:{team_id}:task:{task_id}:round_flags:{round}').execute() try: flag_id = int(secrets.choice(list(flags))) except (ValueError, IndexError, TypeError): return None return get_flag_by_id(flag_id, current_round)
def start_game(): """Starts game Sets `game_running` in DB """ logger.info('Starting game') with storage.get_redis_storage().pipeline(transaction=True) as pipeline: with locking.acquire_redis_lock(pipeline, 'game_starting_lock'): already_started = storage.game.get_game_running() if already_started: logger.info('Game already started') return storage.game.set_round_start(round=0) storage.game.set_game_running(True) game_state = storage.game.construct_game_state_from_db(round=0) if not game_state: logger.warning('Initial game_state missing') else: logger.info(f"Initializing game_state with {game_state.to_dict()}") pipeline.set('game_state', game_state.to_json()) pipeline.execute() storage.get_wro_sio_manager().emit( event='update_scoreboard', data={'data': game_state.to_json()}, namespace='/game_events', )
def add_flag(flag: helplib.models.Flag) -> helplib.models.Flag: """Inserts a newly generated flag into the database and cache :param flag: Flag model instance to be inserted :return: flag with set "id" field """ with storage.db_cursor() as (conn, curs): curs.execute(_INSERT_FLAG_QUERY, ( flag.flag, flag.team_id, flag.task_id, flag.round, flag.flag_data, flag.vuln_number, )) flag.id, = curs.fetchone() conn.commit() game_config = storage.game.get_current_global_config() expires = game_config.flag_lifetime * game_config.round_time * 2 # can be smaller with storage.get_redis_storage().pipeline(transaction=True) as pipeline: round_flags_key = f'team:{flag.team_id}:task:{flag.task_id}:round_flags:{flag.round}' pipeline.sadd(round_flags_key, flag.id) pipeline.expire(round_flags_key, expires) pipeline.set(f'flag:id:{flag.id}', flag.to_json(), ex=expires) pipeline.set(f'flag:str:{flag.flag}', flag.to_json(), ex=expires) pipeline.execute() return flag
def startup(**_kwargs): """Task to run on start of celery, schedules game start""" game_config = storage.game.get_current_global_config() logger.info(f'Received game config: {game_config}') with storage.get_redis_storage().pipeline(transaction=True) as pipeline: with locking.acquire_redis_lock(pipeline, 'game_starting_lock'): already_started = storage.game.get_game_running() if not already_started: logger.info("Game is not running, initializing...") start_game.apply_async(eta=game_config.start_time, ) game_state = storage.game.construct_game_state_from_db(round=0) if not game_state: logger.warning('Initial game_state missing') else: logger.info( f"Initializing game_state with {game_state.to_dict()}") pipeline.set('game_state', game_state.to_json()) pipeline.execute() storage.get_wro_sio_manager().emit( event='update_scoreboard', data={'data': game_state.to_json()}, namespace='/game_events', )
def get_flag_by_field(field_name: str, field_value, round: int) -> helplib.models.Flag: """Get flag by generic field :param field_name: field name to ask cache for :param field_value: value of the field "field_name" to filter on :param round: current round :return: Flag model instance with flag.field_name == field_value :raises: FlagSubmitException if nothing found """ with storage.get_redis_storage().pipeline(transaction=True) as pipeline: cached, = pipeline.exists('flags:cached').execute() if not cached: cache_helper( pipeline=pipeline, cache_key='flags:cached', cache_func=caching.cache_last_flags, cache_args=(round, pipeline), ) pipeline.exists(f'flag:{field_name}:{field_value}') pipeline.get(f'flag:{field_name}:{field_value}') flag_exists, flag_json = pipeline.execute() if not flag_exists: raise helplib.exceptions.FlagSubmitException( 'Flag is invalid or too old') flag = helplib.models.Flag.from_json(flag_json) return flag
def update_attack_data(current_round): logger.info(f'Updating attack data contents for round {current_round}') tasks = storage.tasks.get_tasks() tasks = list(filter(lambda x: x.checker_provides_public_flag_data, tasks)) flags = storage.flags.get_attack_data(current_round, tasks) with storage.get_redis_storage().pipeline(transaction=True) as pipeline: pipeline.set('attack_data', json.dumps(flags)).execute()
def process_round(): """Process new round Updates current round variable, then processes all teams. This function also caches previous state and notifies frontend of a new round. Only one instance of process_round could be run! """ game_running = storage.game.get_game_running() if not game_running: logger.info('Game is not running, exiting') return current_round = storage.game.get_real_round_from_db() finished_round = current_round new_round = current_round + 1 storage.game.update_real_round_in_db(new_round=new_round) storage.tasks.initialize_teamtasks(round=new_round) logger.info(f'Processing round {new_round}') with storage.get_redis_storage().pipeline(transaction=True) as pipeline: pipeline.set('round', finished_round) pipeline.set('real_round', new_round) pipeline.execute() if new_round > 1: storage.caching.cache_teamtasks(round=finished_round) game_state = storage.game.get_game_state(round=finished_round) if not game_state: logger.warning( f'Game state is missing for round {finished_round}, skipping') else: logger.info(f'Publishing scoreboard for round {finished_round}') with storage.get_redis_storage().pipeline( transaction=True) as pipeline: pipeline.publish('scoreboard', game_state.to_json()) pipeline.set('game_state', game_state.to_json()) pipeline.execute() teams = storage.teams.get_teams() for team in teams: process_team.delay(team.to_json(), new_round)
def get_round_start(round: int) -> int: """Get start time for round as unix timestamp""" with storage.get_redis_storage().pipeline(transaction=False) as pipeline: start_time, = pipeline.get(f'round:{round}:start_time').execute() try: start_time = int(start_time) except (ValueError, TypeError): start_time = 0 return start_time
def run(): reset_db.run() init_db.run() r = storage.get_redis_storage() r.flushall() print('New team tokens:') print_tokens.run()
def get_current_round() -> int: """Get current round, returns -1 if round not in cache""" with storage.get_redis_storage().pipeline(transaction=False) as pipeline: round, = pipeline.get('round').execute() try: round = int(round.decode()) except (AttributeError, ValueError): return -1 else: return round
def update_round(finished_round): logger.info(f'Updating round to {finished_round + 1}') storage.game.set_round_start(round=finished_round + 1) storage.game.update_real_round_in_db(new_round=finished_round + 1) # Might think there's a RC here (I thought so too) # But all teamtasks with round >= real_round are updated in the attack handler # So both old and new teamtasks will be updated properly with storage.get_redis_storage().pipeline(transaction=True) as pipeline: pipeline.set('real_round', finished_round + 1).execute()
def get_real_round() -> int: """Get real round of system (for flag submitting), returns -1 if round not in cache """ with storage.get_redis_storage().pipeline(transaction=False) as pipeline: round, = pipeline.get('real_round').execute() try: round = int(round) except (ValueError, TypeError): return -1 return round
def run(self, *args, **kwargs): """Process new round Updates current round variable, then processes all teams. This function also caches previous state and notifies frontend of a new round (for classic game mode). Only one instance of process_round that updates round is to be be run! """ game_running = storage.game.get_game_running() if not game_running: logger.info('Game is not running, exiting') return with storage.get_redis_storage().pipeline( transaction=True) as pipeline: with locking.acquire_redis_lock(pipeline, 'round_update:lock'): current_round = storage.game.get_real_round_from_db() round_to_check = current_round if self.should_update_round(): self.update_round(current_round) round_to_check = current_round + 1 self.update_attack_data(round_to_check) if not round_to_check: logger.info("Not processing, round is 0") return if self.should_update_game_state(): self.update_game_state(current_round) teams = storage.teams.get_teams() random.shuffle(teams) tasks = storage.tasks.get_tasks() random.shuffle(tasks) args = itertools.product(teams, tasks, [round_to_check]) if self.round_type == 'full': logger.info("Running full round") celery_tasks.modes.run_full_round.starmap(args).apply_async() elif self.round_type == 'check_gets': logger.info("Running check_gets round") celery_tasks.modes.run_check_gets_round.starmap(args).apply_async() elif self.round_type == 'puts': logger.info("Running puts round") celery_tasks.modes.run_puts_round.starmap(args).apply_async() else: logger.critical( f"Invalid round type supplied: {self.round_type}, falling back to full" ) celery_tasks.modes.run_full_round.starmap(args).apply_async()
def get_current_global_config() -> models.GlobalConfig: """Get global config from cache is cached, otherwise cache it""" with storage.get_redis_storage().pipeline(transaction=True) as pipeline: cache_helper( pipeline=pipeline, cache_key='global_config:cached', cache_func=storage.caching.cache_global_config, cache_args=(pipeline,), ) result, = pipeline.get('global_config').execute() global_config = models.GlobalConfig.from_json(result) return global_config
def get_tasks() -> List[models.Task]: """Get list of tasks registered in database""" with storage.get_redis_storage().pipeline(transaction=True) as pipeline: cache_helper( pipeline=pipeline, cache_key='tasks:cached', cache_func=caching.cache_tasks, cache_args=(pipeline, ), ) tasks, = pipeline.smembers('tasks').execute() tasks = list(models.Task.from_json(task) for task in tasks) return tasks
def get_teams() -> List[models.Team]: """Get list of teams registered in the database""" with storage.get_redis_storage().pipeline(transaction=True) as pipeline: cache_helper( pipeline=pipeline, cache_key='teams:cached', cache_func=caching.cache_teams, cache_args=(pipeline, ), ) teams, = pipeline.smembers('teams').execute() teams = list(models.Team.from_json(team) for team in teams) return teams
def run(): conn = storage.get_db_pool().getconn() curs = conn.cursor() create_query_path = os.path.join(SCRIPTS_DIR, 'drop_query.sql') create_query = open(create_query_path).read() try: curs.execute(create_query) except psycopg2.errors.UndefinedTable: pass else: conn.commit() finally: curs.close() while True: try: storage.get_redis_storage().flushall() except (redis.exceptions.ConnectionError, redis.exceptions.BusyLoadingError): print('[*] Redis isn\'t running, waiting...') time.sleep(5) else: break
def get_user(): sess = request.cookies.get('session') if not sess: abort(401) with storage.get_redis_storage().pipeline() as pipeline: user_exists, user = pipeline.exists(sess).get(sess).execute() if not user_exists: abort(401) try: user = json.loads(user.decode()) except (UnicodeDecodeError, json.JSONDecodeError): abort(401) return user
def get_teamtasks_of_team(team_id: int) -> List[dict]: """Fetch teamtasks for team for all tasks""" tasks = get_tasks() with storage.get_redis_storage().pipeline(transaction=True) as pipeline: for task in tasks: pipeline.xrevrange(f'teamtasks:{team_id}:{task.id}') data = pipeline.execute() data = sum(data, []) results = [] for timestamp, record in data: record['timestamp'] = timestamp results.append(record) return results
def update_game_state(current_round): if not current_round: return game_state = storage.game.construct_latest_game_state( round=current_round) logger.info(f'Publishing scoreboard for round {current_round}') with storage.get_redis_storage().pipeline( transaction=True) as pipeline: pipeline.set('game_state', game_state.to_json()) pipeline.execute() storage.get_wro_sio_manager().emit( event='update_scoreboard', data={'data': game_state.to_json()}, namespace='/game_events', )
def cache_teamtasks(round: int): """Put "teamtasks" table data for the specified round from database to cache :param round: round to cache This function caches full game state for specified round """ conn = storage.get_db_pool().getconn() curs = conn.cursor(cursor_factory=extras.RealDictCursor) curs.execute(_SELECT_TEAMTASKS_BY_ROUND_QUERY, (round, )) results = curs.fetchall() curs.close() storage.get_db_pool().putconn(conn) data = json.dumps(results) with storage.get_redis_storage().pipeline(transaction=True) as pipeline: pipeline.set(f'teamtasks:{round}', data) pipeline.set(f'teamtasks:{round}:cached', 1) pipeline.execute()
def get_team_id_by_token(token: str) -> Optional[int]: """Get team by token :param token: token string :return: team id """ with storage.get_redis_storage().pipeline(transaction=True) as pipeline: cache_helper( pipeline=pipeline, cache_key='teams:cached', cache_func=caching.cache_teams, cache_args=(pipeline, ), ) team_id, = pipeline.get(f'team:token:{token}').execute() try: team_id = int(team_id) except (ValueError, TypeError): return None else: return team_id
def login(): name = request.json.get('name') password = request.json.get('password') if not name or not password: return get_error('Specify both name and password', 400) with storage.db_cursor(dict_cursor=True) as (conn, curs): query = 'SELECT * FROM users WHERE name=%s AND password=%s' curs.execute(query, (name, password)) user = curs.fetchone() if not user: return get_error('Invalid credentials', 403) session = secrets.token_hex(32) with storage.get_redis_storage().pipeline() as pipeline: pipeline.set(session, json.dumps(dict(user))).execute() resp = make_response(jsonify('ok')) resp.set_cookie('session', session) return resp
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_last_teamtasks() -> List[dict]: """Fetch team tasks, last for each team for each task :return: dictionary of team tasks or None """ teams = storage.teams.get_teams() tasks = storage.tasks.get_tasks() with storage.get_redis_storage().pipeline(transaction=True) as pipeline: for team in teams: for task in tasks: pipeline.xrevrange(f'teamtasks:{team.id}:{task.id}', count=1) data = pipeline.execute() data = sum(data, []) results = [] for timestamp, record in data: record['timestamp'] = timestamp results.append(record) process_teamtasks(results) return results
def try_add_stolen_flag(flag: helplib.models.Flag, attacker: int, round: int): """Check that flag is valid for current round, add it to cache, then add to db :param flag: Flag model instance :param attacker: attacker team id :param round: current round :raises: an instance of FlagSubmitException on validation error """ game_config = storage.game.get_current_global_config() if round - flag.round > game_config.flag_lifetime: raise helplib.exceptions.FlagSubmitException('Flag is too old') if flag.team_id == attacker: raise helplib.exceptions.FlagSubmitException('Flag is your own') with storage.get_redis_storage().pipeline(transaction=True) as pipeline: # optimization of redis request count cached_stolen = pipeline.exists( f'team:{attacker}:stolen_flags:cached').execute() if not cached_stolen: cache_helper( pipeline=pipeline, cache_key=f'team:{attacker}:stolen_flags:cached', cache_func=caching.cache_last_stolen, cache_args=(attacker, round, pipeline), ) is_new, = pipeline.sadd(f'team:{attacker}:stolen_flags', flag.id).execute() if not is_new: raise helplib.exceptions.FlagSubmitException('Flag already stolen') pipeline.incr(f'team:{attacker}:task:{flag.task_id}:stolen') pipeline.incr(f'team:{flag.team_id}:task:{flag.task_id}:lost') pipeline.execute()
def run(): conf_path = os.path.join(CONFIG_DIR, CONFIG_FILENAME) with open(conf_path) as f: file_config = yaml.load(f, Loader=yaml.FullLoader) with storage.db_cursor() as (conn, curs): create_tables_path = os.path.join(SCRIPTS_DIR, 'create_tables.sql') with open(create_tables_path) as f: create_tables_query = f.read() curs.execute(create_tables_query) create_functions_path = os.path.join(SCRIPTS_DIR, 'create_functions.sql') with open(create_functions_path) as f: create_functions_query = f.read() curs.execute(create_functions_query) teams_config = file_config['teams'] 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 = file_config['tasks'] tasks = [] global_defaults = { 'checkers_path': '/checkers/', 'env_path': '/checkers/bin/', 'default_score': 2000.0, 'game_hardness': 3000.0, 'inflation': True, 'flag_lifetime': 5, 'round_time': 60, 'timezone': 'UTC', 'game_mode': 'classic', } global_config = file_config['global'] for k, v in global_defaults.items(): if k not in global_config: global_defaults[k] = v task_defaults = { 'env_path': global_config['env_path'], 'default_score': global_config['default_score'], 'get_period': global_config.get('get_period', global_config['round_time']), 'checker_returns_flag_id': True, 'gevent_optimized': False, } for task_conf in tasks_config: for k, v in task_defaults.items(): if k not in task_conf: task_conf[k] = v task_conf['checker'] = os.path.join(global_config['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, task.checker_returns_flag_id, task.gevent_optimized, task.get_period, )) task.id, = curs.fetchone() tasks.append(task) data = [(task.id, team.id, task.default_score, -1) for team in teams for task in tasks] curs.executemany(_TEAMTASK_INSERT_QUERY, data) global_config.pop('env_path', None) global_config.pop('default_score', None) global_config.pop('checkers_path', None) global_config.pop('get_period', None) tz = pytz.timezone(global_config['timezone']) global_config['start_time'] = tz.localize(global_config['start_time']) keys = global_config.keys() columns = ','.join(keys) values = ','.join(f'%({key})s' for key in keys) curs.execute( _CONFIG_INITIALIZATION_QUERY.format(columns=columns, values=values), global_config, ) conn.commit() game_state = storage.game.construct_game_state_from_db(round=0) with storage.get_redis_storage().pipeline(transaction=True) as pipeline: pipeline.set('game_state', game_state.to_json()) pipeline.execute() storage.get_wro_sio_manager().emit( event='update_scoreboard', data={'data': game_state.to_json()}, namespace='/game_events', )
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()