/atlassian-plugins-osgi/src/main/java/com/atlassian/plugin/osgi/factory/transform/DefaultPluginTransformer.java

package com.atlassian.plugin.osgi.factory.transform;

import com.atlassian.plugin.Application;
import com.atlassian.plugin.JarPluginArtifact;
import com.atlassian.plugin.PluginArtifact;
import com.atlassian.plugin.osgi.container.OsgiContainerManager;
import com.atlassian.plugin.osgi.container.OsgiPersistentCache;
import com.atlassian.plugin.osgi.factory.transform.model.SystemExports;
import com.atlassian.plugin.osgi.factory.transform.stage.AddBundleOverridesStage;
import com.atlassian.plugin.osgi.factory.transform.stage.ComponentImportSpringStage;
import com.atlassian.plugin.osgi.factory.transform.stage.ComponentSpringStage;
import com.atlassian.plugin.osgi.factory.transform.stage.GenerateManifestStage;
import com.atlassian.plugin.osgi.factory.transform.stage.HostComponentSpringStage;
import com.atlassian.plugin.osgi.factory.transform.stage.ModuleTypeSpringStage;
import com.atlassian.plugin.osgi.factory.transform.stage.ScanDescriptorForHostClassesStage;
import com.atlassian.plugin.osgi.factory.transform.stage.ScanInnerJarsStage;
import com.atlassian.plugin.osgi.hostcomponents.HostComponentRegistration;
import com.google.common.collect.ImmutableList;
import org.apache.commons.io.IOUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.BufferedOutputStream;
import java.io.ByteArrayInputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.SortedMap;
import java.util.TreeMap;
import java.util.zip.Deflater;
import java.util.zip.ZipEntry;
import java.util.zip.ZipInputStream;
import java.util.zip.ZipOutputStream;

import static com.google.common.base.Preconditions.checkNotNull;

/**
 * Default implementation of plugin transformation that uses stages to convert a plain JAR into an OSGi bundle.
 */
public class DefaultPluginTransformer implements PluginTransformer {
    private static final Logger log = LoggerFactory.getLogger(DefaultPluginTransformer.class);

    public static final String TRANSFORM_COMPRESSION_LEVEL = "atlassian.plugins.plugin.transformer.compression";

    private final String pluginDescriptorPath;
    private final List<TransformStage> stages;
    private final File bundleCacheDir;
    private final SystemExports systemExports;
    private final Set<Application> applications;
    private final OsgiContainerManager osgiContainerManager;

    /**
     * Gets the default list of transform stages performed by the transformer. Clients wishing to add stages to the
     * transformation process should use this list as a template rather than creating their own from scratch.
     */
    public static ArrayList<TransformStage> getDefaultTransformStages() {
        return new ArrayList<TransformStage>(Arrays.asList(
                new AddBundleOverridesStage(),
                new ScanInnerJarsStage(),
                new ComponentImportSpringStage(),
                new ComponentSpringStage(),
                new ScanDescriptorForHostClassesStage(),
                new ModuleTypeSpringStage(),
                new HostComponentSpringStage(),
                new GenerateManifestStage()
        ));
    }

    /**
     * Constructs a transformer with the default stages
     *
     * @param cache                The OSGi cache configuration for transformed plugins
     * @param systemExports        The packages the system bundle exports
     * @param pluginDescriptorPath The path to the plugin descriptor
     * @since 2.2.0
     */
    public DefaultPluginTransformer(OsgiPersistentCache cache, SystemExports systemExports, Set<Application> applications, String pluginDescriptorPath, OsgiContainerManager osgiContainerManager) {
        this(cache, systemExports, applications, pluginDescriptorPath, osgiContainerManager, getDefaultTransformStages());
    }

    /**
     * Constructs a transformer and its stages
     *
     * @param cache                The OSGi cache configuration for transformed plugins
     * @param systemExports        The packages the system bundle exports
     * @param pluginDescriptorPath The descriptor path
     * @param stages               A set of stages
     * @since 2.2.0
     */
    public DefaultPluginTransformer(OsgiPersistentCache cache, SystemExports systemExports, Set<Application> applications, String pluginDescriptorPath, OsgiContainerManager osgiContainerManager, List<TransformStage> stages) {
        this.pluginDescriptorPath = checkNotNull(pluginDescriptorPath, "The plugin descriptor path is required");
        this.osgiContainerManager = checkNotNull(osgiContainerManager);
        this.stages = ImmutableList.copyOf(checkNotNull(stages, "A list of stages is required"));
        this.bundleCacheDir = checkNotNull(cache).getTransformedPluginCache();
        this.systemExports = systemExports;
        this.applications = applications;
    }

    /**
     * Transforms the file into an OSGi bundle
     *
     * @param pluginJar The plugin jar
     * @param regs      The list of registered host components
     * @return The new OSGi-enabled plugin jar
     * @throws PluginTransformationException If anything goes wrong
     */
    public File transform(File pluginJar, List<HostComponentRegistration> regs) throws PluginTransformationException {
        return transform(new JarPluginArtifact(pluginJar), regs);
    }

    /**
     * Transforms the file into an OSGi bundle
     *
     * @param pluginArtifact The plugin artifact, usually a jar
     * @param regs           The list of registered host components
     * @return The new OSGi-enabled plugin jar
     * @throws PluginTransformationException If anything goes wrong
     */
    public File transform(PluginArtifact pluginArtifact, List<HostComponentRegistration> regs) throws PluginTransformationException {
        checkNotNull(pluginArtifact, "The plugin artifact is required");
        checkNotNull(regs, "The host component registrations are required");

        File artifactFile = pluginArtifact.toFile();

        // Look in cache first
        File cachedPlugin = getFromCache(artifactFile);
        if (cachedPlugin != null) {
            return cachedPlugin;
        }

        final TransformContext context = new TransformContext(regs, systemExports, pluginArtifact, applications, pluginDescriptorPath, osgiContainerManager);
        for (TransformStage stage : stages) {
            stage.execute(context);
        }

        // Create a new jar by overriding the specified files
        try {
            if (log.isDebugEnabled()) {
                StringBuilder sb = new StringBuilder();
                sb.append("Overriding files in ").append(pluginArtifact.toString()).append(":\n");
                for (Map.Entry<String, byte[]> entry : context.getFileOverrides().entrySet()) {
                    sb.append("==").append(entry.getKey()).append("==\n");

                    // Yes, this doesn't take into account encoding, but since only text files are overridden, that
                    // should be fine
                    sb.append(new String(entry.getValue()));
                }
                log.debug(sb.toString());
            }
            return addFilesToExistingZip(artifactFile, context.getFileOverrides());
        } catch (IOException e) {
            throw new PluginTransformationException("Unable to add files to plugin jar", e);
        }
    }

    private File getFromCache(File artifact) {
        String name = generateCacheName(artifact);
        for (File child : bundleCacheDir.listFiles()) {
            if (child.getName().equals(name))
                return child;
        }
        return null;
    }

    /**
     * Generate a cache name that incorporates the timestap and preserves the extension
     *
     * @param file The original file to cache
     * @return The new file name
     */
    static String generateCacheName(File file) {
        int dotPos = file.getName().lastIndexOf('.');
        if (dotPos > 0 && file.getName().length() - 1 > dotPos) {
            return file.getName().substring(0, dotPos) + "_" + file.lastModified() + file.getName().substring(dotPos);
        } else {
            return file.getName() + "_" + file.lastModified();
        }
    }


    /**
     * Creates a new jar by overriding the specified files in the existing one
     *
     * @param zipFile The existing zip file
     * @param files   The files to override
     * @return The new zip
     * @throws IOException If there are any problems processing the streams
     */
    File addFilesToExistingZip(File zipFile,
                               Map<String, byte[]> files) throws IOException {
        // get a temp file
        File tempFile = new File(bundleCacheDir, generateCacheName(zipFile));

        ZipInputStream zin = null;
        ZipOutputStream out = null;
        try {
            zin = new ZipInputStream(new FileInputStream(zipFile));
            out = new ZipOutputStream(new BufferedOutputStream(new FileOutputStream(tempFile)));
            final int requestedCompressionLevel = Integer.getInteger(TRANSFORM_COMPRESSION_LEVEL, Deflater.NO_COMPRESSION);
            final int clampedCompressionLevel =
                    Math.max(Deflater.NO_COMPRESSION, Math.min(requestedCompressionLevel, Deflater.BEST_COMPRESSION));
            out.setLevel(clampedCompressionLevel);

            // Copy any entries that are unmodified from the existing JAR to the new JAR
            ZipEntry entry = zin.getNextEntry();
            while (entry != null) {
                String name = entry.getName();
                if (!files.containsKey(name)) {
                    // Add ZIP entry to output stream, preserve existing modification time
                    final ZipEntry newEntry = new ZipEntry(name);
                    newEntry.setTime(entry.getTime());
                    out.putNextEntry(newEntry);

                    // Transfer bytes from the ZIP file to the output file
                    IOUtils.copyLarge(zin, out);
                }
                entry = zin.getNextEntry();
            }
            // Close the streams
            zin.close();

            // Override/append any entries present in files
            final SortedMap<String, byte[]> sortedFiles = new TreeMap(files);
            for (Map.Entry<String, byte[]> fentry : sortedFiles.entrySet()) {
                InputStream in = null;
                try {
                    in = new ByteArrayInputStream(fentry.getValue());

                    // Add ZIP entry to output stream, use the last modified time of the original jar
                    // for any new files
                    final ZipEntry newEntry = new ZipEntry(fentry.getKey());
                    newEntry.setTime(zipFile.lastModified());
                    out.putNextEntry(newEntry);

                    // Transfer bytes from the file to the ZIP file
                    IOUtils.copyLarge(in, out);

                    // Complete the entry
                    out.closeEntry();
                } finally {
                    IOUtils.closeQuietly(in);
                }
            }
            // Complete the ZIP file
            out.close();
        } finally {
            // Close just in case
            IOUtils.closeQuietly(zin);
            IOUtils.closeQuietly(out);
        }
        return tempFile;
    }

}