// 加密算法库
package util

import (
	"crypto/cipher"
	"crypto/md5"
)

const (
	XOR_KEY_LEN = 32
	XOR_IV_LEN  = 32
)

// 定义一些加密方法
var CipherMethod = map[string]*CipherInfo{
	"plain": {0, 0, NewPlainStream, NewPlainStream},
	"xor":   {XOR_KEY_LEN, XOR_IV_LEN, NewXorStream, NewXorStream},
}

// 普通的xor加密，没有破解的难度，只是简单的让明文不可见而已
type XorStreamCipher struct {
	key   [XOR_KEY_LEN]byte
	index byte // 加解密的位置指示
}

func (c *XorStreamCipher) XORKeyStream(dst, src []byte) {
	l := len(src)
	for i := range l {
		dst[i] = src[i] ^ c.key[c.index]
		c.index++
		if c.index >= XOR_KEY_LEN {
			c.index = 0
		}
	}
}

func NewXorStream(key, iv []byte) (cipher.Stream, error) {
	var c XorStreamCipher
	copy(c.key[:], key[:XOR_KEY_LEN])
	for i := range XOR_KEY_LEN {
		c.key[i] = c.key[i] ^ iv[i]
	}
	return &c, nil
}

// 没有加密模式，考虑到部分通讯管道已经是加密的状态，故增加不加密的方式，提高速度
type PlainStreamCipher struct{}

func (c *PlainStreamCipher) XORKeyStream(dst, src []byte) {
	if &dst != &src {
		copy(dst, src)
	}
}

func NewPlainStream(key, iv []byte) (cipher.Stream, error) {
	var c PlainStreamCipher
	return &c, nil
}

func Md5Sum(d []byte) []byte {
	h := md5.New()
	h.Write(d)
	return h.Sum(nil)
}

// 转字符串密码为byte数组类型
func EvpBytesToKey(password string, keyLen int) (key []byte) {
	if keyLen <= 1 {
		return
	}
	const md5Len = 16

	cnt := (keyLen-1)/md5Len + 1
	m := make([]byte, cnt*md5Len)
	copy(m, Md5Sum([]byte(password)))

	// Repeatedly call md5 until bytes generated is enough.
	// Each call to md5 uses data: prev md5 sum + password.
	d := make([]byte, md5Len+len(password))
	start := 0
	for i := 1; i < cnt; i++ {
		start += md5Len
		copy(d, m[start-md5Len:start])
		copy(d[md5Len:], password)
		copy(m[start:], Md5Sum(d))
	}
	return m[:keyLen]
}

// 预先准备好的加密算法
type CipherInfo struct {
	KeyLen           int
	IvLen            int
	NewEncryptStream func(key, iv []byte) (cipher.Stream, error)
	NewDecryptStream func(key, iv []byte) (cipher.Stream, error)
}

// 建立一个cipher对加密算法统一管理
type Cipher struct {
	enc cipher.Stream
	dec cipher.Stream
}

func (c *Cipher) Encrypt(dst, src []byte) {
	c.enc.XORKeyStream(dst, src)
}

func (c *Cipher) Decrypt(dst, src []byte) {
	c.dec.XORKeyStream(dst, src)
}

// 新建 cipher
func NewCipher(ci *CipherInfo, pwd string, eiv []byte, div []byte) (c *Cipher, err error) {
	key := EvpBytesToKey(pwd, ci.KeyLen)
	enc, err := ci.NewEncryptStream(key, eiv)
	if err != nil {
		return nil, err
	}
	dec, err := ci.NewDecryptStream(key, div)
	if err != nil {
		return nil, err
	}
	c = &Cipher{
		enc: enc,
		dec: dec,
	}
	return
}