package encryption

import "bytes"

type PaddingMode int

const (
	NoPadding PaddingMode = iota
	ZeroPadding
	PKCS7Padding
)

// 对给定的数据进行填充，默认采用填充零的方式（ZeroPadding），也可以采用填充PKCS#7的方式（PKCS7Padding）或者不填充（NoPadding）。
func Padding(data []byte, blockSize int, padding ...PaddingMode) []byte {
	paddingMethod := append(padding, ZeroPadding)[0]
	n := blockSize - len(data)%blockSize
	switch paddingMethod {
	case PKCS7Padding:
		pad := bytes.Repeat([]byte{byte(n)}, n)
		return append(data, pad...)
	case NoPadding:
		return data
	case ZeroPadding:
		fallthrough
	default:
		if n == 0 {
			return data
		}
		return append(data, bytes.Repeat([]byte{0}, n)...)
	}
}

// 对给定的数据进行去填充，默认采用去填充零的方式（对应ZeroPadding），也可以采用去填充PKCS#7的方式（对应PKCS7Padding）或者不填充（对应NoPadding）。
func Unpadding(data []byte, padding ...PaddingMode) []byte {
	paddingMethod := append(padding, ZeroPadding)[0]
	switch paddingMethod {
	case PKCS7Padding:
		length := len(data)
		unpadding := int(data[length-1])
		if length-unpadding < 0 {
			return make([]byte, 0)
		}
		return data[:(length - unpadding)]
	case NoPadding:
		return data
	case ZeroPadding:
		fallthrough
	default:
		return bytes.TrimRightFunc(data, func(r rune) bool {
			return r == rune(0)
		})
	}
}