//! How a file in an encrypted area looks on disk (see the vision's
//! Encryption note): a header with the file's salt, then its contents in
//! chunks, each sealed with XChaCha20-Poly1305 under a fresh random nonce,
//! bound to its place in the file and to whether it is the last one.
//! Changing, swapping or cutting chunks off is detected; restoring an old
//! copy of a whole file is not. Pure: no I/O.
//!
//! Keys: the master key → an area's key (`area_key`, keyed BLAKE2s over
//! the area's name) → each file's key (over its salt).

#![no_std]
#![forbid(unsafe_code)]

mod drbg;

#[cfg(test)]
mod tests;

pub use drbg::Drbg;

use blake2::Blake2sMac256;
use blake2::digest::Mac;
use chacha20poly1305::aead::{AeadInPlace, KeyInit};
use chacha20poly1305::{Tag, XChaCha20Poly1305, XNonce};
use zeroize::Zeroize;

pub const KEY_SIZE: usize = 32;
const MAGIC: &[u8; 8] = b"KOSFILE1";
const SALT_SIZE: usize = 24;
/// The header: magic and salt.
pub const HEADER: usize = 8 + SALT_SIZE;
const NONCE: usize = 24;
const TAG: usize = 16;
/// What a chunk adds to what it holds.
pub const OVERHEAD: usize = NONCE + TAG;
/// A sealed chunk on disk, the last one possibly shorter.
pub const CHUNK: usize = 4096;
/// What a full chunk holds.
pub const PLAIN_CHUNK: usize = CHUNK - OVERHEAD;

/// A key, wiped from memory when dropped.
#[derive(Clone)]
pub struct Key(pub [u8; KEY_SIZE]);

impl Drop for Key {
    fn drop(&mut self) {
        self.0.zeroize();
    }
}

fn mac(key: &[u8; KEY_SIZE], parts: &[&[u8]]) -> Key {
    let mut m =
        <Blake2sMac256 as KeyInit>::new_from_slice(key).expect("BLAKE2s takes a 32-byte key");
    for p in parts {
        m.update(p);
    }
    Key(m.finalize().into_bytes().into())
}

/// The key of the area named `area` (`home`, `data/kos.notes`, ...):
/// each app's data is noise under another app's key.
pub fn area_key(master: &Key, area: &[u8]) -> Key {
    mac(&master.0, &[b"kos/storage/v1/area/", area])
}

/// Randomness for salts and nonces.
pub trait Random {
    fn fill(&mut self, buf: &mut [u8]);
}

/// A file's header: its salt, which with its area's key makes its key.
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub struct Header {
    salt: [u8; SALT_SIZE],
}

impl Header {
    pub fn new(rng: &mut impl Random) -> Self {
        let mut salt = [0u8; SALT_SIZE];
        rng.fill(&mut salt);
        Self { salt }
    }

    pub fn encode(&self) -> [u8; HEADER] {
        let mut out = [0u8; HEADER];
        out[..8].copy_from_slice(MAGIC);
        out[8..].copy_from_slice(&self.salt);
        out
    }

    /// `None` unless it begins an encrypted file (a plaintext file, say).
    pub fn decode(bytes: &[u8]) -> Option<Self> {
        let bytes = bytes.get(..HEADER)?;
        if &bytes[..8] != MAGIC {
            return None;
        }
        let mut salt = [0u8; SALT_SIZE];
        salt.copy_from_slice(&bytes[8..]);
        Some(Self { salt })
    }

    pub fn file_key(&self, area: &Key) -> Key {
        mac(&area.0, &[b"kos/storage/v1/file/", &self.salt])
    }
}

/// How many plaintext bytes a file holding `disk` bytes has; `None` if no
/// encrypted file is that long (a chunk cut short).
pub fn plain_len(disk: usize) -> Option<usize> {
    let body = disk.checked_sub(HEADER)?;
    let (full, rest) = (body / CHUNK, body % CHUNK);
    match rest {
        // The last chunk is full.
        0 if full > 0 => Some(full * PLAIN_CHUNK),
        r if r >= OVERHEAD => Some(full * PLAIN_CHUNK + r - OVERHEAD),
        _ => None,
    }
}

/// How many chunks hold `plain` bytes: always at least one, so even an
/// empty file has a last chunk to say it is complete.
pub fn chunks(plain: usize) -> usize {
    plain.div_ceil(PLAIN_CHUNK).max(1)
}

/// What a file holding `plain` bytes takes on disk.
pub fn disk_len(plain: usize) -> usize {
    HEADER + plain + chunks(plain) * OVERHEAD
}

/// Where chunk `index` starts on disk.
pub fn chunk_offset(index: usize) -> usize {
    HEADER + index * CHUNK
}

/// Binds a chunk to its place, and to being the last or not.
fn associated(index: u64, last: bool) -> [u8; 9] {
    let mut ad = [0u8; 9];
    ad[..8].copy_from_slice(&index.to_le_bytes());
    ad[8] = u8::from(last);
    ad
}

/// Seals `plain` (at most `PLAIN_CHUNK` bytes) as chunk `index` into `out`
/// (room for `plain.len() + OVERHEAD`); its length.
pub fn seal(
    key: &Key,
    index: u64,
    last: bool,
    plain: &[u8],
    out: &mut [u8],
    rng: &mut impl Random,
) -> usize {
    let n = plain.len().min(PLAIN_CHUNK);
    let (nonce, rest) = out.split_at_mut(NONCE);
    rng.fill(nonce);
    let (body, tag) = rest.split_at_mut(n);
    body.copy_from_slice(&plain[..n]);
    let t = XChaCha20Poly1305::new((&key.0).into())
        .encrypt_in_place_detached(XNonce::from_slice(nonce), &associated(index, last), body)
        .expect("a chunk is far below XChaCha20-Poly1305's limits");
    tag[..TAG].copy_from_slice(&t);
    n + OVERHEAD
}

/// Opens chunk `index` (`sealed` as read from disk) into `out`; its length,
/// or `None` if it is not that chunk of that file, unaltered.
pub fn open(key: &Key, index: u64, last: bool, sealed: &[u8], out: &mut [u8]) -> Option<usize> {
    let n = sealed.len().checked_sub(OVERHEAD)?;
    let (nonce, rest) = sealed.split_at(NONCE);
    let (body, tag) = rest.split_at(n);
    let out = out.get_mut(..n)?;
    out.copy_from_slice(body);
    let ok = XChaCha20Poly1305::new((&key.0).into())
        .decrypt_in_place_detached(
            XNonce::from_slice(nonce),
            &associated(index, last),
            out,
            Tag::from_slice(tag),
        )
        .is_ok();
    if !ok {
        out.zeroize();
    }
    ok.then_some(n)
}