def run():
with storage.db_cursor() as (conn, curs):
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))
def get_db_global_config() -> models.GlobalConfig:
"""Get global config from database as it is"""
with storage.db_cursor(dict_cursor=True) as (conn, curs):
curs.execute(_GET_GLOBAL_CONFIG_QUERY)
result = curs.fetchone()
return models.GlobalConfig.from_dict(result)
def handle_attack(attacker_id: int, flag_str: str, round: int) -> float:
"""Check flag, lock team for update, call rating recalculation,
then publish redis message about stolen flag
:param attacker_id: id of the attacking team
:param flag_str: flag to be checked
:param round: round of the attack
:raises FlagSubmitException: when flag check was failed
:return: attacker rating change
"""
flag = flags.try_add_stolen_flag_by_str(flag_str=flag_str, attacker=attacker_id, round=round)
with storage.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()
flag_data = {
'attacker_id': attacker_id,
'victim_id': flag.team_id,
'task_id': flag.task_id,
'attacker_delta': attacker_delta,
'victim_delta': victim_delta,
}
storage.get_wro_sio_manager().emit(
event='flag_stolen',
data={'data': json.dumps(flag_data)},
namespace='/game_events',
)
return attacker_delta
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 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 = storage.game.get_current_global_config()
expires = game_config.flag_lifetime * game_config.round_time * 2 # can be smaller
with storage.db_cursor(dict_cursor=True) as (conn, curs):
curs.execute(_SELECT_ALL_LAST_FLAGS_QUERY,
(round - game_config.flag_lifetime, ))
flags = curs.fetchall()
pipeline.delete('flags:cached')
flag_models = list(helplib.models.Flag.from_dict(data) for data in flags)
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(), ex=expires)
pipeline.set(f'flag:str:{flag.flag}', flag.to_json(), ex=expires)
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('flags:cached', 1)
def get_game_running() -> bool:
"""Update game_running value in db"""
with storage.db_cursor() as (conn, curs):
curs.execute(_GET_GAME_RUNNING_QUERY)
game_running, = curs.fetchone()
return game_running
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 get_real_round_from_db() -> int:
"""Get real round from database
Fully persistent to use with game management"""
with storage.db_cursor() as (conn, curs):
curs.execute(_CURRENT_REAL_ROUND_QUERY)
round, = curs.fetchone()
return round
def get_teamtasks_from_db() -> List[dict]:
"""Fetch current team tasks from database
:return: dictionary of team tasks or None
"""
with storage.db_cursor(dict_cursor=True) as (conn, curs):
curs.execute(_SELECT_TEAMTASKS_QUERY)
data = curs.fetchall()
return data
def cache_tasks(pipeline):
"""Put "tasks" table data from database to cache
Just adds commands to pipeline stack (to support aioredis), don't forget to execute afterwards
"""
with storage.db_cursor(dict_cursor=True) as (conn, curs):
curs.execute(_SELECT_ALL_TASKS_QUERY)
tasks = curs.fetchall()
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 register():
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)
name = re.sub('[\\x00-\\x1f!@#$%^&*()]', '', name)
with storage.db_cursor(dict_cursor=True) as (conn, curs):
query = 'SELECT * FROM users WHERE name=%s'
curs.execute(query, (name, ))
user = curs.fetchone()
if user:
return get_error('name taken', 400)
with storage.db_cursor(dict_cursor=True) as (conn, curs):
query = 'INSERT INTO users (name, password) VALUES (%s, %s)'
curs.execute(query, (name, password))
conn.commit()
return jsonify('ok')
def cache_teams(pipeline):
"""Put "teams" table data from database to cache
Just adds commands to pipeline stack, don't forget to execute afterwards
"""
with storage.db_cursor(dict_cursor=True) as (conn, curs):
curs.execute(_SELECT_ALL_TEAMS_QUERY)
teams = curs.fetchall()
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 init_db():
users_query = '''CREATE TABLE IF NOT EXISTS users(
id SERIAL PRIMARY KEY,
name VARCHAR(255) UNIQUE NOT NULL,
password VARCHAR(255)
)'''
anime_uploads_query = '''CREATE TABLE IF NOT EXISTS anime_uploads(
id SERIAL PRIMARY KEY,
name VARCHAR(255) NOT NULL,
user_id INTEGER NOT NULL,
token VARCHAR(64) NOT NULL
)'''
with storage.db_cursor() as (conn, curs):
curs.execute(users_query)
curs.execute(anime_uploads_query)
conn.commit()
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 = storage.game.get_current_global_config()
with storage.db_cursor() as (conn, curs):
curs.execute(_SELECT_LAST_STOLEN_TEAM_FLAGS_QUERY,
(round - game_config.flag_lifetime, team_id))
flags = curs.fetchall()
pipeline.delete(f'team:{team_id}:stolen_flags:cached',
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}:stolen_flags:cached', 1)
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 get_attack_data(
current_round: int,
tasks: List[helplib.models.Task]) -> Dict[str, Dict[int, List[str]]]:
"""Get unexpired flags for round 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 = storage.game.get_current_global_config()
need_round = current_round - config.flag_lifetime
if task_ids:
with storage.db_cursor() as (conn, curs):
curs.execute(_GET_UNEXPIRED_FLAGS_QUERY, (need_round, task_ids))
flags = curs.fetchall()
else:
flags = []
data = {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 update_real_round_in_db(new_round: int):
"""Update real round stored in DB"""
with storage.db_cursor() as (conn, curs):
curs.execute(_UPDATE_REAL_ROUND_QUERY, (new_round,))
conn.commit()
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 set_game_running(new_value: bool):
"""Update game_running value in db"""
with storage.db_cursor() as (conn, curs):
curs.execute(_SET_GAME_RUNNING_QUERY, (new_value,))
conn.commit()
