From 02cb0e834873b27d822e7a57fd0bd1a0c5550b1f Mon Sep 17 00:00:00 2001
From: informaniac <informaniac@shodan.local>
Date: Fri, 1 Dec 2017 17:36:36 +0100
Subject: [PATCH] added blup balance util

---
 blup/balance_util.py | 138 +++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 138 insertions(+)
 create mode 100644 blup/balance_util.py

diff --git a/blup/balance_util.py b/blup/balance_util.py
new file mode 100644
index 0000000..8fa8626
--- /dev/null
+++ b/blup/balance_util.py
@@ -0,0 +1,138 @@
+"""
+this file is part of blup. it provides an easy way
+to determine the direction the player is leaning
+"""
+from enum import Enum
+import struct
+
+Direction = Enum('BalanceDirection',
+                 ['LEFT', 'RIGHT', 'UP', 'DOWN',
+                  'CENTER', 'QUIT'])
+
+
+def sanitize_threshold(threshold):
+    """
+    function to make sure the threshold is positive
+    threshold -- the given (untrusted) treshold
+    """
+    if threshold < 0:
+        threshold = abs(threshold)
+    return threshold
+
+
+class BalanceSocket(object):
+    def __init__(self, balance_server_socket):
+        self.socket = balance_server_socket
+
+    def close(self):
+        self.socket.close()
+
+    def _get_raw_data(self):
+        self.socket.send(b'a')
+        data = self.socket.recv(4)
+        p0x, p0y, p1x, p1y = struct.unpack('bbbb', data)
+        return p0x, p0y, p1x, p1y
+
+    def get_player1_ready(self):
+        """
+        check if someone is on board 1
+        :return: true if someone is on the board
+        """
+        p0x, p0y, p1x, p1y = self._get_raw_data()
+        return p0x > -128 or p0y > -128
+
+    def get_player2_ready(self):
+        """
+        check if someone is on board 2
+        :return: true if someone is on the board
+        """
+        p0x, p0y, p1x, p1y = self._get_raw_data()
+        return p1x > -128 or p1y > -128
+
+
+class BalanceUtil(BalanceSocket):
+    def __init__(self, balance_server_socket, player_id):
+        """
+        BalanceUtil object
+        balance_socket -- (ip, port) of the balance server
+        player_id -- id of the player
+        """
+        super(BalanceUtil, self).__init__(balance_server_socket)
+        self.player_id = player_id
+
+    def _get_raw_data(self):
+        self.socket.send(b'a')
+        data = self.socket.recv(4)
+        p0x, p0y, p1x, p1y = struct.unpack('bbbb', data)
+        if self.player_id == 0:
+            return p0x, p0y
+        else:
+            return p1x, p1y
+
+    def get_raw_2dir_y(self):
+        """
+        Get raw 2-directional data from the y axis of the board.
+        e.g. determine how fast the paddle is supposed to move
+        in balance pong
+        :return: raw y axis balance value
+        """
+        x_bal, y_bal = self._get_raw_data()
+        return y_bal
+
+    def get_2dir_y(self, threshold):
+        threshold = sanitize_threshold(threshold)
+        x_bal, y_bal = self._get_raw_data()
+        if x_bal == -128 or y_bal == -128:
+            return Direction.QUIT
+        elif abs(y_bal) < threshold:
+            return Direction.CENTER
+        elif y_bal > 0:
+            return Direction.UP
+        else:
+            return Direction.DOWN
+
+    def get_raw_2dir_x(self):
+        x_bal, y_bal = self._get_raw_data()
+        return x_bal
+
+    def get_2dir_x(self, threshold):
+        threshold = sanitize_threshold(threshold)
+        x_bal, y_bal = self._get_raw_data()
+        if x_bal == -128 or y_bal == -128:
+            return Direction.QUIT
+        elif abs(x_bal) < threshold:
+            return Direction.CENTER
+        elif x_bal > 0:
+            return Direction.RIGHT
+        else:
+            return Direction.LEFT
+
+    def get_raw_4dir(self):
+        return self._get_raw_data()
+
+    def get_4dir(self, threshold):
+        """
+        Get evaluated 4-directional data of the board
+        e.g. determine the direction and rotation or drop move
+        in balance tetris
+
+        threshold -- deadzone of the board where nothing is
+        supposed to happen
+        :return: Enum Direction value of direction or event
+        """
+        threshold = sanitize_threshold(threshold)
+        x_bal, y_bal = self._get_raw_data()
+        if x_bal == -128 or y_bal == -128:
+            return Direction.QUIT
+        elif abs(x_bal) < threshold and abs(y_bal) < threshold:
+            return Direction.CENTER
+        elif abs(x_bal) < abs(y_bal):
+            if y_bal > 0:
+                return Direction.UP
+            else:
+                return Direction.DOWN
+        else:
+            if x_bal > 0:
+                return Direction.RIGHT
+            else:
+                return Direction.LEFT