penis 0.1.1

use crate::consts::INITIAL_STATE;
use alloc::vec::Vec;

/// Security Parameter
///
/// Controls the security level of the Penis Protocol by specifying the round number (n)
/// where hash computation is paused/resumed. This parameter is critical for balancing
/// security and functionality.
///
/// # Security Considerations
///
/// - n=40: Maximum security level (recommended for production)
/// - n=32: Minimum recommended for production
/// - n<32: NOT recommended for sensitive data
/// - n<16: UNSAFE, reveals unhashed data
/// - n≥48: Invalid, causes implementation errors
///
/// The best range for n is 24-40. Other values may weaken the protocol's security.
#[derive(Debug)]
pub struct SecsParam {
    /// Nth-round: indicates where to start off
    /// It's the position of round where the calculation paused/resumed
    pub(crate) n: u8,
}

impl SecsParam {
    /// Creates a new Security Parameter with given Nth-round
    ///
    /// # Arguments
    ///
    /// * `n` - The round number where hash computation will pause/resume
    ///
    /// # Panics
    ///
    /// Panics if n is not between 16 and 48
    ///
    /// # Security Notes
    ///
    /// - N must not be equal to 48 (bypasses schedule array recovery protection)
    /// - N must be less than 48 (prevents out of range errors)
    /// - N must be greater than 16 (prevents exposure of unhashed values)
    /// - Best security is achieved with n between 24-40
    #[inline(always)]
    pub const fn new(n: u8) -> Self {
        assert!(n >= 16 && n < 48, "n must be between 16 and 48");
        SecsParam { n }
    }
}

/// Data range
///
/// Specifies a range of bytes in the input data, typically used to mark
/// confidential portions that should be protected by the Penis Protocol.
///
/// # Requirements
///
/// - Ranges must not overlap
/// - Must be in ascending order
/// - End must not exceed data length
/// - Keep ranges as small as necessary
#[derive(Debug)]
pub struct DataRange(
    /// start of byte range
    usize,
    /// end of byte range (exclusive)
    usize,
);

impl DataRange {
    /// Creates a new DataRange
    ///
    /// # Arguments
    ///
    /// * `start` - Starting byte position
    /// * `end` - Ending byte position (exclusive)
    ///
    /// # Panics
    ///
    /// Panics if start is not less than end
    #[inline(always)]
    pub const fn new(start: usize, end: usize) -> Self {
        assert!(start < end, "start must be less than end");
        DataRange(start, end)
    }

    /// Returns the length of the DataRange
    #[inline(always)]
    pub const fn length(&self) -> usize {
        self.1 - self.0
    }

    /// Returns the start of the DataRange
    #[inline(always)]
    pub const fn start(&self) -> usize {
        self.0
    }

    /// Returns the end of the DataRange
    #[inline(always)]
    pub const fn end(&self) -> usize {
        self.1
    }
}

/// ToBeHashed: Input data and confidential ranges
///
/// Contains the data to be hashed along with specifications of which portions
/// should be kept confidential through the Penis Protocol.
///
/// # Security Requirements
///
/// - Confidential ranges must not overlap
/// - Ranges must be in ascending order
/// - Ranges must not exceed data length
/// - Validate ranges before processing
#[derive(Debug)]
pub struct ToBeHashed {
    /// Data string, that must be padded in advance.
    pub data: Vec<u8>,
    /// A list of confidential part. Never overlap and must be ascending.
    pub confidentials: Vec<DataRange>,
}

impl ToBeHashed {
    /// Checks the data ranges within the `confidentials` field of the `ToBeHashed` struct.
    ///
    /// This function iterates over the `confidentials` vector and ensures that each data range
    /// is valid and does not overlap with the previous one. It also ensures that the total
    /// range does not exceed the length of the `data` field.
    ///
    /// # Returns
    ///
    /// * `Ok(())` if all ranges are valid
    /// * `Err(&str)` with an error message if any validation fails
    pub fn check_datarange(&self) -> Result<(), &'static str> {
        let mut last_end = 0;
        for range in &self.confidentials {
            if range.start() < last_end {
                return Err("Invalid range: ranges must not overlap and must be ascending");
            }
            if range.end() > self.data.len() {
                return Err("Invalid range: end exceeds data length");
            }
            last_end = range.end();
        }
        Ok(())
    }

    /// Returns true if the given block range contains any confidential parts.
    ///
    /// This method checks if any confidential range overlaps with the specified block range.
    ///
    /// # Arguments
    ///
    /// * `start` - Start of the block range
    /// * `end` - End of the block range
    ///
    /// # Returns
    ///
    /// true if the block range contains confidential data, false otherwise
    #[inline(always)]
    pub fn is_in_confidentials(&self, start: usize, end: usize) -> bool {
        self.confidentials.iter().any(|range| {
            (range.start() < end && range.end() > start)
                || (range.start() >= start && range.start() < end)
        })
    }
}

/// Represents either a private (confidential) or public block in the hash computation
#[cfg_attr(test, derive(Debug, PartialEq))]
pub enum PenisBlock {
    /// Private block containing intermediate state and partial schedule array
    Private(PrivateBlock),
    /// Public block containing raw data
    Public(RawBlock),
}

impl PenisBlock {
    /// Returns the type flag for the block (1 for private, 0 for public)
    #[inline(always)]
    pub(crate) const fn type_flag(&self) -> u8 {
        match self {
            PenisBlock::Private(_) => 1,
            PenisBlock::Public(_) => 0,
        }
    }
    /// Returns the size of the block in bytes
    #[inline(always)]
    pub(crate) const fn encoded_size(&self) -> usize {
        match self {
            PenisBlock::Private(_) => 0x60,
            PenisBlock::Public(_) => 0x40,
        }
    }
}

/// SchedVariable: An element of Scheduling Variable
pub(crate) type SchedVariable = u32;
/// SchedArray: An array of Scheduling Variable
pub(crate) type SchedArray = [SchedVariable; 64];
/// PartialSchedArray: A minimum subset of SchedArray
pub(crate) type PartialSchedArray = [SchedVariable; 16];

/// Raw block type for SHA-256 (64 bytes)
pub(crate) type RawBlock = [u8; 64];
/// Reference to a raw block
pub(crate) type RawBlockRef<'a> = &'a [u8];

/// Type alias for round numbers
pub(crate) type NRound = usize;

/// Private block containing intermediate state and partial schedule array
///
/// This structure represents a block of confidential data in its partially
/// processed form, containing both the intermediate state and the partial
/// schedule array needed to complete the hash computation.
#[derive(PartialEq, Debug)]
pub struct PrivateBlock {
    /// Partially Extracted array of scheduled variable
    /// It does not contain the raw value. It will be used to complete the scheduled variable
    pub psa: PartialSchedArray,
    /// Intermediate State: partially executed state variables
    /// It's SHA-256 state variable in the middle.
    pub state: State,
}

/// SHA-256 internal state
///
/// Represents the eight 32-bit values that make up the internal state
/// of the SHA-256 hash function during computation.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct State {
    pub a: u32,
    pub b: u32,
    pub c: u32,
    pub d: u32,
    pub e: u32,
    pub f: u32,
    pub g: u32,
    pub h: u32,
}

impl State {
    /// Converts the state into its final 32-byte hash representation
    ///
    /// # Returns
    ///
    /// A 32-byte array containing the final hash value
    #[inline(always)]
    pub fn finalize(&self) -> [u8; 32] {
        let mut final_hash = [0u8; 32];

        final_hash[0..4].clone_from_slice(&self.a.to_be_bytes());
        final_hash[4..8].clone_from_slice(&self.b.to_be_bytes());
        final_hash[8..12].clone_from_slice(&self.c.to_be_bytes());
        final_hash[12..16].clone_from_slice(&self.d.to_be_bytes());
        final_hash[16..20].clone_from_slice(&self.e.to_be_bytes());
        final_hash[20..24].clone_from_slice(&self.f.to_be_bytes());
        final_hash[24..28].clone_from_slice(&self.g.to_be_bytes());
        final_hash[28..32].clone_from_slice(&self.h.to_be_bytes());

        final_hash
    }

    /// Converts the state into an array of eight 32-bit values
    ///
    /// # Returns
    ///
    /// An array containing the eight state values
    #[inline(always)]
    pub fn finalize_u32(&self) -> [u32; 8] {
        [
            self.a, self.b, self.c, self.d, self.e, self.f, self.g, self.h,
        ]
    }

    /// Creates a new State from a 32-byte array
    ///
    /// # Arguments
    ///
    /// * `bytes` - A 32-byte array containing the state values
    ///
    /// # Panics
    ///
    /// Panics if the input array is not exactly 32 bytes
    pub fn from_bytes(bytes: &[u8]) -> Self {
        assert_eq!(bytes.len(), 32);
        State {
            a: u32::from_be_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]),
            b: u32::from_be_bytes([bytes[4], bytes[5], bytes[6], bytes[7]]),
            c: u32::from_be_bytes([bytes[8], bytes[9], bytes[10], bytes[11]]),
            d: u32::from_be_bytes([bytes[12], bytes[13], bytes[14], bytes[15]]),
            e: u32::from_be_bytes([bytes[16], bytes[17], bytes[18], bytes[19]]),
            f: u32::from_be_bytes([bytes[20], bytes[21], bytes[22], bytes[23]]),
            g: u32::from_be_bytes([bytes[24], bytes[25], bytes[26], bytes[27]]),
            h: u32::from_be_bytes([bytes[28], bytes[29], bytes[30], bytes[31]]),
        }
    }
}

impl Default for State {
    #[inline(always)]
    fn default() -> Self {
        INITIAL_STATE
    }
}

#[cfg(test)]
mod tests {
    use alloc::vec;

    use super::*;

    #[test]
    fn test_is_in_confidentials() {
        // First set of confidentials
        let to_be_hashed_1 = ToBeHashed {
            data: vec![],
            confidentials: vec![
                DataRange::new(0, 10),
                DataRange::new(20, 30),
                DataRange::new(40, 50),
            ],
        };

        assert!(to_be_hashed_1.is_in_confidentials(5, 15)); // Overlaps with first range
        assert!(!to_be_hashed_1.is_in_confidentials(10, 20)); // Between first and second range
        assert!(to_be_hashed_1.is_in_confidentials(0, 10)); // Exact match with first range
        assert!(to_be_hashed_1.is_in_confidentials(9, 20)); // Overlaps end of first range

        // Second set of confidentials
        let to_be_hashed_2 = ToBeHashed {
            data: vec![],
            confidentials: vec![
                DataRange::new(5, 15),
                DataRange::new(25, 35),
                DataRange::new(45, 55),
            ],
        };

        assert!(to_be_hashed_2.is_in_confidentials(10, 20)); // Overlaps with first range
        assert!(!to_be_hashed_2.is_in_confidentials(20, 25)); // Between first and second range
        assert!(to_be_hashed_2.is_in_confidentials(5, 15)); // Exact match with first range
        assert!(to_be_hashed_2.is_in_confidentials(14, 25)); // Overlaps end of first range

        // Third set of confidentials
        let to_be_hashed_3 = ToBeHashed {
            data: vec![],
            confidentials: vec![
                DataRange::new(0, 5),
                DataRange::new(10, 15),
                DataRange::new(20, 25),
            ],
        };

        assert!(to_be_hashed_3.is_in_confidentials(3, 12)); // Overlaps with first and second ranges
        assert!(!to_be_hashed_3.is_in_confidentials(5, 10)); // Between first and second range
        assert!(to_be_hashed_3.is_in_confidentials(0, 5)); // Exact match with first range
        assert!(to_be_hashed_3.is_in_confidentials(4, 10)); // Overlaps end of first range
    }
}