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 set_game_running(new_value: int):
"""Update game_running value in db"""
conn = storage.get_db_pool().getconn()
curs = conn.cursor()
curs.execute(_SET_GAME_RUNNING_QUERY, (new_value, ))
conn.commit()
curs.close()
storage.get_db_pool().putconn(conn)
def update_real_round_in_db(new_round: int):
"""Update real round stored in DB"""
conn = storage.get_db_pool().getconn()
curs = conn.cursor()
curs.execute(_UPDATE_REAL_ROUND_QUERY, (new_round, ))
conn.commit()
curs.close()
storage.get_db_pool().putconn(conn)
def run():
conn = storage.get_db_pool().getconn()
curs = conn.cursor()
curs.execute(_SELECT_TEAMS_NAME_TOKEN_QUERY)
result = curs.fetchall()
print('\n'.join("{name}:{token}".format(name=name, token=token)
for name, token in result))
curs.close()
storage.get_db_pool().putconn(conn)
def get_game_running() -> int:
"""Update game_running value in db"""
conn = storage.get_db_pool().getconn()
curs = conn.cursor()
curs.execute(_GET_GAME_RUNNING_QUERY)
game_running, = curs.fetchone()
curs.close()
storage.get_db_pool().putconn(conn)
return game_running
def get_real_round_from_db() -> int:
"""Get real round from database
Fully persistent to use with game management"""
conn = storage.get_db_pool().getconn()
curs = conn.cursor()
curs.execute(_CURRENT_REAL_ROUND_QUERY)
round, = curs.fetchone()
curs.close()
storage.get_db_pool().putconn(conn)
return round
def cache_tasks(pipeline):
"""Put "tasks" table data from database to cache
Just adds commands to pipeline stack, don't forget to execute afterwards
"""
conn = storage.get_db_pool().getconn()
curs = conn.cursor(cursor_factory=extras.DictCursor)
curs.execute(_SELECT_ALL_TASKS_QUERY)
tasks = curs.fetchall()
curs.close()
storage.get_db_pool().putconn(conn)
tasks = list(models.Task.from_dict(task) for task in tasks)
pipeline.delete('tasks', 'tasks:cached')
if tasks:
pipeline.sadd('tasks', *[task.to_json() for task in tasks])
pipeline.set('tasks:cached', 1)
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 cache_teams(pipeline):
"""Put "teams" table data from database to cache
Just adds commands to pipeline stack, don't forget to execute afterwards
"""
conn = storage.get_db_pool().getconn()
curs = conn.cursor(cursor_factory=extras.DictCursor)
curs.execute(_SELECT_ALL_TEAMS_QUERY)
teams = curs.fetchall()
curs.close()
storage.get_db_pool().putconn(conn)
teams = list(models.Team.from_dict(team) for team in teams)
pipeline.delete('teams', 'teams:cached')
if teams:
pipeline.sadd('teams', *[team.to_json() for team in teams])
for team in teams:
pipeline.set(f'team:token:{team.token}', team.id)
pipeline.set('teams:cached', 1)
def cache_teamtasks_for_team(team_id: int, current_round: int, pipeline):
"""Put "teamtasks" for specified team table data for the specified round from database to cache
:param team_id: team id
:param current_round: round to cache
:param pipeline: redis connection to add command to
"""
conn = storage.get_db_pool().getconn()
curs = conn.cursor(cursor_factory=extras.RealDictCursor)
curs.execute(_SELECT_TEAMTASKS_FOR_TEAM_WITH_ROUND_QUERY, (
team_id,
current_round,
))
results = curs.fetchall()
curs.close()
storage.get_db_pool().putconn(conn)
data = json.dumps(results)
pipeline.set(f'teamtasks:team:{team_id}:round:{current_round}', data)
pipeline.set(f'teamtasks:team:{team_id}:round:{current_round}:cached', 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 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 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()
