bitbottle 0.10.0

use cipher::consts::U64;
use cipher::generic_array::GenericArray;
use digest::Digest;
use num_enum::TryFromPrimitive;
use sha2::Sha512;
use std::cell::RefCell;
use std::fmt;
use std::io::{Read, Write};
use std::rc::Rc;
use crate::asymmetric::{BottlePublicKey, BottleSecretKey, Ed25519PublicKey};
use crate::bottle::{BottleReader, BottleStream, BottleWriter};
use crate::bottle_cap::BottleType;
use crate::bottle_error::{BottleError, BottleResult};
use crate::header::Header;
use crate::header_stream::{read_header_stream, write_header_stream};
use crate::streams::{ReadStream, ReadStreamRef, StreamBottle};


// fields in signed bottle header block:
const FIELD_SIGNATURE_TYPE_INT: u8 = 0;
const FIELD_SIGNER_NAME_STRING: u8 = 0;
const FIELD_PUBLIC_KEY_BYTES: u8 = 0;

// fields in signature header block:
const FIELD_SIGNATURE_BYTES: u8 = 0;

const DEFAULT_BLOCK_SIZE_BITS: usize = 20;


#[allow(non_camel_case_types)]
#[derive(Clone, Copy, Debug, Eq, PartialEq, TryFromPrimitive)]
#[repr(u8)]
pub enum SignatureAlgorithm {
    SHA_512_ED25519 = 0,
}

/// Pick your poison for secure hashing. The blake3 variant stores like 2KB
/// on the stack, but it's only created by the file scanner, which would need
/// to create it anyway.
#[allow(clippy::large_enum_variant)]
enum SignatureHashing {
    Sha512(Sha512),
}

impl fmt::Debug for SignatureHashing {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "SignatureHashing({})", match self {
            SignatureHashing::Sha512(_) => "sha512",
        })
    }
}

type HashingOutput = GenericArray<u8, U64>;

impl SignatureHashing {
    pub fn new(algorithm: SignatureAlgorithm) -> SignatureHashing {
        match algorithm {
            SignatureAlgorithm::SHA_512_ED25519 => SignatureHashing::Sha512(Sha512::new()),
        }
    }

    pub fn update(&mut self, data: &[u8]) {
        match self {
            SignatureHashing::Sha512(h) => { h.update(data); },
        }
    }

    pub fn finalize_reset(&mut self) -> HashingOutput {
        match self {
            SignatureHashing::Sha512(h) => h.finalize_reset(),
        }
    }
}

fn public_key_from_bytes(
    algorithm: SignatureAlgorithm,
    data: &[u8],
    name: String
) -> Option<Box<dyn BottlePublicKey>> {
    match algorithm {
        SignatureAlgorithm::SHA_512_ED25519 => Some(Box::new(Ed25519PublicKey::from_bytes(data, name).ok()?)),
        // _ => None
    }
}


/// A `Write` that writes all data into a bottle of type `Signed`.
pub struct SignedBottleWriter<W: Write, K: BottleSecretKey> {
    bottle_writer: BottleWriter<W>,
    hashing: SignatureHashing,
    secret_key: K,
}

impl<W: Write, K: BottleSecretKey> SignedBottleWriter<W, K> {
    pub fn new(writer: W, algorithm: SignatureAlgorithm, secret_key: K) -> BottleResult<SignedBottleWriter<W, K>> {
        let mut header = Header::new();
        header.add_int(FIELD_SIGNATURE_TYPE_INT, (algorithm as u8) as u64)?;
        header.add_string(FIELD_SIGNER_NAME_STRING, secret_key.name())?;
        header.add_bytes(FIELD_PUBLIC_KEY_BYTES, secret_key.public_key().as_bytes(false))?;
        let mut bottle_writer = BottleWriter::new(writer, BottleType::Signed, header, DEFAULT_BLOCK_SIZE_BITS)?;

        bottle_writer.write_data_stream()?;
        let hashing = SignatureHashing::new(algorithm);
        Ok(SignedBottleWriter { bottle_writer, hashing, secret_key })
    }

    pub fn close(mut self) -> BottleResult<W> {
        self.bottle_writer.close_stream()?;

        // make a header to sign
        let mut message = self.bottle_writer.bottle_cap().to_bytes();
        message.extend_from_slice(&self.hashing.finalize_reset());
        let signature = self.secret_key.sign(&message)?;

        let mut header = Header::new();
        header.add_bytes(FIELD_SIGNATURE_BYTES, &signature)?;
        write_header_stream(&mut self.bottle_writer, &header)?;

        self.bottle_writer.close()
    }
}

impl<W: Write, K: BottleSecretKey> Write for SignedBottleWriter<W, K> {
    fn write(&mut self, data: &[u8]) -> std::io::Result<usize> {
        let len = self.bottle_writer.write(data)?;
        self.hashing.update(&data[..len]);
        Ok(len)
    }

    fn flush(&mut self) -> std::io::Result<()> {
        self.bottle_writer.flush()
    }
}


#[derive(Debug)]
pub struct SignedBottleReader<R: Read> {
    bottle_reader: BottleReader<R>,
    hashing: SignatureHashing,

    /// for readers: decoded info from the bottle header
    pub algorithm: SignatureAlgorithm,
    pub public_key: Box<dyn BottlePublicKey>,
}

impl<R: Read> SignedBottleReader<R> {
    pub fn new(mut bottle_reader: BottleReader<R>) -> BottleResult<SignedBottleReader<R>> {
        let cap = &bottle_reader.bottle_cap;
        bottle_reader.expect_type(BottleType::Signed)?;

        let id = cap.header.get_int(FIELD_SIGNATURE_TYPE_INT).ok_or(BottleError::UnknownSignatureAlgorithm)?;
        let algorithm: SignatureAlgorithm = (id as u8).try_into().map_err(|_| BottleError::UnknownSignatureAlgorithm)?;
        let signer_name = cap.header.get_string(FIELD_SIGNER_NAME_STRING).unwrap_or("?").to_string();
        let public_key_bytes = cap.header.get_bytes(FIELD_PUBLIC_KEY_BYTES).ok_or(BottleError::NoSignaturePublicKey)?;
        let public_key = public_key_from_bytes(algorithm, public_key_bytes, signer_name)
            .ok_or(BottleError::CorruptPublicKey)?;
        let hashing = SignatureHashing::new(algorithm);

        bottle_reader.expect_next_stream(BottleStream::Data)?;
        Ok(SignedBottleReader { bottle_reader, hashing, algorithm, public_key })
    }

    fn finish_stream(&mut self) -> BottleResult<()> {
        let hash = self.hashing.finalize_reset();
        self.bottle_reader.close_stream()?;
        let header = read_header_stream(&mut self.bottle_reader)?;
        let signature = header.get_bytes(FIELD_SIGNATURE_BYTES).ok_or(BottleError::MissingHeader)?;
        let mut message = self.bottle_reader.bottle_cap.to_bytes();
        message.extend_from_slice(&hash);
        if !self.public_key.verify(&message, signature) {
            return Err(BottleError::BadSignature);
        }

        self.bottle_reader.expect_end_and_finish()
    }
}

impl<R: fmt::Debug + Read + 'static> StreamBottle for SignedBottleReader<R> {
    fn stream(self) -> ReadStreamRef {
        ReadStreamRef::new(Rc::new(RefCell::new(self)))
    }
}

impl<R: Read> Read for SignedBottleReader<R> {
    fn read(&mut self, buffer: &mut [u8]) -> std::io::Result<usize> {
        let len = self.bottle_reader.data_stream().map_err(|e| e.to_io_error())?.read(buffer)?;
        self.hashing.update(&buffer[..len]);
        Ok(len)
    }
}

impl<R: fmt::Debug + Read> ReadStream for SignedBottleReader<R> {
    fn finish(&mut self) -> BottleResult<()> {
        self.finish_stream()
    }
}


#[cfg(test)]
mod test {
    use std::io::{Read, Write};
    use crate::asymmetric::{Ed25519PublicKey, Ed25519SecretKey};
    use crate::bottle::BottleReader;
    use crate::bottle_cap::BottleType;
    use crate::BottlePublicKey;
    use super::{SignatureAlgorithm, SignedBottleReader, SignedBottleWriter};


    #[test]
    fn write_signature() {
        const PK_HEX: &str = "112905df7c79c726b6250ec6e87aaa35d9127e2f48736a65c7352e8768d77865";
        const SK_HEX: &str = "bc29877fa28b3e03ea6b58d35818e65a1fb4870dae35c4a452de880205108efb";
        let pk = Ed25519PublicKey::from_ssh_bytes(&hex::decode(PK_HEX).unwrap(), String::from("eeyore")).unwrap();
        let sk = Ed25519SecretKey::from_ssh_bytes(&hex::decode(SK_HEX).unwrap(), String::from("eeyore")).unwrap();

        let data = b"i'm daddy, and this is valentine's day!";
        let mut buffer = Vec::new();

        {
            let mut bottle_writer = SignedBottleWriter::new(&mut buffer, SignatureAlgorithm::SHA_512_ED25519, sk).unwrap();
            bottle_writer.write_all(data).unwrap();
            bottle_writer.close().unwrap();
        }

        let bottle_reader = BottleReader::new(&buffer[..]).unwrap();
        assert_eq!(bottle_reader.bottle_cap.bottle_type, BottleType::Signed);
        {
            let mut signed_bottle_reader = SignedBottleReader::new(bottle_reader).unwrap();
            assert_eq!(signed_bottle_reader.algorithm, SignatureAlgorithm::SHA_512_ED25519);
            assert_eq!(signed_bottle_reader.public_key.name(), "eeyore");
            assert_eq!(signed_bottle_reader.public_key.as_bytes(false), pk.as_bytes(false));

            let mut data_buffer = Vec::new();
            let len = signed_bottle_reader.read_to_end(&mut data_buffer).unwrap();
            // signature is ok?
            signed_bottle_reader.finish_stream().unwrap();

            assert_eq!(len, data.len());
            assert_eq!(&data_buffer[0 .. len], data);
        }
    }
}