ceal 0.1.0

use crate::{
	address::Address,
	dns::{query_record, sha256_truncated},
	message::{is_encrypted, plaintext_message, raw_header, split_content_headers}
};
use futures_util::future;
use hickory_resolver::{
	proto::rr::{Name, RData, RecordType},
	recursor::RecursorError
};
use log::{error, info};
use mail_parser::MessageParser;
use rand::{RngExt, distr::Alphanumeric};
use sequoia_openpgp::{
	Cert,
	parse::Parse,
	policy::StandardPolicy,
	serialize::stream::{Armorer, Encryptor, LiteralWriter, Message},
	types::KeyFlags
};
use std::{
	borrow::Cow,
	io::{BufRead as _, Write as _}
};
use thiserror::Error;

#[derive(Debug, Error)]
pub enum Error {
	#[error("DNS Error: {0}")]
	DnsError(#[from] RecursorError),

	#[error("The OPENPGPKEY DNS record(s) contained no usable OpenPGP encryption keys")]
	NoUsableKeys
}

pub type Result<T, E = Error> = std::result::Result<T, E>;

pub async fn find_certs_for_recipient<A>(recipient: &A) -> Result<Vec<Cert>>
where
	A: Address
{
	let mut name: Name = recipient.domain().parse().unwrap();
	name.set_fqdn(true);
	name = name.prepend_label("_openpgpkey").unwrap();
	name = name
		.prepend_label(sha256_truncated(recipient.local().as_bytes()))
		.unwrap();

	let raw_certs = query_record(&name, RecordType::OPENPGPKEY, |record| {
		let RData::OPENPGPKEY(openpgpkey) = record.into_data() else {
			return None;
		};
		Some(openpgpkey.public_key)
	})
	.await?;

	let mut certs = Vec::with_capacity(raw_certs.len());
	let policy = StandardPolicy::new();
	for raw_cert in raw_certs {
		let cert = match Cert::from_bytes(&raw_cert) {
			Ok(cert) => cert,
			Err(err) => {
				error!("Unable parse OPENPGPKEY record for {name}: {err}");
				continue;
			}
		};
		if cert
			.keys()
			.with_policy(&policy, None)
			.supported()
			.alive()
			.revoked(false)
			.next()
			.is_none()
		{
			error!(
				"Found OPENPGPKEY record for {name} with fingerprint {} that contains no usable encryption keys.",
				cert.fingerprint()
			);
			continue;
		}
		certs.push(cert);
	}

	if certs.is_empty() {
		return Err(Error::NoUsableKeys);
	}
	Ok(certs)
}

pub async fn find_certs_for_all_recipients<A>(to: &[A]) -> Result<Vec<Cert>>
where
	A: Address
{
	let certs = future::try_join_all(to.iter().map(find_certs_for_recipient)).await?;
	Ok(certs.into_iter().flatten().collect())
}

pub fn encrypt(msg_bytes: &[u8], recipient_certs: Vec<Cert>) -> Cow<'_, [u8]> {
	let Some(mut msg) = MessageParser::new().parse(msg_bytes) else {
		error!("Unable to parse message, not encrypting");
		return msg_bytes.into();
	};

	// refuse encryption if the message is already encrypted
	if is_encrypted(&msg) {
		info!("Found encrypted message, not encrypting twice");
		return msg_bytes.into();
	}

	// Split the message for encryption
	let content_headers = split_content_headers(&mut msg);
	let plaintext = plaintext_message(&msg, &content_headers);

	// Extract recipient keys from certificates
	let mut recipients = Vec::new();
	let policy = StandardPolicy::new();
	for cert in &recipient_certs {
		for key in cert
			.keys()
			.with_policy(&policy, None)
			.supported()
			.alive()
			.revoked(false)
		{
			if !key.has_any_key_flag(KeyFlags::storage_encryption())
				&& !key.has_any_key_flag(KeyFlags::transport_encryption())
			{
				// key not for encryption, skip
				continue;
			}
			recipients.push(key);
		}
	}

	// Write the remaining headers into the encrypted mail
	let mut encrypted = Vec::<u8>::new();
	for header in msg.root_part().headers() {
		encrypted.extend_from_slice(raw_header(&msg, header));
	}

	// Create a random boundary. We will add an `=` to ensure it is never present in
	// base64-encoded data
	let rng = rand::rng();
	let boundary: Vec<u8> = rng.sample_iter(Alphanumeric).take(16).collect();
	let boundary_begin = b"--=-"
		.iter()
		.chain(&boundary)
		.chain(b"\r\n")
		.copied()
		.collect::<Vec<u8>>();
	let boundary_end = b"--=-"
		.iter()
		.chain(&boundary)
		.chain(b"--\r\n")
		.copied()
		.collect::<Vec<u8>>();

	// Write the OpenPGP content-type header
	encrypted.extend_from_slice(
		b"Content-Type: multipart/encrypted; protocol=\"application/pgp-encrypted\";\r\n"
	);
	encrypted.extend_from_slice(b" boundary=\"=-");
	encrypted.extend_from_slice(&boundary);
	encrypted.extend_from_slice(b"\"\r\n");

	// End of headers
	encrypted.extend_from_slice(b"\r\n");

	// Write the first OpenPGP MIME boilerplate
	encrypted.extend_from_slice(&boundary_begin);
	encrypted.extend_from_slice(b"Content-Type: application/pgp-encrypted\r\n");
	encrypted.extend_from_slice(b"Content-Transfer-Encoding: 7bit\r\n");
	encrypted.extend_from_slice(b"\r\n");
	encrypted.extend_from_slice(b"Version: 1\r\n");
	encrypted.extend_from_slice(b"\r\n");
	encrypted.extend_from_slice(&boundary_begin);
	encrypted.extend_from_slice(
		b"Content-Type: application/octet-stream; name=\"encrypted.asc\"\r\n"
	);
	encrypted.extend_from_slice(b"Content-Description: OpenPGP Encrypted Message\r\n");
	encrypted.extend_from_slice(b"Content-Transfer-Encoding: 7bit\r\n");
	encrypted.extend_from_slice(b"\r\n");

	// Encrypt the message. We must buffer this because sq doesn't like \r\n:
	// https://docs.rs/sequoia-openpgp/2.2.0/src/sequoia_openpgp/armor.rs.html#71
	let mut encrypted_message = Vec::<u8>::new();
	let message = Message::new(&mut encrypted_message);
	let message = Armorer::new(message).build().unwrap();
	let message = Encryptor::for_recipients(message, recipients)
		.build()
		.unwrap();
	let mut message = LiteralWriter::new(message).build().unwrap();
	message.write_all(&plaintext).unwrap();
	message.finalize().unwrap();
	for line in encrypted_message.lines() {
		encrypted.extend_from_slice(line.unwrap().as_bytes());
		encrypted.extend_from_slice(b"\r\n");
	}

	// Write the remaining boilerplate
	encrypted.extend_from_slice(b"\r\n");
	encrypted.extend_from_slice(&boundary_end);

	encrypted.into()
}