From 54ffb6fa7aacf5b7b241b919e7bc48c40de74e17 Mon Sep 17 00:00:00 2001
From: Tim Jensen <tim.l.jensen@gmail.com>
Date: Sat, 25 Sep 2021 11:59:31 -0500
Subject: [PATCH] Adding optional binary option to a2s_rules function

---
 steam/game_servers.py      | 17 +++++-----
 tests/test_game_servers.py | 64 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 74 insertions(+), 7 deletions(-)
 create mode 100644 tests/test_game_servers.py

diff --git a/steam/game_servers.py b/steam/game_servers.py
index 07b09f1..650715b 100644
--- a/steam/game_servers.py
+++ b/steam/game_servers.py
@@ -142,8 +142,11 @@ def _u(data):
 
 
 class StructReader(_StructReader):
-    def read_cstring(self):
-        return _u(super(StructReader, self).read_cstring())
+    def read_cstring(self, binary=False):
+        raw = super(StructReader, self).read_cstring()
+        if binary:
+            return raw
+        return _u(raw)
 
 
 class MSRegion(IntEnum):
@@ -526,7 +529,7 @@ def a2s_players(server_addr, timeout=2, challenge=0):
     return players
 
 
-def a2s_rules(server_addr, timeout=2, challenge=0):
+def a2s_rules(server_addr, timeout=2, challenge=0, binary=False):
     """Get rules from server
 
     :param server_addr: (ip, port) for the server
@@ -571,12 +574,12 @@ def a2s_rules(server_addr, timeout=2, challenge=0):
     rules = {}
 
     while len(rules) != num_rules:
-        name = data.read_cstring()
-        value = data.read_cstring()
+        name = data.read_cstring(binary=binary)
+        value = data.read_cstring(binary=binary)
 
-        if _re_match(r'^\-?[0-9]+$', value):
+        if not binary and _re_match(r'^\-?[0-9]+$', value):
             value = int(value)
-        elif _re_match(r'^\-?[0-9]+\.[0-9]+$', value):
+        elif not binary and _re_match(r'^\-?[0-9]+\.[0-9]+$', value):
             value = float(value)
 
         rules[name] = value
diff --git a/tests/test_game_servers.py b/tests/test_game_servers.py
new file mode 100644
index 0000000..c98ca91
--- /dev/null
+++ b/tests/test_game_servers.py
@@ -0,0 +1,64 @@
+import mock
+import socket
+import unittest
+
+from steam.game_servers import a2s_rules
+
+
+class TestA2SRules(unittest.TestCase):
+    @mock.patch("socket.socket")
+    def test_returns_rules_with_default_arguments(self, mock_socket_class):
+        mock_socket = mock_socket_class.return_value
+        mock_socket.recv.side_effect = [
+            b"\xff\xff\xff\xffA\x01\x02\x03\x04",
+            b"\xff\xff\xff\xffE\x03\0text\0b\x99r\0int\x0042\0float\x0021.12\0"
+        ]
+
+        rules = a2s_rules(("addr", 1234))
+
+        self.assertEqual(
+            {
+                "text": "b\ufffdr",
+                "int": 42,
+                "float": 21.12
+            },
+            rules)
+
+        mock_socket_class.assert_called_once_with(
+            socket.AF_INET, socket.SOCK_DGRAM)
+
+        mock_socket.connect.assert_called_once_with(("addr", 1234))
+        mock_socket.settimeout.assert_called_once_with(2)
+
+        self.assertEqual(2, mock_socket.send.call_count)
+        mock_socket.send.assert_has_calls([
+            mock.call(b"\xff\xff\xff\xffV\0\0\0\0"),
+            mock.call(b"\xff\xff\xff\xffV\x01\x02\x03\x04")
+        ])
+
+        self.assertEqual(2, mock_socket.recv.call_count)
+        mock_socket.recv.assert_has_calls([
+            mock.call(512),
+            mock.call(2048)
+        ])
+
+        mock_socket.close.assert_called_once_with()
+
+    @mock.patch("socket.socket")
+    def test_returns_rules_as_bytes_when_binary_is_true(
+            self, mock_socket_class):
+        mock_socket = mock_socket_class.return_value
+        mock_socket.recv.side_effect = [
+            b"\xff\xff\xff\xffA\x01\x02\x03\x04",
+            b"\xff\xff\xff\xffE\x03\0text\0b\x99r\0int\x0042\0float\x0021.12\0"
+        ]
+
+        rules = a2s_rules(("addr", 1234), binary=True)
+
+        self.assertEqual(
+            {
+                b"text": b"b\x99r",
+                b"int": b"42",
+                b"float": b"21.12"
+            },
+            rules)