Add optional "binary" argument to a2s_rules function (#359)

* Adding optional binary option to a2s_rules function
* Adding "u" prefix to expected string value to keep legacy Python compatibility
* Setting upper limit on protobuf version to keep legacy Python compatibility because protobuf version 3.18 dropped 2.7 support
* Adding docs for new a2s_rules optional binary argument
* Correcting return type for a2s_rules function
* Setting upper limit on protobuf version to keep legacy Python compatibility because protobuf version 3.18 dropped 2.7 support
* Lifting protobuf upper version limit for Python 3, per PR comment
* Removing duplicate binary check, per PR comment
* Lifting protobuf upper version limit for Python 3, per PR comment

requirements.txt

@@ -3,7 +3,8 @@ pycryptodomex>=3.7.0
 requests>=2.9.1
 vdf>=3.3
 gevent>=1.3.0
-protobuf>=3.0.0
+protobuf>~3.0; python_version >= '3'
+protobuf<3.18.0; python_version < '3'
 gevent-eventemitter~=2.1
 cachetools>=3.0.0
 enum34==1.1.2; python_version < '3.4'

setup.py

@@ -26,7 +26,8 @@ install_requires = [
 install_extras = {
     'client': [
         'gevent>=1.3.0',
-        'protobuf>=3.0.0',
+        'protobuf>~3.0; python_version >= "3"',
+        'protobuf<3.18.0; python_version < "3"',
         'gevent-eventemitter~=2.1',
     ],
 }

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
@@ -535,9 +538,11 @@ def a2s_rules(server_addr, timeout=2, challenge=0):
     :type  timeout: float
     :param challenge: (optional) challenge number
     :type  challenge: int
+    :param binary: (optional) return rules as raw bytes
+    :type  binary: bool
     :raises: :class:`RuntimeError`, :class:`socket.timeout`
     :returns: a list of rules
-    :rtype: :class:`list`
+    :rtype: :class:`dict`
     """
     ss = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
     ss.connect(server_addr)
@@ -571,13 +576,14 @@ def a2s_rules(server_addr, timeout=2, challenge=0):
     rules = {}
 
     while len(rules) != num_rules:
-        name = data.read_cstring()
-        value = data.read_cstring()
-
-        if _re_match(r'^\-?[0-9]+$', value):
-            value = int(value)
-        elif _re_match(r'^\-?[0-9]+\.[0-9]+$', value):
-            value = float(value)
+        name = data.read_cstring(binary=binary)
+        value = data.read_cstring(binary=binary)
+
+        if not binary:
+            if _re_match(r'^\-?[0-9]+$', value):
+                value = int(value)
+            elif _re_match(r'^\-?[0-9]+\.[0-9]+$', value):
+                value = float(value)
 
         rules[name] = value
 

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": u"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)