From 7fe32d7a5d84dbf3036ce805495e9e9b71cf74ec Mon Sep 17 00:00:00 2001
From: Frank <frankli0324@hotmail.com>
Date: Sat, 30 Jan 2021 04:09:09 +0800
Subject: [PATCH] field based schema validation (#1789)

* Clean up Page and Challenges schema validation
---
 CTFd/schemas/challenges.py | 55 +++++++++++++++++++-------------
 CTFd/schemas/pages.py      | 64 ++++++++++++++++++++++++--------------
 2 files changed, 73 insertions(+), 46 deletions(-)

diff --git a/CTFd/schemas/challenges.py b/CTFd/schemas/challenges.py
index 312347dc..c9ad6f4a 100644
--- a/CTFd/schemas/challenges.py
+++ b/CTFd/schemas/challenges.py
@@ -1,4 +1,5 @@
-from marshmallow import ValidationError, pre_load
+from marshmallow import validate
+from marshmallow_sqlalchemy import field_for
 
 from CTFd.models import Challenges, ma
 
@@ -9,29 +10,39 @@ class ChallengeSchema(ma.ModelSchema):
         include_fk = True
         dump_only = ("id",)
 
-    @pre_load
-    def validate_name(self, data):
-        name = data.get("name", "")
-        if len(name) > 80:
-            raise ValidationError(
-                "Challenge could not be saved. Challenge name too long",
-                field_names=["name"],
+    name = field_for(
+        Challenges,
+        "name",
+        validate=[
+            validate.Length(
+                min=0,
+                max=80,
+                error="Challenge could not be saved. Challenge name too long",
             )
+        ],
+    )
 
-    @pre_load
-    def validate_category(self, data):
-        category = data.get("category", "")
-        if len(category) > 80:
-            raise ValidationError(
-                "Challenge could not be saved. Challenge category too long",
-                field_names=["category"],
+    category = field_for(
+        Challenges,
+        "category",
+        validate=[
+            validate.Length(
+                min=0,
+                max=80,
+                error="Challenge could not be saved. Challenge category too long",
             )
+        ],
+    )
 
-    @pre_load
-    def validate_description(self, data):
-        description = data.get("description", "")
-        if len(description) >= 65536:
-            raise ValidationError(
-                "Challenge could not be saved. Challenge description is too long.",
-                field_names=["description"],
+    description = field_for(
+        Challenges,
+        "description",
+        allow_none=True,
+        validate=[
+            validate.Length(
+                min=0,
+                max=65535,
+                error="Challenge could not be saved. Challenge description too long",
             )
+        ],
+    )
diff --git a/CTFd/schemas/pages.py b/CTFd/schemas/pages.py
index 8c889c73..f71a784c 100644
--- a/CTFd/schemas/pages.py
+++ b/CTFd/schemas/pages.py
@@ -1,4 +1,5 @@
-from marshmallow import ValidationError, pre_load
+from marshmallow import pre_load, validate
+from marshmallow_sqlalchemy import field_for
 
 from CTFd.models import Pages, ma
 from CTFd.utils import string_types
@@ -10,34 +11,49 @@ class PageSchema(ma.ModelSchema):
         include_fk = True
         dump_only = ("id",)
 
-    @pre_load
-    def validate_title(self, data):
-        title = data.get("title", "")
-        if len(title) > 128:
-            raise ValidationError(
-                "Page could not be saved. Your title is too long.",
-                field_names=["title"],
+    title = field_for(
+        Pages,
+        "title",
+        validate=[
+            validate.Length(
+                min=0,
+                max=128,
+                error="Page could not be saved. Your title is too long.",
             )
+        ],
+    )
+
+    route = field_for(
+        Pages,
+        "route",
+        allow_none=True,
+        validate=[
+            validate.Length(
+                min=0,
+                max=128,
+                error="Page could not be saved. Your route is too long.",
+            )
+        ],
+    )
+
+    content = field_for(
+        Pages,
+        "content",
+        allow_none=True,
+        validate=[
+            validate.Length(
+                min=0,
+                max=65535,
+                error="Page could not be saved. Your content is too long.",
+            )
+        ],
+    )
 
     @pre_load
     def validate_route(self, data):
-        route = data.get("route", "")
-        if route.startswith("/"):
+        route = data.get("route")
+        if route and route.startswith("/"):
             data["route"] = route.strip("/")
-        if len(route) > 128:
-            raise ValidationError(
-                "Page could not be saved. Your route is too long.",
-                field_names=["route"],
-            )
-
-    @pre_load
-    def validate_content(self, data):
-        content = data.get("content", "")
-        if len(content) >= 65536:
-            raise ValidationError(
-                "Page could not be saved. Your content is too long.",
-                field_names=["content"],
-            )
 
     def __init__(self, view=None, *args, **kwargs):
         if view: