package udp

import (
	"crypto/rc4"
)

type Crypto struct {
	cipher    *rc4.Cipher
	key       []byte
	encrypted bool
}

func NewCrypto() *Crypto {
	return &Crypto{
		encrypted: false,
	}
}

func (c *Crypto) SetKey(key []byte) error {
	cipher, err := rc4.NewCipher(key)
	if err != nil {
		return err
	}

	c.cipher = cipher
	c.key = make([]byte, len(key))
	copy(c.key, key)
	c.encrypted = true

	return nil
}

func (c *Crypto) IsEncrypted() bool {
	return c.encrypted
}

func (c *Crypto) Encrypt(data []byte) []byte {
	if !c.encrypted {
		return data
	}

	encrypted := make([]byte, len(data))
	copy(encrypted, data)
	c.cipher.XORKeyStream(encrypted, encrypted)
	return encrypted
}

func (c *Crypto) Decrypt(data []byte) []byte {
	if !c.encrypted {
		return data
	}

	decrypted := make([]byte, len(data))
	copy(decrypted, data)
	c.cipher.XORKeyStream(decrypted, decrypted)
	return decrypted
}