From 86e40475822792c26ed9fb7a2ee2cc10fe484585 Mon Sep 17 00:00:00 2001
From: Rossen Georgiev <rossen@rgp.io>
Date: Wed, 16 Nov 2016 22:37:02 +0200
Subject: [PATCH] StructReader improvements

* added read() to get raw bytes
* renamed read_format() to unpack()
* added parameter in read_cstring() to specify terminator
---
 docs/api/steam.util.rst   | 16 ++++++++--------
 steam/core/msg/structs.py | 12 ++++++------
 steam/util/binary.py      | 23 +++++++++++++++++------
 3 files changed, 31 insertions(+), 20 deletions(-)

diff --git a/docs/api/steam.util.rst b/docs/api/steam.util.rst
index a2585f3..759d9cb 100644
--- a/docs/api/steam.util.rst
+++ b/docs/api/steam.util.rst
@@ -6,6 +6,14 @@ util
     :undoc-members:
     :show-inheritance:
 
+util.binary
+-----------
+
+.. automodule:: steam.util.binary
+    :members:
+    :undoc-members:
+    :show-inheritance:
+
 util.throttle
 -------------
 
@@ -21,11 +29,3 @@ util.web
     :members:
     :undoc-members:
     :show-inheritance:
-
-util.binary
------------
-
-.. automodule:: steam.util.binary
-    :members:
-    :undoc-members:
-    :show-inheritance:
diff --git a/steam/core/msg/structs.py b/steam/core/msg/structs.py
index 9a88655..9322fb8 100644
--- a/steam/core/msg/structs.py
+++ b/steam/core/msg/structs.py
@@ -139,13 +139,13 @@ class ClientVACBanStatus(StructMessage):
 
     def load(self, data):
         buf = StructReader(data)
-        numBans, = buf.read_format("<I")
+        numBans, = buf.unpack("<I")
 
         for _ in range(numBans):
             m = self.VACBanRange()
             self.ranges.append(m)
 
-            m.start, m.end, _ = buf.read_format("<III")
+            m.start, m.end, _ = buf.unpack("<III")
 
             if m.start > m.end:
                 m.start, m.end = m.end, m.start
@@ -181,7 +181,7 @@ class ClientChatMsg(StructMessage):
 
     def load(self, data):
         buf = StructReader(data)
-        self.steamIdChatter, self.steamIdChatRoom, self.ChatMsgType = buf.read_format("<QQI")
+        self.steamIdChatter, self.steamIdChatRoom, self.ChatMsgType = buf.unpack("<QQI")
         self.text = buf.read_cstring().decode('utf-8')
 
     def __str__(self):
@@ -269,15 +269,15 @@ class ClientMarketingMessageUpdate2(StructMessage):
 
     def load(self, data):
         buf = StructReader(data)
-        self.time, count = buf.read_format("<II")
+        self.time, count = buf.unpack("<II")
 
         for _ in range(count):
             m = self.MarketingMessage()
             self.messages.append(m)
 
-            length, m.id = buf.read_format("<IQ")
+            length, m.id = buf.unpack("<IQ")
             m.url = buf.read_cstring().decode('utf-8')
-            m.flags = buf.read_format("<I")
+            m.flags = buf.unpack("<I")
 
     def __str__(self):
         text = ["time: %s" % self.time,
diff --git a/steam/util/binary.py b/steam/util/binary.py
index 56043ad..9050d2c 100644
--- a/steam/util/binary.py
+++ b/steam/util/binary.py
@@ -12,18 +12,29 @@ class StructReader(object):
         self.data = data
         self.offset = 0
 
-    def read_cstring(self):
+    def read(self, n=1):
+        """Return n bytes
+
+        :param n: number of bytes to return
+        :type  n: :class:`int`
+        :return: bytes
+        :rtype: :class:`bytes`
+        """
+        self.offset += n
+        return self.data[self.offset - n:self.offset]
+
+    def read_cstring(self, terminator=b'\x00'):
         """Reads a single null termianted string
 
         :return: string without bytes
         :rtype: :class:`bytes`
         """
-        null_index = self.data.find(b'\x00', self.offset)
-        text = self.data[self.offset:null_index]  # bytes without the null
-        self.offset = null_index + 1  # advanced past null
-        return text
+        null_index = self.data.find(terminator, self.offset)
+        result = self.data[self.offset:null_index]    # bytes without the terminator
+        self.offset = null_index + len(terminator)  # advance offset past terminator
+        return result
 
-    def read_format(self, format_text):
+    def unpack(self, format_text):
         """Unpack bytes using struct modules format
 
         :param format_text: struct's module format