cnsprcy 0.3.2

use std::fmt;

use crypto_box::aead::{
	generic_array::typenum::Unsigned,
	AeadCore,
	AeadInPlace
};
use crypto_box::{
	ChaChaBox,
	PublicKey,
	SecretKey,
	KEY_SIZE as PUBLIC_KEY_SIZE
};
use data_encoding::BASE64;
use serde::{
	de::{self, Visitor},
	Deserialize, Serialize,
};

use crate::control::KeyArg;
use crate::peer::{
	Address,
	NodeID
};
use crate::util::{
	bincode_deserialize,
	bincode_serialize_into
};

const NONCE_SIZE: usize = <ChaChaBox as AeadCore>::NonceSize::USIZE;
const TAG_SIZE: usize = <ChaChaBox as AeadCore>::TagSize::USIZE;


#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Invitation {
	/// Inductor public key
	pub from: [u8; PUBLIC_KEY_SIZE],
	/// [`crypto_box`] nonce
	nonce: [u8; NONCE_SIZE],
	/// [`crypto_box`] nonce
	tag: [u8; TAG_SIZE],
	/// [`crypto_box`] ciphertext
	ciphertext: Vec<u8>
}

#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
pub struct Content {
	/// Inductor ID
	pub id: NodeID,
	/// Inductor name
	pub name: String,
	/// CNSPRCY encryption key
	pub key: KeyArg,
	/// Addresses the inductor is listening on
	pub addresses: Vec<Address>
}

impl Invitation {
	pub fn encrypt(key: &PublicKey, data: &Content) -> Result<Self, String> {
		//generate ephemeral DH keypair
		let our_private_key = SecretKey::from_bytes(rand::random());
		let our_public_key = our_private_key.public_key();
		//derive shared secret
		let cipher = ChaChaBox::new(key, &our_private_key);
		// generate random nonce
		let nonce: [u8; NONCE_SIZE] = rand::random();

		//serialize invitation data
		let mut buf = Vec::with_capacity(128);
		bincode_serialize_into(&mut buf, data)
			.map_err(|err| err.to_string())?;
		//TODO why not just encrypt_in_place ?
		let tag = cipher
			.encrypt_in_place_detached(&nonce.into(), &[], &mut buf)
			.map_err(|err| err.to_string())?
			.into();
		let from = our_public_key.to_bytes();
		let ciphertext = buf;

		Ok(Self { from, nonce, tag, ciphertext })
	}
	pub fn decrypt(mut self, key: &SecretKey) -> Result<Content, String> {
		let their_public_key = PublicKey::from_bytes(self.from);
		let cipher = ChaChaBox::new(&their_public_key, key);

		cipher.decrypt_in_place_detached(
			&self.nonce.into(),
			&[],
			&mut self.ciphertext,
			&self.tag.into()
		).map_err(|err| err.to_string())?;
		bincode_deserialize(&self.ciphertext)
			.map_err(|err| err.to_string())
	}

	pub fn into_bytes(self) -> Vec<u8> {
		let mut vec = Vec::new();
		vec.extend_from_slice(&self.from);
		vec.extend_from_slice(&self.nonce);
		vec.extend_from_slice(&self.tag);
		vec.extend_from_slice(&self.ciphertext);
		vec
	}

	#[cfg(test)]
	pub(crate) fn random() -> Self {
		let rcpt_key = PublicKey::from_bytes(rand::random());
		Self::encrypt(&rcpt_key, &Content::random()).unwrap()
	}
}

impl Content {
	#[cfg(test)]
	fn random() -> Self {
		Self {
			id: NodeID::random(),
			name: "Test123".to_string(),
			key: KeyArg::random(),
			addresses: Vec::new()
		}
	}
}


impl std::str::FromStr for Invitation {
	type Err = String;

	fn from_str(s: &str) -> Result<Self, Self::Err> {
		let vec = BASE64.decode(s.as_bytes()).map_err(|err| err.to_string())?;
		//TODO
		Self::try_from(vec.as_slice()).map_err(|_| "too short!".to_owned())
	}
}

impl fmt::Display for Invitation {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
		// it's possible to serialize directly into a fmt::Write, but that requires a contiguous byte slice
		let mut output = String::new();
		let encoding = BASE64;
		let mut encoder = encoding.new_encoder(&mut output);
		encoder.append(&self.from);
		encoder.append(&self.nonce);
		encoder.append(&self.tag);
		encoder.append(&self.ciphertext);
		encoder.finalize();
		f.write_str(&output)
	}
}

pub struct TooShort;

impl TryFrom<&[u8]> for Invitation {
	type Error = TooShort;
	fn try_from(value: &[u8]) -> Result<Self, Self::Error> {
		let (&from, rest) = value.split_first_chunk().ok_or(TooShort)?;
		let (&nonce, rest) = rest.split_first_chunk().ok_or(TooShort)?;
		let (&tag, rest) = rest.split_first_chunk().ok_or(TooShort)?;
		let ciphertext = Vec::from(rest);
		Ok(Self {from, nonce, tag, ciphertext})
	}
}

impl Serialize for Invitation {
	fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
	where
		S: serde::Serializer,
	{
		if serializer.is_human_readable() {
			serializer.serialize_str(&self.to_string())
		}
		else {
			let bytes = &[
				&self.from[..],
				&self.nonce,
				&self.tag,
				&self.ciphertext
			];
			// unfortunately it does not seem possible to do this without allocating a new Vec
			let vec: Vec<u8> = bytes.iter()
				.copied()
				.flatten()
				.copied()
				.collect();
			serializer.serialize_bytes(&vec)
		}
	}
}

impl<'de> Deserialize<'de> for Invitation {
	fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
	where
		D: serde::Deserializer<'de>,
	{
		let human_readable = deserializer.is_human_readable();
		let visitor = InvitationVisitor { human_readable };
		if deserializer.is_human_readable() {
			deserializer.deserialize_str(visitor)
		}
		else {
			deserializer.deserialize_byte_buf(visitor)
		}
	}
}

struct InvitationVisitor {
	human_readable: bool,
}

impl Visitor<'_> for InvitationVisitor {
	type Value = Invitation;

	fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
		if self.human_readable {
			formatter.write_str("a base64-encoded invitation")
		}
		else {
			formatter.write_str("an invitation as raw bytes")
		}
	}
	fn visit_str<E: de::Error>(self, value: &str) -> Result<Self::Value, E> {
		value.parse::<Invitation>().map_err(E::custom)
	}
	fn visit_bytes<E: de::Error>(self, v: &[u8]) -> Result<Self::Value, E> {
		Invitation::try_from(v)
			.map_err(|_| E::custom("byte sequence too short"))
	}
}

#[test]
fn human_readable() {
	use crate::control::Command;

	let invitation = Invitation::random();
	let serialized = serde_json::to_string(&invitation).unwrap();
	let deserialized = serde_json::from_str(&serialized).unwrap();
	assert_eq!(invitation, deserialized);

	let serialized = serde_json::to_value(&invitation).unwrap();
	dbg!(&serialized);
	assert!(serialized.is_string());
	let deserialized = serde_json::from_value(serialized).unwrap();
	assert_eq!(invitation, deserialized);

	let cmd = Command::Accept(invitation.clone());
	let serialized = serde_json::to_string(&cmd).unwrap();
	dbg!(&serialized);
	let deserialized = serde_json::from_str(&serialized).unwrap();
	let Command::Accept(inv_deserialized) = deserialized else {panic!();};
	assert_eq!(invitation, inv_deserialized);

	let example = serde_json::json!(
		"qoNFKuxatD48taSHyuPBYdACi4dtoHUG74yOH2MdvVGFxBH0\
		nbc9PD2RBu6+4F0/xhJoVK1+GUNlQjy+2PLXG+WLO0KtNnlU"
	);
	let _ = serde_json::from_value::<Invitation>(example).unwrap();
	const EXAMPLE: &str = "\"qoNFKuxatD48taSHyuPBYdACi4dtoHUG74yOH2MdvVGFxBH0\
		nbc9PD2RBu6+4F0/xhJoVK1+GUNlQjy+2PLXG+WLO0KtNnlU\"";
	let _ = serde_json::from_str::<Invitation>(EXAMPLE).unwrap();
}

#[test]
fn round_trip() {
	use crate::util::bincode_serialize;

	let rx_privkey = SecretKey::from_bytes(rand::random());
	let rx_pubkey = rx_privkey.public_key();

	let content = Content {
		id: NodeID::random(),
		name: "Test123".to_string(),
		key: KeyArg::random(),
		addresses: Vec::new()
	};

	// Node → Daemon
	let invitation = Invitation::encrypt(&rx_pubkey, &content).unwrap();

	// Daemon → Client
	let serialized = serde_json::to_string(&invitation).unwrap();
	let deserialized = serde_json::from_str(&serialized).unwrap();
	assert_eq!(invitation, deserialized);

	// Client → Client
	let serialized = bincode_serialize(&deserialized).unwrap();
	let deserialized = bincode_deserialize(&serialized).unwrap();
	assert_eq!(invitation, deserialized);

	// Client → Daemon
	let serialized = serde_json::to_string(&deserialized).unwrap();
	let deserialized = serde_json::from_str(&serialized).unwrap();
	assert_eq!(invitation, deserialized);

	// Daemon → Node
	let decrypted = deserialized.decrypt(&rx_privkey).unwrap();
	assert_eq!(content, decrypted);
}