From 2a8c8d2d2a3a745d32215c6430c3b280edbc285d Mon Sep 17 00:00:00 2001
From: "brian m. carlson" <sandals@crustytoothpaste.net>
Date: Sat, 11 Jan 2014 00:58:54 +0000
Subject: [PATCH] Support SHA-384 and SHA-512.

Turn the SHA256Transaction store into a HashTransactionStore and support
SHA-384 and SHA-512, which provide better performance on 64-bit systems.
Don't support SHA-224 since it is weaker and is only useful where the hash
would be truncated anyway.

Signed-off-by: brian m. carlson <sandals@crustytoothpaste.net>
---
 lib/newfol/exception.py |  9 +++++++++
 lib/newfol/filemanip.py | 21 +++++++++++++--------
 test/testfilemanip.py   |  4 ++--
 3 files changed, 24 insertions(+), 10 deletions(-)

diff --git a/lib/newfol/exception.py b/lib/newfol/exception.py
index 8a65bc7..d88ab59 100644
--- a/lib/newfol/exception.py
+++ b/lib/newfol/exception.py
@@ -18,3 +18,12 @@ class DataError(FilemanipError):
 
 class RepairError(FilemanipError):
     pass
+
+class HashError(FilemanipError):
+    pass
+
+class InvalidHashError(HashError):
+    pass
+
+class HashMismatchError(HashError):
+    pass
diff --git a/lib/newfol/filemanip.py b/lib/newfol/filemanip.py
index 7ee0c91..4f589d2 100644
--- a/lib/newfol/filemanip.py
+++ b/lib/newfol/filemanip.py
@@ -159,27 +159,30 @@ class GitTransactionStore(TransactionStore):
                 message += "{:%Y-%m-%d %H:%M:%S} UTC".format(now)
                 self._call_git("commit", "-m", message)
 
-class SHA256TransactionStore(TransactionStore):
-    def __init__(self, options=None):
+class HashTransactionStore(TransactionStore):
+    def __init__(self, hashname="sha256", options=None):
         self._filename = None
         self._mode = None
         self._options = options or {}
         if "forgiving" not in self._options:
             self._options["forgiving"] = True
+        if hashname not in ("sha256", "sha384", "sha512"):
+            raise newfol.exception.InvalidHashError("Invalid hash: " + hashname)
+        self._hashname = hashname
     @staticmethod
-    def _hash_file(filename):
-        h = hashlib.sha256()
+    def _hash_file(hashname, filename):
+        h = hashlib.__dict__[hashname]()
         fp = open(filename, "rb")
         data = fp.read(-1)
         totallen = len(data)
         h.update(data)
         fp.close()
-        return "sha256:" + str(totallen) + ":" + h.hexdigest()
+        return "%s:%s:%s" % (hashname, str(totallen), h.hexdigest())
     def prepare_open(self, filename, mode):
         self._filename = filename
         self._mode = mode
         if "r" in mode:
-            result = self._hash_file(filename)
+            result = self._hash_file(self._hashname, filename)
             try:
                 fp = open(filename + ".checksum", "r")
                 line = fp.readline()[:-1]
@@ -193,7 +196,7 @@ class SHA256TransactionStore(TransactionStore):
                             "checksum")
     def commit_close(self):
         if "w" in self._mode:
-            result = self._hash_file(self._filename)
+            result = self._hash_file(self._hashname, self._filename)
             fp = open(self._filename + ".checksum", "w")
             fp.write(result + "\n")
             fp.close()
@@ -462,7 +465,9 @@ class FileStorage:
     def _make_transaction_store(items, options):
         txntypes = {
             "git": GitTransactionStore,
-            "sha256": SHA256TransactionStore
+            "sha256": lambda *a, **kw: HashTransactionStore("sha256", *a, **kw),
+            "sha384": lambda *a, **kw: HashTransactionStore("sha384", *a, **kw),
+            "sha512": lambda *a, **kw: HashTransactionStore("sha512", *a, **kw),
         }
         if items is None or items == "":
             return TransactionStore()
diff --git a/test/testfilemanip.py b/test/testfilemanip.py
index 9d21663..e2d35d6 100755
--- a/test/testfilemanip.py
+++ b/test/testfilemanip.py
@@ -1,6 +1,6 @@
 #!/usr/bin/python3
 
-from newfol.filemanip import Record, FileStorage, SHA256TransactionStore
+from newfol.filemanip import Record, FileStorage, HashTransactionStore
 import newfol.exception
 import tempfile
 import unittest
@@ -69,7 +69,7 @@ class SHA256TransactionTest(unittest.TestCase):
             wfp = open(temp, "w")
             wfp.write(k)
             wfp.close()
-            self.assertEqual(SHA256TransactionStore._hash_file(temp),
+            self.assertEqual(HashTransactionStore._hash_file("sha256", temp),
                     "sha256:" + v)
         tempdir.cleanup()
     def create_small_database(self, tempdir):