#!/usr/bin/python2

"""mkshar

Generate a shar (with a zip file payload) for a Switch Light installer.
"""

import sys, os, shutil
import subprocess
import hashlib
import optparse

HELP = """\
%prog [OPTIONS] SHAR-OUTPUT SFX-FRAGMENT INSTALL-SCRIPT [ARG ...]
"""

parser = optparse.OptionParser(usage=HELP)
parser.add_option('--blocksize',
                  type=int, default=512,
                  help="Set block size for e.g. 'dd'")
parser.add_option('--bzip',
                  action='store_true',
                  help="Enable bzip2")
parser.add_option('--lazy',
                  action='store_true',
                  help="Enable lazy unpacking")
parser.add_option('--unzip-sfx',
                  action='store_true',
                  help="Use 'unzip' to natively handle SFX preamble")
parser.add_option('--unzip-pipe',
                  action='store_true',
                  help="Enable 'dd' to handle SFX preamble")
parser.add_option('--unzip-loop',
                  action='store_true',
                  help="Enable 'losetup' to handle SFX preamble")
parser.add_option('--unzip-pad',
                  action='store_true',
                  help="Special pad options for deficient unzip")
parser.add_option('--inplace',
                  action='store_true',
                  help="Perform fixups in-place")
parser.add_option('--fixup-perms',
                  type=str,
                  help="Post-unpack shell script to fix permissions")

opts, args = parser.parse_args()

ZIP = "/usr/bin/zip"
ZIP_OPTS = ['-v',
            '-y',
            '-n', '.zip:.loader:.swi:.jar:.gz:.bz:.bz2:.xz',]

# enable (experimental) bzip2, make sure that
# busybox unzip supports it
# Ha ha, using '-9' not only cranks up the compression,
# but *ignores* the suffix list for 'store-only'.
if opts.bzip:
    ZIP_COMPRESS_OPTS = ['-Z', 'bzip2', '-8',]
else:
    ZIP_COMPRESS_OPTS = ['-8',]

shar = args.pop(0)
sfx = args.pop(0)
install = args.pop(0)

import logging
logging.basicConfig()
logger = logging.getLogger("mkshar")
logger.setLevel(logging.DEBUG)

logger.info("initializing ZIP")
zipFile = os.path.abspath(shar + ".zip.in")
if os.path.exists(zipFile):
    logger.debug("+ /bin/rm %s", zipFile)
    os.unlink(zipFile)

def _addfile(path, *opts):
    logger.info("adding file %s with opts %s",
                path, " ".join(opts))
    cmd = [ZIP,] + ZIP_OPTS + list(opts) + ['-u', zipFile, path,]
    subprocess.check_call(cmd)

def _addfiles(paths, *opts):
    logger.info("adding %d files with opts %s",
                len(paths), " ".join(opts))
    cmd = [ZIP,] + ZIP_OPTS + list(opts) + [zipFile,] + paths
    subprocess.check_call(cmd)

def _zipopts(path):
    ext = os.path.splitext(path)[1].lower()
    if ext in ('.loader', '.zip', '.swi',):
        return ['-0',]
    else:
        return ZIP_COMPRESS_OPTS

# block to 512 bytes
buf = open(sfx).read()
sz = len(buf)
logger.info("script size is %d bytes", sz)
rem = sz % opts.blocksize
pad = opts.blocksize - rem

# extra block for house keeping
if opts.unzip_pad:
    pad = pad + opts.blocksize

if pad == 1:
    buf += "\n"
elif pad > 1:
    pad -= 1
    buf = buf + ('#' * pad) + "\n"

logger.info("padded script size is %d bytes", len(buf))

if opts.unzip_pad:
    logger.info("adding SFX pad to beginning as archive item")
    fd = open("pad.bin", "w")
    fd.write(('#' * (len(buf)-1)))
    fd.write('\n')
    fd.close()
    _addfile("pad.bin", '-0', '-j')
    os.unlink("pad.bin")

_addfile(install, '-0')

if opts.fixup_perms:
    _addfile(opts.fixup_perms, '-0')

logger.info("adding other payload items")
o = ['-r',] + ZIP_COMPRESS_OPTS
_addfiles(args, *o)

def _splice(tag, val):
    global buf
    pat = tag + '='
    p = buf.find(pat)
    q = buf.find("\n", p+1)
    if p < 0:
        raise ValueError("cannot find tag %s" % repr(tag))
    llen = q - p
    line = "%s=%s #" % (tag, str(val),)
    if len(line) > llen:
        raise SystemExit("cannot insert byte count marker")
    llen -= len(line)
    line = line + ('#' * llen)
    buf = buf[:p] + line + buf[q:]

def _spliceMaybe(tag, val):
    val = "${%s-\"%s\"}" % (tag, val,)
    _splice(tag, val)

logger.info("prepping SFX")

_splice('SFX_BYTES', len(buf))
_splice('SFX_BLOCKSIZE', opts.blocksize)
_splice('SFX_INSTALL', os.path.split(install)[1])
_splice('SFX_CHECKSUM', '')
if opts.lazy:
    _splice('SFX_LAZY', '1')
else:
    _splice('SFX_LAZY', '')
if opts.unzip_sfx:
    _spliceMaybe('SFX_UNZIP', '1')
else:
    _spliceMaybe('SFX_UNZIP', '')
if opts.unzip_pipe:
    _spliceMaybe('SFX_PIPE', '1')
else:
    _spliceMaybe('SFX_PIPE', '')
if opts.unzip_loop:
    _spliceMaybe('SFX_LOOP', '1')
else:
    _spliceMaybe('SFX_LOOP', '')
if opts.unzip_pad:
    _splice('SFX_PAD', 'pad.bin')
else:
    _splice('SFX_PAD', '')
if opts.fixup_perms:
    _splice('SFX_PERMS', opts.fixup_perms)
if opts.inplace:
    _spliceMaybe('SFX_INPLACE', '1')
else:
    _spliceMaybe('SFX_INPLACE', '')

# remember the checksum offset
ckStart = buf.find("SFX_CHECKSUM=")
ckEnd = buf.find("\n", ckStart+1)

logger.info("generating shar")

if opts.unzip_pad:

    wfd = open(shar, "w+")
    rfd = open(zipFile, "r")
    shutil.copyfileobj(rfd, wfd)
    rfd.close()

    logger.info("capturing first block")
    wfd.seek(0, 0)
    magic = wfd.read(opts.blocksize)

    logger.info("splicing in SFX")
    wfd.seek(0, 0)
    wfd.write(buf)

    logger.info("saving first block")
    wfd.seek(len(buf)-opts.blocksize, 0)
    wfd.write(magic)

    wfd.close()

    # adjust buf to include magic
    buf = buf[:-opts.blocksize] + magic

else:

    wfd = open(shar, "w")
    wfd.write(buf)
    rfd = open(zipFile, "r")
    shutil.copyfileobj(rfd, wfd)
    rfd.close()
    wfd.close()

    logger.info("adjusting SFX")
    cmd = [ZIP,] + ZIP_OPTS + ['-A', shar,]
    subprocess.check_call(cmd)

# compute the actual checksum
ckf = hashlib.md5()
ckf.update(buf[:ckStart])
ckf.update(buf[ckEnd+1:])

rfd = open(shar, "r")
rfd.seek(len(buf), 0)
while True:
    buf = rfd.read(10*1024)
    if not buf: break
    ckf.update(buf)
rfd.close()

logger.info("adjusting checksum")
wfd = open(shar, "r+")
wfd.seek(ckStart, 0)
wfd.write("SFX_CHECKSUM=%s #" % ckf.hexdigest())
wfd.close()

logger.info("+ /bin/chmod +x %s", shar)
os.chmod(shar, 0755)

logger.debug("+ /bin/rm %s", zipFile)
os.unlink(zipFile)