From d78d7518f9a53a69008154e58315f33af45a1c90 Mon Sep 17 00:00:00 2001
From: Kevin Chung <kchung@nyu.edu>
Date: Thu, 20 Aug 2020 03:20:21 -0400
Subject: [PATCH] Add initial schema changes

---
 CTFd/schemas/fields.py | 16 +++++++++++++++-
 CTFd/schemas/teams.py  | 11 ++++++++++-
 2 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/CTFd/schemas/fields.py b/CTFd/schemas/fields.py
index a6c8841a..0199b9d7 100644
--- a/CTFd/schemas/fields.py
+++ b/CTFd/schemas/fields.py
@@ -1,6 +1,6 @@
 from marshmallow import fields
 
-from CTFd.models import Fields, UserFieldEntries, db, ma
+from CTFd.models import Fields, TeamFieldEntries, UserFieldEntries, db, ma
 
 
 class FieldSchema(ma.ModelSchema):
@@ -22,3 +22,17 @@ class UserFieldEntriesSchema(ma.ModelSchema):
     name = fields.Nested(FieldSchema, only=("name"), attribute="field")
     description = fields.Nested(FieldSchema, only=("description"), attribute="field")
     type = fields.Nested(FieldSchema, only=("field_type"), attribute="field")
+
+
+class TeamFieldEntriesSchema(ma.ModelSchema):
+    class Meta:
+        model = TeamFieldEntries
+        sqla_session = db.session
+        include_fk = True
+        load_only = ("id",)
+        exclude = ("field", "team", "team_id")
+        dump_only = ("team_id", "name", "description", "type")
+
+    name = fields.Nested(FieldSchema, only=("name"), attribute="field")
+    description = fields.Nested(FieldSchema, only=("description"), attribute="field")
+    type = fields.Nested(FieldSchema, only=("field_type"), attribute="field")
diff --git a/CTFd/schemas/teams.py b/CTFd/schemas/teams.py
index 669a6ed0..2b8b428f 100644
--- a/CTFd/schemas/teams.py
+++ b/CTFd/schemas/teams.py
@@ -1,7 +1,9 @@
 from marshmallow import ValidationError, pre_load, validate
+from marshmallow.fields import Nested
 from marshmallow_sqlalchemy import field_for
 
-from CTFd.models import Teams, Users, ma
+from CTFd.models import TeamFieldEntries, TeamFields, Teams, Users, ma
+from CTFd.schemas.fields import TeamFieldEntriesSchema
 from CTFd.utils import get_config, string_types
 from CTFd.utils.crypto import verify_password
 from CTFd.utils.user import get_current_team, get_current_user, is_admin
@@ -44,6 +46,9 @@ class TeamSchema(ma.ModelSchema):
         ],
     )
     country = field_for(Teams, "country", validate=[validate_country_code])
+    fields = Nested(
+        TeamFieldEntriesSchema, partial=True, many=True, attribute="field_entries"
+    )
 
     @pre_load
     def validate_name(self, data):
@@ -197,6 +202,7 @@ class TeamSchema(ma.ModelSchema):
             "id",
             "oauth_id",
             "captain_id",
+            "fields",
         ],
         "self": [
             "website",
@@ -210,6 +216,7 @@ class TeamSchema(ma.ModelSchema):
             "oauth_id",
             "password",
             "captain_id",
+            "fields",
         ],
         "admin": [
             "website",
@@ -227,6 +234,7 @@ class TeamSchema(ma.ModelSchema):
             "oauth_id",
             "password",
             "captain_id",
+            "fields",
         ],
     }
 
@@ -236,5 +244,6 @@ class TeamSchema(ma.ModelSchema):
                 kwargs["only"] = self.views[view]
             elif isinstance(view, list):
                 kwargs["only"] = view
+        self.view = view
 
         super(TeamSchema, self).__init__(*args, **kwargs)