#57593 Begin adding overloaded WorkbookFactory.create methods which take the spreadsheet password

git-svn-id: https://svn.apache.org/repos/asf/poi/trunk@1676843 13f79535-47bb-0310-9956-ffa450edef68

2 files changed, 149 insertions, 2 deletions

diff --git a/src/ooxml/java/org/apache/poi/ss/usermodel/WorkbookFactory.java b/src/ooxml/java/org/apache/poi/ss/usermodel/WorkbookFactory.java
index 747d88c263..b4f341533d 100644
--- a/src/ooxml/java/org/apache/poi/ss/usermodel/WorkbookFactory.java
+++ b/src/ooxml/java/org/apache/poi/ss/usermodel/WorkbookFactory.java
@@ -21,12 +21,17 @@ import java.io.FileNotFoundException;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.PushbackInputStream;
+import java.security.GeneralSecurityException;
 
 import org.apache.poi.EncryptedDocumentException;
 import org.apache.poi.POIXMLDocument;
+import org.apache.poi.hssf.record.crypto.Biff8EncryptionKey;
 import org.apache.poi.hssf.usermodel.HSSFWorkbook;
 import org.apache.poi.openxml4j.exceptions.InvalidFormatException;
 import org.apache.poi.openxml4j.opc.OPCPackage;
+import org.apache.poi.poifs.crypt.Decryptor;
+import org.apache.poi.poifs.crypt.EncryptionInfo;
+import org.apache.poi.poifs.filesystem.DirectoryNode;
 import org.apache.poi.poifs.filesystem.NPOIFSFileSystem;
 import org.apache.poi.poifs.filesystem.OfficeXmlFileException;
 import org.apache.poi.poifs.filesystem.POIFSFileSystem;
@@ -53,7 +58,50 @@ public class WorkbookFactory {
      *  Workbook should be closed after use.
      */
     public static Workbook create(NPOIFSFileSystem fs) throws IOException {
-        return new HSSFWorkbook(fs.getRoot(), true);
+        try {
+            return create(fs, null);
+        } catch (InvalidFormatException e) {
+            // Special case of OOXML-in-POIFS which is broken
+            throw new IOException(e);
+        }
+    }
+    
+    /**
+     * Creates a Workbook from the given NPOIFSFileSystem, which may
+     *  be password protected
+     */
+    private static Workbook create(NPOIFSFileSystem fs, String password) throws IOException, InvalidFormatException {
+        DirectoryNode root = fs.getRoot();
+        if (root.hasEntry(Decryptor.DEFAULT_POIFS_ENTRY)) {
+            if (password == null) {
+                throw new EncryptedDocumentException("The supplied spreadsheet is protected, but no password was supplied");
+            } else {
+                EncryptionInfo info = new EncryptionInfo(fs);
+                Decryptor d = Decryptor.getInstance(info);
+                
+                boolean passwordCorrect = false;
+                InputStream stream = null;
+                try {
+                    if (d.verifyPassword(password)) {
+                        passwordCorrect = true;
+                        stream = d.getDataStream(root);
+                    }
+                } catch (GeneralSecurityException e) {}
+                
+                if (! passwordCorrect) 
+                    throw new EncryptedDocumentException("Password incorrect");
+                
+                OPCPackage pkg = OPCPackage.open(stream);
+                return create(pkg);
+            }
+        }
+        
+        if (password != null) {
+            Biff8EncryptionKey.setCurrentUserPassword(password);
+        }
+        Workbook wb = new HSSFWorkbook(root, true);
+        Biff8EncryptionKey.setCurrentUserPassword(null);
+        return wb;
     }
     
     /**
@@ -77,13 +125,29 @@ public class WorkbookFactory {
      * @throws EncryptedDocumentException If the workbook given is password protected
      */
     public static Workbook create(InputStream inp) throws IOException, InvalidFormatException, EncryptedDocumentException {
+        return create(inp, null);
+    }
+    
+    /**
+     * Creates the appropriate HSSFWorkbook / XSSFWorkbook from
+     *  the given InputStream, which may be password protected.
+     * <p>Your input stream MUST either support mark/reset, or
+     *  be wrapped as a {@link PushbackInputStream}! Note that 
+     *  using an {@link InputStream} has a higher memory footprint 
+     *  than using a {@link File}.</p> 
+     * <p>Note that in order to properly release resources the 
+     *  Workbook should be closed after use.
+     * @throws EncryptedDocumentException If the wrong password is given for a protected file
+     */
+    public static Workbook create(InputStream inp, String password) throws IOException, InvalidFormatException, EncryptedDocumentException {
         // If clearly doesn't do mark/reset, wrap up
         if (! inp.markSupported()) {
             inp = new PushbackInputStream(inp, 8);
         }
 
         if (POIFSFileSystem.hasPOIFSHeader(inp)) {
-            return new HSSFWorkbook(inp);
+            NPOIFSFileSystem fs = new NPOIFSFileSystem(inp);
+            return create(fs, password);
         }
         if (POIXMLDocument.hasOOXMLHeader(inp)) {
             return new XSSFWorkbook(OPCPackage.open(inp));
@@ -91,6 +155,7 @@ public class WorkbookFactory {
         throw new IllegalArgumentException("Your InputStream was neither an OLE2 stream, nor an OOXML stream");
     }
     
+    // TODO file+password
     /**
      * Creates the appropriate HSSFWorkbook / XSSFWorkbook from
      *  the given File, which must exist and be readable.
diff --git a/src/ooxml/testcases/org/apache/poi/ss/TestWorkbookFactory.java b/src/ooxml/testcases/org/apache/poi/ss/TestWorkbookFactory.java
index 28983ff6fa..1c7e4b7a87 100644
--- a/src/ooxml/testcases/org/apache/poi/ss/TestWorkbookFactory.java
+++ b/src/ooxml/testcases/org/apache/poi/ss/TestWorkbookFactory.java
@@ -19,6 +19,7 @@ package org.apache.poi.ss;
 
 import java.io.InputStream;
 
+import org.apache.poi.EncryptedDocumentException;
 import org.apache.poi.hssf.usermodel.HSSFWorkbook;
 import org.apache.poi.hssf.HSSFTestDataSamples;
 import org.apache.poi.poifs.filesystem.POIFSFileSystem;
@@ -32,11 +33,15 @@ import junit.framework.TestCase;
 public final class TestWorkbookFactory extends TestCase {
 	private String xls;
 	private String xlsx;
+    private String[] xls_prot;
+    private String[] xlsx_prot;
 	private String txt;
 
 	protected void setUp() {
 		xls = "SampleSS.xls";
 		xlsx = "SampleSS.xlsx";
+		xls_prot = new String[] {"password.xls", "password"};
+		xlsx_prot = new String[]{"protected_passtika.xlsx", "tika"};
 		txt = "SampleSS.txt";
 	}
 
@@ -114,4 +119,81 @@ public final class TestWorkbookFactory extends TestCase {
 			// Good
 		}
 	}
+	
+	/**
+	 * Check that the overloaded stream methods which take passwords work properly
+	 */
+	public void testCreateWithPasswordFromStream() throws Exception {
+        Workbook wb;
+
+        
+	    // Unprotected, no password given, opens normally
+        wb = WorkbookFactory.create(
+                HSSFTestDataSamples.openSampleFileStream(xls), null
+        );
+        assertNotNull(wb);
+        assertTrue(wb instanceof HSSFWorkbook);
+        wb.close();
+	    
+        wb = WorkbookFactory.create(
+                HSSFTestDataSamples.openSampleFileStream(xlsx), null
+        );
+        assertNotNull(wb);
+        assertTrue(wb instanceof XSSFWorkbook);
+        
+        
+        // Unprotected, wrong password, opens normally
+        wb = WorkbookFactory.create(
+                HSSFTestDataSamples.openSampleFileStream(xls), "wrong"
+        );
+        assertNotNull(wb);
+        assertTrue(wb instanceof HSSFWorkbook);
+        wb.close();
+        
+        wb = WorkbookFactory.create(
+                HSSFTestDataSamples.openSampleFileStream(xlsx), "wrong"
+        );
+        assertNotNull(wb);
+        assertTrue(wb instanceof XSSFWorkbook);
+        
+        
+        // Protected, correct password, opens fine
+        wb = WorkbookFactory.create(
+                HSSFTestDataSamples.openSampleFileStream(xls_prot[0]), xls_prot[1]
+        );
+        assertNotNull(wb);
+        assertTrue(wb instanceof HSSFWorkbook);
+        wb.close();
+        
+        wb = WorkbookFactory.create(
+                HSSFTestDataSamples.openSampleFileStream(xlsx_prot[0]), xlsx_prot[1]
+        );
+        assertNotNull(wb);
+        assertTrue(wb instanceof XSSFWorkbook);
+        
+        
+        // Protected, wrong password, throws Exception
+        try {
+            wb = WorkbookFactory.create(
+                    HSSFTestDataSamples.openSampleFileStream(xls_prot[0]), "wrong"
+            );
+            fail("Shouldn't be able to open with the wrong password");
+        } catch (EncryptedDocumentException e) {}
+
+        try {
+            wb = WorkbookFactory.create(
+                    HSSFTestDataSamples.openSampleFileStream(xlsx_prot[0]), "wrong"
+            );
+            fail("Shouldn't be able to open with the wrong password");
+        } catch (EncryptedDocumentException e) {}
+	}
+	
+    /**
+     * Check that the overloaded file methods which take passwords work properly
+     */
+    public void testCreateWithPasswordFromFile() throws Exception {
+        Workbook wb;
+        
+        // TODO
+    }
 }