import uuid
import wtforms
from flask import abort, request, session

CSRF_SESSION_KEY = '_token'


def _generate_csrf_token():
    return str(uuid.uuid4())


class Form(wtforms.Form):
    """
    Subclass of WTForms `Form` class. Flask `request.form` is passed
    as `formdata` argument to the constructor so can handle request
    data implicitly. In addition this `Form` implementation has
    automatic CSRF handling.
    """

    # token field
    csrf = wtforms.fields.HiddenField()

    method = 'POST'
    action = ''

    def __init__(self, formdata=None, *args, **kwargs):
        # set token
        token = session.get(CSRF_SESSION_KEY, None)
        if token is None:
            token = _generate_csrf_token()
        session[CSRF_SESSION_KEY] = token
        super(Form, self).__init__(formdata, csrf=token, *args, **kwargs)

    def process(self, formdata=None, obj=None, **kwargs):
        if request.method in ('PUT', 'POST'):
            if formdata is None:
                formdata = request.form
            # handle the case where the POST data is empty
            if not formdata:
                kwargs['csrf'] = None
        super(Form, self).process(formdata, obj, **kwargs)

    def omit(self, *args):
        for field in args:
            delattr(self, field)
        return self

    def update(self, obj):
        for name, field in self._fields.iteritems():
            try:
                field.populate_obj(obj, name)
            except:
                pass

    def validate_csrf(self, field):
        token = session.get(CSRF_SESSION_KEY, None)
        if not token or field.data != token:
            abort(403)