game socket code audit done

2 files changed, 72 insertions, 27 deletions

diff --git a/api/game/socket.py b/api/game/socket.py
index 8bbc951..6038586 100644
--- a/api/game/socket.py
+++ b/api/game/socket.py
@@ -1,15 +1,42 @@
 from flask import Blueprint, request, make_response
 from flask_socketio import SocketIO, emit, Namespace, join_room
 from game.voerbak_connector import bord
-from auth.login_token import token_login
 from db import cursor, connection
+from socket_io import io
+from hierarchy import io_auth_required
+from auth.login_token import token_login
 import time
 import json
-from socket_io import io
 
 games = {}
 
 
+def participants_only(func):
+	'''
+	listener should have two parameters:
+	listener(data: socket.io.data, user_id: str, game: game)
+
+	listener should only be executed if the request comes from one of
+	the game participants (player_1_id | player_2_id)
+	'''
+	def wrapper(data, user_id):
+		game_id = data["game_id"]
+
+		if not game_id or \
+           not game_id in games:
+			return
+
+		game = games[game_id]
+		if game.player_1_id != user_id and \
+           game.player_2_id != user_id:
+			return
+
+		return func(data, user_id, game)
+
+	wrapper.__name__ = func.__name__
+	return wrapper
+
+
 class game:
 	def __init__(self, game_id, io, player_1_id, player_2_id):
 		self.game_id = game_id
@@ -21,7 +48,9 @@ class game:
 
 	# drop a disc in `column`
 	def move(self, user_id, column):
-		if user_id != self.player_1_id and user_id != self.player_2_id: return
+		if len(self.board.win_positions) > 0: return
+		if self.board.board_full: return
+
 		move = self.player_1_id if self.board.player_1 else self.player_2_id
 		if user_id != move: return
 
@@ -79,34 +108,20 @@ class game:
 
 
 @io.on("newMove")
-def new_move(data):
-	if not data["game_id"] or \
-                                  not data["move"] or \
-                                  not data["token"]:
-		return
-	if not data["game_id"] in games: return
+@io_auth_required("none")
+@participants_only
+def new_move(data, user_id, game):
+	move = data.get("move")
+	if not move: return
 
-	game = games[data["game_id"]]
-	if (len(game.board.win_positions) > 0 or game.board.board_full): return
-	user_id = token_login(data["token"])
-	game.move(user_id, data["move"])
+	game.move(user_id, move)
 
 
 @io.on("resign")
-def resign(data):
-	if not data["game_id"] or \
-                                  not request.cookies.get("token"):
-		return
-	if not data["game_id"] in games: return
-
-	user_id = token_login(request.cookies.get("token"))
-	if not user_id: return
-
-	if games[data["game_id"]].player_1_id != user_id and \
-                                  games[data["game_id"]].player_2_id != user_id:
-		return
-
-	games[data["game_id"]].resign()
+@io_auth_required("none")
+@participants_only
+def resign(data, user_id, game):
+	game.resign()
 
 
 @io.on("registerGameListener")
diff --git a/api/hierarchy.py b/api/hierarchy.py
index 4e065eb..2e7db66 100644
--- a/api/hierarchy.py
+++ b/api/hierarchy.py
@@ -105,3 +105,33 @@ def auth_required(level):
 		return wrapper
 
 	return decorator
+
+
+def io_auth_required(level):
+	'''
+	level = "none" | "user" | "moderator" | "admin" | "bot"
+	endpoint should have two parameters:
+	endpoint(data: socket.io.data, user_id: str) # `user_id` can only be `None` when `level == "none"`
+
+	uses the @auth_required decorator, but is only for use with
+	@io.on decorators
+	'''
+	def decorator(func):
+		# data is the original @io.on data
+		def wrapper(data):
+
+			token = request.cookies.get("token") or ""
+			user_id = token_login(token)
+
+			if not user_id:
+				if level == ranks[0]:
+					return func(data, None)
+				else:
+					return
+
+			return func(data, user_id)
+
+		wrapper.__name__ = func.__name__
+		return wrapper
+
+	return decorator