Extract extra property into a shared function

CTFd/forms/auth.py

@@ -3,8 +3,8 @@ from wtforms.fields.html5 import EmailField
 from wtforms.validators import InputRequired
 
 from CTFd.forms import BaseForm
-from CTFd.forms.users import attach_custom_user_fields
 from CTFd.forms.fields import SubmitField
+from CTFd.forms.users import attach_custom_user_fields, build_custom_user_fields
 from CTFd.models import UserFields
 
 
@@ -17,12 +17,7 @@ def RegistrationForm(*args, **kwargs):
 
         @property
         def extra(self):
-            fields = []
-            new_fields = UserFields.query.all()
-            for field in new_fields:
-                entry = (field.name, getattr(self, f"fields[{field.id}]"))
-                fields.append(entry)
-            return fields
+            return build_custom_user_fields(self, include_entries=False)
 
     attach_custom_user_fields(_RegistrationForm)
 

CTFd/forms/self.py

@@ -3,8 +3,8 @@ from wtforms import PasswordField, SelectField, StringField
 from wtforms.fields.html5 import DateField, URLField
 
 from CTFd.forms import BaseForm
-from CTFd.forms.users import attach_custom_user_fields
 from CTFd.forms.fields import SubmitField
+from CTFd.forms.users import attach_custom_user_fields, build_custom_user_fields
 from CTFd.models import FieldEntries, UserFields
 from CTFd.utils.countries import SELECT_COUNTRIES_LIST
 
@@ -22,24 +22,12 @@ def SettingsForm(*args, **kwargs):
 
         @property
         def extra(self):
-            fields = []
-            new_fields = UserFields.query.filter_by(editable=True).all()
-            user_fields = {}
-
-            for f in FieldEntries.query.filter_by(user_id=session["id"]).all():
-                user_fields[f.field_id] = f.value
-
-            for field in new_fields:
-                form_field = getattr(self, f"fields[{field.id}]")
-                initial = user_fields.get(field.id, "")
-                form_field.data = initial
-                if form_field.render_kw:
-                    form_field.render_kw["data-initial"] = initial
-                else:
-                    form_field.render_kw = {"data-initial": initial}
-                entry = (field.name, form_field)
-                fields.append(entry)
-            return fields
+            return build_custom_user_fields(
+                self,
+                include_entries=True,
+                fields_kwargs={"editable": True},
+                field_entries_kwargs={"user_id": session["id"]},
+            )
 
     attach_custom_user_fields(_SettingsForm, editable=True)
 

CTFd/forms/users.py

@@ -7,12 +7,52 @@ from CTFd.forms.fields import SubmitField
 from CTFd.models import FieldEntries, UserFields
 from CTFd.utils.countries import SELECT_COUNTRIES_LIST
 
-def build_custom_user_fields():
-    pass
+
+def build_custom_user_fields(
+    form_cls, include_entries=False, fields_kwargs=None, field_entries_kwargs=None
+):
+    """
+    Function used to reinject values back into forms for accessing by themes
+    """
+    if fields_kwargs is None:
+        fields_kwargs = {}
+    if field_entries_kwargs is None:
+        field_entries_kwargs = {}
+
+    fields = []
+    new_fields = UserFields.query.filter_by(**fields_kwargs).all()
+    user_fields = {}
+
+    # Only include preexisting values if asked
+    if include_entries is True:
+        for f in FieldEntries.query.filter_by(**field_entries_kwargs).all():
+            user_fields[f.field_id] = f.value
+
+    for field in new_fields:
+        form_field = getattr(form_cls, f"fields[{field.id}]")
+
+        # Only include preexisting values if asked
+        if include_entries is True:
+            initial = user_fields.get(field.id, "")
+            form_field.data = initial
+            if form_field.render_kw:
+                form_field.render_kw["data-initial"] = initial
+            else:
+                form_field.render_kw = {"data-initial": initial}
+
+        entry = (field.name, form_field)
+        fields.append(entry)
+    return fields
 
 
 def attach_custom_user_fields(form_cls, **kwargs):
-    new_fields = UserFields.filter_by(**kwargs).query.all()
+    """
+    Function used to attach form fields to wtforms.
+    Not really a great solution but is approved by wtforms.
+
+    https://wtforms.readthedocs.io/en/2.3.x/specific_problems/#dynamic-form-composition
+    """
+    new_fields = UserFields.query.filter_by(**kwargs).all()
     for field in new_fields:
         validators = []
         if field.required:
@@ -80,24 +120,12 @@ def UserEditForm(*args, **kwargs):
 
         @property
         def extra(self):
-            fields = []
-            new_fields = UserFields.query.all()
-            user_fields = {}
-
-            for f in FieldEntries.query.filter_by(user_id=self.obj.id).all():
-                user_fields[f.field_id] = f.value
-
-            for field in new_fields:
-                form_field = getattr(self, f"fields[{field.id}]")
-                initial = user_fields.get(field.id, "")
-                form_field.data = initial
-                if form_field.render_kw:
-                    form_field.render_kw["initial"] = initial
-                else:
-                    form_field.render_kw = {"data-initial": initial}
-                entry = (field.name, form_field)
-                fields.append(entry)
-            return fields
+            return build_custom_user_fields(
+                self,
+                include_entries=True,
+                fields_kwargs=None,
+                field_entries_kwargs={"user_id": self.obj.id},
+            )
 
         def __init__(self, *args, **kwargs):
             """
@@ -119,14 +147,7 @@ def UserCreateForm(*args, **kwargs):
 
         @property
         def extra(self):
-            fields = []
-            new_fields = UserFields.query.all()
-
-            for field in new_fields:
-                form_field = getattr(self, f"fields[{field.id}]")
-                entry = (field.name, form_field)
-                fields.append(entry)
-            return fields
+            return build_custom_user_fields(self, include_entries=False)
 
     attach_custom_user_fields(_UserCreateForm)