btc-keygen 0.2.0

use crate::Error;
use crate::entropy::{EntropyError, EntropySource};
use secp256k1::SecretKey;
use zeroize::{Zeroize, ZeroizeOnDrop};

/// A validated secp256k1 private key that zeroizes its bytes on drop.
///
/// Created by [`generate`](crate::generate). The key is guaranteed to be a
/// valid scalar in the range `[1, n-1]` where `n` is the secp256k1 curve order.
///
/// When this value goes out of scope, the underlying bytes are securely
/// overwritten with zeros to prevent secrets from lingering in memory.
#[derive(Zeroize, ZeroizeOnDrop)]
pub struct PrivateKey {
    bytes: [u8; 32],
}

impl PrivateKey {
    /// Returns a reference to the raw 32-byte private key.
    pub fn as_bytes(&self) -> &[u8; 32] {
        &self.bytes
    }

    /// Converts the private key into a [`secp256k1::SecretKey`] for use with
    /// the `secp256k1` crate directly.
    pub fn to_secret_key(&self) -> SecretKey {
        SecretKey::from_byte_array(self.bytes).expect("PrivateKey always holds a validated scalar")
    }

    /// Creates a `PrivateKey` from 32 raw bytes, validating that they form a
    /// valid secp256k1 scalar.
    ///
    /// Use this when you have your own source of private key material (for
    /// example, physical entropy like dice rolls converted to hex) and want
    /// to skip OS entropy generation.
    ///
    /// # Errors
    ///
    /// Returns [`Error`](crate::Error) if `bytes` is zero or greater than or
    /// equal to the secp256k1 curve order `n`.
    ///
    /// # Example
    ///
    /// ```
    /// let mut bytes = [0u8; 32];
    /// bytes[31] = 0x01;
    /// let key = btc_keygen::PrivateKey::from_bytes(bytes)?;
    /// # Ok::<(), btc_keygen::Error>(())
    /// ```
    pub fn from_bytes(bytes: [u8; 32]) -> Result<PrivateKey, Error> {
        if !is_valid_key(bytes) {
            return Err(Error("not a valid secp256k1 scalar".into()));
        }

        Ok(PrivateKey { bytes })
    }

    /// Creates a `PrivateKey` from a 64-character hexadecimal string,
    /// validating that the decoded bytes form a valid secp256k1 scalar.
    ///
    /// Convenience wrapper around [`from_bytes`](Self::from_bytes) for callers
    /// that have the key material as a hex string (for example, from a CLI
    /// argument or a text file).
    ///
    /// # Errors
    ///
    /// Returns [`Error`](crate::Error) if:
    ///
    /// - `hex` is not exactly 64 characters long.
    /// - `hex` contains a character that is not a valid hexadecimal digit.
    /// - The decoded bytes are zero or greater than or equal to the secp256k1
    ///   curve order `n`.
    ///
    /// # Example
    ///
    /// ```
    /// let hex = "0000000000000000000000000000000000000000000000000000000000000001";
    /// let key = btc_keygen::PrivateKey::from_hex(hex)?;
    /// # Ok::<(), btc_keygen::Error>(())
    /// ```
    pub fn from_hex(hex: &str) -> Result<PrivateKey, Error> {
        if hex.len() != 64 {
            return Err(Error(format!(
                "expected 64 hex characters, got {}",
                hex.len()
            )));
        }

        let mut bytes = [0u8; 32];
        for i in 0..32 {
            bytes[i] = u8::from_str_radix(&hex[i * 2..i * 2 + 2], 16)
                .map_err(|_| Error(format!("invalid hex at position {}", i * 2)))?;
        }
        Self::from_bytes(bytes)
    }
}

/// Checks whether 32 bytes represent a valid secp256k1 private key.
///
/// A valid key is a scalar in `[1, n-1]` where `n` is the curve order.
pub fn is_valid_key(bytes: [u8; 32]) -> bool {
    SecretKey::from_byte_array(bytes).is_ok()
}

/// Generates a new private key using the provided entropy source.
///
/// Retries up to `MAX_RETRIES` times if the random bytes fall outside the
/// valid secp256k1 scalar range. This is astronomically unlikely but handled
/// for correctness.
pub(crate) fn generate_with_entropy(
    entropy: &dyn EntropySource,
) -> Result<PrivateKey, EntropyError> {
    for _ in 0..MAX_RETRIES {
        let mut bytes = [0u8; 32];
        entropy.fill_bytes(&mut bytes)?;

        if is_valid_key(bytes) {
            return Ok(PrivateKey { bytes });
        }
        // Invalid scalar: zeroize and retry.
        bytes.zeroize();
    }

    Err(EntropyError(
        "failed to generate valid key after maximum retries".into(),
    ))
}

/// Generates a new Bitcoin private key using OS-provided cryptographic randomness.
///
/// Returns a [`PrivateKey`] containing a validated secp256k1 scalar. The
/// entropy comes from the operating system's CSPRNG (`getrandom` syscall on
/// Linux, `getentropy` on macOS, `BCryptGenRandom` on Windows).
///
/// # Errors
///
/// Returns [`Error`](crate::Error) if the OS random number generator fails.
///
/// # Example
///
/// ```no_run
/// let key = btc_keygen::generate().expect("key generation failed");
/// ```
pub fn generate() -> Result<PrivateKey, crate::Error> {
    generate_with_entropy(&crate::entropy::OsEntropy).map_err(crate::Error::from)
}

/// Maximum retry attempts for key generation. A safety net against infinite
/// loops: the probability of needing even one retry is ~10^-38.
const MAX_RETRIES: u32 = 32;

#[cfg(test)]
mod tests {
    use super::*;
    use crate::entropy::{FailingEntropy, FixedEntropy};

    /// secp256k1 curve order n.
    const CURVE_ORDER: [u8; 32] = [
        0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
        0xFE, 0xBA, 0xAE, 0xDC, 0xE6, 0xAF, 0x48, 0xA0, 0x3B, 0xBF, 0xD2, 0x5E, 0x8C, 0xD0, 0x36,
        0x41, 0x41,
    ];

    /// n - 1: the maximum valid private key.
    fn curve_order_minus_one() -> [u8; 32] {
        let mut bytes = CURVE_ORDER;
        bytes[31] -= 1;
        bytes
    }

    /// n + 1: one above the curve order.
    fn curve_order_plus_one() -> [u8; 32] {
        let mut bytes = CURVE_ORDER;
        bytes[31] += 1;
        bytes
    }

    // ---------------------------------------------------------------
    // 6.1: Private key boundary validation
    // ---------------------------------------------------------------

    #[test]
    fn test_zero_key_rejected() {
        let zero = [0u8; 32];
        assert!(!is_valid_key(zero), "zero must not be a valid private key");
    }

    #[test]
    fn test_one_key_valid() {
        let mut one = [0u8; 32];
        one[31] = 1;
        assert!(is_valid_key(one), "scalar 1 must be a valid private key");
    }

    #[test]
    fn test_curve_order_minus_one_valid() {
        let n_minus_1 = curve_order_minus_one();
        assert!(
            is_valid_key(n_minus_1),
            "n-1 must be a valid private key (maximum scalar)"
        );
    }

    #[test]
    fn test_curve_order_rejected() {
        assert!(
            !is_valid_key(CURVE_ORDER),
            "the curve order n itself must not be a valid private key"
        );
    }

    #[test]
    fn test_curve_order_plus_one_rejected() {
        let n_plus_1 = curve_order_plus_one();
        assert!(
            !is_valid_key(n_plus_1),
            "n+1 must not be a valid private key"
        );
    }

    #[test]
    fn test_all_ff_rejected() {
        let all_ff = [0xFF; 32];
        assert!(
            !is_valid_key(all_ff),
            "all 0xFF bytes exceed curve order and must be rejected"
        );
    }

    #[test]
    fn test_valid_midrange_key() {
        // A known midrange value well within [1, n-1].
        let mut key = [0u8; 32];
        key[0] = 0x0A;
        key[31] = 0x0B;
        assert!(is_valid_key(key));
    }

    // ---------------------------------------------------------------
    // 6.2: Deterministic key generation with injectable entropy
    // ---------------------------------------------------------------

    #[test]
    fn test_fixed_entropy_produces_expected_key() {
        let mut key_bytes = [0u8; 32];
        key_bytes[31] = 0x01; // scalar = 1, valid
        let entropy = FixedEntropy::new(key_bytes.to_vec());

        let key = generate_with_entropy(&entropy).expect("generation should succeed");
        assert_eq!(key.as_bytes(), &key_bytes);
    }

    #[test]
    fn test_different_entropy_produces_different_keys() {
        let mut bytes_a = [0u8; 32];
        bytes_a[31] = 0x01;
        let mut bytes_b = [0u8; 32];
        bytes_b[31] = 0x02;

        let key_a = generate_with_entropy(&FixedEntropy::new(bytes_a.to_vec())).unwrap();
        let key_b = generate_with_entropy(&FixedEntropy::new(bytes_b.to_vec())).unwrap();

        assert_ne!(key_a.as_bytes(), key_b.as_bytes());
    }

    #[test]
    fn test_same_entropy_produces_same_key() {
        let mut key_bytes = [0u8; 32];
        key_bytes[31] = 0x05;

        let key1 = generate_with_entropy(&FixedEntropy::new(key_bytes.to_vec())).unwrap();
        let key2 = generate_with_entropy(&FixedEntropy::new(key_bytes.to_vec())).unwrap();

        assert_eq!(key1.as_bytes(), key2.as_bytes());
    }

    #[test]
    fn test_invalid_entropy_triggers_retry() {
        // First 32 bytes: the curve order (invalid).
        // Next 32 bytes: scalar 1 (valid).
        let mut data = CURVE_ORDER.to_vec();
        let mut valid = [0u8; 32];
        valid[31] = 0x01;
        data.extend_from_slice(&valid);

        let entropy = FixedEntropy::new(data);
        let key = generate_with_entropy(&entropy).expect("should succeed after retry");
        assert_eq!(key.as_bytes(), &valid);
    }

    #[test]
    fn test_entropy_failure_propagates() {
        let result = generate_with_entropy(&FailingEntropy);
        assert!(result.is_err(), "entropy failure must propagate as error");
    }

    #[test]
    fn test_generated_key_converts_to_secret_key() {
        let mut key_bytes = [0u8; 32];
        key_bytes[31] = 0x01;
        let key = generate_with_entropy(&FixedEntropy::new(key_bytes.to_vec())).unwrap();

        // Must not panic: validates the internal invariant.
        let _sk = key.to_secret_key();
    }

    // ---------------------------------------------------------------
    // 6.12: PrivateKey::from_bytes validation
    // ---------------------------------------------------------------

    #[test]
    fn test_from_bytes_accepts_valid_scalar() {
        let mut bytes = [0u8; 32];
        bytes[31] = 0x01;
        let key = PrivateKey::from_bytes(bytes).expect("scalar 1 must be accepted");
        assert_eq!(key.as_bytes(), &bytes);
    }

    #[test]
    fn test_from_bytes_rejects_invalid_scalar() {
        let zero = [0u8; 32];
        assert!(
            PrivateKey::from_bytes(zero).is_err(),
            "invalid scalar must be rejected"
        );
    }

    // ---------------------------------------------------------------
    // 6.13: PrivateKey::from_hex parsing and validation
    // ---------------------------------------------------------------

    #[test]
    fn test_from_hex_accepts_valid_scalar_one() {
        let hex = "0000000000000000000000000000000000000000000000000000000000000001";
        let key = PrivateKey::from_hex(hex).expect("scalar 1 must be accepted");
        let mut expected = [0u8; 32];
        expected[31] = 0x01;
        assert_eq!(key.as_bytes(), &expected);
    }

    #[test]
    fn test_from_hex_rejects_wrong_length() {
        assert!(
            PrivateKey::from_hex("01").is_err(),
            "hex not exactly 64 chars must be rejected"
        );
    }

    #[test]
    fn test_from_hex_rejects_non_hex_characters() {
        let hex = "zz00000000000000000000000000000000000000000000000000000000000001";
        assert!(
            PrivateKey::from_hex(hex).is_err(),
            "non-hex characters must be rejected"
        );
    }

    #[test]
    fn test_from_hex_propagates_invalid_scalar() {
        let hex = "0000000000000000000000000000000000000000000000000000000000000000";
        assert!(
            PrivateKey::from_hex(hex).is_err(),
            "invalid scalar from from_bytes must propagate"
        );
    }
}