// Copyright 2016 - 2020 The excelize Authors. All rights reserved. Use of
// this source code is governed by a BSD-style license that can be found in
// the LICENSE file.
//
// Package excelize providing a set of functions that allow you to write to
// and read from XLSX files. Support reads and writes XLSX file generated by
// Microsoft Excel™ 2007 and later. Support save file without losing original
// charts of XLSX. This library needs Go version 1.10 or later.

package excelize

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/md5"
	"crypto/sha1"
	"crypto/sha256"
	"crypto/sha512"
	"encoding/base64"
	"encoding/binary"
	"encoding/xml"
	"hash"
	"strings"

	"github.com/richardlehane/mscfb"
	"golang.org/x/crypto/md4"
	"golang.org/x/crypto/ripemd160"
	"golang.org/x/text/encoding/unicode"
)

var (
	blockKey                   = []byte{0x14, 0x6e, 0x0b, 0xe7, 0xab, 0xac, 0xd0, 0xd6} // Block keys used for encryption
	packageOffset              = 8                                                      // First 8 bytes are the size of the stream
	packageEncryptionChunkSize = 4096
	cryptoIdentifier           = []byte{ // checking protect workbook by [MS-OFFCRYPTO] - v20181211 3.1 FeatureIdentifier
		0x3c, 0x00, 0x00, 0x00, 0x4d, 0x00, 0x69, 0x00, 0x63, 0x00, 0x72, 0x00, 0x6f, 0x00, 0x73, 0x00,
		0x6f, 0x00, 0x66, 0x00, 0x74, 0x00, 0x2e, 0x00, 0x43, 0x00, 0x6f, 0x00, 0x6e, 0x00, 0x74, 0x00,
		0x61, 0x00, 0x69, 0x00, 0x6e, 0x00, 0x65, 0x00, 0x72, 0x00, 0x2e, 0x00, 0x44, 0x00, 0x61, 0x00,
		0x74, 0x00, 0x61, 0x00, 0x53, 0x00, 0x70, 0x00, 0x61, 0x00, 0x63, 0x00, 0x65, 0x00, 0x73, 0x00,
		0x01, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
	}
)

// Encryption specifies the encryption structure, streams, and storages are
// required when encrypting ECMA-376 documents.
type Encryption struct {
	KeyData       KeyData       `xml:"keyData"`
	DataIntegrity DataIntegrity `xml:"dataIntegrity"`
	KeyEncryptors KeyEncryptors `xml:"keyEncryptors"`
}

// KeyData specifies the cryptographic attributes used to encrypt the data.
type KeyData struct {
	SaltSize        int    `xml:"saltSize,attr"`
	BlockSize       int    `xml:"blockSize,attr"`
	KeyBits         int    `xml:"keyBits,attr"`
	HashSize        int    `xml:"hashSize,attr"`
	CipherAlgorithm string `xml:"cipherAlgorithm,attr"`
	CipherChaining  string `xml:"cipherChaining,attr"`
	HashAlgorithm   string `xml:"hashAlgorithm,attr"`
	SaltValue       string `xml:"saltValue,attr"`
}

// DataIntegrity specifies the encrypted copies of the salt and hash values
// used to help ensure that the integrity of the encrypted data has not been
// compromised.
type DataIntegrity struct {
	EncryptedHmacKey   string `xml:"encryptedHmacKey,attr"`
	EncryptedHmacValue string `xml:"encryptedHmacValue,attr"`
}

// KeyEncryptors specifies the key encryptors used to encrypt the data.
type KeyEncryptors struct {
	KeyEncryptor []KeyEncryptor `xml:"keyEncryptor"`
}

// KeyEncryptor specifies that the schema used by this encryptor is the schema
// specified for password-based encryptors.
type KeyEncryptor struct {
	XMLName      xml.Name     `xml:"keyEncryptor"`
	URI          string       `xml:"uri,attr"`
	EncryptedKey EncryptedKey `xml:"encryptedKey"`
}

// EncryptedKey used to generate the encrypting key.
type EncryptedKey struct {
	XMLName                    xml.Name `xml:"http://schemas.microsoft.com/office/2006/keyEncryptor/password encryptedKey"`
	SpinCount                  int      `xml:"spinCount,attr"`
	EncryptedVerifierHashInput string   `xml:"encryptedVerifierHashInput,attr"`
	EncryptedVerifierHashValue string   `xml:"encryptedVerifierHashValue,attr"`
	EncryptedKeyValue          string   `xml:"encryptedKeyValue,attr"`
	KeyData
}

// Decrypt API decrypt the CFB file format with Agile Encryption. Support
// cryptographic algorithm: MD4, MD5, RIPEMD-160, SHA1, SHA256, SHA384 and
// SHA512.
func Decrypt(raw []byte, opt *Options) (packageBuf []byte, err error) {
	doc, err := mscfb.New(bytes.NewReader(raw))
	if err != nil {
		return
	}
	encryptionInfoBuf, encryptedPackageBuf := extractPart(doc)
	var encryptionInfo Encryption
	if encryptionInfo, err = parseEncryptionInfo(encryptionInfoBuf[8:]); err != nil {
		return
	}
	// Convert the password into an encryption key.
	key, err := convertPasswdToKey(opt.Password, encryptionInfo)
	if err != nil {
		return
	}
	// Use the key to decrypt the package key.
	encryptedKey := encryptionInfo.KeyEncryptors.KeyEncryptor[0].EncryptedKey
	saltValue, err := base64.StdEncoding.DecodeString(encryptedKey.SaltValue)
	if err != nil {
		return
	}
	encryptedKeyValue, err := base64.StdEncoding.DecodeString(encryptedKey.EncryptedKeyValue)
	if err != nil {
		return
	}
	packageKey, err := crypt(false, encryptedKey.CipherAlgorithm, encryptedKey.CipherChaining, key, saltValue, encryptedKeyValue)
	// Use the package key to decrypt the package.
	return cryptPackage(false, packageKey, encryptedPackageBuf, encryptionInfo)
}

// extractPart extract data from storage by specified part name.
func extractPart(doc *mscfb.Reader) (encryptionInfoBuf, encryptedPackageBuf []byte) {
	for entry, err := doc.Next(); err == nil; entry, err = doc.Next() {
		switch entry.Name {
		case "EncryptionInfo":
			buf := make([]byte, entry.Size)
			i, _ := doc.Read(buf)
			if i > 0 {
				encryptionInfoBuf = buf
				break
			}
		case "EncryptedPackage":
			buf := make([]byte, entry.Size)
			i, _ := doc.Read(buf)
			if i > 0 {
				encryptedPackageBuf = buf
				break
			}
		}
	}
	return
}

// convertPasswdToKey convert the password into an encryption key.
func convertPasswdToKey(passwd string, encryption Encryption) (key []byte, err error) {
	var b bytes.Buffer
	saltValue, err := base64.StdEncoding.DecodeString(encryption.KeyEncryptors.KeyEncryptor[0].EncryptedKey.SaltValue)
	if err != nil {
		return
	}
	b.Write(saltValue)
	encoder := unicode.UTF16(unicode.LittleEndian, unicode.IgnoreBOM).NewEncoder()
	passwordBuffer, err := encoder.Bytes([]byte(passwd))
	if err != nil {
		return
	}
	b.Write(passwordBuffer)
	// Generate the initial hash.
	key = hashing(encryption.KeyData.HashAlgorithm, b.Bytes())
	// Now regenerate until spin count.
	for i := 0; i < encryption.KeyEncryptors.KeyEncryptor[0].EncryptedKey.SpinCount; i++ {
		iterator := createUInt32LEBuffer(i)
		key = hashing(encryption.KeyData.HashAlgorithm, iterator, key)
	}
	// Now generate the final hash.
	key = hashing(encryption.KeyData.HashAlgorithm, key, blockKey)
	// Truncate or pad as needed to get to length of keyBits.
	keyBytes := encryption.KeyEncryptors.KeyEncryptor[0].EncryptedKey.KeyBits / 8
	if len(key) < keyBytes {
		tmp := make([]byte, 0x36)
		key = append(key, tmp...)
		key = tmp
	} else if len(key) > keyBytes {
		key = key[:keyBytes]
	}
	return
}

// hashing data by specified hash algorithm.
func hashing(hashAlgorithm string, buffer ...[]byte) (key []byte) {
	var hashMap = map[string]hash.Hash{
		"md4":        md4.New(),
		"md5":        md5.New(),
		"ripemd-160": ripemd160.New(),
		"sha1":       sha1.New(),
		"sha256":     sha256.New(),
		"sha384":     sha512.New384(),
		"sha512":     sha512.New(),
	}
	handler, ok := hashMap[strings.ToLower(hashAlgorithm)]
	if !ok {
		return key
	}
	for _, buf := range buffer {
		handler.Write(buf)
	}
	key = handler.Sum(nil)
	return key
}

// createUInt32LEBuffer create buffer with little endian 32-bit unsigned
// integer.
func createUInt32LEBuffer(value int) []byte {
	buf := make([]byte, 4)
	binary.LittleEndian.PutUint32(buf, uint32(value))
	return buf
}

// parseEncryptionInfo parse the encryption info XML into an object.
func parseEncryptionInfo(encryptionInfo []byte) (encryption Encryption, err error) {
	err = xml.Unmarshal(encryptionInfo, &encryption)
	return
}

// crypt encrypt / decrypt input by given cipher algorithm, cipher chaining,
// key and initialization vector.
func crypt(encrypt bool, cipherAlgorithm, cipherChaining string, key, iv, input []byte) (packageKey []byte, err error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return input, err
	}
	stream := cipher.NewCBCDecrypter(block, iv)
	stream.CryptBlocks(input, input)
	return input, nil
}

// cryptPackage encrypt / decrypt package by given packageKey and encryption
// info.
func cryptPackage(encrypt bool, packageKey, input []byte, encryption Encryption) (outputChunks []byte, err error) {
	encryptedKey := encryption.KeyData
	var offset = packageOffset
	if encrypt {
		offset = 0
	}
	var i, start, end int
	var iv, outputChunk []byte
	for end < len(input) {
		start = end
		end = start + packageEncryptionChunkSize

		if end > len(input) {
			end = len(input)
		}
		// Grab the next chunk
		var inputChunk []byte
		if (end + offset) < len(input) {
			inputChunk = input[start+offset : end+offset]
		} else {
			inputChunk = input[start+offset : end]
		}

		// Pad the chunk if it is not an integer multiple of the block size
		remainder := len(inputChunk) % encryptedKey.BlockSize
		if remainder != 0 {
			inputChunk = append(inputChunk, make([]byte, encryptedKey.BlockSize-remainder)...)
		}
		// Create the initialization vector
		iv, err = createIV(encrypt, i, encryption)
		if err != nil {
			return
		}
		// Encrypt/decrypt the chunk and add it to the array
		outputChunk, err = crypt(encrypt, encryptedKey.CipherAlgorithm, encryptedKey.CipherChaining, packageKey, iv, inputChunk)
		if err != nil {
			return
		}
		outputChunks = append(outputChunks, outputChunk...)
		i++
	}
	return
}

// createIV create an initialization vector (IV).
func createIV(encrypt bool, blockKey int, encryption Encryption) ([]byte, error) {
	encryptedKey := encryption.KeyData
	// Create the block key from the current index
	blockKeyBuf := createUInt32LEBuffer(blockKey)
	var b bytes.Buffer
	saltValue, err := base64.StdEncoding.DecodeString(encryptedKey.SaltValue)
	if err != nil {
		return nil, err
	}
	b.Write(saltValue)
	b.Write(blockKeyBuf)
	// Create the initialization vector by hashing the salt with the block key.
	// Truncate or pad as needed to meet the block size.
	iv := hashing(encryptedKey.HashAlgorithm, b.Bytes())
	if len(iv) < encryptedKey.BlockSize {
		tmp := make([]byte, 0x36)
		iv = append(iv, tmp...)
		iv = tmp
	} else if len(iv) > encryptedKey.BlockSize {
		iv = iv[0:encryptedKey.BlockSize]
	}
	return iv, nil
}