1   package com.atlassian.maven.plugins.amps.util;
2   
3   import org.apache.commons.io.IOUtils;
4   import org.apache.commons.lang.StringUtils;
5   
6   import com.google.common.collect.Lists;
7   
8   
9   import java.io.File;
10  import java.io.FileInputStream;
11  import java.io.FileOutputStream;
12  import java.io.IOException;
13  import java.io.InputStream;
14  import java.io.OutputStream;
15  import java.util.Enumeration;
16  import java.util.List;
17  import java.util.zip.ZipEntry;
18  import java.util.zip.ZipException;
19  import java.util.zip.ZipFile;
20  import java.util.zip.ZipOutputStream;
21  
22  public class ZipUtils
23  {
24      public static void unzip(final File zipFile, final String destDir) throws IOException
25      {
26          unzip(zipFile, destDir, 0);
27      }
28  
29      /**
30       * Unzips a file
31       *
32       * @param zipFile
33       *            the Zip file
34       * @param destDir
35       *            the destination folder
36       * @param leadingPathSegmentsToTrim
37       *            number of root folders to skip. Example: If all files are in generated-resources/home/*,
38       *            then you may want to skip 2 folders.
39       * @throws IOException
40       */
41      public static void unzip(final File zipFile, final String destDir, int leadingPathSegmentsToTrim) throws IOException
42      {
43          final ZipFile zip = new ZipFile(zipFile);
44          try
45          {
46              final Enumeration<? extends ZipEntry> entries = zip.entries();
47              while (entries.hasMoreElements())
48              {
49                  final ZipEntry zipEntry = entries.nextElement();
50                  String zipPath = trimPathSegments(zipEntry.getName(), leadingPathSegmentsToTrim);
51                  final File file = new File(destDir + "/" + zipPath);
52                  if (zipEntry.isDirectory())
53                  {
54                      file.mkdirs();
55                      continue;
56                  }
57                  // make sure our parent exists in case zipentries are out of order
58                  if (!file.getParentFile().exists())
59                  {
60                      file.getParentFile().mkdirs();
61                  }
62  
63                  InputStream is = null;
64                  OutputStream fos = null;
65                  try
66                  {
67                      is = zip.getInputStream(zipEntry);
68                      fos = new FileOutputStream(file);
69                      IOUtils.copy(is, fos);
70                  }
71                  finally
72                  {
73                      IOUtils.closeQuietly(is);
74                      IOUtils.closeQuietly(fos);
75                  }
76                  file.setLastModified(zipEntry.getTime());
77              }
78          }
79          finally
80          {
81              try
82              {
83                  zip.close();
84              }
85              catch (IOException e)
86              {
87                  // ignore
88              }
89          }
90      }
91  
92      /**
93       * Count the number of nested root folders. A root folder is a folder which contains 0 or 1 file or folder.
94       *
95       * Example: A zip with only "generated-resources/home/database.log" has 2 root folders.
96       *
97       * @param zip the zip file
98       * @return the number of root folders.
99       */
100     public static int countNestingLevel(File zip) throws ZipException, IOException
101     {
102         ZipFile zipFile = null;
103         try
104         {
105             zipFile = new ZipFile(zip);
106             List<String> filenames = toList(zipFile.entries());
107             return countNestingLevel(filenames);
108         }
109         finally
110         {
111             if (zipFile != null)
112             {
113                 try
114                 {
115                     zipFile.close();
116                 }
117                 catch (IOException e)
118                 {
119                     // ignore
120                 }
121             }
122         }
123     }
124 
125     /**
126      * Count the number of nested root directories in the filenames.
127      *
128      * A root directory is a directory that has no sibling.
129      * @param filenames the list of filenames, using / as a separator. Must be a mutable copy,
130      * as it will be modified.
131      */
132     static int countNestingLevel(List<String> filenames)
133     {
134         String prefix = StringUtils.getCommonPrefix(filenames.toArray(new String[filenames.size()]));
135         if (!prefix.endsWith("/"))
136         {
137             prefix = prefix.substring(0, prefix.lastIndexOf("/") + 1);
138         }
139 
140         // The first prefix may be wrong, example:
141         // root/ <- to be discarded
142         // root/nested/ <- to be discarded
143         // root/nested/folder1/file.txt <- the root "root/nested/" will be detected properly
144         // root/nested/folder2/file.txt
145         if (filenames.remove(prefix))
146         {
147             return countNestingLevel(filenames);
148         }
149 
150         // The client can't use these filenames anymore.
151         filenames.clear();
152         return StringUtils.countMatches(prefix, "/");
153     }
154 
155     private static List<String> toList(final Enumeration<? extends ZipEntry> entries)
156     {
157         List<String> filenamesList = Lists.newArrayList();
158         while (entries.hasMoreElements())
159         {
160             final ZipEntry zipEntry = entries.nextElement();
161             filenamesList.add(zipEntry.getName());
162         }
163         return filenamesList;
164     }
165 
166     /**
167      * @param prefix the prefix. If empty, uses the srcDir's name. That means you can't create a zip with no
168      * root folder.
169      */
170     public static void zipDir(final File zipFile, final File srcDir, final String prefix) throws IOException
171     {
172         ZipOutputStream out = new ZipOutputStream(new FileOutputStream(zipFile));
173         try
174         {
175             addZipPrefixes(srcDir, out, prefix);
176             addZipDir(srcDir, out, prefix);
177         }
178         finally
179         {
180             // Complete the ZIP file
181             IOUtils.closeQuietly(out);
182         }
183     }
184 
185     private static void addZipPrefixes(File dirObj, ZipOutputStream out, String prefix) throws IOException
186     {
187         // need to manually add the prefix folders
188         String entryPrefix = ensurePrefixWithSlash(dirObj, prefix);
189 
190         String[] prefixes = entryPrefix.split("/");
191         String lastPrefix = "";
192         for (int i = 0; i < prefixes.length; i++)
193         {
194             ZipEntry entry = new ZipEntry(lastPrefix + prefixes[i] + "/");
195             out.putNextEntry(entry);
196             out.closeEntry();
197 
198             lastPrefix = prefixes[i] + "/";
199         }
200     }
201 
202     private static void addZipDir(File dirObj, ZipOutputStream out, String prefix) throws IOException
203     {
204         File[] files = dirObj.listFiles();
205         byte[] tmpBuf = new byte[1024];
206         File currentFile;
207         String entryPrefix = ensurePrefixWithSlash(dirObj, prefix);
208         String entryName = "";
209 
210         for (int i = 0; i < files.length; i++)
211         {
212             currentFile = files[i];
213             if (currentFile.isDirectory())
214             {
215                 entryName = entryPrefix + currentFile.getName() + "/";
216 
217                 // need to manually add folders so entries are in order
218                 ZipEntry entry = new ZipEntry(entryName);
219                 out.putNextEntry(entry);
220                 out.closeEntry();
221 
222                 // add the files in the folder
223                 addZipDir(currentFile, out, entryName);
224             }
225             else if (currentFile.isFile())
226             {
227 
228                 entryName = entryPrefix + currentFile.getName();
229                 FileInputStream in = new FileInputStream(currentFile.getAbsolutePath());
230                 try
231                 {
232                     out.putNextEntry(new ZipEntry(entryName));
233                     // Transfer from the file to the ZIP file
234                     int len;
235                     while ((len = in.read(tmpBuf)) > 0)
236                     {
237                         out.write(tmpBuf, 0, len);
238                     }
239 
240                     // Complete the entry
241                     out.closeEntry();
242                 }
243                 finally
244                 {
245                     IOUtils.closeQuietly(in);
246                 }
247             }
248         }
249     }
250 
251     /**
252      * Make sure 'prefix' is in format 'entry/' or, by default, 'rootDir/'
253      * (not '', '/', '/entry', or 'entry').
254      */
255     private static String ensurePrefixWithSlash(File rootDir, String prefix)
256     {
257         String entryPrefix = prefix;
258 
259         if (StringUtils.isNotBlank(entryPrefix) && !entryPrefix.equals("/"))
260         {
261             // strip leading '/'
262             if (entryPrefix.charAt(0) == '/')
263             {
264                 entryPrefix = entryPrefix.substring(1);
265             }
266             // ensure trailing '/'
267             if (entryPrefix.charAt(entryPrefix.length() - 1) != '/')
268             {
269                 entryPrefix = entryPrefix + "/";
270             }
271         }
272         else
273         {
274             entryPrefix = rootDir.getName() + "/";
275         }
276 
277         return entryPrefix;
278     }
279 
280     private static String trimPathSegments(String zipPath, final int trimLeadingPathSegments)
281     {
282         int startIndex = 0;
283         for (int i = 0; i < trimLeadingPathSegments; i++)
284         {
285             int nextSlash = zipPath.indexOf("/", startIndex);
286             if (nextSlash == -1)
287             {
288                 break;
289             }
290             else
291             {
292                 startIndex = nextSlash + 1;
293             }
294         }
295 
296         return zipPath.substring(startIndex);
297     }
298 
299 }