diff options
| -rw-r--r-- | squaredeal/__init__.py | 68 |
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) |