Merge pull request #49 from python-discord/ks123/role-assigning

Implement Discord roles assigning for form submissions

4 files changed, 78 insertions, 14 deletions

diff --git a/backend/constants.py b/backend/constants.py
index d90fd9a..7ea4519 100644
--- a/backend/constants.py
+++ b/backend/constants.py
@@ -28,6 +28,8 @@ FORMS_BACKEND_DSN = os.getenv("FORMS_BACKEND_DSN")
 DOCS_PASSWORD = os.getenv("DOCS_PASSWORD")
 
 SECRET_KEY = os.getenv("SECRET_KEY", binascii.hexlify(os.urandom(30)).decode())
+DISCORD_BOT_TOKEN = os.getenv("DISCORD_BOT_TOKEN")
+DISCORD_GUILD = os.getenv("DISCORD_GUILD", 267624335836053506)
 
 HCAPTCHA_API_SECRET = os.getenv("HCAPTCHA_API_SECRET")
 
@@ -60,6 +62,8 @@ REQUIRED_QUESTION_TYPE_DATA = {
     },
 }
 
+DISCORD_API_BASE_URL = "https://discord.com/api/v8"
+
 
 class FormFeatures(Enum):
     """Lists form features. Read more in SCHEMA.md."""
@@ -70,6 +74,7 @@ class FormFeatures(Enum):
     COLLECT_EMAIL = "COLLECT_EMAIL"
     DISABLE_ANTISPAM = "DISABLE_ANTISPAM"
     WEBHOOK_ENABLED = "WEBHOOK_ENABLED"
+    ASSIGN_ROLE = "ASSIGN_ROLE"
 
 
 class WebHook(Enum):
diff --git a/backend/discord.py b/backend/discord.py
index 8cb602c..e5c7f8f 100644
--- a/backend/discord.py
+++ b/backend/discord.py
@@ -2,11 +2,9 @@
 import httpx
 
 from backend.constants import (
-    OAUTH2_CLIENT_ID, OAUTH2_CLIENT_SECRET
+    DISCORD_API_BASE_URL, OAUTH2_CLIENT_ID, OAUTH2_CLIENT_SECRET
 )
 
-API_BASE_URL = "https://discord.com/api/v8"
-
 
 async def fetch_bearer_token(code: str, redirect: str, *, refresh: bool) -> dict:
     async with httpx.AsyncClient() as client:
@@ -23,7 +21,7 @@ async def fetch_bearer_token(code: str, redirect: str, *, refresh: bool) -> dict
             data["grant_type"] = "authorization_code"
             data["code"] = code
 
-        r = await client.post(f"{API_BASE_URL}/oauth2/token", headers={
+        r = await client.post(f"{DISCORD_API_BASE_URL}/oauth2/token", headers={
             "Content-Type": "application/x-www-form-urlencoded"
         }, data=data)
 
@@ -34,7 +32,7 @@ async def fetch_bearer_token(code: str, redirect: str, *, refresh: bool) -> dict
 
 async def fetch_user_details(bearer_token: str) -> dict:
     async with httpx.AsyncClient() as client:
-        r = await client.get(f"{API_BASE_URL}/users/@me", headers={
+        r = await client.get(f"{DISCORD_API_BASE_URL}/users/@me", headers={
             "Authorization": f"Bearer {bearer_token}"
         })
 
diff --git a/backend/models/form.py b/backend/models/form.py
index eac0b63..30ae0e7 100644
--- a/backend/models/form.py
+++ b/backend/models/form.py
@@ -1,13 +1,21 @@
 import typing as t
 
 import httpx
-from pydantic import BaseModel, Field, validator
+from pydantic import BaseModel, Field, root_validator, validator
 from pydantic.error_wrappers import ErrorWrapper, ValidationError
 
 from backend.constants import FormFeatures, WebHook
 from .question import Question
 
-PUBLIC_FIELDS = ["id", "features", "questions", "name", "description", "submitted_text"]
+PUBLIC_FIELDS = [
+    "id",
+    "features",
+    "questions",
+    "name",
+    "description",
+    "submitted_text",
+    "discord_role"
+]
 
 
 class _WebHook(BaseModel):
@@ -34,6 +42,7 @@ class Form(BaseModel):
     description: str
     submitted_text: t.Optional[str] = None
     webhook: _WebHook = None
+    discord_role: t.Optional[str]
 
     class Config:
         allow_population_by_field_name = True
@@ -47,11 +56,30 @@ class Form(BaseModel):
         if any(v not in allowed_values for v in value):
             raise ValueError("Form features list contains one or more invalid values.")
 
-        if FormFeatures.COLLECT_EMAIL in value and FormFeatures.REQUIRES_LOGIN not in value:
-            raise ValueError("COLLECT_EMAIL feature require REQUIRES_LOGIN feature.")
+        if FormFeatures.REQUIRES_LOGIN.value not in value:
+            if FormFeatures.COLLECT_EMAIL.value in value:
+                raise ValueError(
+                    "COLLECT_EMAIL feature require REQUIRES_LOGIN feature."
+                )
+
+            if FormFeatures.ASSIGN_ROLE.value in value:
+                raise ValueError("ASSIGN_ROLE feature require REQUIRES_LOGIN feature.")
 
         return value
 
+    @root_validator
+    def validate_role(cls, values: dict[str, t.Any]) -> t.Optional[dict[str, t.Any]]:
+        """Validates does Discord role provided when flag provided."""
+        if (
+            FormFeatures.ASSIGN_ROLE.value in values.get("features", [])
+            and not values.get("discord_role")
+        ):
+            raise ValueError(
+                "discord_role field is required when ASSIGN_ROLE flag is provided."
+            )
+
+        return values
+
     def dict(self, admin: bool = True, **kwargs: t.Any) -> dict[str, t.Any]:
         """Wrapper for original function to exclude private data for public access."""
         data = super().dict(**kwargs)
diff --git a/backend/routes/forms/submit.py b/backend/routes/forms/submit.py
index 2624c98..4d15ab7 100644
--- a/backend/routes/forms/submit.py
+++ b/backend/routes/forms/submit.py
@@ -2,6 +2,7 @@
 Submit a form.
 """
 
+import asyncio
 import binascii
 import datetime
 import hashlib
@@ -12,7 +13,7 @@ import httpx
 from pydantic import ValidationError
 from pydantic.main import BaseModel
 from spectree import Response
-from starlette.background import BackgroundTask
+from starlette.background import BackgroundTasks
 from starlette.requests import Request
 from starlette.responses import JSONResponse
 
@@ -29,6 +30,10 @@ HCAPTCHA_HEADERS = {
     "Content-Type": "application/x-www-form-urlencoded"
 }
 
+DISCORD_HEADERS = {
+    "Authorization": f"Bot {constants.DISCORD_BOT_TOKEN}"
+}
+
 
 class SubmissionResponse(BaseModel):
     form: Form
@@ -180,19 +185,26 @@ class SubmitForm(Route):
                 response_obj.dict(by_alias=True)
             )
 
-            send_webhook = None
+            tasks = BackgroundTasks()
             if constants.FormFeatures.WEBHOOK_ENABLED.value in form.features:
-                send_webhook = BackgroundTask(
+                tasks.add_task(
                     self.send_submission_webhook,
                     form=form,
                     response=response_obj,
                     request_user=request.user
                 )
 
+            if constants.FormFeatures.ASSIGN_ROLE.value in form.features:
+                tasks.add_task(
+                    self.assign_role,
+                    form=form,
+                    request_user=request.user
+                )
+
             return JSONResponse({
                 "form": form.dict(admin=False),
                 "response": response_obj.dict()
-            }, background=send_webhook)
+            }, background=tasks)
 
         else:
             return JSONResponse({
@@ -203,7 +215,7 @@ class SubmitForm(Route):
     async def send_submission_webhook(
             form: Form,
             response: FormResponse,
-            request_user: Request.user
+            request_user: User
     ) -> None:
         """Helper to send a submission message to a discord webhook."""
         # Stop if webhook is not available
@@ -262,3 +274,24 @@ class SubmitForm(Route):
         async with httpx.AsyncClient() as client:
             r = await client.post(form.webhook.url, json=hook)
             r.raise_for_status()
+
+    @staticmethod
+    async def assign_role(form: Form, request_user: User) -> None:
+        """Assigns Discord role to user when user submitted response."""
+        if not form.discord_role:
+            raise ValueError("Got empty Discord role ID.")
+
+        url = (
+            f"{constants.DISCORD_API_BASE_URL}/guilds/{constants.DISCORD_GUILD}"
+            f"/members/{request_user.payload['id']}/roles/{form.discord_role}"
+        )
+
+        async with httpx.AsyncClient() as client:
+            resp = await client.put(url, headers=DISCORD_HEADERS)
+            # Handle Rate Limits
+            while resp.status_code == 429:
+                retry_after = float(resp.headers["X-Ratelimit-Reset-After"])
+                await asyncio.sleep(retry_after)
+                resp = await client.put(url, headers=DISCORD_HEADERS)
+
+            resp.raise_for_status()