From 875d1d61e05191308547d7d10cea2f35e1f128f0 Mon Sep 17 00:00:00 2001
From: Davide Casale <davi0kprogramsthings@gmail.com>
Date: Fri, 28 Jul 2023 17:37:15 +0200
Subject: [PATCH] Add order book checksum handling in
 /examples/websocket/public/raw_order_book.py.

---
 examples/websocket/public/raw_order_book.py | 46 ++++++++++++++++++++-
 1 file changed, 45 insertions(+), 1 deletion(-)

diff --git a/examples/websocket/public/raw_order_book.py b/examples/websocket/public/raw_order_book.py
index 33ef321..264c3f0 100644
--- a/examples/websocket/public/raw_order_book.py
+++ b/examples/websocket/public/raw_order_book.py
@@ -2,7 +2,9 @@
 
 from collections import OrderedDict
 
-from typing import List
+from typing import List, Dict
+
+import crcmod
 
 from bfxapi import Client, PUB_WSS_HOST
 
@@ -18,6 +20,9 @@ class RawOrderBook:
             } for symbol in symbols
         }
 
+        self.cooldown: Dict[str, bool] = \
+            { symbol: False for symbol in symbols }
+
     def update(self, symbol: str, data: TradingPairRawBook) -> None:
         order_id, price, amount = data.order_id, data.price, data.amount
 
@@ -34,6 +39,31 @@ class RawOrderBook:
             if order_id in self.__raw_order_book[symbol][kind]:
                 del self.__raw_order_book[symbol][kind][order_id]
 
+    def verify(self, symbol: str, checksum: int) -> bool:
+        values: List[int] = [ ]
+
+        bids = sorted([ (data["order_id"], data["price"], data["amount"]) \
+            for _, data in self.__raw_order_book[symbol]["bids"].items() ],
+                key=lambda data: (-data[1], data[0]))
+
+        asks = sorted([ (data["order_id"], data["price"], data["amount"]) \
+            for _, data in self.__raw_order_book[symbol]["asks"].items() ],
+                key=lambda data: (data[1], data[0]))
+
+        if len(bids) < 25 or len(asks) < 25:
+            raise AssertionError("Not enough bids or asks (need at least 25).")
+
+        for _i in range(25):
+            bid, ask = bids[_i], asks[_i]
+            values.extend([ bid[0], bid[2] ])
+            values.extend([ ask[0], ask[2] ])
+
+        local = ":".join(str(value) for value in values).encode("UTF-8")
+
+        crc32 = crcmod.mkCrcFun(0x104C11DB7, initCrc=0, xorOut=0xFFFFFFFF)
+
+        return crc32(local) == checksum
+
 SYMBOLS = [ "tBTCUSD", "tLTCUSD", "tLTCBTC", "tETHUSD", "tETHBTC" ]
 
 raw_order_book = RawOrderBook(symbols=SYMBOLS)
@@ -62,4 +92,18 @@ def on_t_raw_book_snapshot(subscription: Book, snapshot: List[TradingPairRawBook
 def on_t_raw_book_update(subscription: Book, data: TradingPairRawBook):
     raw_order_book.update(subscription["symbol"], data)
 
+@bfx.wss.on("checksum")
+async def on_checksum(subscription: Book, value: int):
+    symbol = subscription["symbol"]
+
+    if raw_order_book.verify(symbol, value):
+        raw_order_book.cooldown[symbol] = False
+    elif not raw_order_book.cooldown[symbol]:
+        print("Mismatch between local and remote checksums: "
+            f"restarting book for symbol <{symbol}>...")
+
+        await bfx.wss.resubscribe(sub_id=subscription["subId"])
+
+        raw_order_book.cooldown[symbol] = True
+
 bfx.wss.run()