Add support for workbook protection (#1431)

crypt.go

@@ -37,16 +37,17 @@ import (
 )
 
 var (
-	blockKey                   = []byte{0x14, 0x6e, 0x0b, 0xe7, 0xab, 0xac, 0xd0, 0xd6} // Block keys used for encryption
-	oleIdentifier              = []byte{0xd0, 0xcf, 0x11, 0xe0, 0xa1, 0xb1, 0x1a, 0xe1}
-	headerCLSID                = make([]byte, 16)
-	difSect                    = -4
-	endOfChain                 = -2
-	fatSect                    = -3
-	iterCount                  = 50000
-	packageEncryptionChunkSize = 4096
-	packageOffset              = 8 // First 8 bytes are the size of the stream
-	sheetProtectionSpinCount   = 1e5
+	blockKey                    = []byte{0x14, 0x6e, 0x0b, 0xe7, 0xab, 0xac, 0xd0, 0xd6} // Block keys used for encryption
+	oleIdentifier               = []byte{0xd0, 0xcf, 0x11, 0xe0, 0xa1, 0xb1, 0x1a, 0xe1}
+	headerCLSID                 = make([]byte, 16)
+	difSect                     = -4
+	endOfChain                  = -2
+	fatSect                     = -3
+	iterCount                   = 50000
+	packageEncryptionChunkSize  = 4096
+	packageOffset               = 8 // First 8 bytes are the size of the stream
+	sheetProtectionSpinCount    = 1e5
+	workbookProtectionSpinCount = 1e5
 )
 
 // Encryption specifies the encryption structure, streams, and storages are

errors.go

@@ -230,4 +230,10 @@ var (
 	// ErrSheetNameLength defined the error message on receiving the sheet
 	// name length exceeds the limit.
 	ErrSheetNameLength = fmt.Errorf("the sheet name length exceeds the %d characters limit", MaxSheetNameLength)
+	// ErrUnprotectWorkbook defined the error message on workbook has set no
+	// protection.
+	ErrUnprotectWorkbook = errors.New("workbook has set no protect")
+	// ErrUnprotectWorkbookPassword defined the error message on remove workbook
+	// protection with password verification failed.
+	ErrUnprotectWorkbookPassword = errors.New("workbook protect password not match")
 )

excelize_test.go

@@ -1329,6 +1329,61 @@ func TestUnprotectSheet(t *testing.T) {
 	assert.EqualError(t, f.UnprotectSheet(sheetName, "wrongPassword"), "illegal base64 data at input byte 8")
 }
 
+func TestProtectWorkbook(t *testing.T) {
+	f := NewFile()
+	assert.NoError(t, f.ProtectWorkbook(nil))
+	// Test protect workbook with default hash algorithm
+	assert.NoError(t, f.ProtectWorkbook(&WorkbookProtectionOptions{
+		Password:      "password",
+		LockStructure: true,
+	}))
+	wb, err := f.workbookReader()
+	assert.NoError(t, err)
+	assert.Equal(t, "SHA-512", wb.WorkbookProtection.WorkbookAlgorithmName)
+	assert.Equal(t, 24, len(wb.WorkbookProtection.WorkbookSaltValue))
+	assert.Equal(t, 88, len(wb.WorkbookProtection.WorkbookHashValue))
+	assert.Equal(t, int(workbookProtectionSpinCount), wb.WorkbookProtection.WorkbookSpinCount)
+	assert.NoError(t, f.SaveAs(filepath.Join("test", "TestProtectWorkbook.xlsx")))
+	// Test protect workbook with password exceeds the limit length
+	assert.EqualError(t, f.ProtectWorkbook(&WorkbookProtectionOptions{
+		AlgorithmName: "MD4",
+		Password:      strings.Repeat("s", MaxFieldLength+1),
+	}), ErrPasswordLengthInvalid.Error())
+	// Test protect workbook with unsupported hash algorithm
+	assert.EqualError(t, f.ProtectWorkbook(&WorkbookProtectionOptions{
+		AlgorithmName: "RIPEMD-160",
+		Password:      "password",
+	}), ErrUnsupportedHashAlgorithm.Error())
+}
+
+func TestUnprotectWorkbook(t *testing.T) {
+	f, err := OpenFile(filepath.Join("test", "Book1.xlsx"))
+	if !assert.NoError(t, err) {
+		t.FailNow()
+	}
+
+	assert.NoError(t, f.UnprotectWorkbook())
+	assert.EqualError(t, f.UnprotectWorkbook("password"), ErrUnprotectWorkbook.Error())
+	assert.NoError(t, f.SaveAs(filepath.Join("test", "TestUnprotectWorkbook.xlsx")))
+	assert.NoError(t, f.Close())
+
+	f = NewFile()
+	assert.NoError(t, f.ProtectWorkbook(&WorkbookProtectionOptions{Password: "password"}))
+	// Test remove workbook protection with an incorrect password
+	assert.EqualError(t, f.UnprotectWorkbook("wrongPassword"), ErrUnprotectWorkbookPassword.Error())
+	// Test remove workbook protection with password verification
+	assert.NoError(t, f.UnprotectWorkbook("password"))
+	// Test with invalid salt value
+	assert.NoError(t, f.ProtectWorkbook(&WorkbookProtectionOptions{
+		AlgorithmName: "SHA-512",
+		Password:      "password",
+	}))
+	wb, err := f.workbookReader()
+	assert.NoError(t, err)
+	wb.WorkbookProtection.WorkbookSaltValue = "YWJjZA====="
+	assert.EqualError(t, f.UnprotectWorkbook("wrongPassword"), "illegal base64 data at input byte 8")
+}
+
 func TestSetDefaultTimeStyle(t *testing.T) {
 	f := NewFile()
 	// Test set default time style on not exists worksheet.

workbook.go

@@ -59,6 +59,67 @@ func (f *File) GetWorkbookProps() (WorkbookPropsOptions, error) {
 	return opts, err
 }
 
+// ProtectWorkbook provides a function to prevent other users from accidentally or
+// deliberately changing, moving, or deleting data in a workbook.
+func (f *File) ProtectWorkbook(opts *WorkbookProtectionOptions) error {
+	wb, err := f.workbookReader()
+	if err != nil {
+		return err
+	}
+	if wb.WorkbookProtection == nil {
+		wb.WorkbookProtection = new(xlsxWorkbookProtection)
+	}
+	if opts == nil {
+		opts = &WorkbookProtectionOptions{}
+	}
+	wb.WorkbookProtection = &xlsxWorkbookProtection{
+		LockStructure: opts.LockStructure,
+		LockWindows:   opts.LockWindows,
+	}
+	if opts.Password != "" {
+		if opts.AlgorithmName == "" {
+			opts.AlgorithmName = "SHA-512"
+		}
+		hashValue, saltValue, err := genISOPasswdHash(opts.Password, opts.AlgorithmName, "", int(workbookProtectionSpinCount))
+		if err != nil {
+			return err
+		}
+		wb.WorkbookProtection.WorkbookAlgorithmName = opts.AlgorithmName
+		wb.WorkbookProtection.WorkbookSaltValue = saltValue
+		wb.WorkbookProtection.WorkbookHashValue = hashValue
+		wb.WorkbookProtection.WorkbookSpinCount = int(workbookProtectionSpinCount)
+	}
+	return nil
+}
+
+// UnprotectWorkbook provides a function to remove protection for workbook,
+// specified the second optional password parameter to remove workbook
+// protection with password verification.
+func (f *File) UnprotectWorkbook(password ...string) error {
+	wb, err := f.workbookReader()
+	if err != nil {
+		return err
+	}
+	// password verification
+	if len(password) > 0 {
+		if wb.WorkbookProtection == nil {
+			return ErrUnprotectWorkbook
+		}
+		if wb.WorkbookProtection.WorkbookAlgorithmName != "" {
+			// check with given salt value
+			hashValue, _, err := genISOPasswdHash(password[0], wb.WorkbookProtection.WorkbookAlgorithmName, wb.WorkbookProtection.WorkbookSaltValue, wb.WorkbookProtection.WorkbookSpinCount)
+			if err != nil {
+				return err
+			}
+			if wb.WorkbookProtection.WorkbookHashValue != hashValue {
+				return ErrUnprotectWorkbookPassword
+			}
+		}
+	}
+	wb.WorkbookProtection = nil
+	return err
+}
+
 // setWorkbook update workbook property of the spreadsheet. Maximum 31
 // characters are allowed in sheet title.
 func (f *File) setWorkbook(name string, sheetID, rid int) {

xmlWorkbook.go

@@ -320,3 +320,11 @@ type WorkbookPropsOptions struct {
 	FilterPrivacy *bool   `json:"filter_privacy,omitempty"`
 	CodeName      *string `json:"code_name,omitempty"`
 }
+
+// WorkbookProtectionOptions directly maps the settings of workbook protection.
+type WorkbookProtectionOptions struct {
+	AlgorithmName string `json:"algorithmName,omitempty"`
+	Password      string `json:"password,omitempty"`
+	LockStructure bool   `json:"lockStructure,omitempty"`
+	LockWindows   bool   `json:"lockWindows,omitempty"`
+}