//! What the gatekeeper does, over its three ports: the keystore device,
//! randomness, and the owner (`Prompt`). No kernel, no hardware: tested on
//! the host.

use keystore::{Account, Block, Copies, KdfParams, Key, Name, Outcome, RECORD_SIZE, Random, Record, Refused, Role, SLOTS, Session, User, UserId};
use zeroize::Zeroize;

/// The keystore device: the record's two copies.
pub trait Keystore {
    fn read(&mut self, copy: usize, out: &mut [u8; RECORD_SIZE]) -> bool;
    fn write(&mut self, copy: usize, data: &[u8; RECORD_SIZE]) -> bool;
}

/// A passcode as typed, wiped from memory when dropped.
pub struct Secret {
    bytes: [u8; MAX_SECRET],
    len: usize,
}

pub const MAX_SECRET: usize = 64;

impl Secret {
    pub const fn new() -> Self {
        Self { bytes: [0; MAX_SECRET], len: 0 }
    }

    pub fn set(&mut self, typed: &[u8]) {
        self.bytes.zeroize();
        let n = typed.len().min(MAX_SECRET);
        self.bytes[..n].copy_from_slice(&typed[..n]);
        self.len = n;
    }

    pub fn as_bytes(&self) -> &[u8] {
        &self.bytes[..self.len]
    }

    /// The same bytes, compared without stopping at the first difference.
    fn same(&self, other: &Secret) -> bool {
        let diff = self.bytes.iter().zip(&other.bytes).fold(0u8, |d, (a, b)| d | (a ^ b));
        (diff | (self.len ^ other.len) as u8) == 0 && self.len == other.len
    }
}

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

impl Drop for Secret {
    fn drop(&mut self) {
        self.bytes.zeroize();
    }
}

/// The owner, on the screen and keyboard.
pub trait Prompt {
    fn say(&mut self, text: &str);
    /// Blanks the screen.
    fn clear(&mut self);
    /// Shows `question` and reads a passcode without showing it; false if
    /// nothing could be read.
    fn secret(&mut self, question: &str, out: &mut Secret) -> bool;
    /// Shows `question` and reads an ordinary line (a name, a choice).
    fn line(&mut self, question: &str, out: &mut Secret) -> bool;
}

/// Why the gatekeeper could not hand over a key.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Failure {
    /// The keystore could not be read or written.
    Device,
    /// No randomness for keys and salts.
    Random,
    /// The owner could not be asked.
    Prompt,
}

/// Wrong passcodes in a row (for any name) that wipe every user.
pub const THRESHOLD: u8 = 10;
/// Shortest passcode accepted.
pub const MIN_PASSCODE: usize = 4;

/// Someone is in.
pub struct Unlocked {
    /// Who, or `None` for a guest (whose key is made up for this session).
    pub user: Option<User>,
    pub key: Key,
    /// The keystore is new (first boot, or just wiped): nothing sealed
    /// under an earlier key exists any more.
    pub fresh: bool,
    /// Users removed just now, whose files are to go.
    pub removed: [Option<UserId>; SLOTS],
}

// Each step keeps its records and buffers (KiBs) to itself: inlined into
// `run` together they would overflow the 16 KiB stack init gives a program.

/// The keystore's newest record, and where the next write goes.
#[inline(never)]
fn load(ks: &mut impl Keystore) -> Result<(Option<Record>, usize), Failure> {
    let (mut a, mut b) = (alloc::boxed::Box::new([0u8; RECORD_SIZE]), alloc::boxed::Box::new([0u8; RECORD_SIZE]));
    if !ks.read(0, &mut a) || !ks.read(1, &mut b) {
        return Err(Failure::Device);
    }
    let copies = Copies { copies: [&a[..], &b[..]] };
    Ok((copies.newest().map(|(r, _)| r), copies.next_copy()))
}

#[inline(never)]
fn store(ks: &mut impl Keystore, copy: usize, r: &Record) -> Result<(), Failure> {
    if ks.write(copy, &r.encode()) { Ok(()) } else { Err(Failure::Device) }
}

/// `r` written as the newest record.
fn save(ks: &mut impl Keystore, r: &Record) -> Result<(), Failure> {
    let (_, next) = load(ks)?;
    store(ks, next, r)
}

/// Destroys every slot and the directory: noise over both copies, twice,
/// then blank (the device then looks never set up). Marked as wiping
/// first, so a power cut part way is finished on the next boot.
#[inline(never)]
fn wipe(ks: &mut impl Keystore, record: &Record, next: usize, rng: &mut impl Random) -> Result<(), Failure> {
    let first = record.wiping(rng).ok_or(Failure::Random)?;
    let second = first.wiping(rng).ok_or(Failure::Random)?;
    for r in [&first, &second] {
        store(ks, next, r)?;
        store(ks, 1 - next, r)?;
    }
    let blank = [0u8; RECORD_SIZE];
    if ks.write(next, &blank) && ks.write(1 - next, &blank) { Ok(()) } else { Err(Failure::Device) }
}

/// Asks for a new passcode twice until both match; `optional` allows none.
fn choose(prompt: &mut impl Prompt, question: &str, optional: bool, out: &mut Secret) -> Result<(), Failure> {
    loop {
        let mut again = Secret::new();
        if !prompt.secret(question, out) || !prompt.secret("Again: ", &mut again) {
            return Err(Failure::Prompt);
        }
        if !out.same(&again) {
            prompt.say("They differ; once more.\n");
        } else if out.len == 0 && optional {
            return Ok(());
        } else if out.len < MIN_PASSCODE {
            prompt.say("At least 4 characters.\n");
        } else {
            return Ok(());
        }
    }
}

/// A name, asked until it is a valid one.
fn ask_name(prompt: &mut impl Prompt, question: &str) -> Result<Name, Failure> {
    loop {
        let mut typed = Secret::new();
        if !prompt.line(question, &mut typed) {
            return Err(Failure::Prompt);
        }
        match Name::new(typed.as_bytes()) {
            Some(n) => return Ok(n),
            None => prompt.say("A name is 1 to 16 letters, digits or signs, no spaces.\n"),
        }
    }
}

/// A passcode and an optional duress passcode, different from it.
fn ask_passcodes(prompt: &mut impl Prompt, real: &mut Secret, duress: &mut Secret) -> Result<(), Failure> {
    choose(prompt, "New passcode: ", false, real)?;
    prompt.say("A duress passcode wipes the device when it is entered instead. Leave it empty for none.\n");
    loop {
        choose(prompt, "Duress passcode: ", true, duress)?;
        if duress.len == 0 || !duress.same(real) {
            return Ok(());
        }
        prompt.say("It must differ from the passcode.\n");
    }
}

/// A device never set up (or wiped): its first user, an owner.
#[inline(never)]
fn setup(ks: &mut impl Keystore, rng: &mut impl Random, prompt: &mut impl Prompt, params: KdfParams, memory: &mut [Block]) -> Result<Unlocked, Failure> {
    // After a wipe nothing earlier may show: it must look like a new device.
    prompt.clear();
    prompt.say("This device has no users yet. Set up its owner; the passcode protects everything on it.\n");
    let name = ask_name(prompt, "Your name: ")?;
    let (mut real, mut duress) = (Secret::new(), Secret::new());
    ask_passcodes(prompt, &mut real, &mut duress)?;
    let duress = (duress.len > 0).then_some(duress.as_bytes());
    let (record, session) = Record::setup(name, real.as_bytes(), duress, THRESHOLD, params, memory, rng).ok_or(Failure::Random)?;
    store(ks, 0, &record)?;
    store(ks, 1, &record)?;
    Ok(Unlocked { user: Some(session.user), key: session.key.clone(), fresh: true, removed: [None; SLOTS] })
}

/// What the owner may change at the lock screen, after signing in.
#[inline(never)]
fn manage(ks: &mut impl Keystore, rng: &mut impl Random, prompt: &mut impl Prompt, memory: &mut [Block], session: &Session, removed: &mut [Option<UserId>; SLOTS]) -> Result<(), Failure> {
    loop {
        let (Some(record), _) = load(ks)? else { return Err(Failure::Device) };
        let users = record.users(session).map_err(|_| Failure::Device)?;
        prompt.say("Users:");
        for u in users.iter().flatten() {
            prompt.say(" ");
            prompt.say(core::str::from_utf8(u.name.as_bytes()).unwrap_or("?"));
            if u.role == Role::Owner {
                prompt.say(" (owner)");
            }
        }
        prompt.say(if record.guests { ". Guests: allowed.\n" } else { ". Guests: not allowed.\n" });
        let mut choice = Secret::new();
        if !prompt.line("[a] add, [r] remove, [g] guests on/off, [Enter] done: ", &mut choice) {
            return Err(Failure::Prompt);
        }
        let changed = match choice.as_bytes() {
            b"" => return Ok(()),
            b"g" => Ok(record.next(|r| r.guests = !r.guests)),
            b"a" => {
                let name = ask_name(prompt, "Name: ")?;
                let mut role = Secret::new();
                if !prompt.line("Owner too? [y/N] ", &mut role) {
                    return Err(Failure::Prompt);
                }
                let role = if role.as_bytes() == b"y" { Role::Owner } else { Role::Standard };
                let (mut real, mut duress) = (Secret::new(), Secret::new());
                ask_passcodes(prompt, &mut real, &mut duress)?;
                let duress = (duress.len > 0).then_some(duress.as_bytes());
                record.add_user(session, Account { name, role, passcode: real.as_bytes(), duress }, memory, rng).map(|(r, _)| r)
            }
            b"r" => {
                let name = ask_name(prompt, "Remove whom: ")?;
                record.remove_user(session, &name, rng).map(|(r, id)| {
                    if let Some(slot) = removed.iter_mut().find(|s| s.is_none()) {
                        *slot = Some(id);
                    }
                    r
                })
            }
            _ => continue,
        };
        match changed {
            Ok(r) => save(ks, &r)?,
            Err(Refused::NameTaken) => prompt.say("That name is taken.\n"),
            Err(Refused::Full) => prompt.say("No room for another user.\n"),
            Err(Refused::NoSuchUser) => prompt.say("No such user.\n"),
            Err(Refused::LastOwner) => prompt.say("The device needs an owner.\n"),
            Err(_) => return Err(Failure::Random),
        }
    }
}

/// Authenticates someone and returns their key: setting up a fresh
/// device's owner, wiping everyone on a duress passcode or too many
/// failures (the device then looks fresh, and is set up anew).
pub fn run(ks: &mut impl Keystore, rng: &mut impl Random, prompt: &mut impl Prompt, params: KdfParams, memory: &mut [Block]) -> Result<Unlocked, Failure> {
    loop {
        let (record, next) = load(ks)?;
        let Some(record) = record else {
            return setup(ks, rng, prompt, params, memory);
        };
        if record.must_wipe() {
            wipe(ks, &record, next, rng)?;
            continue;
        }
        let mut name = Secret::new();
        if !prompt.line("Name: ", &mut name) {
            return Err(Failure::Prompt);
        }
        if name.len == 0 {
            if !record.guests {
                prompt.say("Guests are not allowed here.\n");
                continue;
            }
            // A key nobody keeps: what a guest leaves behind is noise.
            let mut key = Key([0; keystore::KEY_SIZE]);
            if !rng.fill(&mut key.0) {
                return Err(Failure::Random);
            }
            return Ok(Unlocked { user: None, key, fresh: false, removed: [None; SLOTS] });
        }
        let mut typed = Secret::new();
        if !prompt.secret("Passcode: ", &mut typed) {
            return Err(Failure::Prompt);
        }
        // Counted before it is checked: pulling the power cannot save a guess.
        let counted = record.next(|r| r.failures = r.failures.saturating_add(1));
        store(ks, next, &counted)?;
        match record.unlock(name.as_bytes(), typed.as_bytes(), memory) {
            Outcome::User(session) => {
                store(ks, 1 - next, &counted.next(|r| r.failures = 0))?;
                let mut removed = [None; SLOTS];
                loop {
                    let mut choice = Secret::new();
                    let menu = if session.is_owner() { "[Enter] continue, [p] passcode, [u] users: " } else { "[Enter] continue, [p] passcode: " };
                    if !prompt.line(menu, &mut choice) {
                        return Err(Failure::Prompt);
                    }
                    match choice.as_bytes() {
                        b"" => break,
                        b"p" => {
                            let mut new = Secret::new();
                            choose(prompt, "New passcode: ", false, &mut new)?;
                            let (Some(now), _) = load(ks)? else { return Err(Failure::Device) };
                            let changed = now.change_passcode(&session, new.as_bytes(), memory, rng).map_err(|_| Failure::Random)?;
                            save(ks, &changed)?;
                            prompt.say("Passcode changed.\n");
                        }
                        b"u" if session.is_owner() => manage(ks, rng, prompt, memory, &session, &mut removed)?,
                        _ => {}
                    }
                }
                return Ok(Unlocked { user: Some(session.user), key: session.key.clone(), fresh: false, removed });
            }
            Outcome::Duress => {
                // Recorded as worn out first: a wipe cut short resumes.
                store(ks, 1 - next, &counted.next(|r| r.failures = r.threshold))?;
                let (worn, next) = load(ks)?;
                wipe(ks, &worn.ok_or(Failure::Device)?, next, rng)?;
            }
            Outcome::Wrong => {
                let left = counted.threshold.saturating_sub(counted.failures);
                if left > 0 {
                    prompt.say(if left == 1 { "Wrong name or passcode. One try left before the device is wiped.\n" } else { "Wrong name or passcode.\n" });
                }
            }
        }
    }
}