package sbproxy

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/sha256"
)

type cfb8 struct {
	block   cipher.Block
	iv      []byte
	tmp     []byte
	decrypt bool
}

func NewCFB8(block cipher.Block, iv []byte, decrypt bool) cipher.Stream {
	bs := block.BlockSize()
	if len(iv) != bs {
		panic("cfb8: IV length must equal block size")
	}
	return &cfb8{
		block:   block,
		iv:      append([]byte(nil), iv...),
		tmp:     make([]byte, bs),
		decrypt: decrypt,
	}
}

func (x *cfb8) XORKeyStream(dst, src []byte) {
	for i := range src {
		x.block.Encrypt(x.tmp, x.iv)
		val := src[i]
		dst[i] = val ^ x.tmp[0]
		if x.decrypt {
			copy(x.iv, x.iv[1:])
			x.iv[len(x.iv)-1] = val
		} else {
			copy(x.iv, x.iv[1:])
			x.iv[len(x.iv)-1] = dst[i]
		}
	}
}

func NewCipherStream(password string, decrypt bool) (cipher.Stream, error) {
	key := sha256.Sum256([]byte(password))
	block, err := aes.NewCipher(key[:])
	if err != nil {
		return nil, err
	}
	// Use first 16 bytes of Hash as IV
	return NewCFB8(block, key[:16], decrypt), nil
}