From 2299098884833ccf7017edb3b5d718305c93991c Mon Sep 17 00:00:00 2001
From: Victor Shyba <victor1984@riseup.net>
Date: Mon, 7 May 2018 15:10:19 -0300
Subject: [PATCH] add integrity check for the headers file

---
 lbrynet/conf.py        |  1 +
 lbrynet/core/Wallet.py | 24 +++++++++++++++++++++++-
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/lbrynet/conf.py b/lbrynet/conf.py
index 0be7a423e..9c596216e 100644
--- a/lbrynet/conf.py
+++ b/lbrynet/conf.py
@@ -236,6 +236,7 @@ FIXED_SETTINGS = {
     'SLACK_WEBHOOK': ('nUE0pUZ6Yl9bo29epl5moTSwnl5wo20ip2IlqzywMKZiIQSFZR5'
                       'AHx4mY0VmF0WQZ1ESEP9kMHZlp1WzJwWOoKN3ImR1M2yUAaMyqGZ='),
     'WALLET_TYPES': [LBRYUM_WALLET, LBRYCRD_WALLET],
+    'HEADERS_FILE_SHA256_CHECKSUM': (366295, 'b0c8197153a33ccbc52fb81a279588b6015b68b7726f73f6a2b81f7e25bfe4b9')
 }
 
 ADJUSTABLE_SETTINGS = {
diff --git a/lbrynet/core/Wallet.py b/lbrynet/core/Wallet.py
index 577e6e3d4..00f1511a9 100644
--- a/lbrynet/core/Wallet.py
+++ b/lbrynet/core/Wallet.py
@@ -10,6 +10,7 @@ from twisted.internet import threads, reactor, defer, task
 from twisted.python.failure import Failure
 from twisted.internet.error import ConnectionAborted
 
+from hashlib import sha256
 from lbryum import wallet as lbryum_wallet
 from lbryum.network import Network
 from lbryum.simple_config import SimpleConfig
@@ -127,7 +128,6 @@ class Wallet(object):
             return os.stat(headers_path).st_size
         return 0
 
-
     @defer.inlineCallbacks
     def get_remote_height(self, server, port):
         connected = defer.Deferred()
@@ -143,6 +143,7 @@ class Wallet(object):
         from lbrynet import conf
         if conf.settings['blockchain_name'] != "lbrycrd_main":
             defer.returnValue(False)
+        self._check_header_file_integrity(conf)
         s3_headers_depth = conf.settings['s3_headers_depth']
         if not s3_headers_depth:
             defer.returnValue(False)
@@ -158,6 +159,27 @@ class Wallet(object):
                 log.warning("error requesting remote height from %s:%i - %s", server_url, port, err)
         defer.returnValue(False)
 
+    def _check_header_file_integrity(self, conf):
+        # TODO: temporary workaround for usability. move to txlbryum and check headers instead of file integrity
+        hashsum = sha256()
+        checksum_height, checksum = conf.settings['HEADERS_FILE_SHA256_CHECKSUM']
+        checksum_length_in_bytes = checksum_height * HEADER_SIZE
+        if self.local_header_file_size() < checksum_length_in_bytes:
+            return
+        headers_path = os.path.join(self.config.path, "blockchain_headers")
+        with open(headers_path, "rb") as headers_file:
+            hashsum.update(headers_file.read(checksum_length_in_bytes))
+        current_checksum = hashsum.hexdigest()
+        if current_checksum != checksum:
+            msg = "Expected checksum {}, got {}".format(checksum, current_checksum)
+            log.warning("Wallet file corrupted, checksum mismatch. " + msg)
+            log.warning("Deleting header file so it can be downloaded again.")
+            os.unlink(headers_path)
+        elif (self.local_header_file_size() % HEADER_SIZE) != 0:
+            log.warning("Header file is good up to checkpoint height, but incomplete. Truncating to checkpoint.")
+            with open(headers_path, "rb+") as headers_file:
+                headers_file.truncate(checksum_length_in_bytes)
+
     @defer.inlineCallbacks
     def start(self):
         should_download_headers = yield self.should_download_headers_from_s3()