moth/lib/python/mothballer.py

116 lines
3.4 KiB
Python
Executable File

#!/usr/bin/env python3
import argparse
import binascii
import hashlib
import io
import json
import logging
import moth
import os
import shutil
import tempfile
import zipfile
import random
SEEDFN = "SEED"
def write_kv_pairs(ziphandle, filename, kv):
""" Write out a sorted map to file
:param ziphandle: a zipfile object
:param filename: The filename to write within the zipfile object
:param kv: the map to write out
:return:
"""
filehandle = io.StringIO()
for key in sorted(kv.keys()):
if isinstance(kv[key], list):
for val in kv[key]:
filehandle.write("%s %s\n" % (key, val))
else:
filehandle.write("%s %s\n" % (key, kv[key]))
filehandle.seek(0)
ziphandle.writestr(filename, filehandle.read())
def escape(s):
return s.replace('&', '&amp;').replace('<', '&lt;').replace('>', '&gt;')
def build_category(categorydir, outdir):
category_seed = random.getrandbits(32)
categoryname = os.path.basename(categorydir.strip(os.sep))
zipfilename = os.path.join(outdir, "%s.mb" % categoryname)
logging.info("Building {} from {}".format(zipfilename, categorydir))
if os.path.exists(zipfilename):
# open and gather some state
existing = zipfile.ZipFile(zipfilename, 'r')
try:
category_seed = int(existing.open(SEEDFN).read().strip())
except Exception:
pass
existing.close()
logging.debug("Using PRNG seed {}".format(category_seed))
zipfileraw = tempfile.NamedTemporaryFile(delete=False)
mothball = package(categoryname, categorydir, category_seed)
shutil.copyfileobj(mothball, zipfileraw)
zipfileraw.close()
shutil.move(zipfileraw.name, zipfilename)
# Returns a file-like object containing the contents of the new zip file
def package(categoryname, categorydir, seed):
zfraw = io.BytesIO()
zf = zipfile.ZipFile(zfraw, 'x')
zf.writestr("category_seed.txt", str(seed))
cat = moth.Category(categorydir, seed)
mapping = {}
answers = {}
summary = {}
for puzzle in cat:
logging.info("Processing point value {}".format(puzzle.points))
hashmap = hashlib.sha1(str(seed).encode('utf-8'))
hashmap.update(str(puzzle.points).encode('utf-8'))
puzzlehash = hashmap.hexdigest()
mapping[puzzle.points] = puzzlehash
answers[puzzle.points] = puzzle.answers
summary[puzzle.points] = puzzle.summary
puzzledir = os.path.join("content", puzzlehash)
for fn, f in puzzle.files.items():
payload = f.stream.read()
zf.writestr(os.path.join(puzzledir, fn), payload)
obj = puzzle.package()
zf.writestr(os.path.join(puzzledir, 'puzzle.json'), json.dumps(obj))
write_kv_pairs(zf, 'map.txt', mapping)
write_kv_pairs(zf, 'answers.txt', answers)
write_kv_pairs(zf, 'summaries.txt', summary)
# clean up
zf.close()
zfraw.seek(0)
return zfraw
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Build a category package')
parser.add_argument('outdir', help='Output directory')
parser.add_argument('categorydirs', nargs='+', help='Directory of category source')
args = parser.parse_args()
logging.basicConfig(level=logging.DEBUG)
outdir = os.path.abspath(args.outdir)
for categorydir in args.categorydirs:
categorydir = os.path.abspath(categorydir)
build_category(categorydir, outdir)