#!/usr/bin/env python
"""%prog [options] -x|-t|-c marfile [files]

Utility for managing mar files"""

### The canonical location for this file is
###   https://hg.mozilla.org/build/tools/file/default/buildfarm/utils/mar.py
###
### Please update the copy in puppet to deploy new changes to
### stage.mozilla.org, see
# https://wiki.mozilla.org/ReleaseEngineering/How_To/Modify_scripts_on_stage

import struct
import os
import bz2
import hashlib
import tempfile
from subprocess import Popen, PIPE


def rsa_sign(digest, keyfile):
    proc = Popen(['openssl', 'pkeyutl', '-sign', '-inkey', keyfile],
                 stdin=PIPE, stdout=PIPE)
    proc.stdin.write(digest)
    proc.stdin.close()
    sig = proc.stdout.read()
    return sig


def rsa_verify(digest, signature, keyfile):
    tmp = tempfile.NamedTemporaryFile()
    tmp.write(signature)
    tmp.flush()
    proc = Popen(['openssl', 'pkeyutl', '-pubin', '-verify', '-sigfile',
                 tmp.name, '-inkey', keyfile], stdin=PIPE, stdout=PIPE)
    proc.stdin.write(digest)
    proc.stdin.close()
    data = proc.stdout.read()
    return "Signature Verified Successfully" in data


def packint(i):
    return struct.pack(">L", i)


def unpackint(s):
    return struct.unpack(">L", s)[0]


def generate_signature(fp, updatefunc):
    fp.seek(0)
    # Magic
    updatefunc(fp.read(4))
    # index_offset
    updatefunc(fp.read(4))
    # file size
    updatefunc(fp.read(8))
    # number of signatures
    num_sigs = fp.read(4)
    updatefunc(num_sigs)
    num_sigs = unpackint(num_sigs)
    for i in range(num_sigs):
        # signature algo
        updatefunc(fp.read(4))

        # signature size
        sigsize = fp.read(4)
        updatefunc(sigsize)
        sigsize = unpackint(sigsize)

        # Read this, but don't update the signature with it
        fp.read(sigsize)

    # Read the rest of the file
    while True:
        block = fp.read(512 * 1024)
        if not block:
            break
        updatefunc(block)


class MarSignature:
    """Represents a signature"""
    size = None
    sigsize = None
    algo_id = None
    signature = None
    _offset = None  # where in the file this signature is located
    keyfile = None  # what key to use

    @classmethod
    def from_fileobj(cls, fp):
        _offset = fp.tell()
        algo_id = unpackint(fp.read(4))
        self = cls(algo_id)
        self._offset = _offset
        sigsize = unpackint(fp.read(4))
        assert sigsize == self.sigsize
        self.signature = fp.read(self.sigsize)
        return self

    def __init__(self, algo_id, keyfile=None):
        self.algo_id = algo_id
        self.keyfile = keyfile
        if self.algo_id == 1:
            self.sigsize = 256
            self.size = self.sigsize + 4 + 4
            self._hsh = hashlib.new('sha1')
        else:
            raise ValueError("Unsupported signature algorithm: %s" % algo_id)

    def update(self, data):
        self._hsh.update(data)

    def verify_signature(self):
        if self.algo_id == 1:
            print "digest is", self._hsh.hexdigest()
            assert self.keyfile
            return rsa_verify(self._hsh.digest(), self.signature, self.keyfile)

    def write_signature(self, fp):
        assert self.keyfile
        print "digest is", self._hsh.hexdigest()
        self.signature = rsa_sign(self._hsh.digest(), self.keyfile)
        assert len(self.signature) == self.sigsize
        fp.seek(self._offset)
        fp.write("%s%s%s" % (
            packint(self.algo_id),
            packint(self.sigsize),
            self.signature))
        print "wrote signature %s" % self.algo_id


class MarInfo:
    """Represents information about a member of a MAR file. The following
    attributes are supported:
        `size`:     the file's size
        `name`:     the file's name
        `flags`:    file permission flags
        `_offset`:  where in the MAR file this member exists
    """
    size = None
    name = None
    flags = None
    _offset = None

    # The member info is serialized as a sequence of 4-byte integers
    # representing the offset, size, and flags of the file followed by a null
    # terminated string representing the filename
    _member_fmt = ">LLL"

    @classmethod
    def from_fileobj(cls, fp):
        """Return a MarInfo object by reading open file object `fp`"""
        self = cls()
        data = fp.read(12)
        if not data:
            # EOF
            return None
        if len(data) != 12:
            raise ValueError("Malformed mar?")
        self._offset, self.size, self.flags = struct.unpack(
            cls._member_fmt, data)
        name = ""
        while True:
            c = fp.read(1)
            if c is None:
                raise ValueError("Malformed mar?")
            if c == "\x00":
                break
            name += c
        self.name = name
        return self

    def __repr__(self):
        return "<%s %o %s bytes starting at %i>" % (
            self.name, self.flags, self.size, self._offset)

    def to_bytes(self):
        return struct.pack(self._member_fmt, self._offset, self.size, self.flags) + \
            self.name + "\x00"


class MarFile:
    """Represents a MAR file on disk.

    `name`:     filename of MAR file
    `mode`:     either 'r' or 'w', depending on if you're reading or writing.
                defaults to 'r'
    """

    _longint_fmt = ">L"

    def __init__(self, name, mode="r", signature_versions=[]):
        if mode not in "rw":
            raise ValueError("Mode must be either 'r' or 'w'")

        self.name = name
        self.mode = mode
        if mode == 'w':
            self.fileobj = open(name, 'wb')
        else:
            self.fileobj = open(name, 'rb')

        self.members = []

        # Current offset of our index in the file. This gets updated as we add
        # files to the MAR. This also refers to the end of the file until we've
        # actually written the index
        self.index_offset = 8

        # Flag to indicate that we need to re-write the index
        self.rewrite_index = False

        # Signatures, if any
        self.signatures = []
        self.signature_versions = signature_versions

        if mode == "r":
            # Read the file's index
            self._read_index()
        elif mode == "w":
            # Create placeholder signatures
            if signature_versions:
                # Space for num_signatures and file size
                self.index_offset += 4 + 8

            # Write the magic and placeholder for the index
            self.fileobj.write("MAR1" + packint(self.index_offset))

            # Write placeholder for file size
            self.fileobj.write(struct.pack(">Q", 0))

            # Write num_signatures
            self.fileobj.write(packint(len(signature_versions)))

            for algo_id, keyfile in signature_versions:
                sig = MarSignature(algo_id, keyfile)
                sig._offset = self.index_offset
                self.index_offset += sig.size
                self.signatures.append(sig)
                # algoid
                self.fileobj.write(packint(algo_id))
                # numbytes
                self.fileobj.write(packint(sig.sigsize))
                # space for signature
                self.fileobj.write("\0" * sig.sigsize)

    def _read_index(self):
        fp = self.fileobj
        fp.seek(0)
        # Read the header
        header = fp.read(8)
        magic, self.index_offset = struct.unpack(">4sL", header)
        if magic != "MAR1":
            raise ValueError("Bad magic")
        fp.seek(self.index_offset)

        # Read the index_size, we don't use it though
        # We just read all the info sections from here to the end of the file
        fp.read(4)

        self.members = []

        while True:
            info = MarInfo.from_fileobj(fp)
            if not info:
                break
            self.members.append(info)

        # Sort them by where they are in the file
        self.members.sort(key=lambda info: info._offset)

        # Read any signatures
        first_offset = self.members[0]._offset
        signature_offset = 16
        if signature_offset < first_offset:
            fp.seek(signature_offset)
            num_sigs = unpackint(fp.read(4))
            for i in range(num_sigs):
                sig = MarSignature.from_fileobj(fp)
                for algo_id, keyfile in self.signature_versions:
                    if algo_id == sig.algo_id:
                        sig.keyfile = keyfile
                        break
                else:
                    print "no key specified to validate %i signature" % sig.algo_id
                self.signatures.append(sig)

    def verify_signatures(self):
        if not self.signatures:
            return

        fp = self.fileobj

        generate_signature(fp, self._update_signatures)

        for sig in self.signatures:
            if not sig.verify_signature():
                raise IOError("Verification failed")
            else:
                print "Verification OK (%s)" % sig.algo_id

    def _update_signatures(self, data):
        for sig in self.signatures:
            sig.update(data)

    def add(self, path, name=None, fileobj=None, flags=None):
        """Adds `path` to this MAR file.

        If `name` is set, the file is named with `name` in the MAR, otherwise
        use the normalized version of `path`.

        If `fileobj` is set, file data is read from it, otherwise the file
        located at `path` will be opened and read.

        If `flags` is set, it will be used as the permission flags in the MAR,
        otherwise the permissions of `path` will be used.
        """
        if self.mode != "w":
            raise ValueError("File not opened for writing")

        # If path refers to a directory, add all the files inside of it
        if os.path.isdir(path):
            self.add_dir(path)
            return

        info = MarInfo()
        if not fileobj:
            info.name = name or os.path.normpath(path)
            info.size = os.path.getsize(path)
            info.flags = flags or os.stat(path).st_mode & 0777
            info._offset = self.index_offset

            f = open(path, 'rb')
            self.fileobj.seek(self.index_offset)
            while True:
                block = f.read(512 * 1024)
                if not block:
                    break
                self.fileobj.write(block)
        else:
            assert flags
            info.name = name or path
            info.size = 0
            info.flags = flags
            info._offset = self.index_offset
            self.fileobj.seek(self.index_offset)
            while True:
                block = fileobj.read(512 * 1024)
                if not block:
                    break
                info.size += len(block)
                self.fileobj.write(block)

        # Shift our index, and mark that we have to re-write it on close
        self.index_offset += info.size
        self.rewrite_index = True
        self.members.append(info)

    def add_dir(self, path):
        """Add all of the files under `path` to the MAR file"""
        for root, dirs, files in os.walk(path):
            for f in files:
                self.add(os.path.join(root, f))

    def close(self):
        """Close the MAR file, writing out the new index if required.

        Furthur modifications to the file are not allowed."""
        if self.mode == "w" and self.rewrite_index:
            self._write_index()

        # Update file size
        self.fileobj.seek(0, 2)
        totalsize = self.fileobj.tell()
        self.fileobj.seek(8)
        # print "File size is", totalsize, repr(struct.pack(">Q", totalsize))
        self.fileobj.write(struct.pack(">Q", totalsize))

        if self.mode == "w" and self.signatures:
            self.fileobj.flush()
            fileobj = open(self.name, 'rb')
            generate_signature(fileobj, self._update_signatures)
            for sig in self.signatures:
                # print sig._offset
                sig.write_signature(self.fileobj)

        self.fileobj.close()
        self.fileobj = None

    def __del__(self):
        """Close the file when we're garbage collected"""
        if self.fileobj:
            self.close()

    def _write_index(self):
        """Writes the index of all members at the end of the file"""
        self.fileobj.seek(self.index_offset + 4)
        index_size = 0
        for m in self.members:
            member_bytes = m.to_bytes()
            index_size += len(member_bytes)

        for m in self.members:
            member_bytes = m.to_bytes()
            self.fileobj.write(member_bytes)
        self.fileobj.seek(self.index_offset)
        self.fileobj.write(packint(index_size))

        # Update the offset to the index
        self.fileobj.seek(4)
        self.fileobj.write(packint(self.index_offset))

    def extractall(self, path=".", members=None):
        """Extracts members into `path`. If members is None (the default), then
        all members are extracted."""
        if members is None:
            members = self.members
        for m in members:
            self.extract(m, path)

    def extract(self, member, path="."):
        """Extract `member` into `path` which defaults to the current
        directory."""
        dstpath = os.path.join(path, member.name)
        dirname = os.path.dirname(dstpath)
        if not os.path.exists(dirname):
            os.makedirs(dirname)

        self.fileobj.seek(member._offset)
        open(dstpath, "wb").write(self.fileobj.read(member.size))
        os.chmod(dstpath, member.flags)


class BZ2MarFile(MarFile):
    """Subclass of MarFile that compresses/decompresses members using BZ2.

    BZ2 compression is used for most update MARs."""
    def extract(self, member, path="."):
        """Extract and decompress `member` into `path` which defaults to the
        current directory."""
        dstpath = os.path.join(path, member.name)
        dirname = os.path.dirname(dstpath)
        if not os.path.exists(dirname):
            os.makedirs(dirname)

        self.fileobj.seek(member._offset)
        decomp = bz2.BZ2Decompressor()
        output = open(dstpath, "wb")
        toread = member.size
        while True:
            thisblock = min(128 * 1024, toread)
            block = self.fileobj.read(thisblock)
            if not block:
                break
            toread -= len(block)
            output.write(decomp.decompress(block))
        output.close()
        os.chmod(dstpath, member.flags)

    def add(self, path, name=None, fileobj=None, mode=None):
        """Adds `path` compressed with BZ2 to this MAR file.

        If `name` is set, the file is named with `name` in the MAR, otherwise
        use the normalized version of `path`.

        If `fileobj` is set, file data is read from it, otherwise the file
        located at `path` will be opened and read.

        If `flags` is set, it will be used as the permission flags in the MAR,
        otherwise the permissions of `path` will be used.
        """
        if self.mode != "w":
            raise ValueError("File not opened for writing")
        if os.path.isdir(path):
            self.add_dir(path)
            return
        info = MarInfo()
        info.name = name or os.path.normpath(path)
        info.size = 0
        if not fileobj:
            info.flags = os.stat(path).st_mode & 0777
        else:
            info.flags = mode
        info._offset = self.index_offset

        if not fileobj:
            f = open(path, 'rb')
        else:
            f = fileobj
        comp = bz2.BZ2Compressor(9)
        self.fileobj.seek(self.index_offset)
        while True:
            block = f.read(512 * 1024)
            if not block:
                break
            block = comp.compress(block)
            info.size += len(block)
            self.fileobj.write(block)
        block = comp.flush()
        info.size += len(block)
        self.fileobj.write(block)

        self.index_offset += info.size
        self.rewrite_index = True
        self.members.append(info)

if __name__ == "__main__":
    from optparse import OptionParser

    parser = OptionParser(__doc__)
    parser.set_defaults(
        action=None,
        bz2=False,
        chdir=None,
        keyfile=None,
        verify=False,
    )
    parser.add_option("-x", "--extract", action="store_const", const="extract",
                      dest="action", help="extract MAR")
    parser.add_option("-t", "--list", action="store_const", const="list",
                      dest="action", help="print out MAR contents")
    parser.add_option("-c", "--create", action="store_const", const="create",
                      dest="action", help="create MAR")
    parser.add_option("-j", "--bzip2", action="store_true", dest="bz2",
                      help="compress/decompress members with BZ2")
    parser.add_option("-k", "--keyfile", dest="keyfile",
                      help="sign/verify with given key")
    parser.add_option("-v", "--verify", dest="verify", action="store_true",
                      help="verify the marfile")
    parser.add_option("-C", "--chdir", dest="chdir",
                      help="chdir to this directory before creating or extracing; location of marfile isn't affected by this option.")

    options, args = parser.parse_args()

    if not options.action:
        parser.error("Must specify something to do (one of -x, -t, -c)")

    if not args:
        parser.error("You must specify at least a marfile to work with")

    marfile, files = args[0], args[1:]
    marfile = os.path.abspath(marfile)

    if options.bz2:
        mar_class = BZ2MarFile
    else:
        mar_class = MarFile

    signatures = []
    if options.keyfile:
        signatures.append((1, options.keyfile))

    # Move into the directory requested
    if options.chdir:
        os.chdir(options.chdir)

    if options.action == "extract":
        m = mar_class(marfile)
        m.extractall()

    elif options.action == "list":
        m = mar_class(marfile, signature_versions=signatures)
        if options.verify:
            m.verify_signatures()
        print "%-7s %-7s %-7s" % ("SIZE", "MODE", "NAME")
        for m in m.members:
            print "%-7i %04o    %s" % (m.size, m.flags, m.name)

    elif options.action == "create":
        if not files:
            parser.error("Must specify at least one file to add to marfile")
        m = mar_class(marfile, "w", signature_versions=signatures)
        for f in files:
            m.add(f)
        m.close()