ctfbot.py 文件源码

def advance_round(teams, services):
    """Advance the round: update results, truncate the active_flags table and
    store new flags in the database for each team and service."""

    with db_conn.cursor() as cur:
        # advance the round and clear the flag tables
        try:
            cur.execute('SELECT * FROM switch_round()')
            rnd = cur.fetchone()['switch_round']
        except psycopg2.Error as e:
            logger.critical(("Error while incrementing the round, "
                              "aborting: {}").format(e))
            abort()
    # commit the stored procedure operations (probably not needed)
    db_conn.commit()
    logger.info("Round {} started".format(rnd))

    # generate and insert the new flags to the database
    cur = db_conn.cursor()
    for service in services:
        for team in teams:
            inserted = False
            while not inserted:
                flag = utils.generate_flag(config['FLAG_PREFIX'], config['FLAG_SUFFIX'],
                                           config['FLAG_CHARS'], config['FLAG_LENGTH'])
                try:
                    cur.execute((
                        'INSERT INTO flags (flag, team_id, service_id, round) '
                        'VALUES (%s, %s, %s, %s)'),
                        (flag, team.id, service.id, rnd))
                except psycopg2.IntegrityError:
                    logger.warning('Duplicate flag, generating a new one')
                except psycopg2.Error as e:
                    logger.critical(('Error while adding a new flag to the '
                                      'database, aborting: {}').format(e))
                    abort()
                else:
                    inserted = True
                    logger.debug(('New flag just added to the database: {}').format(flag))
    db_conn.commit()
    cur.close()