#!/usr/bin/env python
# -*- coding: utf-8 -*-
# 
# Copyright 2004, 2005 Zuza Software Foundation
# 
# This file is part of translate.
#
# translate is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
# 
# translate is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with translate; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA

"""module for accessing mozilla xpi packages"""

from __future__ import generators
import zipfile
import os.path
from translate import __version__
import StringIO
import re

# we have some enhancements to zipfile in a file called zipfileext
# hopefully they will be included in a future version of python
from translate.misc import zipfileext
ZipFileBase = zipfileext.ZipFileExt

from translate.misc import wStringIO
# this is a fix to the StringIO in Python 2.3.3
# submitted as patch 951915 on sourceforge
class FixedStringIO(wStringIO.StringIO):
    def truncate(self, size=None):
        StringIO.StringIO.truncate(self, size)
        self.len = len(self.buf)

NamedStringInput = wStringIO.StringIO
NamedStringOutput = wStringIO.StringIO

def _commonprefix(itemlist):
    def cp(a, b):
        l = min(len(a), len(b))
        for n in range(l):
            if a[n] != b[n]: return a[:n]
        return a[:l]
    if itemlist:
        return reduce(cp, itemlist)
    else:
        return ''

def rememberchanged(self, method):
    def changed(*args, **kwargs):
        self.changed = True
        method(*args, **kwargs)
    return changed

class CatchPotentialOutput(NamedStringInput, object):
    """catches output if there has been, before closing"""
    def __init__(self, contents, onclose):
        """Set up the output stream, and remember a method to call on closing"""
        NamedStringInput.__init__(self, contents)
        self.onclose = onclose
        self.changed = False
        s = super(CatchPotentialOutput, self)
        self.write = rememberchanged(self, s.write)
        self.writelines = rememberchanged(self, s.writelines)
        self.truncate = rememberchanged(self, s.truncate)

    def close(self):
        """wrap the underlying close method, to pass the value to onclose before it goes"""
        if self.changed:
            value = self.getvalue()
            self.onclose(value)
        NamedStringInput.close(self)

    def flush(self):
        """zip files call flush, not close, on file-like objects"""
        value = self.getvalue()
        self.onclose(value)
        NamedStringInput.flush(self)

    def slam(self):
        """use this method to force the closing of the stream if it isn't closed yet"""
        if not self.closed:
            self.close()

class ZipFileCatcher(ZipFileBase, object):
    """a ZipFile that calls any methods its instructed to before closing (useful for catching stream output)"""
    def __init__(self, *args, **kwargs):
        """initialize the ZipFileCatcher"""
        # storing oldclose as attribute, since if close is called from __del__ it has no access to external variables
        self.oldclose = super(ZipFileCatcher, self).close
        super(ZipFileCatcher, self).__init__(*args, **kwargs)

    def addcatcher(self, pendingsave):
        """remember to call the given method before closing"""
        if hasattr(self, "pendingsaves"):
            if not pendingsave in self.pendingsaves:
                self.pendingsaves.append(pendingsave)
        else:
            self.pendingsaves = [pendingsave]

    def close(self):
        """close the stream, remembering to call any addcatcher methods first"""
        if hasattr(self, "pendingsaves"):
            for pendingsave in self.pendingsaves:
                pendingsave()
        # if close is called from __del__, it somehow can't see ZipFileCatcher, so we've cached oldclose...
        if ZipFileCatcher is None:
            self.oldclose()
        else:
            super(ZipFileCatcher, self).close()

    def overwritestr(self, zinfo_or_arcname, bytes):
        """writes the string into the archive, overwriting the file if it exists...""" 
        if isinstance(zinfo_or_arcname, zipfile.ZipInfo):
            filename = zinfo_or_arcname.filename
        else:
            filename = zinfo_or_arcname
        if filename in self.NameToInfo:
            self.delete(filename)
        self.writestr(zinfo_or_arcname, bytes)
        self.writeendrec()

class XpiFile(ZipFileCatcher):
    def __init__(self, *args, **kwargs):
        """sets up the xpi file"""
        self.includenonloc = kwargs.get("includenonloc", True)
        if "includenonloc" in kwargs:
            del kwargs["includenonloc"]
        if "compression" not in kwargs:
            kwargs["compression"] = zipfile.ZIP_DEFLATED
        self.locale = kwargs.pop("locale", None)
        self.region = kwargs.pop("region", None)
        super(XpiFile, self).__init__(*args, **kwargs)
        self.jarfiles = {}
        self.findlangreg()
        self.jarprefixes = self.findjarprefixes()
        self.reverseprefixes = dict([
            (prefix,jarfilename) for jarfilename, prefix in self.jarprefixes.iteritems() if prefix])
        self.reverseprefixes["package/"] = None

    def iterjars(self):
        """iterate through the jar files in the xpi as ZipFile objects"""
        for filename in self.namelist():
            if filename.lower().endswith('.jar'):
                if filename not in self.jarfiles:
                    jarstream = self.openinputstream(None, filename)
                    jarfile = ZipFileCatcher(jarstream, mode=self.mode)
                    self.jarfiles[filename] = jarfile
                else:
                    jarfile = self.jarfiles[filename]
                yield filename, jarfile

    def islocfile(self, filename):
        """returns whether the given file is needed for localization (basically .dtd and .properties)"""
        base, ext = os.path.splitext(filename)
        return ext in (os.extsep + "dtd", os.extsep + "properties")

    def findlangreg(self):
        """finds the common prefix of all the files stored in the jar files"""
        dirstructure = {}
        locale = self.locale
        region = self.region
        localematch = re.compile("^[a-z]{2,3}(-[a-zA-Z]{2,3}|)$")
        regionmatch = re.compile("^[a-zA-Z]{2,3}$")
        # exclude en-mac, en-win, en-unix for seamonkey
        osmatch = re.compile("^[a-z]{2,3}-(mac|unix|win)$")
        for jarfilename, jarfile in self.iterjars():
            jarname = "".join(jarfilename.split('/')[-1:]).replace(".jar", "", 1)
            if localematch.match(jarname) and not osmatch.match(jarname):
                if locale is None:
                    locale = jarname
                elif locale != jarname:
                    locale = 0
            elif regionmatch.match(jarname):
                if region is None:
                    region = jarname
                elif region != jarname:
                    region = 0
            for filename in jarfile.namelist():
                if filename.endswith('/'): continue
                if not self.islocfile(filename) and not self.includenonloc: continue
                parts = filename.split('/')[:-1]
                treepoint = dirstructure
                for partnum in range(len(parts)):
                    part = parts[partnum]
                    if part in treepoint:
                        treepoint = treepoint[part]
                    else:
                        treepoint[part] = {}
                        treepoint = treepoint[part]
        localeentries = {}
        if 'locale' in dirstructure:
            for dirname in dirstructure['locale']:
                localeentries[dirname] = 1
                if localematch.match(dirname) and not osmatch.match(dirname):
                    if locale is None:
                        locale = dirname
                    elif locale != dirname:
                        print "locale dir mismatch - ", dirname, "but locale is", locale, "setting to 0"
                        locale = 0
                elif regionmatch.match(dirname):
                    if region is None:
                        region = dirname
                    elif region != dirname:
                        region = 0
        if locale and locale in localeentries:
            del localeentries[locale]
        if region and region in localeentries:
            del localeentries[region]
        if locale and not region:
            if "-" in locale:
                region = locale.split("-", 1)[1]
            else:
                region = ""
        self.setlangreg(locale, region)

    def setlangreg(self, locale, region):
        """set the locale and region of this xpi"""
        if locale == 0 or locale is None:
            raise ValueError("unable to determine locale")
        self.locale = locale
        self.region = region
        self.dirmap = {}
        if self.locale is not None:
            self.dirmap[('locale', self.locale)] = ('lang-reg',)
        if self.region:
            self.dirmap[('locale', self.region)] = ('reg',)

    def findjarprefixes(self):
        """checks the uniqueness of the jar files contents"""
        uniquenames = {}
        jarprefixes = {}
        for jarfilename, jarfile in self.iterjars():
            jarprefixes[jarfilename] = ""
            for filename in jarfile.namelist():
                if filename.endswith('/'): continue
                if filename in uniquenames:
                    jarprefixes[jarfilename] = True
                    jarprefixes[uniquenames[filename]] = True
                else:
                    uniquenames[filename] = jarfilename
        for jarfilename, hasconflicts in jarprefixes.items():
            if hasconflicts:
                shortjarfilename = os.path.split(jarfilename)[1]
                shortjarfilename = os.path.splitext(shortjarfilename)[0]
                jarprefixes[jarfilename] = shortjarfilename+'/'
        # this is a clever trick that will e.g. remove zu- from zu-win, zu-mac, zu-unix
        commonjarprefix = _commonprefix([prefix for prefix in jarprefixes.itervalues() if prefix])
        if commonjarprefix:
            for jarfilename, prefix in jarprefixes.items():
                if prefix:
                    jarprefixes[jarfilename] = prefix.replace(commonjarprefix, '', 1)
        return jarprefixes

    def ziptoospath(self, zippath):
        """converts a zipfile filepath to an os-style filepath"""
        return os.path.join(*zippath.split('/'))

    def ostozippath(self, ospath):
        """converts an os-style filepath to a zipfile filepath"""
        return '/'.join(ospath.split(os.sep))

    def mapfilename(self, filename):
        """uses a map to simplify the directory structure"""
        parts = tuple(filename.split('/'))
        possiblematch = None
        for prefix, mapto in self.dirmap.iteritems():
            if parts[:len(prefix)] == prefix:
                if possiblematch is None or len(possiblematch[0]) < len(prefix):
                    possiblematch = prefix, mapto
        if possiblematch is not None:
            prefix, mapto = possiblematch
            mapped = mapto + parts[len(prefix):]
            return '/'.join(mapped)
        return filename

    def mapxpifilename(self, filename):
        """uses a map to rename files that occur straight in the xpi"""
        if filename.startswith('bin/chrome/') and filename.endswith(".manifest"):
            return 'bin/chrome/lang-reg.manifest'
        return filename

    def reversemapfile(self, filename):
        """unmaps the filename..."""
        possiblematch = None
        parts = tuple(filename.split('/'))
        for prefix, mapto in self.dirmap.iteritems():
            if parts[:len(mapto)] == mapto:
                if possiblematch is None or len(possiblematch[0]) < len(mapto):
                    possiblematch = (mapto, prefix)
        if possiblematch is None:
            return filename
        mapto, prefix = possiblematch
        reversemapped = prefix + parts[len(mapto):]
        return '/'.join(reversemapped)

    def reversemapxpifilename(self, filename):
        """uses a map to rename files that occur straight in the xpi"""
        if filename == 'bin/chrome/lang-reg.manifest':
            if self.locale:
                return '/'.join(('bin', 'chrome', self.locale + '.manifest'))
            else:
                for otherfilename in self.namelist():
                    if otherfilename.startswith("bin/chrome/") and otherfilename.endswith(".manifest"):
                        return otherfilename
        return filename

    def jartoospath(self, jarfilename, filename):
        """converts a filename from within a jarfile to an os-style filepath"""
        if jarfilename:
            jarprefix = self.jarprefixes[jarfilename]
            return self.ziptoospath(jarprefix+self.mapfilename(filename))
        else:
            return self.ziptoospath(os.path.join("package", self.mapxpifilename(filename)))

    def ostojarpath(self, ospath):
        """converts an extracted os-style filepath to a jarfilename and filename"""
        zipparts = ospath.split(os.sep)
        prefix = zipparts[0] + '/'
        if prefix in self.reverseprefixes:
            jarfilename = self.reverseprefixes[prefix]
            filename = self.reversemapfile('/'.join(zipparts[1:]))
            if jarfilename is None:
                filename = self.reversemapxpifilename(filename)
            return jarfilename, filename
        else:
            filename = self.ostozippath(ospath)
            if filename in self.namelist():
                return None, filename
            filename = self.reversemapfile('/'.join(zipparts))
            possiblejarfilenames = [jarfilename for jarfilename, prefix in self.jarprefixes.iteritems() if not prefix]
            for jarfilename in possiblejarfilenames:
                jarfile = self.jarfiles[jarfilename]
                if filename in jarfile.namelist():
                    return jarfilename, filename
            raise IndexError("ospath not found in xpi file, could not guess location: %r" % ospath)

    def jarfileexists(self, jarfilename, filename):
        """checks whether the given file exists inside the xpi"""
        if jarfilename is None:
            return filename in self.namelist()
        else:
            jarfile = self.jarfiles[jarfilename]
            return filename in jarfile.namelist()

    def ospathexists(self, ospath):
        """checks whether the given file exists inside the xpi"""
        jarfilename, filename = self.ostojarpath(ospath)
        if jarfilename is None:
            return filename in self.namelist()
        else:
            jarfile = self.jarfiles[jarfilename]
            return filename in jarfile.namelist()

    def openinputstream(self, jarfilename, filename):
        """opens a file (possibly inside a jarfile as a StringIO"""
        if jarfilename is None:
            contents = self.read(filename)
            def onclose(contents):
                if contents != self.read(filename):
                    self.overwritestr(filename, contents)
            inputstream = CatchPotentialOutput(contents, onclose)
            self.addcatcher(inputstream.slam)
        else:
            jarfile = self.jarfiles[jarfilename]
            contents = jarfile.read(filename)
            inputstream = NamedStringInput(contents)
        inputstream.name = self.jartoospath(jarfilename, filename)
        if hasattr(self.fp, 'name'):
            inputstream.name = "%s:%s" % (self.fp.name, inputstream.name)
        return inputstream

    def openoutputstream(self, jarfilename, filename):
        """opens a file for writing (possibly inside a jarfile as a StringIO"""
        if jarfilename is None:
            def onclose(contents):
                self.overwritestr(filename, contents)
        else:
            if jarfilename in self.jarfiles:
                jarfile = self.jarfiles[jarfilename]
            else:
                jarstream = self.openoutputstream(None, jarfilename)
                jarfile = ZipFileCatcher(jarstream, "w")
                self.jarfiles[jarfilename] = jarfile
                self.addcatcher(jarstream.slam)
            def onclose(contents):
                jarfile.overwritestr(filename, contents)
        outputstream = wStringIO.CatchStringOutput(onclose)
        outputstream.name = "%s %s" % (jarfilename, filename)
        if jarfilename is None:
            self.addcatcher(outputstream.slam)
        else:
            jarfile.addcatcher(outputstream.slam)
        return outputstream

    def close(self):
        """Close the file, and for mode "w" and "a" write the ending records."""
        for jarfile in self.jarfiles.itervalues():
            jarfile.close()
        super(XpiFile, self).close()

    def testzip(self):
        """test the xpi zipfile and all enclosed jar files..."""
        for jarfile in self.jarfiles.itervalues():
            jarfile.testzip()
        super(XpiFile, self).testzip()

    def restructurejar(self, origjarfilename, newjarfilename, otherxpi, newlang, newregion):
        """Create a new .jar file with the same contents as the given name, but rename directories, write to outputstream"""
        jarfile = self.jarfiles[origjarfilename]
        origlang = self.locale[:self.locale.find("-")]
        if newregion:
            newlocale = "%s-%s" % (newlang, newregion)
        else:
            newlocale = newlang
        for filename in jarfile.namelist():
            filenameparts = filename.split("/")
            for i in range(len(filenameparts)):
                part = filenameparts[i]
                if part == origlang:
                    filenameparts[i] = newlang
                elif part == self.locale:
                    filenameparts[i] = newlocale
                elif part == self.region:
                    filenameparts[i] = newregion
            newfilename = '/'.join(filenameparts)
            fileoutputstream = otherxpi.openoutputstream(newjarfilename, newfilename)
            fileinputstream = self.openinputstream(origjarfilename, filename)
            fileoutputstream.write(fileinputstream.read())
            fileinputstream.close()
            fileoutputstream.close()

    def clone(self, newfilename, newmode=None, newlang=None, newregion=None):
        """Create a new .xpi file with the same contents as this one..."""
        other = XpiFile(newfilename, "w", locale=newlang, region=newregion)
        origlang = self.locale[:self.locale.find("-")]
        # TODO: check if this language replacement code is still neccessary
        if newlang is None:
            newlang = origlang
        if newregion is None:
            newregion = self.region
        if newregion:
            newlocale = "%s-%s" % (newlang, newregion)
        else:
            newlocale = newlang
        for filename in self.namelist():
            filenameparts = filename.split('/')
            basename = filenameparts[-1]
            if basename.startswith(self.locale):
                newbasename = basename.replace(self.locale, newlocale)
            elif basename.startswith(origlang):
                newbasename = basename.replace(origlang, newlang)
            elif basename.startswith(self.region):
                newbasename = basename.replace(self.region, newregion)
            else:
                newbasename = basename
            if newbasename != basename:
                filenameparts[-1] = newbasename
                renamefilename = "/".join(filenameparts)
                print "cloning", filename, "and renaming to", renamefilename
            else:
                print "cloning", filename
                renamefilename = filename
            if filename.lower().endswith(".jar"):
                self.restructurejar(filename, renamefilename, other, newlang, newregion)
            else:
                inputstream = self.openinputstream(None, filename)
                outputstream = other.openoutputstream(None, renamefilename)
                outputstream.write(inputstream.read())
                inputstream.close()
                outputstream.close()
        other.close()
        if newmode is None: newmode = self.mode
        if newmode == "w": newmode = "a"
        other = XpiFile(newfilename, newmode)
        other.setlangreg(newlocale, newregion)
        return other

    def iterextractnames(self, includenonjars=False, includedirs=False):
        """iterates through all the localization files with the common prefix stripped and a jarfile name added if neccessary"""
        if includenonjars:
            for filename in self.namelist():
                if filename.endswith('/') and not includedirs: continue
                if not self.islocfile(filename) and not self.includenonloc: continue
                if not filename.lower().endswith(".jar"):
                    yield self.jartoospath(None, filename)
        for jarfilename, jarfile in self.iterjars():
            for filename in jarfile.namelist():
                if filename.endswith('/'):
                    if not includedirs: continue
                if not self.islocfile(filename) and not self.includenonloc: continue
                yield self.jartoospath(jarfilename, filename)

    # the following methods are required by translate.convert.ArchiveConvertOptionParser #
    def __iter__(self):
        """iterates through all the files. this is the method use by the converters"""
        for inputpath in self.iterextractnames(includenonjars=True):
            yield inputpath

    def __contains__(self, fullpath):
        """returns whether the given pathname exists in the archive"""
        try:
            jarfilename, filename = self.ostojarpath(fullpath)
        except IndexError:
            return False
        return self.jarfileexists(jarfilename, filename)

    def openinputfile(self, fullpath):
        """opens an input file given the full pathname"""
        jarfilename, filename = self.ostojarpath(fullpath)
        return self.openinputstream(jarfilename, filename)

    def openoutputfile(self, fullpath):
        """opens an output file given the full pathname"""
        try:
            jarfilename, filename = self.ostojarpath(fullpath)
        except IndexError:
            return None
        return self.openoutputstream(jarfilename, filename)

if __name__ == '__main__':
    import optparse
    optparser = optparse.OptionParser(version="%prog "+__version__.sver)
    optparser.usage = "%prog [-l|-x] [options] file.xpi"
    optparser.add_option("-l", "--list", help="list files", \
        action="store_true", dest="listfiles", default=False)
    optparser.add_option("-p", "--prefix", help="show common prefix", \
        action="store_true", dest="showprefix", default=False)
    optparser.add_option("-x", "--extract", help="extract files", \
        action="store_true", dest="extractfiles", default=False)
    optparser.add_option("-d", "--extractdir", help="extract into EXTRACTDIR", \
        default=".", metavar="EXTRACTDIR")
    (options, args) = optparser.parse_args()
    if len(args) < 1:
        optparser.error("need at least one argument")
    xpifile = XpiFile(args[0])
    if options.showprefix:
        for prefix, mapto in xpifile.dirmap.iteritems():
            print "/".join(prefix), "->", "/".join(mapto)
    if options.listfiles:
        for name in xpifile.iterextractnames(includenonjars=True, includedirs=True):
            print name #, xpifile.ostojarpath(name)
    if options.extractfiles:
        if options.extractdir and not os.path.isdir(options.extractdir):
            os.mkdir(options.extractdir)
        for name in xpifile.iterextractnames(includenonjars=True, includedirs=False):
            abspath = os.path.join(options.extractdir, name)
            # check neccessary directories exist - this way we don't create empty directories
            currentpath = options.extractdir
            subparts = os.path.dirname(name).split(os.sep)
            for part in subparts:
                currentpath = os.path.join(currentpath, part)
                if not os.path.isdir(currentpath):
                    os.mkdir(currentpath)
            outputstream = open(abspath, 'w')
            jarfilename, filename = xpifile.ostojarpath(name)
            inputstream = xpifile.openinputstream(jarfilename, filename)
            outputstream.write(inputstream.read())
            outputstream.close()


#!/usr/bin/env python
# -*- coding: utf-8 -*-
# 
# Copyright 2004, 2005 Zuza Software Foundation
# 
# This file is part of translate.
#
# translate is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
# 
# translate is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with translate; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA

"""module for accessing mozilla xpi packages"""

from __future__ import generators
import zipfile
import os.path
from translate import __version__
import StringIO
import re

# we have some enhancements to zipfile in a file called zipfileext
# hopefully they will be included in a future version of python
from translate.misc import zipfileext
ZipFileBase = zipfileext.ZipFileExt

from translate.misc import wStringIO
# this is a fix to the StringIO in Python 2.3.3
# submitted as patch 951915 on sourceforge
class FixedStringIO(wStringIO.StringIO):
    def truncate(self, size=None):
        StringIO.StringIO.truncate(self, size)
        self.len = len(self.buf)

NamedStringInput = wStringIO.StringIO
NamedStringOutput = wStringIO.StringIO

def _commonprefix(itemlist):
    def cp(a, b):
        l = min(len(a), len(b))
        for n in range(l):
            if a[n] != b[n]: return a[:n]
        return a[:l]
    if itemlist:
        return reduce(cp, itemlist)
    else:
        return ''

def rememberchanged(self, method):
    def changed(*args, **kwargs):
        self.changed = True
        method(*args, **kwargs)
    return changed

class CatchPotentialOutput(NamedStringInput, object):
    """catches output if there has been, before closing"""
    def __init__(self, contents, onclose):
        """Set up the output stream, and remember a method to call on closing"""
        NamedStringInput.__init__(self, contents)
        self.onclose = onclose
        self.changed = False
        s = super(CatchPotentialOutput, self)
        self.write = rememberchanged(self, s.write)
        self.writelines = rememberchanged(self, s.writelines)
        self.truncate = rememberchanged(self, s.truncate)

    def close(self):
        """wrap the underlying close method, to pass the value to onclose before it goes"""
        if self.changed:
            value = self.getvalue()
            self.onclose(value)
        NamedStringInput.close(self)

    def flush(self):
        """zip files call flush, not close, on file-like objects"""
        value = self.getvalue()
        self.onclose(value)
        NamedStringInput.flush(self)

    def slam(self):
        """use this method to force the closing of the stream if it isn't closed yet"""
        if not self.closed:
            self.close()

class ZipFileCatcher(ZipFileBase, object):
    """a ZipFile that calls any methods its instructed to before closing (useful for catching stream output)"""
    def __init__(self, *args, **kwargs):
        """initialize the ZipFileCatcher"""
        # storing oldclose as attribute, since if close is called from __del__ it has no access to external variables
        self.oldclose = super(ZipFileCatcher, self).close
        super(ZipFileCatcher, self).__init__(*args, **kwargs)

    def addcatcher(self, pendingsave):
        """remember to call the given method before closing"""
        if hasattr(self, "pendingsaves"):
            if not pendingsave in self.pendingsaves:
                self.pendingsaves.append(pendingsave)
        else:
            self.pendingsaves = [pendingsave]

    def close(self):
        """close the stream, remembering to call any addcatcher methods first"""
        if hasattr(self, "pendingsaves"):
            for pendingsave in self.pendingsaves:
                pendingsave()
        # if close is called from __del__, it somehow can't see ZipFileCatcher, so we've cached oldclose...
        if ZipFileCatcher is None:
            self.oldclose()
        else:
            super(ZipFileCatcher, self).close()

    def overwritestr(self, zinfo_or_arcname, bytes):
        """writes the string into the archive, overwriting the file if it exists...""" 
        if isinstance(zinfo_or_arcname, zipfile.ZipInfo):
            filename = zinfo_or_arcname.filename
        else:
            filename = zinfo_or_arcname
        if filename in self.NameToInfo:
            self.delete(filename)
        self.writestr(zinfo_or_arcname, bytes)
        self.writeendrec()

class XpiFile(ZipFileCatcher):
    def __init__(self, *args, **kwargs):
        """sets up the xpi file"""
        self.includenonloc = kwargs.get("includenonloc", True)
        if "includenonloc" in kwargs:
            del kwargs["includenonloc"]
        if "compression" not in kwargs:
            kwargs["compression"] = zipfile.ZIP_DEFLATED
        self.locale = kwargs.pop("locale", None)
        self.region = kwargs.pop("region", None)
        super(XpiFile, self).__init__(*args, **kwargs)
        self.jarfiles = {}
        self.findlangreg()
        self.jarprefixes = self.findjarprefixes()
        self.reverseprefixes = dict([
            (prefix,jarfilename) for jarfilename, prefix in self.jarprefixes.iteritems() if prefix])
        self.reverseprefixes["package/"] = None

    def iterjars(self):
        """iterate through the jar files in the xpi as ZipFile objects"""
        for filename in self.namelist():
            if filename.lower().endswith('.jar'):
                if filename not in self.jarfiles:
                    jarstream = self.openinputstream(None, filename)
                    jarfile = ZipFileCatcher(jarstream, mode=self.mode)
                    self.jarfiles[filename] = jarfile
                else:
                    jarfile = self.jarfiles[filename]
                yield filename, jarfile

    def islocfile(self, filename):
        """returns whether the given file is needed for localization (basically .dtd and .properties)"""
        base, ext = os.path.splitext(filename)
        return ext in (os.extsep + "dtd", os.extsep + "properties")

    def findlangreg(self):
        """finds the common prefix of all the files stored in the jar files"""
        dirstructure = {}
        locale = self.locale
        region = self.region
        localematch = re.compile("^[a-z]{2,3}(-[a-zA-Z]{2,3}|)$")
        regionmatch = re.compile("^[a-zA-Z]{2,3}$")
        # exclude en-mac, en-win, en-unix for seamonkey
        osmatch = re.compile("^[a-z]{2,3}-(mac|unix|win)$")
        for jarfilename, jarfile in self.iterjars():
            jarname = "".join(jarfilename.split('/')[-1:]).replace(".jar", "", 1)
            if localematch.match(jarname) and not osmatch.match(jarname):
                if locale is None:
                    locale = jarname
                elif locale != jarname:
                    locale = 0
            elif regionmatch.match(jarname):
                if region is None:
                    region = jarname
                elif region != jarname:
                    region = 0
            for filename in jarfile.namelist():
                if filename.endswith('/'): continue
                if not self.islocfile(filename) and not self.includenonloc: continue
                parts = filename.split('/')[:-1]
                treepoint = dirstructure
                for partnum in range(len(parts)):
                    part = parts[partnum]
                    if part in treepoint:
                        treepoint = treepoint[part]
                    else:
                        treepoint[part] = {}
                        treepoint = treepoint[part]
        localeentries = {}
        if 'locale' in dirstructure:
            for dirname in dirstructure['locale']:
                localeentries[dirname] = 1
                if localematch.match(dirname) and not osmatch.match(dirname):
                    if locale is None:
                        locale = dirname
                    elif locale != dirname:
                        print "locale dir mismatch - ", dirname, "but locale is", locale, "setting to 0"
                        locale = 0
                elif regionmatch.match(dirname):
                    if region is None:
                        region = dirname
                    elif region != dirname:
                        region = 0
        if locale and locale in localeentries:
            del localeentries[locale]
        if region and region in localeentries:
            del localeentries[region]
        if locale and not region:
            if "-" in locale:
                region = locale.split("-", 1)[1]
            else:
                region = ""
        self.setlangreg(locale, region)

    def setlangreg(self, locale, region):
        """set the locale and region of this xpi"""
        if locale == 0 or locale is None:
            raise ValueError("unable to determine locale")
        self.locale = locale
        self.region = region
        self.dirmap = {}
        if self.locale is not None:
            self.dirmap[('locale', self.locale)] = ('lang-reg',)
        if self.region:
            self.dirmap[('locale', self.region)] = ('reg',)

    def findjarprefixes(self):
        """checks the uniqueness of the jar files contents"""
        uniquenames = {}
        jarprefixes = {}
        for jarfilename, jarfile in self.iterjars():
            jarprefixes[jarfilename] = ""
            for filename in jarfile.namelist():
                if filename.endswith('/'): continue
                if filename in uniquenames:
                    jarprefixes[jarfilename] = True
                    jarprefixes[uniquenames[filename]] = True
                else:
                    uniquenames[filename] = jarfilename
        for jarfilename, hasconflicts in jarprefixes.items():
            if hasconflicts:
                shortjarfilename = os.path.split(jarfilename)[1]
                shortjarfilename = os.path.splitext(shortjarfilename)[0]
                jarprefixes[jarfilename] = shortjarfilename+'/'
        # this is a clever trick that will e.g. remove zu- from zu-win, zu-mac, zu-unix
        commonjarprefix = _commonprefix([prefix for prefix in jarprefixes.itervalues() if prefix])
        if commonjarprefix:
            for jarfilename, prefix in jarprefixes.items():
                if prefix:
                    jarprefixes[jarfilename] = prefix.replace(commonjarprefix, '', 1)
        return jarprefixes

    def ziptoospath(self, zippath):
        """converts a zipfile filepath to an os-style filepath"""
        return os.path.join(*zippath.split('/'))

    def ostozippath(self, ospath):
        """converts an os-style filepath to a zipfile filepath"""
        return '/'.join(ospath.split(os.sep))

    def mapfilename(self, filename):
        """uses a map to simplify the directory structure"""
        parts = tuple(filename.split('/'))
        possiblematch = None
        for prefix, mapto in self.dirmap.iteritems():
            if parts[:len(prefix)] == prefix:
                if possiblematch is None or len(possiblematch[0]) < len(prefix):
                    possiblematch = prefix, mapto
        if possiblematch is not None:
            prefix, mapto = possiblematch
            mapped = mapto + parts[len(prefix):]
            return '/'.join(mapped)
        return filename

    def mapxpifilename(self, filename):
        """uses a map to rename files that occur straight in the xpi"""
        if filename.startswith('bin/chrome/') and filename.endswith(".manifest"):
            return 'bin/chrome/lang-reg.manifest'
        return filename

    def reversemapfile(self, filename):
        """unmaps the filename..."""
        possiblematch = None
        parts = tuple(filename.split('/'))
        for prefix, mapto in self.dirmap.iteritems():
            if parts[:len(mapto)] == mapto:
                if possiblematch is None or len(possiblematch[0]) < len(mapto):
                    possiblematch = (mapto, prefix)
        if possiblematch is None:
            return filename
        mapto, prefix = possiblematch
        reversemapped = prefix + parts[len(mapto):]
        return '/'.join(reversemapped)

    def reversemapxpifilename(self, filename):
        """uses a map to rename files that occur straight in the xpi"""
        if filename == 'bin/chrome/lang-reg.manifest':
            if self.locale:
                return '/'.join(('bin', 'chrome', self.locale + '.manifest'))
            else:
                for otherfilename in self.namelist():
                    if otherfilename.startswith("bin/chrome/") and otherfilename.endswith(".manifest"):
                        return otherfilename
        return filename

    def jartoospath(self, jarfilename, filename):
        """converts a filename from within a jarfile to an os-style filepath"""
        if jarfilename:
            jarprefix = self.jarprefixes[jarfilename]
            return self.ziptoospath(jarprefix+self.mapfilename(filename))
        else:
            return self.ziptoospath(os.path.join("package", self.mapxpifilename(filename)))

    def ostojarpath(self, ospath):
        """converts an extracted os-style filepath to a jarfilename and filename"""
        zipparts = ospath.split(os.sep)
        prefix = zipparts[0] + '/'
        if prefix in self.reverseprefixes:
            jarfilename = self.reverseprefixes[prefix]
            filename = self.reversemapfile('/'.join(zipparts[1:]))
            if jarfilename is None:
                filename = self.reversemapxpifilename(filename)
            return jarfilename, filename
        else:
            filename = self.ostozippath(ospath)
            if filename in self.namelist():
                return None, filename
            filename = self.reversemapfile('/'.join(zipparts))
            possiblejarfilenames = [jarfilename for jarfilename, prefix in self.jarprefixes.iteritems() if not prefix]
            for jarfilename in possiblejarfilenames:
                jarfile = self.jarfiles[jarfilename]
                if filename in jarfile.namelist():
                    return jarfilename, filename
            raise IndexError("ospath not found in xpi file, could not guess location: %r" % ospath)

    def jarfileexists(self, jarfilename, filename):
        """checks whether the given file exists inside the xpi"""
        if jarfilename is None:
            return filename in self.namelist()
        else:
            jarfile = self.jarfiles[jarfilename]
            return filename in jarfile.namelist()

    def ospathexists(self, ospath):
        """checks whether the given file exists inside the xpi"""
        jarfilename, filename = self.ostojarpath(ospath)
        if jarfilename is None:
            return filename in self.namelist()
        else:
            jarfile = self.jarfiles[jarfilename]
            return filename in jarfile.namelist()

    def openinputstream(self, jarfilename, filename):
        """opens a file (possibly inside a jarfile as a StringIO"""
        if jarfilename is None:
            contents = self.read(filename)
            def onclose(contents):
                if contents != self.read(filename):
                    self.overwritestr(filename, contents)
            inputstream = CatchPotentialOutput(contents, onclose)
            self.addcatcher(inputstream.slam)
        else:
            jarfile = self.jarfiles[jarfilename]
            contents = jarfile.read(filename)
            inputstream = NamedStringInput(contents)
        inputstream.name = self.jartoospath(jarfilename, filename)
        if hasattr(self.fp, 'name'):
            inputstream.name = "%s:%s" % (self.fp.name, inputstream.name)
        return inputstream

    def openoutputstream(self, jarfilename, filename):
        """opens a file for writing (possibly inside a jarfile as a StringIO"""
        if jarfilename is None:
            def onclose(contents):
                self.overwritestr(filename, contents)
        else:
            if jarfilename in self.jarfiles:
                jarfile = self.jarfiles[jarfilename]
            else:
                jarstream = self.openoutputstream(None, jarfilename)
                jarfile = ZipFileCatcher(jarstream, "w")
                self.jarfiles[jarfilename] = jarfile
                self.addcatcher(jarstream.slam)
            def onclose(contents):
                jarfile.overwritestr(filename, contents)
        outputstream = wStringIO.CatchStringOutput(onclose)
        outputstream.name = "%s %s" % (jarfilename, filename)
        if jarfilename is None:
            self.addcatcher(outputstream.slam)
        else:
            jarfile.addcatcher(outputstream.slam)
        return outputstream

    def close(self):
        """Close the file, and for mode "w" and "a" write the ending records."""
        for jarfile in self.jarfiles.itervalues():
            jarfile.close()
        super(XpiFile, self).close()

    def testzip(self):
        """test the xpi zipfile and all enclosed jar files..."""
        for jarfile in self.jarfiles.itervalues():
            jarfile.testzip()
        super(XpiFile, self).testzip()

    def restructurejar(self, origjarfilename, newjarfilename, otherxpi, newlang, newregion):
        """Create a new .jar file with the same contents as the given name, but rename directories, write to outputstream"""
        jarfile = self.jarfiles[origjarfilename]
        origlang = self.locale[:self.locale.find("-")]
        if newregion:
            newlocale = "%s-%s" % (newlang, newregion)
        else:
            newlocale = newlang
        for filename in jarfile.namelist():
            filenameparts = filename.split("/")
            for i in range(len(filenameparts)):
                part = filenameparts[i]
                if part == origlang:
                    filenameparts[i] = newlang
                elif part == self.locale:
                    filenameparts[i] = newlocale
                elif part == self.region:
                    filenameparts[i] = newregion
            newfilename = '/'.join(filenameparts)
            fileoutputstream = otherxpi.openoutputstream(newjarfilename, newfilename)
            fileinputstream = self.openinputstream(origjarfilename, filename)
            fileoutputstream.write(fileinputstream.read())
            fileinputstream.close()
            fileoutputstream.close()

    def clone(self, newfilename, newmode=None, newlang=None, newregion=None):
        """Create a new .xpi file with the same contents as this one..."""
        other = XpiFile(newfilename, "w", locale=newlang, region=newregion)
        origlang = self.locale[:self.locale.find("-")]
        # TODO: check if this language replacement code is still neccessary
        if newlang is None:
            newlang = origlang
        if newregion is None:
            newregion = self.region
        if newregion:
            newlocale = "%s-%s" % (newlang, newregion)
        else:
            newlocale = newlang
        for filename in self.namelist():
            filenameparts = filename.split('/')
            basename = filenameparts[-1]
            if basename.startswith(self.locale):
                newbasename = basename.replace(self.locale, newlocale)
            elif basename.startswith(origlang):
                newbasename = basename.replace(origlang, newlang)
            elif basename.startswith(self.region):
                newbasename = basename.replace(self.region, newregion)
            else:
                newbasename = basename
            if newbasename != basename:
                filenameparts[-1] = newbasename
                renamefilename = "/".join(filenameparts)
                print "cloning", filename, "and renaming to", renamefilename
            else:
                print "cloning", filename
                renamefilename = filename
            if filename.lower().endswith(".jar"):
                self.restructurejar(filename, renamefilename, other, newlang, newregion)
            else:
                inputstream = self.openinputstream(None, filename)
                outputstream = other.openoutputstream(None, renamefilename)
                outputstream.write(inputstream.read())
                inputstream.close()
                outputstream.close()
        other.close()
        if newmode is None: newmode = self.mode
        if newmode == "w": newmode = "a"
        other = XpiFile(newfilename, newmode)
        other.setlangreg(newlocale, newregion)
        return other

    def iterextractnames(self, includenonjars=False, includedirs=False):
        """iterates through all the localization files with the common prefix stripped and a jarfile name added if neccessary"""
        if includenonjars:
            for filename in self.namelist():
                if filename.endswith('/') and not includedirs: continue
                if not self.islocfile(filename) and not self.includenonloc: continue
                if not filename.lower().endswith(".jar"):
                    yield self.jartoospath(None, filename)
        for jarfilename, jarfile in self.iterjars():
            for filename in jarfile.namelist():
                if filename.endswith('/'):
                    if not includedirs: continue
                if not self.islocfile(filename) and not self.includenonloc: continue
                yield self.jartoospath(jarfilename, filename)

    # the following methods are required by translate.convert.ArchiveConvertOptionParser #
    def __iter__(self):
        """iterates through all the files. this is the method use by the converters"""
        for inputpath in self.iterextractnames(includenonjars=True):
            yield inputpath

    def __contains__(self, fullpath):
        """returns whether the given pathname exists in the archive"""
        try:
            jarfilename, filename = self.ostojarpath(fullpath)
        except IndexError:
            return False
        return self.jarfileexists(jarfilename, filename)

    def openinputfile(self, fullpath):
        """opens an input file given the full pathname"""
        jarfilename, filename = self.ostojarpath(fullpath)
        return self.openinputstream(jarfilename, filename)

    def openoutputfile(self, fullpath):
        """opens an output file given the full pathname"""
        try:
            jarfilename, filename = self.ostojarpath(fullpath)
        except IndexError:
            return None
        return self.openoutputstream(jarfilename, filename)

if __name__ == '__main__':
    import optparse
    optparser = optparse.OptionParser(version="%prog "+__version__.sver)
    optparser.usage = "%prog [-l|-x] [options] file.xpi"
    optparser.add_option("-l", "--list", help="list files", \
        action="store_true", dest="listfiles", default=False)
    optparser.add_option("-p", "--prefix", help="show common prefix", \
        action="store_true", dest="showprefix", default=False)
    optparser.add_option("-x", "--extract", help="extract files", \
        action="store_true", dest="extractfiles", default=False)
    optparser.add_option("-d", "--extractdir", help="extract into EXTRACTDIR", \
        default=".", metavar="EXTRACTDIR")
    (options, args) = optparser.parse_args()
    if len(args) < 1:
        optparser.error("need at least one argument")
    xpifile = XpiFile(args[0])
    if options.showprefix:
        for prefix, mapto in xpifile.dirmap.iteritems():
            print "/".join(prefix), "->", "/".join(mapto)
    if options.listfiles:
        for name in xpifile.iterextractnames(includenonjars=True, includedirs=True):
            print name #, xpifile.ostojarpath(name)
    if options.extractfiles:
        if options.extractdir and not os.path.isdir(options.extractdir):
            os.mkdir(options.extractdir)
        for name in xpifile.iterextractnames(includenonjars=True, includedirs=False):
            abspath = os.path.join(options.extractdir, name)
            # check neccessary directories exist - this way we don't create empty directories
            currentpath = options.extractdir
            subparts = os.path.dirname(name).split(os.sep)
            for part in subparts:
                currentpath = os.path.join(currentpath, part)
                if not os.path.isdir(currentpath):
                    os.mkdir(currentpath)
            outputstream = open(abspath, 'w')
            jarfilename, filename = xpifile.ostojarpath(name)
            inputstream = xpifile.openinputstream(jarfilename, filename)
            outputstream.write(inputstream.read())
            outputstream.close()


#!/usr/bin/env python # -*- coding: utf-8 -*- # # Copyright 2004, 2005 Zuza Software Foundation # # This file is part of translate. # # translate is free software; you can redistribute it and/or modify # it under the terms of the GNU General Public License as published by # the Free Software Foundation; either version 2 of the License, or # (at your option) any later version. # # translate is distributed in the hope that it will be useful, # but WITHOUT ANY WARRANTY; without even the implied warranty of # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the # GNU General Public License for more details. # # You should have received a copy of the GNU General Public License # along with translate; if not, write to the Free Software # Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA module for accessing mozilla xpi packages # we have some enhancements to zipfile in a file called zipfileext # hopefully they will be included in a future version of python # this is a fix to the StringIO in Python 2.3.3 # submitted as patch 951915 on sourceforge catches output if there has been, before closing Set up the output stream, and remember a method to call on closing wrap the underlying close method, to pass the value to onclose before it goes zip files call flush, not close, on file-like objects use this method to force the closing of the stream if it isn't closed yet a ZipFile that calls any methods its instructed to before closing (useful for catching stream output) initialize the ZipFileCatcher # storing oldclose as attribute, since if close is called from __del__ it has no access to external variables remember to call the given method before closing close the stream, remembering to call any addcatcher methods first # if close is called from __del__, it somehow can't see ZipFileCatcher, so we've cached oldclose... writes the string into the archive, overwriting the file if it exists... sets up the xpi file iterate through the jar files in the xpi as ZipFile objects returns whether the given file is needed for localization (basically .dtd and .properties) finds the common prefix of all the files stored in the jar files # exclude en-mac, en-win, en-unix for seamonkey set the locale and region of this xpi checks the uniqueness of the jar files contents # this is a clever trick that will e.g. remove zu- from zu-win, zu-mac, zu-unix converts a zipfile filepath to an os-style filepath converts an os-style filepath to a zipfile filepath uses a map to simplify the directory structure uses a map to rename files that occur straight in the xpi unmaps the filename... uses a map to rename files that occur straight in the xpi converts a filename from within a jarfile to an os-style filepath converts an extracted os-style filepath to a jarfilename and filename checks whether the given file exists inside the xpi checks whether the given file exists inside the xpi opens a file (possibly inside a jarfile as a StringIO opens a file for writing (possibly inside a jarfile as a StringIO Close the file, and for mode "w" and "a" write the ending records. test the xpi zipfile and all enclosed jar files... Create a new .jar file with the same contents as the given name, but rename directories, write to outputstream Create a new .xpi file with the same contents as this one... # TODO: check if this language replacement code is still neccessary iterates through all the localization files with the common prefix stripped and a jarfile name added if neccessary # the following methods are required by translate.convert.ArchiveConvertOptionParser # iterates through all the files. this is the method use by the converters returns whether the given pathname exists in the archive opens an input file given the full pathname opens an output file given the full pathname #, xpifile.ostojarpath(name) # check neccessary directories exist - this way we don't create empty directories