SQD/SQK writing routine + bunch of refactoring

1 files changed, 55 insertions, 13 deletions

diff --git a/squaredeal/__init__.py b/squaredeal/__init__.py
index a036f20..c08f946 100644
--- a/squaredeal/__init__.py
+++ b/squaredeal/__init__.py
@@ -1,4 +1,4 @@
-import hashlib, os
+import hashlib, os, shutil
 
 
 class SquareDealError(Exception):
@@ -23,6 +23,9 @@ class SquareDealPhase(object):
         self.info = parts[3]
         self.s_keys = [None] * self.sessions
 
+    def tostring(self):
+        return ':'.join([str(self.sessions), str(self.boards), self.prefix, self.info or ''])
+
 
 class SquareDeal(object):
     def __init__(self):
@@ -32,7 +35,7 @@ class SquareDeal(object):
         self.phases = []
         self.published = False
 
-    def fromfile(self, sqdpath, encoding='utf-8', sqkpath=None):
+    def fromfile(self, sqdpath, sqkpath=None, encoding='utf-8'):
         with open(sqdpath, encoding=encoding) as sqdfile:
             contents = [line.strip() for line in sqdfile.readlines()]
         for idx, line in enumerate(contents):
@@ -52,15 +55,15 @@ class SquareDeal(object):
             else:
                 raise SquareDealError('Unrecognized tag %s on line %d' % (linetype, idx))
         if sqkpath is None:
-            sqkpath = list(os.path.splitext(sqdpath))
-            sqkpath[-1] = '.sqk'
-            sqkpath = ''.join(sqkpath)
+            sqkpath = self._deduce_sqk_path(sqdpath)
         try:
             with open(sqkpath, encoding=encoding) as sqkfile:
                 contents = [line.strip() for line in sqkfile.readlines()]
         except FileNotFoundError:
             raise SquareDealError('Unable to locate SQK file for %s' % (sqdpath))
         for line in contents:
+            if not line.strip():
+                continue
             lineparts = line.split(':')
             if len(lineparts) != 2:
                 raise SquareDealError('Malformed SQK line: %s' % (line))
@@ -77,13 +80,52 @@ class SquareDeal(object):
             for s_idx, s_key in enumerate(phase.s_keys):
                 if s_key is None:
                     raise SquareDealError('Session %d,%d missing a key in SQK' % (ph_idx+1, s_idx+1))
-        with open(sqkpath, 'rb') as sqkfile:
-            sqk_hash = hashlib.sha256()
-            while True:
-                sqk_chunk = sqkfile.read(1024)
-                if not sqk_chunk:
-                    break
-                sqk_hash.update(sqk_chunk)
-            sqk_hash = sqk_hash.hexdigest()
+        sqk_hash = self._get_file_hash(sqkpath)
         if sqk_hash != self.hash:
             raise SquareDealError('SQK hash mismtach: %s in SQD, % actual' % (self.hash, sqk_hash))
+
+    def _deduce_sqk_path(self, sqdpath):
+        sqkpath = list(os.path.splitext(sqdpath))
+        sqkpath[-1] = '.sqk'
+        return ''.join(sqkpath)
+
+    def _get_file_hash(self, path):
+        with open(path, 'rb') as hashed_file:
+            hash = hashlib.sha256()
+            while True:
+                chunk = hashed_file.read(1024)
+                if not chunk:
+                    break
+                hash.update(chunk)
+            return hash.hexdigest()
+
+    def _make_backups(self, sqdpath, sqkpath):
+        for f in [sqdpath, sqkpath]:
+            if os.path.exists(f):
+                shutil.copy(f, f + '.bak')
+
+    def _write_session_keys(self, sqkpath):
+        with open(sqkpath, 'wb') as sqkfile:
+            for ph_idx, phase in enumerate(self.phases):
+                for s_idx, session_key in enumerate(phase.s_keys):
+                    if session_key is None:
+                        raise SquareDealError('Missing session key for session %d,%d' % (ph_idx+1, s_idx+1))
+                    sqkfile.write(('%d,%d:%s\r\n' % (ph_idx+1, s_idx+1, session_key)).encode('utf8'))
+        self.hash = self._get_file_hash(sqkpath)
+
+    def tofile(self, sqdpath, sqkpath=None, make_backups=True):
+        if sqkpath is None:
+            sqkpath = self._deduce_sqk_path(sqdpath)
+        if make_backups:
+            self._make_backups(sqdpath, sqkpath)
+        self._write_session_keys(sqkpath)
+        sqd_contents = []
+        sqd_contents.append('TN %s\n' % (self.name or ''))
+        sqd_contents.append('KH %s\n' % (self.hash))
+        sqd_contents.append('DI %s\n' % (self.delayed_info or ''))
+        for phase in self.phases:
+            sqd_contents.append('SN %s\n' % (phase.tostring()))
+        if self.published:
+            sqd_contents.append('PU\n')
+        with open(sqdpath, 'w') as sqdfile:
+            sqdfile.writelines(sqd_contents)