\r
import java.io.ByteArrayInputStream;\r
import java.io.ByteArrayOutputStream;\r
+import java.io.Closeable;\r
import java.io.File;\r
import java.io.FileInputStream;\r
import java.io.IOException;\r
import java.util.zip.ZipEntry;\r
import java.util.zip.ZipInputStream;\r
\r
-import org.apache.poi.poifs.eventfilesystem.POIFSReader;\r
-import org.apache.poi.poifs.eventfilesystem.POIFSReaderEvent;\r
-import org.apache.poi.poifs.eventfilesystem.POIFSReaderListener;\r
+import org.apache.poi.poifs.filesystem.DirectoryNode;\r
import org.apache.poi.poifs.filesystem.DocumentInputStream;\r
+import org.apache.poi.poifs.filesystem.DocumentNode;\r
+import org.apache.poi.poifs.filesystem.Entry;\r
import org.apache.poi.poifs.filesystem.NPOIFSFileSystem;\r
import org.apache.poi.poifs.filesystem.OfficeXmlFileException;\r
import org.apache.poi.util.IOUtils;\r
* Finds all VBA Macros in an office file (OLE2/POIFS and OOXML/OPC),\r
* and returns them\r
*/\r
-public class VBAMacroReader {\r
- protected static final String VBA_PROJECT = "xl/vbaProject.bin";\r
+public class VBAMacroReader implements Closeable {\r
+ protected static final String VBA_PROJECT_OOXML = "xl/vbaProject.bin";\r
+ protected static final String VBA_PROJECT_POIFS = "VBA";\r
\r
private NPOIFSFileSystem fs;\r
\r
if (NPOIFSFileSystem.hasPOIFSHeader(header8)) {\r
fs = new NPOIFSFileSystem(stream);\r
} else {\r
- stream.unread(header8);\r
openOOXML(stream);\r
}\r
}\r
ZipInputStream zis = new ZipInputStream(zipFile);\r
ZipEntry zipEntry;\r
while ((zipEntry = zis.getNextEntry()) != null) {\r
- if (VBA_PROJECT.equals(zipEntry.getName())) {\r
+ if (VBA_PROJECT_OOXML.equals(zipEntry.getName())) {\r
try {\r
+ // Make a NPOIFS from the contents, and close the stream\r
this.fs = new NPOIFSFileSystem(zis);\r
- } finally {\r
- zis.closeEntry();\r
+ return;\r
+ } catch (IOException e) {\r
+ // Tidy up\r
+ zis.close();\r
+ \r
+ // Pass on\r
+ throw e;\r
}\r
- zis.close();\r
- return;\r
}\r
}\r
zis.close();\r
throw new IllegalArgumentException("No VBA project found");\r
}\r
+ \r
+ public void close() throws IOException {\r
+ fs.close();\r
+ fs = null;\r
+ }\r
\r
/**\r
* Reads all macros from all modules of the opened office file. \r
+ * @return All the macros and their contents\r
*/\r
public Map<String, String> readMacros() throws IOException {\r
- class Module {\r
- Integer offset;\r
- byte[] buf;\r
+ final ModuleMap modules = new ModuleMap();\r
+ findMacros(fs.getRoot(), modules);\r
+ \r
+ Map<String, String> moduleSources = new HashMap<String, String>();\r
+ for (Map.Entry<String, Module> entry : modules.entrySet()) {\r
+ Module module = entry.getValue();\r
+ if (module.buf != null && module.buf.length > 0) { // Skip empty modules\r
+ moduleSources.put(entry.getKey(), new String(module.buf, modules.charset));\r
+ }\r
}\r
- class ModuleMap extends HashMap<String, Module> {\r
-\r
- Charset charset = Charset.forName("Cp1252"); // default charset\r
+ return moduleSources;\r
+ }\r
+ \r
+ protected static class Module {\r
+ Integer offset;\r
+ byte[] buf;\r
+ }\r
+ protected static class ModuleMap extends HashMap<String, Module> {\r
+ Charset charset = Charset.forName("Cp1252"); // default charset\r
+ }\r
+ \r
+ protected void findMacros(DirectoryNode dir, ModuleMap modules) throws IOException {\r
+ if (VBA_PROJECT_POIFS.equals(dir.getName())) {\r
+ // VBA project directory, process\r
+ readMacros(dir, modules);\r
+ } else {\r
+ // Check children\r
+ for (Entry child : dir) {\r
+ if (child instanceof DirectoryNode) {\r
+ findMacros((DirectoryNode)child, modules);\r
+ }\r
+ }\r
}\r
- try {\r
- final ModuleMap modules = new ModuleMap();\r
- POIFSReader dirReader = new POIFSReader();\r
- dirReader.registerListener(new POIFSReaderListener() {\r
-\r
- public void processPOIFSReaderEvent(POIFSReaderEvent event) {\r
- try {\r
- String name = event.getName();\r
- if (event.getPath().toString().endsWith("\\VBA")) {\r
- if ("dir".equals(name)) {\r
- // process DIR\r
- RLEDecompressingInputStream in = new RLEDecompressingInputStream(event.getStream());\r
- String streamName = null;\r
- while (true) {\r
- int id = in.readShort();\r
- if (id == -1 || id == 0x0010) {\r
- break; // EOF or TERMINATOR\r
- }\r
- int len = in.readInt();\r
- switch (id) {\r
- case 0x0009: // PROJECTVERSION\r
- in.skip(6);\r
- break;\r
- case 0x0003: // PROJECTCODEPAGE\r
- int codepage = in.readShort();\r
- modules.charset = Charset.forName("Cp" + codepage);\r
- break;\r
- case 0x001A: // STREAMNAME\r
- byte[] streamNameBuf = new byte[len];\r
- int count = in.read(streamNameBuf);\r
- streamName = new String(streamNameBuf, 0, count, modules.charset);\r
- break;\r
- case 0x0031: // MODULEOFFSET\r
- int moduleOffset = in.readInt();\r
- Module module = modules.get(streamName);\r
- if (module != null) {\r
- ByteArrayOutputStream out = new ByteArrayOutputStream();\r
- RLEDecompressingInputStream stream = new RLEDecompressingInputStream(new ByteArrayInputStream(\r
- module.buf, moduleOffset, module.buf.length - moduleOffset));\r
- IOUtils.copy(stream, out);\r
- stream.close();\r
- out.close();\r
- module.buf = out.toByteArray();\r
- } else {\r
- module = new Module();\r
- module.offset = moduleOffset;\r
- modules.put(streamName, module);\r
- }\r
- break;\r
- default:\r
- in.skip(len);\r
- break;\r
- }\r
- }\r
- } else if (!name.startsWith("__SRP") && !name.startsWith("_VBA_PROJECT")) {\r
- // process module, skip __SRP and _VBA_PROJECT since these do not contain macros\r
- Module module = modules.get(name);\r
- final DocumentInputStream stream = event.getStream();\r
- final InputStream in;\r
- if (module == null) {\r
- // no DIR stream with offsets yet, so store the compressed bytes for later\r
- module = new Module();\r
- modules.put(name, module);\r
- in = stream;\r
- } else {\r
- // we know the offset already, so decompress immediately on-the-fly\r
- stream.skip(module.offset);\r
- in = new RLEDecompressingInputStream(stream);\r
- }\r
- final ByteArrayOutputStream out = new ByteArrayOutputStream();\r
- IOUtils.copy(in, out);\r
- in.close();\r
- out.close();\r
- module.buf = out.toByteArray();\r
- }\r
+ }\r
+ protected void readMacros(DirectoryNode macroDir, ModuleMap modules) throws IOException {\r
+ for (Entry entry : macroDir) {\r
+ if (! (entry instanceof DocumentNode)) { continue; }\r
+ \r
+ String name = entry.getName();\r
+ DocumentNode document = (DocumentNode)entry;\r
+ DocumentInputStream dis = new DocumentInputStream(document);\r
+ if ("dir".equals(name)) {\r
+ // process DIR\r
+ RLEDecompressingInputStream in = new RLEDecompressingInputStream(dis);\r
+ String streamName = null;\r
+ while (true) {\r
+ int id = in.readShort();\r
+ if (id == -1 || id == 0x0010) {\r
+ break; // EOF or TERMINATOR\r
+ }\r
+ int len = in.readInt();\r
+ switch (id) {\r
+ case 0x0009: // PROJECTVERSION\r
+ in.skip(6);\r
+ break;\r
+ case 0x0003: // PROJECTCODEPAGE\r
+ int codepage = in.readShort();\r
+ modules.charset = Charset.forName("Cp" + codepage);\r
+ break;\r
+ case 0x001A: // STREAMNAME\r
+ byte[] streamNameBuf = new byte[len];\r
+ int count = in.read(streamNameBuf);\r
+ streamName = new String(streamNameBuf, 0, count, modules.charset);\r
+ break;\r
+ case 0x0031: // MODULEOFFSET\r
+ int moduleOffset = in.readInt();\r
+ Module module = modules.get(streamName);\r
+ if (module != null) {\r
+ ByteArrayOutputStream out = new ByteArrayOutputStream();\r
+ RLEDecompressingInputStream stream = new RLEDecompressingInputStream(new ByteArrayInputStream(\r
+ module.buf, moduleOffset, module.buf.length - moduleOffset));\r
+ IOUtils.copy(stream, out);\r
+ stream.close();\r
+ out.close();\r
+ module.buf = out.toByteArray();\r
+ } else {\r
+ module = new Module();\r
+ module.offset = moduleOffset;\r
+ modules.put(streamName, module);\r
}\r
- } catch (IOException e) {\r
- throw new RuntimeException(e);\r
+ break;\r
+ default:\r
+ in.skip(len);\r
+ break;\r
}\r
}\r
- });\r
- dirReader.read(null); // TODO\r
- Map<String, String> moduleSources = new HashMap<String, String>();\r
- for (Map.Entry<String, Module> entry : modules.entrySet()) {\r
- Module module = entry.getValue();\r
- if (module.buf != null && module.buf.length > 0) { // Skip empty modules\r
- moduleSources.put(entry.getKey(), new String(module.buf, modules.charset));\r
+ in.close();\r
+ } else if (!name.startsWith("__SRP") && !name.startsWith("_VBA_PROJECT")) {\r
+ // process module, skip __SRP and _VBA_PROJECT since these do not contain macros\r
+ Module module = modules.get(name);\r
+ final InputStream in;\r
+ // TODO Refactor this to fetch dir then do the rest\r
+ if (module == null) {\r
+ // no DIR stream with offsets yet, so store the compressed bytes for later\r
+ module = new Module();\r
+ modules.put(name, module);\r
+ in = dis;\r
+ } else {\r
+ // we know the offset already, so decompress immediately on-the-fly\r
+ dis.skip(module.offset);\r
+ in = new RLEDecompressingInputStream(dis);\r
}\r
+ final ByteArrayOutputStream out = new ByteArrayOutputStream();\r
+ IOUtils.copy(in, out);\r
+ in.close();\r
+ out.close();\r
+ module.buf = out.toByteArray();\r
}\r
- return moduleSources;\r
- } catch (IOException e) {\r
- e.printStackTrace();\r
- throw e;\r
}\r
}\r
}\r
--- /dev/null
+/* ====================================================================
+ Licensed to the Apache Software Foundation (ASF) under one or more
+ contributor license agreements. See the NOTICE file distributed with
+ this work for additional information regarding copyright ownership.
+ The ASF licenses this file to You under the Apache License, Version 2.0
+ (the "License"); you may not use this file except in compliance with
+ the License. You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
+==================================================================== */
+
+package org.apache.poi.poifs.macros;
+
+import static org.apache.poi.POITestCase.assertContains;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+import java.io.File;
+import java.io.FileInputStream;
+import java.util.Map;
+
+import org.apache.poi.hssf.HSSFTestDataSamples;
+import org.apache.poi.poifs.filesystem.NPOIFSFileSystem;
+import org.apache.poi.util.IOUtils;
+import org.apache.poi.util.StringUtil;
+import org.junit.Test;
+
+public class TestVBAMacroReader {
+ private final String testMacroContents;
+ private final String testMacroNoSub;
+ public TestVBAMacroReader() throws Exception {
+ File macro = HSSFTestDataSamples.getSampleFile("SimpleMacro.vba");
+ testMacroContents = new String(
+ IOUtils.toByteArray(new FileInputStream(macro)),
+ StringUtil.UTF8
+ );
+
+ if (! testMacroContents.startsWith("Sub ")) {
+ throw new IllegalArgumentException("Not a macro");
+ }
+ testMacroNoSub = testMacroContents.substring(testMacroContents.indexOf("()")+3);
+ }
+
+ @Test
+ public void fromStream() throws Exception {
+ VBAMacroReader r;
+
+ r = new VBAMacroReader(HSSFTestDataSamples.openSampleFileStream("SimpleMacro.xls"));
+ assertMacroContents(r);
+ r.close();
+
+ r = new VBAMacroReader(HSSFTestDataSamples.openSampleFileStream("SimpleMacro.xlsm"));
+ assertMacroContents(r);
+ r.close();
+ }
+ @Test
+ public void fromFile() throws Exception {
+ VBAMacroReader r;
+
+ r = new VBAMacroReader(HSSFTestDataSamples.getSampleFile("SimpleMacro.xls"));
+ assertMacroContents(r);
+ r.close();
+
+ r = new VBAMacroReader(HSSFTestDataSamples.getSampleFile("SimpleMacro.xlsm"));
+ assertMacroContents(r);
+ r.close();
+ }
+ @Test
+ public void fromNPOIFS() throws Exception {
+ NPOIFSFileSystem fs = new NPOIFSFileSystem(
+ HSSFTestDataSamples.getSampleFile("SimpleMacro.xls"));
+ VBAMacroReader r = new VBAMacroReader(fs);
+ assertMacroContents(r);
+ r.close();
+ }
+
+ protected void assertMacroContents(VBAMacroReader r) throws Exception {
+ Map<String,String> contents = r.readMacros();
+
+ assertFalse(contents.isEmpty());
+ assertEquals(5, contents.size());
+
+ // Check the ones without scripts
+ String[] noScripts = new String[] { "ThisWorkbook",
+ "Sheet1", "Sheet2", "Sheet3" };
+ for (String entry : noScripts) {
+ assertTrue(entry, contents.containsKey(entry));
+
+ String content = contents.get(entry);
+ assertContains(content, "Attribute VB_Exposed = True");
+ assertContains(content, "Attribute VB_Customizable = True");
+ assertContains(content, "Attribute VB_TemplateDerived = False");
+ assertContains(content, "Attribute VB_GlobalNameSpace = False");
+ assertContains(content, "Attribute VB_Exposed = True");
+ }
+
+ // Check the script one
+ String content = contents.get("Module1");
+ assertContains(content, "Attribute VB_Name = \"Module1\"");
+ assertContains(content, "Attribute TestMacro.VB_Description = \"This is a test macro\"");
+
+ // And the macro itself
+ assertContains(content, testMacroNoSub);
+ }
+}
projects / poi.git / commitdiff