//! The encrypted areas (see the vision's Encryption and Users notes), over
//! any file system. Each user's folder `users/<id>/` holds their areas:
//! `home/`, `permissions/`, `settings/` (init's record of their choices,
//! such as the theme), and per app `apps/<app>/` and `data/<app>/`,
//! each kept as `filecrypt` files under its own key, derived from that
//! user's master key. Only the signed-in user's areas open; anyone else's
//! are refused (and noise to their key anyway). The rest (programs,
//! `factory/`) is plaintext. Signed out, every area refuses everything, and
//! a plaintext file in one is refused too: nothing there is trusted unless
//! it was sealed.

use filecrypt::{CHUNK, Drbg, HEADER, Header, Key, OVERHEAD, PLAIN_CHUNK, area_key, chunk_offset, chunks, plain_len};

use crate::filesystem::{DirEntry, FileHandle, FileSystem};
use crate::slots::Slots;

/// Where a path falls.
#[derive(Debug, PartialEq, Eq)]
enum Area<'a> {
    Plain,
    /// Encrypted under the key of the area with this name.
    Sealed(&'a [u8]),
    /// `data/` itself, or an app's folder in it: folders, never files.
    Neither,
}

/// Where `path` falls, and whose it is: the user folder's name.
fn area(path: &[u8]) -> (Area<'_>, Option<&[u8]>) {
    let path = path.strip_prefix(b"/").unwrap_or(path);
    let mut parts = path.splitn(5, |b| *b == b'/');
    let (users, user, kind, app, rest) = (parts.next(), parts.next(), parts.next(), parts.next(), parts.next());
    if users != Some(b"users") {
        return (Area::Plain, None);
    }
    // `users/` itself holds only the folders, which anyone may list.
    let Some(user) = user else { return (Area::Plain, None) };
    if user.is_empty() {
        return (Area::Neither, None);
    }
    // `users/<id>/` up to and including `kind`, and the app's folder.
    let upto = |n: usize| &path[..n];
    let base = 6 + user.len() + 1;
    let area = match (kind, app, rest) {
        (Some(k @ (b"home" | b"permissions" | b"settings")), _, _) => Area::Sealed(upto(base + k.len())),
        (Some(k @ (b"apps" | b"data")), Some(a), Some(r)) if !a.is_empty() && !r.is_empty() => Area::Sealed(upto(base + k.len() + 1 + a.len())),
        _ => Area::Neither,
    };
    (area, Some(user))
}

/// An open encrypted file.
struct Open {
    key: Key,
    /// Plaintext bytes in it.
    size: usize,
    /// Where the next read or write starts, in plaintext.
    pos: usize,
}

pub struct Encrypted<F: FileSystem> {
    inner: F,
    /// The signed-in user's master key and folder name, and randomness.
    master: Option<Key>,
    user: [u8; USER_NAME],
    user_len: usize,
    rng: Option<Drbg>,
    /// Sealed files open, by the inner file system's handle.
    open: Slots<Open>,
}

/// A user folder's name: 16 hex digits, or `guest`.
pub const USER_NAME: usize = 16;

impl<F: FileSystem> Encrypted<F> {
    pub fn new(inner: F) -> Self {
        Self { inner, master: None, user: [0; USER_NAME], user_len: 0, rng: None, open: Slots::new() }
    }

    pub fn inner(&self) -> &F {
        &self.inner
    }

    #[cfg(test)]
    pub fn inner_mut_for_test(&mut self) -> &mut F {
        &mut self.inner
    }

    pub fn into_inner(self) -> F {
        self.inner
    }

    /// Whether `user`'s folder is the signed-in user's.
    fn signed_in(&self, user: Option<&[u8]>) -> bool {
        self.master.is_some() && user == Some(&self.user[..self.user_len])
    }

    /// The key of `path`'s area, if it is encrypted and its user signed in:
    /// `Ok(None)` for a plaintext path, `Err` for one it must refuse.
    fn key_for(&self, path: &[u8]) -> Result<Option<Key>, ()> {
        match area(path) {
            (Area::Plain, _) => Ok(None),
            (Area::Sealed(name), user) if self.signed_in(user) => Ok(self.master.as_ref().map(|m| area_key(m, name))),
            _ => Err(()),
        }
    }

    /// Whether `path` may be used as a directory (made, listed, removed):
    /// anywhere plaintext, and in the signed-in user's own folder.
    fn dir_allowed(&self, path: &[u8]) -> bool {
        match area(path) {
            (Area::Plain, _) => true,
            (_, user) => self.signed_in(user),
        }
    }

    /// Chunk `index` of a file of `size` bytes, opened into `out`; its length.
    fn read_chunk(inner: &mut F, h: FileHandle, key: &Key, size: usize, index: usize, out: &mut [u8; PLAIN_CHUNK]) -> Option<usize> {
        let last = index + 1 == chunks(size);
        let plain = if last { size - index * PLAIN_CHUNK } else { PLAIN_CHUNK };
        let mut sealed = [0u8; CHUNK];
        let sealed = &mut sealed[..plain + OVERHEAD];
        if !inner.seek(h, chunk_offset(index)) || inner.read(h, sealed)? != sealed.len() {
            return None;
        }
        filecrypt::open(key, index as u64, last, sealed, out)
    }

    fn write_chunk(&mut self, h: FileHandle, key: &Key, index: usize, last: bool, plain: &[u8]) -> bool {
        let Some(rng) = self.rng.as_mut() else { return false };
        let mut sealed = [0u8; CHUNK];
        let n = filecrypt::seal(key, index as u64, last, plain, &mut sealed, rng);
        self.inner.seek(h, chunk_offset(index)) && self.inner.write(h, &sealed[..n]) == Some(n)
    }
}

impl<F: FileSystem> FileSystem for Encrypted<F> {
    fn writable(&self) -> bool {
        self.inner.writable()
    }

    fn unlock(&mut self, master: &[u8; 32], seed: &[u8; 32], user: &[u8]) {
        if self.master.is_none() && !user.is_empty() && user.len() <= USER_NAME {
            self.master = Some(Key(*master));
            self.rng = Some(Drbg::new(*seed));
            self.user[..user.len()].copy_from_slice(user);
            self.user_len = user.len();
        }
    }

    fn lock(&mut self) {
        let open: alloc::vec::Vec<usize> = self.open.iter().map(|(i, _)| i).collect();
        for i in open {
            self.open.take(i);
            self.inner.close(FileHandle::new(i as u16));
        }
        self.master = None;
        self.rng = None;
        self.user = [0; USER_NAME];
        self.user_len = 0;
    }

    /// Removes a whole tree without opening anything in it: a removed
    /// user's or a guest's files go though no key for them is held.
    fn purge(&mut self, path: &[u8]) -> bool {
        crate::filesystem::remove_tree(&mut self.inner, path)
    }

    fn open(&mut self, path: &[u8]) -> Option<FileHandle> {
        let key = self.key_for(path).ok()?;
        let h = self.inner.open(path)?;
        let Some(area) = key else { return Some(h) };
        let mut head = [0u8; HEADER];
        let header = (self.inner.read(h, &mut head) == Some(HEADER)).then(|| Header::decode(&head)).flatten();
        let size = self.inner.file_size(h).and_then(plain_len);
        let (Some(header), Some(size)) = (header, size) else {
            // Not sealed (planted, or cut short): not one of ours.
            self.inner.close(h);
            return None;
        };
        let key = header.file_key(&area);
        self.open.put(h.raw() as usize, Open { key, size, pos: 0 });
        Some(h)
    }

    fn create(&mut self, path: &[u8]) -> Option<FileHandle> {
        let key = self.key_for(path).ok()?;
        if key.is_some() && self.rng.is_none() {
            return None;
        }
        let h = self.inner.create(path)?;
        let Some(area) = key else { return Some(h) };
        let header = Header::new(self.rng.as_mut()?);
        let file_key = header.file_key(&area);
        // A header and an empty last chunk: a whole, empty file.
        let ok = self.inner.write(h, &header.encode()) == Some(HEADER) && self.write_chunk(h, &file_key, 0, true, &[]);
        if !ok {
            self.inner.close(h);
            let _ = self.inner.delete(path);
            return None;
        }
        self.open.put(h.raw() as usize, Open { key: file_key, size: 0, pos: 0 });
        Some(h)
    }

    fn read(&mut self, h: FileHandle, buf: &mut [u8]) -> Option<usize> {
        let Some(f) = self.open.get(h.raw() as usize) else { return self.inner.read(h, buf) };
        let (key, size, mut pos) = (f.key.clone(), f.size, f.pos);
        let mut done = 0;
        let mut chunk = [0u8; PLAIN_CHUNK];
        while done < buf.len() && pos < size {
            let index = pos / PLAIN_CHUNK;
            let n = Self::read_chunk(&mut self.inner, h, &key, size, index, &mut chunk)?;
            let from = pos % PLAIN_CHUNK;
            let take = (n - from).min(buf.len() - done);
            buf[done..done + take].copy_from_slice(&chunk[from..from + take]);
            done += take;
            pos += take;
        }
        chunk.fill(0);
        self.open.get_mut(h.raw() as usize)?.pos = pos;
        Some(done)
    }

    fn write(&mut self, h: FileHandle, buf: &[u8]) -> Option<usize> {
        let Some(f) = self.open.get(h.raw() as usize) else { return self.inner.write(h, buf) };
        let (key, old_size, pos) = (f.key.clone(), f.size, f.pos);
        let new_size = old_size.max(pos + buf.len());
        let (old_last, new_last) = (chunks(old_size) - 1, chunks(new_size) - 1);
        let first = pos / PLAIN_CHUNK;
        // The old last chunk is no longer last: resealed so, even when the
        // write itself does not touch it.
        let resealed = (old_last < first && old_last != new_last).then_some(old_last);
        let mut chunk = [0u8; PLAIN_CHUNK];
        for index in resealed.into_iter().chain(first..=new_last) {
            let start = index * PLAIN_CHUNK;
            let mut len = match index <= old_last {
                true => Self::read_chunk(&mut self.inner, h, &key, old_size, index, &mut chunk)?,
                false => 0,
            };
            // The part of `buf` that falls in this chunk.
            let (from, to) = (pos.max(start), (pos + buf.len()).min(start + PLAIN_CHUNK));
            if from < to {
                chunk[from - start..to - start].copy_from_slice(&buf[from - pos..to - pos]);
                len = len.max(to - start);
            }
            if !self.write_chunk(h, &key, index, index == new_last, &chunk[..len]) {
                return None;
            }
        }
        chunk.fill(0);
        let f = self.open.get_mut(h.raw() as usize)?;
        f.size = new_size;
        f.pos = pos + buf.len();
        Some(buf.len())
    }

    fn seek(&mut self, h: FileHandle, offset: usize) -> bool {
        match self.open.get_mut(h.raw() as usize) {
            Some(f) if offset <= f.size => {
                f.pos = offset;
                true
            }
            Some(_) => false,
            None => self.inner.seek(h, offset),
        }
    }

    fn close(&mut self, h: FileHandle) {
        self.open.take(h.raw() as usize);
        self.inner.close(h);
    }

    fn file_size(&mut self, h: FileHandle) -> Option<usize> {
        match self.open.get(h.raw() as usize) {
            Some(f) => Some(f.size),
            None => self.inner.file_size(h),
        }
    }

    fn list_dir(&mut self, path: &[u8], entries: &mut [DirEntry]) -> Option<usize> {
        if !self.dir_allowed(path) {
            return None;
        }
        let n = self.inner.list_dir(path, entries)?;
        // In an encrypted area a file's size is what it holds, not its disk size.
        // In a user's folder, a file's size is what it holds.
        let sealed = area(path).0 != Area::Plain;
        for e in entries[..n].iter_mut().filter(|e| sealed && !e.is_dir) {
            e.size = plain_len(e.size).unwrap_or(0);
        }
        Some(n)
    }

    fn delete(&mut self, path: &[u8]) -> bool {
        self.key_for(path).is_ok() && self.inner.delete(path)
    }

    fn mkdir(&mut self, path: &[u8]) -> bool {
        self.dir_allowed(path) && self.inner.mkdir(path)
    }

    fn rmdir(&mut self, path: &[u8]) -> bool {
        self.dir_allowed(path) && self.inner.rmdir(path)
    }
}

#[cfg(test)]
pub(crate) fn area_of(path: &[u8]) -> &'static str {
    match area(path).0 {
        Area::Plain => "plain",
        Area::Sealed(_) => "sealed",
        Area::Neither => "neither",
    }
}