From aeac3c4b1062aa15071b4cfba6b37d594ef2cfd2 Mon Sep 17 00:00:00 2001
From: Ske <voltasalt@gmail.com>
Date: Mon, 11 Mar 2019 21:53:08 +0100
Subject: [PATCH] Add OAuth2 token flow to API

---
 docker-compose.yml      |  3 +++
 src/api_main.py         | 32 ++++++++++++++++++++++++++++++--
 src/pluralkit/system.py |  5 +++++
 3 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/docker-compose.yml b/docker-compose.yml
index b233579c..ccf3f073 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -24,6 +24,9 @@ services:
     - "2939:8080"
     environment:
     - "DATABASE_URI=postgres://postgres:postgres@db:5432/postgres"
+    - "CLIENT_ID"
+    - "CLIENT_SECRET"
+    - "REDIRECT_URI"
   db:
     image: postgres:alpine
     volumes:
diff --git a/src/api_main.py b/src/api_main.py
index 39ebc68d..e7527962 100644
--- a/src/api_main.py
+++ b/src/api_main.py
@@ -2,7 +2,7 @@ import json
 import logging
 import os
 
-from aiohttp import web
+from aiohttp import web, ClientSession
 
 from pluralkit import db, utils
 from pluralkit.errors import PluralKitError
@@ -166,6 +166,33 @@ class Handlers:
 
         return web.json_response(await switch.to_json(hid_getter))
 
+    async def discord_oauth(request):
+        code = await request.text()
+        async with ClientSession() as sess:
+            data = {
+                'client_id': os.environ["CLIENT_ID"],
+                'client_secret': os.environ["CLIENT_SECRET"],
+                'grant_type': 'authorization_code',
+                'code': code,
+                'redirect_uri': os.environ["REDIRECT_URI"],
+                'scope': 'identify'
+            }
+            headers = {
+                'Content-Type': 'application/x-www-form-urlencoded'
+            }
+            res = await sess.post("https://discordapp.com/api/v6/oauth2/token", data=data, headers=headers)
+            if res.status != 200:
+                raise web.HTTPBadRequest()
+
+            access_token = (await res.json())["access_token"]
+            res = await sess.get("https://discordapp.com/api/v6/users/@me", headers={"Authorization": "Bearer " + access_token})
+            user_id = int((await res.json())["id"])
+            
+            system = await System.get_by_account(request["conn"], user_id)
+            if not system:
+                raise web.HTTPUnauthorized()
+            return web.Response(text=await system.get_token(request["conn"]))
+
 async def run():
     app = web.Application(middlewares=[db_middleware, auth_middleware, error_middleware])
 
@@ -179,7 +206,8 @@ async def run():
         web.get("/m/{member}", Handlers.get_member),
         web.post("/m", Handlers.post_member),
         web.patch("/m/{member}", Handlers.patch_member),
-        web.delete("/m/{member}", Handlers.delete_member)
+        web.delete("/m/{member}", Handlers.delete_member),
+        web.post("/discord_oauth", Handlers.discord_oauth)
     ])
     app["pool"] = await db.connect(
         os.environ["DATABASE_URI"]
diff --git a/src/pluralkit/system.py b/src/pluralkit/system.py
index a0ea39b8..2112c650 100644
--- a/src/pluralkit/system.py
+++ b/src/pluralkit/system.py
@@ -116,6 +116,11 @@ class System(namedtuple("System", ["id", "hid", "name", "description", "tag", "a
         await db.update_system_field(conn, self.id, "token", new_token)
         return new_token
 
+    async def get_token(self, conn) -> str:
+        if self.token:
+            return self.token
+        return await self.refresh_token(conn)
+
     async def create_member(self, conn, member_name: str) -> Member:
         # TODO: figure out what to do if this errors out on collision on generate_hid
         new_hid = generate_hid()