From 3f3109f589945f16a25079b02b74027979447745 Mon Sep 17 00:00:00 2001
From: Kevin Chung <kchung@nyu.edu>
Date: Wed, 29 Apr 2020 20:33:51 -0400
Subject: [PATCH] Add team attrs and cache banning SQL query

---
 CTFd/constants/teams.py               | 20 ++++++++++++++++++++
 CTFd/constants/users.py               |  2 +-
 CTFd/utils/initialization/__init__.py | 14 +++++++++++---
 CTFd/utils/user/__init__.py           | 22 +++++++++++++++++++++-
 serve.py                              |  5 +----
 5 files changed, 54 insertions(+), 9 deletions(-)
 create mode 100644 CTFd/constants/teams.py

diff --git a/CTFd/constants/teams.py b/CTFd/constants/teams.py
new file mode 100644
index 0000000..d9de99d
--- /dev/null
+++ b/CTFd/constants/teams.py
@@ -0,0 +1,20 @@
+from collections import namedtuple
+
+TeamAttrs = namedtuple(
+    "TeamAttrs",
+    [
+        "id",
+        "oauth_id",
+        "name",
+        "email",
+        "secret",
+        "website",
+        "affiliation",
+        "country",
+        "bracket",
+        "hidden",
+        "banned",
+        "captain_id",
+        "created",
+    ],
+)
diff --git a/CTFd/constants/users.py b/CTFd/constants/users.py
index 123e39d..8fcf815 100644
--- a/CTFd/constants/users.py
+++ b/CTFd/constants/users.py
@@ -19,4 +19,4 @@ UserAttrs = namedtuple(
         "team_id",
         "created",
     ],
-)
\ No newline at end of file
+)
diff --git a/CTFd/utils/initialization/__init__.py b/CTFd/utils/initialization/__init__.py
index e245503..85c7694 100644
--- a/CTFd/utils/initialization/__init__.py
+++ b/CTFd/utils/initialization/__init__.py
@@ -38,7 +38,15 @@ from CTFd.utils.plugins import (
 )
 from CTFd.utils.security.auth import login_user, logout_user, lookup_user_token
 from CTFd.utils.security.csrf import generate_nonce
-from CTFd.utils.user import authed, get_current_team, get_current_user, get_ip, is_admin
+from CTFd.utils.user import (
+    authed,
+    get_current_team,
+    get_current_user,
+    get_current_user_attrs,
+    get_current_team_attrs,
+    get_ip,
+    is_admin,
+)
 
 
 def init_template_filters(app):
@@ -191,8 +199,8 @@ def init_request_processors(app):
             return
 
         if authed():
-            user = get_current_user()
-            team = get_current_team()
+            user = get_current_user_attrs()
+            team = get_current_team_attrs()
 
             if user and user.banned:
                 return (
diff --git a/CTFd/utils/user/__init__.py b/CTFd/utils/user/__init__.py
index c754b5e..68c2989 100644
--- a/CTFd/utils/user/__init__.py
+++ b/CTFd/utils/user/__init__.py
@@ -6,6 +6,7 @@ from flask import request, session
 
 from CTFd.cache import cache
 from CTFd.constants.users import UserAttrs
+from CTFd.constants.teams import TeamAttrs
 from CTFd.models import Fails, Users, db, Teams
 from CTFd.utils import get_config
 
@@ -33,7 +34,7 @@ def get_user_attrs(user_id):
         for field in UserAttrs._fields:
             d[field] = getattr(user, field)
         return UserAttrs(**d)
-    return user
+    return None
 
 
 def get_current_team():
@@ -44,6 +45,25 @@ def get_current_team():
         return None
 
 
+def get_current_team_attrs():
+    if authed():
+        user = get_user_attrs(user_id=session["id"])
+        if user.team_id:
+            return get_team_attrs(team_id=user.team_id)
+    return None
+
+
+@cache.memoize()
+def get_team_attrs(team_id):
+    team = Teams.query.filter_by(id=team_id).first()
+    if team:
+        d = {}
+        for field in TeamAttrs._fields:
+            d[field] = getattr(user, field)
+        return TeamAttrs(**d)
+    return None
+
+
 def get_current_user_type(fallback=None):
     if authed():
         user = get_current_user_attrs()
diff --git a/serve.py b/serve.py
index 1e3a79a..962e577 100644
--- a/serve.py
+++ b/serve.py
@@ -18,10 +18,7 @@ if args.profile:
         "enabled": app.config["DEBUG"],
         "storage": {"engine": "sqlite"},
         "basicAuth": {"enabled": False},
-        "ignore": [
-	        "^/themes/.*",
-	        "^/events",
-	    ]
+        "ignore": ["^/themes/.*", "^/events"],
     }
     flask_profiler.init_app(app)
     app.config["DEBUG_TB_PROFILER_ENABLED"] = True