1#![allow(non_snake_case)]
19
20use std::sync::Arc;
45use crate::core::{RiResult, RiError};
46
47#[derive(Debug, Clone)]
48#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
49pub struct DilithiumPublicKey(pub Vec<u8>);
50
51#[derive(Debug, Clone)]
52#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
53pub struct DilithiumSecretKey(pub Vec<u8>);
54
55#[derive(Debug, Clone)]
56#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
57pub struct DilithiumSignature(pub Vec<u8>);
58
59#[derive(Debug, Clone)]
60#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
61pub struct DilithiumSigner {
62 algorithm: Arc<std::sync::RwLock<DilithiumAlgorithm>>,
63}
64
65#[derive(Debug, Clone, Copy)]
66enum DilithiumAlgorithm {
67 Dilithium2,
68 Dilithium3,
69 Dilithium5,
70}
71
72impl DilithiumSigner {
73 pub fn new() -> Self {
74 Self {
75 algorithm: Arc::new(std::sync::RwLock::new(DilithiumAlgorithm::Dilithium2)),
76 }
77 }
78
79 pub fn with_algorithm(algorithm: super::RiPostQuantumAlgorithm) -> Self {
80 let algo = match algorithm {
81 super::RiPostQuantumAlgorithm::Dilithium2 => DilithiumAlgorithm::Dilithium2,
82 super::RiPostQuantumAlgorithm::Dilithium3 => DilithiumAlgorithm::Dilithium3,
83 super::RiPostQuantumAlgorithm::Dilithium5 => DilithiumAlgorithm::Dilithium5,
84 _ => DilithiumAlgorithm::Dilithium2,
85 };
86 Self {
87 algorithm: Arc::new(std::sync::RwLock::new(algo)),
88 }
89 }
90
91 #[cfg(feature = "protocol")]
92 pub fn keygen(&self) -> RiResult<(Vec<u8>, Vec<u8>)> {
93 use oqs::sig::Sig;
94
95 let algo = *self.algorithm.read().map_err(|e|
96 RiError::InvalidState(format!("Lock error: {}", e))
97 )?;
98 let sig = match algo {
99 DilithiumAlgorithm::Dilithium2 => Sig::new(oqs::sig::Algorithm::Dilithium2),
100 DilithiumAlgorithm::Dilithium3 => Sig::new(oqs::sig::Algorithm::Dilithium3),
101 DilithiumAlgorithm::Dilithium5 => Sig::new(oqs::sig::Algorithm::Dilithium5),
102 }.map_err(|e| RiError::Other(format!("Failed to initialize Dilithium: {:?}", e)))?;
103
104 let (pk, sk) = sig.keypair()
105 .map_err(|e| RiError::Other(format!("Dilithium keygen failed: {:?}", e)))?;
106 Ok((pk.into_vec(), sk.into_vec()))
107 }
108
109 #[cfg(not(feature = "protocol"))]
110 pub fn keygen(&self) -> RiResult<(Vec<u8>, Vec<u8>)> {
111 Err(RiError::Other(
112 "Post-quantum cryptography requires the 'protocol' feature. \
113 Enable with: cargo build --features protocol".to_string()
114 ))
115 }
116
117 #[cfg(feature = "protocol")]
118 pub fn sign(&self, secret_key: &[u8], message: &[u8]) -> RiResult<Vec<u8>> {
119 use oqs::sig::Sig;
120
121 let algo = *self.algorithm.read().map_err(|e|
122 RiError::InvalidState(format!("Lock error: {}", e))
123 )?;
124 let sig = match algo {
125 DilithiumAlgorithm::Dilithium2 => Sig::new(oqs::sig::Algorithm::Dilithium2),
126 DilithiumAlgorithm::Dilithium3 => Sig::new(oqs::sig::Algorithm::Dilithium3),
127 DilithiumAlgorithm::Dilithium5 => Sig::new(oqs::sig::Algorithm::Dilithium5),
128 }.map_err(|e| RiError::Other(format!("Failed to initialize Dilithium: {:?}", e)))?;
129
130 let sk = sig.secret_key_from_bytes(secret_key)
131 .ok_or_else(|| RiError::Other("Invalid secret key".to_string()))?;
132 let signature = sig.sign(message, &sk)
133 .map_err(|e| RiError::Other(format!("Dilithium sign failed: {:?}", e)))?;
134 Ok(signature.into_vec())
135 }
136
137 #[cfg(not(feature = "protocol"))]
138 pub fn sign(&self, _secret_key: &[u8], _message: &[u8]) -> RiResult<Vec<u8>> {
139 Err(RiError::Other(
140 "Post-quantum cryptography requires the 'protocol' feature. \
141 Enable with: cargo build --features protocol".to_string()
142 ))
143 }
144
145 #[cfg(feature = "protocol")]
146 pub fn verify(&self, public_key: &[u8], message: &[u8], signature: &[u8]) -> RiResult<bool> {
147 use oqs::sig::Sig;
148
149 let algo = *self.algorithm.read().map_err(|e|
150 RiError::InvalidState(format!("Lock error: {}", e))
151 )?;
152 let sig = match algo {
153 DilithiumAlgorithm::Dilithium2 => Sig::new(oqs::sig::Algorithm::Dilithium2),
154 DilithiumAlgorithm::Dilithium3 => Sig::new(oqs::sig::Algorithm::Dilithium3),
155 DilithiumAlgorithm::Dilithium5 => Sig::new(oqs::sig::Algorithm::Dilithium5),
156 }.map_err(|e| RiError::Other(format!("Failed to initialize Dilithium: {:?}", e)))?;
157
158 let pk = sig.public_key_from_bytes(public_key)
159 .ok_or_else(|| RiError::Other("Invalid public key".to_string()))?;
160 let sig_bytes = sig.signature_from_bytes(signature)
161 .ok_or_else(|| RiError::Other("Invalid signature".to_string()))?;
162 let result = sig.verify(message, &sig_bytes, &pk);
163 Ok(result.is_ok())
164 }
165
166 #[cfg(not(feature = "protocol"))]
167 pub fn verify(&self, _public_key: &[u8], _message: &[u8], _signature: &[u8]) -> RiResult<bool> {
168 Err(RiError::Other(
169 "Post-quantum cryptography requires the 'protocol' feature. \
170 Enable with: cargo build --features protocol".to_string()
171 ))
172 }
173}
174
175impl Default for DilithiumSigner {
176 fn default() -> Self {
177 Self::new()
178 }
179}
180
181#[cfg(feature = "pyo3")]
182#[pyo3::prelude::pymethods]
183impl DilithiumSigner {
184 #[new]
185 pub fn new_py() -> Self {
186 Self::new()
187 }
188
189 pub fn keygen_py(&self) -> Option<(Vec<u8>, Vec<u8>)> {
190 self.keygen().ok()
191 }
192
193 pub fn sign_py(&self, secret_key: &[u8], message: &[u8]) -> Option<Vec<u8>> {
194 self.sign(secret_key, message).ok()
195 }
196
197 pub fn verify_py(&self, public_key: &[u8], message: &[u8], signature: &[u8]) -> bool {
198 self.verify(public_key, message, signature).unwrap_or(false)
199 }
200}