//! Read-only disk ciphers and encrypted volume unlocking.
//!
//! An unlocked volume is exposed as a standard `Source`: the rest of Guará
//! (scanning, recovery) sees the plaintext filesystem transparently.
//!

pub mod bitlocker;
pub mod filevault;
pub mod luks;

use crate::source::Source;
use aes::cipher::{Block, BlockCipherDecrypt, BlockCipherEncrypt, KeyInit};
use aes::{Aes128, Aes192, Aes256};
use std::io;
use std::sync::Arc;

/// AES cipher with runtime key size selection.
#[derive(Clone)]
pub enum AesAny {
    A128(Aes128),
    A192(Aes192),
    A256(Aes256),
}

impl AesAny {
    pub fn new(key: &[u8]) -> Option<Self> {
        Some(match key.len() {
            16 => AesAny::A128(Aes128::new_from_slice(key).ok()?),
            24 => AesAny::A192(Aes192::new_from_slice(key).ok()?),
            32 => AesAny::A256(Aes256::new_from_slice(key).ok()?),
            _ => return None,
        })
    }
    pub fn encrypt(&self, b: &mut [u8; 16]) {
        let g: &mut Block<Aes128> = (&mut b[..]).try_into().expect("bloco AES de 16 bytes");
        match self {
            AesAny::A128(c) => c.encrypt_block(g),
            AesAny::A192(c) => c.encrypt_block(g),
            AesAny::A256(c) => c.encrypt_block(g),
        }
    }
    pub fn decrypt(&self, b: &mut [u8; 16]) {
        let g: &mut Block<Aes128> = (&mut b[..]).try_into().expect("bloco AES de 16 bytes");
        match self {
            AesAny::A128(c) => c.decrypt_block(g),
            AesAny::A192(c) => c.decrypt_block(g),
            AesAny::A256(c) => c.decrypt_block(g),
        }
    }
}

fn xor16(a: &mut [u8; 16], b: &[u8]) {
    for i in 0..16 {
        a[i] ^= b[i];
    }
}

/// Sector initialization vector (IV) generation mode.
#[derive(Clone)]
pub enum Iv {
    /// 512-byte sector number as 64-bit little-endian.
    Plain64,
    /// Sector number as 32-bit little-endian.
    Plain32,
    Plain64Be,
    /// ESSIV: AES(SHA-256(key), sector).
    Essiv(AesAny),
    /// Data unit number (`usize`-byte sector) as 64-bit little-endian.
    Unit(u64),
    /// BitLocker (Windows 10 1511+): AES(key, byte offset).
    Eboiv(AesAny),
    /// BitLocker CBC with Elephant diffuser: IV = AES(key, byte offset).
    Bitlocker(AesAny),
}

impl Iv {
    fn make(&self, sector_512: u64, byte_off: u64) -> [u8; 16] {
        let mut iv = [0u8; 16];
        match self {
            Iv::Plain64 => iv[..8].copy_from_slice(&sector_512.to_le_bytes()),
            Iv::Plain32 => iv[..4].copy_from_slice(&(sector_512 as u32).to_le_bytes()),
            Iv::Plain64Be => iv[8..].copy_from_slice(&sector_512.to_be_bytes()),
            Iv::Unit(u) => iv[..8].copy_from_slice(&(byte_off / u).to_le_bytes()),
            Iv::Essiv(k) => {
                iv[..8].copy_from_slice(&sector_512.to_le_bytes());
                k.encrypt(&mut iv);
            }
            Iv::Eboiv(k) | Iv::Bitlocker(k) => {
                iv[..8].copy_from_slice(&byte_off.to_le_bytes());
                k.encrypt(&mut iv);
            }
        }
        iv
    }
}

/// Sector-level disk cipher.
#[derive(Clone)]
pub enum SectorCipher {
    Xts {
        k1: AesAny,
        k2: AesAny,
        iv: Iv,
    },
    Cbc {
        key: AesAny,
        iv: Iv,
    },
    /// BitLocker AES-CBC with Elephant diffuser (data key + tweak key).
    Elephant {
        key: AesAny,
        tweak: AesAny,
    },
}

fn gf_double(t: &mut [u8; 16]) {
    let mut carry = 0u8;
    for b in t.iter_mut() {
        let next = *b >> 7;
        *b = (*b << 1) | carry;
        carry = next;
    }
    if carry != 0 {
        t[0] ^= 0x87;
    }
}

impl SectorCipher {
    /// Decrypts a sector in place. `sector_512` is the 512-byte sector index; `byte_off` is the
    /// byte offset relative to the start of the encrypted payload.
    pub fn decrypt_sector(&self, buf: &mut [u8], sector_512: u64, byte_off: u64) {
        match self {
            SectorCipher::Xts { k1, k2, iv } => {
                let mut t = iv.make(sector_512, byte_off);
                k2.encrypt(&mut t);
                for ch in buf.chunks_exact_mut(16) {
                    let mut b = [0u8; 16];
                    b.copy_from_slice(ch);
                    xor16(&mut b, &t);
                    k1.decrypt(&mut b);
                    xor16(&mut b, &t);
                    ch.copy_from_slice(&b);
                    gf_double(&mut t);
                }
            }
            SectorCipher::Cbc { key, iv } => {
                cbc_decrypt(key, &iv.make(sector_512, byte_off), buf);
            }
            SectorCipher::Elephant { key, tweak } => {
                bitlocker::elephant_decrypt(key, tweak, buf, byte_off);
            }
        }
    }
}

pub fn cbc_decrypt(key: &AesAny, iv: &[u8; 16], buf: &mut [u8]) {
    let mut prev = *iv;
    for ch in buf.chunks_exact_mut(16) {
        let mut c = [0u8; 16];
        c.copy_from_slice(ch);
        let mut b = c;
        key.decrypt(&mut b);
        xor16(&mut b, &prev);
        ch.copy_from_slice(&b);
        prev = c;
    }
}

pub fn cbc_encrypt(key: &AesAny, iv: &[u8; 16], buf: &mut [u8]) {
    let mut prev = *iv;
    for ch in buf.chunks_exact_mut(16) {
        let mut b = [0u8; 16];
        b.copy_from_slice(ch);
        xor16(&mut b, &prev);
        key.encrypt(&mut b);
        ch.copy_from_slice(&b);
        prev = b;
    }
}

/// Plaintext view of an encrypted volume: `len` bytes starting at `data_offset` in the source.
pub struct CryptSource {
    inner: Arc<dyn Source>,
    cipher: SectorCipher,
    data_offset: u64,
    len: u64,
    sector: u64,
    /// Sectors before this threshold are either unencrypted (partially encrypted BitLocker) or
    /// redirected to a backup copy (BitLocker relocated boot sectors).
    patches: Vec<Patch>,
    /// Bytes at or past this payload-relative offset are still unencrypted.
    plain_from: u64,
    /// Added to the sector index when computing the IV (LUKS2 `iv_tweak`).
    iv_tweak: u64,
}

/// Volume range redirected to another offset on the source (BitLocker relocated header sectors).
#[derive(Clone)]
pub struct Patch {
    /// Start offset in the logical plaintext volume.
    pub at: u64,
    pub len: u64,
    /// Source offset where the data resides.
    pub from: u64,
    /// Whether the source bytes must be decrypted (using cipher offset `from_crypt`).
    pub decrypt_from: Option<u64>,
    /// Returns zeros (BitLocker metadata regions masked out by the OS).
    pub zero: bool,
}

impl CryptSource {
    pub fn new(
        inner: Arc<dyn Source>,
        cipher: SectorCipher,
        data_offset: u64,
        len: u64,
        sector: u64,
    ) -> Self {
        let sector = sector.max(512);
        Self {
            inner,
            cipher,
            data_offset,
            len: len / sector * sector,
            sector,
            patches: vec![],
            plain_from: u64::MAX,
            iv_tweak: 0,
        }
    }
    pub fn with_patches(mut self, p: Vec<Patch>) -> Self {
        self.patches = p;
        self
    }
    pub fn with_iv_tweak(mut self, t: u64) -> Self {
        self.iv_tweak = t;
        self
    }
    pub fn with_plain_from(mut self, off: u64) -> Self {
        self.plain_from = off;
        self
    }

    fn read_sector_plain(&self, rel: u64, out: &mut [u8]) -> io::Result<()> {
        // `rel` is sector-aligned; `out.len() == sector`.
        for p in &self.patches {
            if rel >= p.at && rel < p.at + p.len {
                if p.zero {
                    out.fill(0);
                    return Ok(());
                }
                let src_off = p.from + (rel - p.at);
                crate::source::read_exact_at(self.inner.as_ref(), self.data_offset + src_off, out)?;
                if let Some(pos) = p.decrypt_from {
                    let r = pos + (rel - p.at);
                    self.cipher.decrypt_sector(out, r / 512, r);
                }
                return Ok(());
            }
        }
        crate::source::read_exact_at(self.inner.as_ref(), self.data_offset + rel, out)?;
        if rel < self.plain_from {
            self.cipher
                .decrypt_sector(out, rel / 512 + self.iv_tweak, rel);
        }
        Ok(())
    }
}

impl Source for CryptSource {
    fn size(&self) -> u64 {
        self.len
    }
    fn read_at(&self, offset: u64, buf: &mut [u8]) -> io::Result<usize> {
        if offset >= self.len || buf.is_empty() {
            return Ok(0);
        }
        let want = buf.len().min((self.len - offset) as usize);
        let s = self.sector;
        let mut done = 0usize;
        let mut tmp = vec![0u8; s as usize];
        while done < want {
            let pos = offset + done as u64;
            let rel = pos / s * s;
            let within = (pos - rel) as usize;
            let n = (s as usize - within).min(want - done);
            self.read_sector_plain(rel, &mut tmp)?;
            buf[done..done + n].copy_from_slice(&tmp[within..within + n]);
            done += n;
        }
        Ok(done)
    }
}

/// User-supplied credential for unlocking an encrypted volume.
#[derive(Clone)]
pub enum Secret {
    Password(String),
    /// BitLocker numeric recovery password (8 groups of 6 digits).
    RecoveryKey(String),
    /// Raw contents of a key file (BitLocker .BEK, LUKS keyfile).
    KeyFile(Vec<u8>),
    /// Known volume master key (FVEK / LUKS master key).
    VolumeKey(Vec<u8>),
}

impl std::fmt::Debug for Secret {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(match self {
            Secret::Password(_) => "Secret::Password(..)",
            Secret::RecoveryKey(_) => "Secret::RecoveryKey(..)",
            Secret::KeyFile(_) => "Secret::KeyFile(..)",
            Secret::VolumeKey(_) => "Secret::VolumeKey(..)",
        })
    }
}

#[derive(Debug)]
pub enum UnlockError {
    /// None of the supplied passwords or keys unlocked the volume.
    WrongSecret,
    Unsupported(String),
    Corrupt(String),
}

impl std::fmt::Display for UnlockError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            UnlockError::WrongSecret => write!(f, "senha ou chave incorreta"),
            UnlockError::Unsupported(s) => write!(f, "não suportado: {s}"),
            UnlockError::Corrupt(s) => write!(f, "cabeçalho inválido: {s}"),
        }
    }
}

impl From<UnlockError> for crate::Error {
    fn from(e: UnlockError) -> Self {
        crate::Error::msg(e.to_string())
    }
}

/// Unlocked encrypted volume.
pub struct Unlocked {
    pub source: Arc<dyn Source>,
    pub description: String,
    /// Volume master key (can be saved in the session to reopen without a password if permitted).
    pub volume_key: Option<Vec<u8>>,
}

pub type UnlockResult = std::result::Result<Unlocked, UnlockError>;

pub fn to_hex(b: &[u8]) -> String {
    b.iter().map(|x| format!("{x:02x}")).collect()
}

pub fn from_hex(s: &str) -> Option<Vec<u8>> {
    let s = s.trim();
    if s.len() % 2 != 0 {
        return None;
    }
    (0..s.len() / 2)
        .map(|i| u8::from_str_radix(&s[i * 2..i * 2 + 2], 16).ok())
        .collect()
}