From 3b2c91594e2d2a2e41113d9a079d0f5d9ea74b30 Mon Sep 17 00:00:00 2001
From: huaiyu <shunyu.lsy@alibaba-inc.com>
Date: Thu, 10 Oct 2019 15:11:36 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E6=94=B9=E5=8E=9F=E6=9D=A5=E7=9A=84?=
 =?UTF-8?q?=E5=86=99=E6=B3=95?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../java/com/alibaba/excel/ExcelWriter.java   | 20 ++++++++++--------
 .../com/alibaba/excel/write/ExcelBuilder.java |  9 ++++++++
 .../alibaba/excel/write/ExcelBuilderImpl.java | 21 ++++++++++++++-----
 .../builder/ExcelWriterSheetBuilder.java      |  3 +--
 .../easyexcel/test/demo/write/WriteTest.java  |  5 ++---
 5 files changed, 39 insertions(+), 19 deletions(-)

diff --git a/src/main/java/com/alibaba/excel/ExcelWriter.java b/src/main/java/com/alibaba/excel/ExcelWriter.java
index 3a2dde55..4800f840 100644
--- a/src/main/java/com/alibaba/excel/ExcelWriter.java
+++ b/src/main/java/com/alibaba/excel/ExcelWriter.java
@@ -138,7 +138,17 @@ public class ExcelWriter {
      * @return this current writer
      */
     public ExcelWriter write(List data, WriteSheet writeSheet) {
-        return write(data, writeSheet, null);
+        return write(data, writeSheet, (WriteTable)null);
+    }
+
+
+    public ExcelWriter write(List data, WriteSheet writeSheet, String password) {
+        return write(data, writeSheet, null, password);
+    }
+
+    private ExcelWriter write(List data, WriteSheet writeSheet, WriteTable writeTable, String password) {
+        excelBuilder.addContent(data, writeSheet, writeTable, password);
+        return this;
     }
 
     /**
@@ -331,12 +341,4 @@ public class ExcelWriter {
     public WriteContext writeContext() {
         return excelBuilder.writeContext();
     }
-
-    /**
-     * doEncrypt
-     * @param password
-     */
-    public void doEncrypt(String password){
-        excelBuilder.setPassword(password);
-    }
 }
diff --git a/src/main/java/com/alibaba/excel/write/ExcelBuilder.java b/src/main/java/com/alibaba/excel/write/ExcelBuilder.java
index 8a60444e..8a7d57da 100644
--- a/src/main/java/com/alibaba/excel/write/ExcelBuilder.java
+++ b/src/main/java/com/alibaba/excel/write/ExcelBuilder.java
@@ -73,4 +73,13 @@ public interface ExcelBuilder {
      * Close io
      */
     void finish();
+
+    /**
+     * add password
+     * @param data
+     * @param writeSheet
+     * @param writeTable
+     * @param password
+     */
+    void addContent(List data, WriteSheet writeSheet, WriteTable writeTable, String password);
 }
diff --git a/src/main/java/com/alibaba/excel/write/ExcelBuilderImpl.java b/src/main/java/com/alibaba/excel/write/ExcelBuilderImpl.java
index 27f65905..19284475 100644
--- a/src/main/java/com/alibaba/excel/write/ExcelBuilderImpl.java
+++ b/src/main/java/com/alibaba/excel/write/ExcelBuilderImpl.java
@@ -88,11 +88,6 @@ public class ExcelBuilderImpl implements ExcelBuilder {
         }
     }
 
-    @Override
-    public void setPassword(String password) {
-        context.setPassword(password);
-    }
-
     @Override
     public void finish() {
         if (context != null) {
@@ -100,6 +95,22 @@ public class ExcelBuilderImpl implements ExcelBuilder {
         }
     }
 
+    @Override
+    public void addContent(List data, WriteSheet writeSheet, WriteTable writeTable, String password) {
+        try {
+            context.currentSheet(writeSheet);
+            context.currentTable(writeTable);
+            context.setPassword(password);
+            doAddContent(data);
+        } catch (RuntimeException e) {
+            finish();
+            throw e;
+        } catch (Throwable e) {
+            finish();
+            throw new ExcelGenerateException(e);
+        }
+    }
+
     @Override
     public void merge(int firstRow, int lastRow, int firstCol, int lastCol) {
         CellRangeAddress cra = new CellRangeAddress(firstRow, lastRow, firstCol, lastCol);
diff --git a/src/main/java/com/alibaba/excel/write/builder/ExcelWriterSheetBuilder.java b/src/main/java/com/alibaba/excel/write/builder/ExcelWriterSheetBuilder.java
index 893d2180..194fedb0 100644
--- a/src/main/java/com/alibaba/excel/write/builder/ExcelWriterSheetBuilder.java
+++ b/src/main/java/com/alibaba/excel/write/builder/ExcelWriterSheetBuilder.java
@@ -157,8 +157,7 @@ public class ExcelWriterSheetBuilder {
         if (excelWriter == null) {
             throw new ExcelGenerateException("Must use 'EasyExcelFactory.write().sheet()' to call this method");
         }
-        excelWriter.write(data, build());
-        excelWriter.doEncrypt(password);
+        excelWriter.write(data, build(), password);
         excelWriter.finish();
     }
 
diff --git a/src/test/java/com/alibaba/easyexcel/test/demo/write/WriteTest.java b/src/test/java/com/alibaba/easyexcel/test/demo/write/WriteTest.java
index a86f4e52..b2a1722c 100644
--- a/src/test/java/com/alibaba/easyexcel/test/demo/write/WriteTest.java
+++ b/src/test/java/com/alibaba/easyexcel/test/demo/write/WriteTest.java
@@ -57,9 +57,8 @@ public class WriteTest {
         // 这里 需要指定写用哪个class去读
         ExcelWriter excelWriter = EasyExcel.write(fileName, DemoData.class).build();
         WriteSheet writeSheet = EasyExcel.writerSheet("模板").build();
-        excelWriter.write(data(), writeSheet);
-        // doEncrypt 进行加密
-        //excelWriter.doEncrypt("123456");
+        excelWriter.write(data(), writeSheet,"12345");
+        // 进行加密
         /// 千万别忘记finish 会帮忙关闭流
         excelWriter.finish();
     }