// Copyright (c) 2024 blackshirt.
// Use of this source code is governed by an MIT license
// that can be found in the LICENSE file.
//
// Chacha20 symmetric key stream cipher encryption based on RFC 8439
module chacha20

import crypto.internal.subtle

// The size of ChaCha20 key, ie 256 bits size, in bytes
pub const key_size = 32
// The size of standard IETF ChaCha20 nonce, ie 96 bits size, in bytes
pub const nonce_size = 12
// The size of extended variant of standard ChaCha20 (XChaCha20) nonce, 192 bits
pub const x_nonce_size = 24
// The size of original ChaCha20 nonce, 64 bits
pub const orig_nonce_size = 8
// internal block size ChaCha20 operates on, in bytes
const block_size = 64

// four constants of ChaCha20 state.
const cc0 = u32(0x61707865) // expa
const cc1 = u32(0x3320646e) // nd 3
const cc2 = u32(0x79622d32) // 2-by
const cc3 = u32(0x6b206574) // te k

// CipherMode was enumeration of ChaCha20 supported variant.
enum CipherMode {
	// The standard IETF ChaCha20 (and XChaCha20), with 32-bit internal counter.
	standard
	// The original ChaCha20 with 64-bit internal counter.
	original
}

// Configuration options
@[params]
pub struct Options {
pub mut:
	// currently, used for XChaCha20 construct
	use_64bit_counter bool
}

// encrypt encrypts plaintext bytes with ChaCha20 cipher instance with provided key and nonce.
// It was a thin wrapper around two supported nonce size, ChaCha20 with 96 bits
// and XChaCha20 with 192 bits nonce. Internally, encrypt start with 0's counter value.
// If you want more control, use Cipher instance and setup the counter by your self.
pub fn encrypt(key []u8, nonce []u8, plaintext []u8, opt Options) ![]u8 {
	mut stream := new_stream_with_options(key, nonce, opt)!
	mut dst := []u8{len: plaintext.len}
	stream.keystream_full(mut dst, plaintext)!
	unsafe { stream.reset() }
	return dst
}

// decrypt does reverse of encrypt operation by decrypting ciphertext with ChaCha20 cipher
// instance with provided key and nonce.
pub fn decrypt(key []u8, nonce []u8, ciphertext []u8, opt Options) ![]u8 {
	mut stream := new_stream_with_options(key, nonce)!
	mut dst := []u8{len: ciphertext.len}
	stream.keystream_full(mut dst, ciphertext)!
	unsafe { stream.reset() }
	return dst
}

// Cipher represents ChaCha20 stream cipher instances.
@[noinit]
pub struct Cipher {
	Stream
mut:
	// internal buffer for storing key stream results
	block []u8 = []u8{len: block_size}
	// The last length of leftover unprocessed keystream from internal buffer
	length int
}

// new_cipher creates a new ChaCha20 stream cipher with the given 32 bytes key
// and bytes of nonce with supported size, ie, 8, 12 or 24 bytes nonce.
// Standard IETF variant use 12 bytes nonce's, if you want create original ChaCha20 cipher
// with support for 64-bit counter, use 8 bytes length nonce's instead
// If 24 bytes of nonce was provided, the XChaCha20 construction will be used.
// It returns new ChaCha20 cipher instance or an error if key or nonce have any other length.
pub fn new_cipher(key []u8, nonce []u8, opt Options) !&Cipher {
	stream := new_stream_with_options(key, nonce, opt)!
	return &Cipher{
		Stream: stream
	}
}

// xor_key_stream xors each byte in the given slice in the src with a byte from the
// cipher's key stream. It fulfills `cipher.Stream` interface. It encrypts the plaintext message
// in src and stores the ciphertext result in dst in a key stream fashion.
// You must never use the same (key, nonce) pair more than once for encryption.
// This would void any confidentiality guarantees for the messages encrypted with the same nonce and key.
@[direct_array_access]
pub fn (mut c Cipher) xor_key_stream(mut dst []u8, src []u8) {
	if src.len == 0 {
		return
	}
	if dst.len < src.len {
		panic('chacha20/chacha: dst buffer is to small')
	}
	dst = unsafe { dst[..src.len] }
	if subtle.inexact_overlap(dst, src) {
		panic('chacha20: invalid buffer overlap')
	}
	// index of position within src bytes
	mut idx := 0

	// First, try to drain any remaining key stream from internal buffer
	if c.length != 0 {
		// remaining keystream on internal buffer
		mut kstream := c.block[block_size - c.length..]
		if src.len < kstream.len {
			kstream = unsafe { kstream[..src.len] }
		}
		// xors every bytes in src with bytes from key stream and stored into dst
		for i, b in kstream {
			dst[idx + i] = src[idx + i] ^ b
		}
		// updates position and internal buffer length.
		// when c.length reaches the block_size, we reset it for future use.
		c.length -= kstream.len
		idx += kstream.len
		if c.length == block_size {
			unsafe { c.block.reset() }
			c.length = 0
		}
	}
	// process for remaining unprocessed src bytes
	mut remains := unsafe { src[idx..] }
	nr_blocks := remains.len / block_size

	// process for full block_size-d message
	for i := 0; i < nr_blocks; i++ {
		// for every block_sized message, we generates 64-bytes block key stream
		// and then xor-ing this block with generated key stream
		block := unsafe { remains[i * block_size..(i + 1) * block_size] }
		ks := c.keystream() or { panic(err) }
		for j, b in ks {
			dst[idx + j] = block[j] ^ b
		}
		// updates position
		idx += block_size
	}

	// process for remaining partial block
	if remains.len % block_size != 0 {
		last_block := unsafe { remains[nr_blocks * block_size..] }
		// generates one 64-bytes keystream block
		c.block = c.keystream() or { panic(err) }
		for i, b in last_block {
			dst[idx + i] = b ^ c.block[i]
		}
		c.length = block_size - last_block.len
		idx += last_block.len
	}
}

// encrypt encrypts src and stores into dst buffer. It works like `xor_key_stream` except
// its ignore key streaming process by ignoring remaining key stream in the internal buffer,
// so, its works in one shot of fashion.
// Its added to allow `chacha20poly1305` modules to work without key stream fashion.
// TODO: integrates it with the rest
@[direct_array_access]
pub fn (mut c Cipher) encrypt(mut dst []u8, src []u8) ! {
	if src.len == 0 {
		return
	}
	if dst.len < src.len {
		return error('chacha20: dst buffer is to small')
	}
	if subtle.inexact_overlap(dst, src) {
		return error('chacha20: invalid buffer overlap')
	}

	c.Stream.keystream_full(mut dst, src)!
}

// free the resources taken by the Cipher `c`. Dont use cipher after .free call
@[unsafe]
pub fn (mut c Cipher) free() {
	$if prealloc {
		return
	}
	unsafe {
		c.block.free()
	}
}

// reset quickly sets all Cipher's fields to default value.
// This method will be deprecated.
@[deprecated: 'do not use .reset() at all, create a new Cipher instead']
@[deprecated_after: '2025-11-30']
@[unsafe]
pub fn (mut c Cipher) reset() {
	c.Stream.reset()
	unsafe {
		c.block.reset()
	}
	c.length = 0
}

// set_counter sets Cipher's counter
pub fn (mut c Cipher) set_counter(ctr u64) {
	c.Stream.set_ctr(ctr)
}

// counter returns a current underlying counter value, as u64.
pub fn (c Cipher) counter() u64 {
	return c.Stream.ctr()
}

// rekey resets internal Cipher's state and reinitializes state with the provided key and nonce
pub fn (mut c Cipher) rekey(key []u8, nonce []u8) ! {
	// start of .reset() inline
	unsafe {
		c.Stream.reset()
		c.block.reset()
	}
	c.length = 0
	// end of .reset() inline

	// we use c.Stream.mode info to get 64-bit counter capability
	w64 := if c.mode == .original { true } else { false }
	opt := Options{
		use_64bit_counter: w64
	}
	stream := new_stream_with_options(key, nonce, opt)!
	c.Stream = stream
}