# forms.py
class CheckFlagForm(FlaskForm):
    flag = StringField('Flag', validators=[DataRequired()])
    submit = SubmitField('Check')

    def __init__(self, challenge_id, *args, **kwargs):
        super(CheckFlagForm, self).__init__(*args, **kwargs)
        self.challenge_id = challenge_id

    def validate_flag(self, flag):

        challenge = Challenge.query.filter_by(id=self.challenge_id).first()

        if challenge.check_flag(flag.data):
            ## success!
            pass
        else:
            raise ValidationError("Invalid flag")

# routes.py

@app.route('/challenge/<challenge_id>', methods=['GET', 'POST'])
@login_required
def show_challenge(challenge_id):
    challenge = Challenge.query.filter_by(is_deleted=False, is_published=True, id=challenge_id).first()
    form = CheckFlagForm(challenge_id)
    if form.validate_on_submit():
        flag = form.flag.data
        if challenge.check_flag(flag):
            ## SUCCESS
            flash('You solved the challenge!', category="success")
        else:
            ## FAIL
            flash('Invalid flag!', category="danger")
        return redirect(url_for('show_challenge', challenge_id=challenge_id))
    elif request.method == 'GET':
        return render_template('challenge.html', title=f'Challenge: {challenge.name}', user=user, challenge=challenge,
                               form=form)
    return render_template('challenge.html', title=f'Challenge: {challenge.name}', user=user, challenge=challenge,
                           form=form)