From 8fde918d981e9229b43fc6045b259a1db31cf1f4 Mon Sep 17 00:00:00 2001
From: xuri <xuri.me@gmail.com>
Date: Tue, 31 May 2022 11:05:10 +0800
Subject: This update docs and tests for workbook encryption

---
 .gitignore     |  2 +-
 crypt.go       | 27 +++++++++++----------------
 crypt_test.go  | 13 +++++++++++++
 excelize.go    |  1 +
 stream_test.go |  2 ++
 5 files changed, 28 insertions(+), 17 deletions(-)

diff --git a/.gitignore b/.gitignore
index 4dce768..e697544 100644
--- a/.gitignore
+++ b/.gitignore
@@ -5,7 +5,7 @@ test/Test*.xlsx
 test/Test*.xltm
 test/Test*.xltx
 # generated files
-test/BadEncrypt.xlsx
+test/Encryption*.xlsx
 test/BadWorkbook.SaveAsEmptyStruct.xlsx
 test/*.png
 test/excelize-*
diff --git a/crypt.go b/crypt.go
index 65d5dae..239208d 100644
--- a/crypt.go
+++ b/crypt.go
@@ -143,15 +143,10 @@ func Decrypt(raw []byte, opt *Options) (packageBuf []byte, err error) {
 	if err != nil || mechanism == "extensible" {
 		return
 	}
-	switch mechanism {
-	case "agile":
+	if mechanism == "agile" {
 		return agileDecrypt(encryptionInfoBuf, encryptedPackageBuf, opt)
-	case "standard":
-		return standardDecrypt(encryptionInfoBuf, encryptedPackageBuf, opt)
-	default:
-		err = ErrUnsupportedEncryptMechanism
 	}
-	return
+	return standardDecrypt(encryptionInfoBuf, encryptedPackageBuf, opt)
 }
 
 // Encrypt API encrypt data with the password.
@@ -1112,7 +1107,7 @@ func (c *cfb) writeDirectoryEntry(propertyCount, customSectID, size int) []byte
 	return storage.stream
 }
 
-// writeMSAT provides a function to write compound file sector allocation
+// writeMSAT provides a function to write compound file master sector allocation
 // table.
 func (c *cfb) writeMSAT(MSATBlocks, SATBlocks int, MSAT []int) []int {
 	if MSATBlocks > 0 {
@@ -1129,19 +1124,19 @@ func (c *cfb) writeMSAT(MSATBlocks, SATBlocks int, MSAT []int) []int {
 			}
 			MSAT = append(MSAT, -1)
 		}
-	} else {
-		for i := 0; i < 109; i++ {
-			if i < SATBlocks {
-				MSAT = append(MSAT, i)
-				continue
-			}
-			MSAT = append(MSAT, -1)
+		return MSAT
+	}
+	for i := 0; i < 109; i++ {
+		if i < SATBlocks {
+			MSAT = append(MSAT, i)
+			continue
 		}
+		MSAT = append(MSAT, -1)
 	}
 	return MSAT
 }
 
-// writeSAT provides a function to write compound file master sector allocation
+// writeSAT provides a function to write compound file sector allocation
 // table.
 func (c *cfb) writeSAT(MSATBlocks, SATBlocks, SSATBlocks, directoryBlocks, fileBlocks, streamBlocks int, SAT []int) (int, []int) {
 	var blocks int
diff --git a/crypt_test.go b/crypt_test.go
index 2df5af2..d2fba35 100644
--- a/crypt_test.go
+++ b/crypt_test.go
@@ -12,6 +12,7 @@
 package excelize
 
 import (
+	"io/ioutil"
 	"path/filepath"
 	"strings"
 	"testing"
@@ -30,6 +31,13 @@ func TestEncrypt(t *testing.T) {
 	assert.NoError(t, err)
 	assert.Equal(t, "SECRET", cell)
 	assert.NoError(t, f.Close())
+	// Test decrypt spreadsheet with unsupported encrypt mechanism
+	raw, err := ioutil.ReadFile(filepath.Join("test", "encryptAES.xlsx"))
+	assert.NoError(t, err)
+	raw[2050] = 3
+	_, err = Decrypt(raw, &Options{Password: "password"})
+	assert.EqualError(t, err, ErrUnsupportedEncryptMechanism.Error())
+
 	// Test encrypt spreadsheet with invalid password
 	assert.EqualError(t, f.SaveAs(filepath.Join("test", "Encryption.xlsx"), Options{Password: strings.Repeat("*", MaxFieldLength+1)}), ErrPasswordLengthInvalid.Error())
 	// Test encrypt spreadsheet with new password
@@ -51,6 +59,11 @@ func TestEncryptionMechanism(t *testing.T) {
 	assert.EqualError(t, err, ErrUnknownEncryptMechanism.Error())
 }
 
+func TestEncryptionWriteDirectoryEntry(t *testing.T) {
+	cfb := cfb{}
+	assert.Equal(t, 1536, len(cfb.writeDirectoryEntry(0, 0, -1)))
+}
+
 func TestHashing(t *testing.T) {
 	assert.Equal(t, hashing("unsupportedHashAlgorithm", []byte{}), []byte(nil))
 }
diff --git a/excelize.go b/excelize.go
index 8c71b16..aaa4953 100644
--- a/excelize.go
+++ b/excelize.go
@@ -93,6 +93,7 @@ type Options struct {
 //        return
 //    }
 //
+// Close the file by Close function after opening the spreadsheet.
 func OpenFile(filename string, opt ...Options) (*File, error) {
 	file, err := os.Open(filepath.Clean(filename))
 	if err != nil {
diff --git a/stream_test.go b/stream_test.go
index 3df898a..9776b38 100644
--- a/stream_test.go
+++ b/stream_test.go
@@ -129,6 +129,8 @@ func TestStreamWriter(t *testing.T) {
 	}
 	assert.NoError(t, rows.Close())
 	assert.Equal(t, 2559558, cells)
+	// Save spreadsheet with password.
+	assert.NoError(t, file.SaveAs(filepath.Join("test", "EncryptionTestStreamWriter.xlsx"), Options{Password: "password"}))
 	assert.NoError(t, file.Close())
 }
 
-- 
cgit v1.2.1