tlock_age 0.0.8

use std::{
    collections::HashSet,
    io,
    sync::{Arc, Mutex},
};

use age::secrecy::ExposeSecret;
use age_core::format::{FileKey, Stanza};
use zeroize::Zeroize;

pub const STANZA_TAG: &str = "tlock";

// Identity implements the age Identity interface. This is used to decrypt
// data with the age Decrypt API.
pub struct Identity {
    hash: Vec<u8>,
    signature: Vec<u8>,
}

impl Identity {
    pub fn new(hash: &[u8], signature: &[u8]) -> Self {
        Self {
            hash: hash.to_vec(),
            signature: signature.to_vec(),
        }
    }
}

impl age::Identity for Identity {
    // Unwrap is called by the age Decrypt API and is provided the DEK that was time
    // lock encrypted by the Wrap function via the Stanza. Inside of Unwrap we decrypt
    // the DEK and provide back to age.
    fn unwrap_stanza(&self, stanza: &Stanza) -> Option<Result<FileKey, age::DecryptError>> {
        if stanza.tag != STANZA_TAG {
            return None;
        }
        if stanza.args.len() != 2 {
            return Some(Err(age::DecryptError::InvalidHeader));
        }
        let args: [String; 2] = [stanza.args[0].clone(), stanza.args[1].clone()];

        let _round = args[0]
            .parse::<u64>()
            .map_err(|_| age::DecryptError::InvalidHeader)
            .ok()?;

        if self.hash != hex::decode(&args[1]).ok()? {
            return Some(Err(age::DecryptError::InvalidHeader));
        }

        let dst = InMemoryWriter::new();
        let decryption = tlock::decrypt(dst.to_owned(), stanza.body.as_slice(), &self.signature);
        decryption
            .map_err(|_| age::DecryptError::DecryptionFailed)
            .ok()?;
        let mut dst = dst.memory();
        dst.resize(16, 0);
        let file_key: [u8; 16] = dst[..].try_into().ok()?;
        Some(Ok(FileKey::new(Box::new(file_key))))
    }
}

// Identity implements the age Identity interface. This is used to decrypt
// data with the age Decrypt API.
pub struct HeaderIdentity {
    hash: Mutex<Option<Vec<u8>>>,
    round: Mutex<Option<u64>>,
}

impl HeaderIdentity {
    pub fn new() -> Self {
        Self {
            hash: Mutex::new(None),
            round: Mutex::new(None),
        }
    }

    pub fn hash(&self) -> Option<Vec<u8>> {
        self.hash.lock().unwrap().clone()
    }

    pub fn round(&self) -> Option<u64> {
        *self.round.lock().unwrap()
    }
}

impl Default for HeaderIdentity {
    fn default() -> Self {
        Self::new()
    }
}

impl age::Identity for HeaderIdentity {
    // Unwrap is called by the age Decrypt API and is provided the DEK that was time
    // lock encrypted by the Wrap function via the Stanza. Inside of Unwrap we extract
    // tlock header and assign it to the identity.
    fn unwrap_stanza(&self, stanza: &Stanza) -> Option<Result<FileKey, age::DecryptError>> {
        if stanza.tag != STANZA_TAG {
            return None;
        }
        if stanza.args.len() != 2 {
            return Some(Err(age::DecryptError::InvalidHeader));
        }
        let args: [String; 2] = [stanza.args[0].clone(), stanza.args[1].clone()];

        let round = args[0]
            .parse::<u64>()
            .map_err(|_| age::DecryptError::InvalidHeader)
            .ok()?;
        let hash = hex::decode(&args[1])
            .map_err(|_| age::DecryptError::InvalidHeader)
            .ok()?;

        *self.round.lock().unwrap() = Some(round);
        *self.hash.lock().unwrap() = Some(hash);
        None
    }
}

/// Recipient implements the age Recipient interface. This is used to encrypt
/// data with the age Encrypt API.
pub struct Recipient {
    hash: Vec<u8>,
    public_key_bytes: Vec<u8>,
    round: u64,
}

impl Recipient {
    pub fn new(hash: &[u8], public_key_bytes: &[u8], round: u64) -> Self {
        Self {
            hash: hash.to_vec(),
            public_key_bytes: public_key_bytes.to_vec(),
            round,
        }
    }
}

#[derive(Clone)]
struct InMemoryWriter {
    memory: Arc<Mutex<Vec<u8>>>,
}

impl InMemoryWriter {
    pub fn new() -> Self {
        Self {
            memory: Arc::new(Mutex::new(vec![])),
        }
    }

    pub fn memory(&self) -> Vec<u8> {
        self.memory.lock().unwrap().to_owned()
    }
}

impl io::Write for InMemoryWriter {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.memory.lock().unwrap().extend(buf);
        Ok(buf.len())
    }

    fn flush(&mut self) -> io::Result<()> {
        self.memory.lock().unwrap().to_owned().zeroize();
        Ok(())
    }
}

impl age::Recipient for Recipient {
    /// Wrap is called by the age Encrypt API and is provided the DEK generated by
    /// age that is used for encrypting/decrypting data. Inside of Wrap we encrypt
    /// the DEK using time lock encryption.
    fn wrap_file_key(
        &self,
        file_key: &FileKey,
    ) -> Result<(Vec<Stanza>, HashSet<String>), age::EncryptError> {
        let src = file_key.expose_secret().as_slice();
        let dst = InMemoryWriter::new();
        let _ = tlock::encrypt(dst.to_owned(), src, &self.public_key_bytes, self.round);

        Ok((
            vec![Stanza {
                tag: STANZA_TAG.to_string(),
                args: vec![self.round.to_string(), hex::encode(&self.hash)],
                body: dst.memory(),
            }],
            HashSet::new(),
        ))
    }
}

#[cfg(test)]
mod tests {
    use std::{
        io::{Read, Write},
        iter,
    };

    use drand_core::HttpClient;

    use crate::{Identity, Recipient};

    #[test]
    fn it_works() {
        let client: HttpClient =
            "https://api.drand.sh/52db9ba70e0cc0f6eaf7803dd07447a1f5477735fd3f661792ba94600c84e971"
                .try_into()
                .unwrap();
        let info = client.chain_info().unwrap();

        let round = 100;
        let beacon = client.get(round).unwrap();
        let id = Identity::new(&info.hash(), &beacon.signature());
        let recipient = Recipient::new(&info.hash(), &info.public_key(), round);

        let mut plaintext = vec![0u8; 1000];
        plaintext.fill_with(rand::random);
        let encrypted = {
            let encryptor =
                age::Encryptor::with_recipients(iter::once(&recipient as &dyn age::Recipient))
                    .expect("we provided a recipient");

            let mut encrypted = vec![];
            let mut writer = encryptor.wrap_output(&mut encrypted).unwrap();
            writer.write_all(&plaintext).unwrap();
            writer.finish().unwrap();

            encrypted
        };

        let decrypted = {
            let decryptor = age::Decryptor::new(encrypted.as_slice()).unwrap();
            assert!(!decryptor.is_scrypt());

            let mut decrypted = vec![];
            let mut reader = decryptor
                .decrypt(iter::once(&id as &dyn age::Identity))
                .unwrap();
            reader.read_to_end(&mut decrypted).unwrap();

            decrypted
        };

        assert_eq!(decrypted, plaintext);
    }
}