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 KyberPublicKey(pub Vec<u8>);
50
51#[derive(Debug, Clone)]
52#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
53pub struct KyberSecretKey(pub Vec<u8>);
54
55#[derive(Debug, Clone)]
56#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
57pub struct KyberCiphertext(pub Vec<u8>);
58
59#[derive(Debug, Clone)]
60#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
61pub struct KyberKEM {
62 algorithm: Arc<std::sync::RwLock<KyberAlgorithm>>,
63}
64
65#[derive(Debug, Clone, Copy)]
66enum KyberAlgorithm {
67 Kyber512,
68 Kyber768,
69 Kyber1024,
70}
71
72impl KyberKEM {
73 pub fn new() -> Self {
74 Self {
75 algorithm: Arc::new(std::sync::RwLock::new(KyberAlgorithm::Kyber512)),
76 }
77 }
78
79 pub fn with_algorithm(algorithm: super::RiPostQuantumAlgorithm) -> Self {
80 let algo = match algorithm {
81 super::RiPostQuantumAlgorithm::Kyber512 => KyberAlgorithm::Kyber512,
82 super::RiPostQuantumAlgorithm::Kyber768 => KyberAlgorithm::Kyber768,
83 super::RiPostQuantumAlgorithm::Kyber1024 => KyberAlgorithm::Kyber1024,
84 _ => KyberAlgorithm::Kyber512,
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::kem::Kem;
94
95 let algo = *self.algorithm.read().map_err(|e|
96 RiError::InvalidState(format!("Lock error: {}", e))
97 )?;
98 let kem = match algo {
99 KyberAlgorithm::Kyber512 => Kem::new(oqs::kem::Algorithm::Kyber512),
100 KyberAlgorithm::Kyber768 => Kem::new(oqs::kem::Algorithm::Kyber768),
101 KyberAlgorithm::Kyber1024 => Kem::new(oqs::kem::Algorithm::Kyber1024),
102 }.map_err(|e| RiError::Other(format!("Failed to initialize Kyber: {:?}", e)))?;
103
104 let (pk, sk) = kem.keypair()
105 .map_err(|e| RiError::Other(format!("Kyber 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 encapsulate(&self, public_key: &[u8]) -> RiResult<super::KEMResult> {
119 use oqs::kem::Kem;
120
121 let algo = *self.algorithm.read().map_err(|e|
122 RiError::InvalidState(format!("Lock error: {}", e))
123 )?;
124 let kem = match algo {
125 KyberAlgorithm::Kyber512 => Kem::new(oqs::kem::Algorithm::Kyber512),
126 KyberAlgorithm::Kyber768 => Kem::new(oqs::kem::Algorithm::Kyber768),
127 KyberAlgorithm::Kyber1024 => Kem::new(oqs::kem::Algorithm::Kyber1024),
128 }.map_err(|e| RiError::Other(format!("Failed to initialize Kyber: {:?}", e)))?;
129
130 let pk = kem.public_key_from_bytes(public_key)
131 .ok_or_else(|| RiError::Other("Invalid public key".to_string()))?;
132 let (ct, ss) = kem.encapsulate(&pk)
133 .map_err(|e| RiError::Other(format!("Kyber encapsulate failed: {:?}", e)))?;
134 Ok(super::KEMResult {
135 ciphertext: ct.into_vec(),
136 shared_secret: ss.into_vec(),
137 })
138 }
139
140 #[cfg(not(feature = "protocol"))]
141 pub fn encapsulate(&self, _public_key: &[u8]) -> RiResult<super::KEMResult> {
142 Err(RiError::Other(
143 "Post-quantum cryptography requires the 'protocol' feature. \
144 Enable with: cargo build --features protocol".to_string()
145 ))
146 }
147
148 #[cfg(feature = "protocol")]
149 pub fn decapsulate(&self, ciphertext: &[u8], secret_key: &[u8]) -> RiResult<Vec<u8>> {
150 use oqs::kem::Kem;
151
152 let algo = *self.algorithm.read().map_err(|e|
153 RiError::InvalidState(format!("Lock error: {}", e))
154 )?;
155 let kem = match algo {
156 KyberAlgorithm::Kyber512 => Kem::new(oqs::kem::Algorithm::Kyber512),
157 KyberAlgorithm::Kyber768 => Kem::new(oqs::kem::Algorithm::Kyber768),
158 KyberAlgorithm::Kyber1024 => Kem::new(oqs::kem::Algorithm::Kyber1024),
159 }.map_err(|e| RiError::Other(format!("Failed to initialize Kyber: {:?}", e)))?;
160
161 let ct = kem.ciphertext_from_bytes(ciphertext)
162 .ok_or_else(|| RiError::Other("Invalid ciphertext".to_string()))?;
163 let sk = kem.secret_key_from_bytes(secret_key)
164 .ok_or_else(|| RiError::Other("Invalid secret key".to_string()))?;
165 let ss = kem.decapsulate(&sk, &ct)
166 .map_err(|e| RiError::Other(format!("Kyber decapsulate failed: {:?}", e)))?;
167 Ok(ss.into_vec())
168 }
169
170 #[cfg(not(feature = "protocol"))]
171 pub fn decapsulate(&self, _ciphertext: &[u8], _secret_key: &[u8]) -> RiResult<Vec<u8>> {
172 Err(RiError::Other(
173 "Post-quantum cryptography requires the 'protocol' feature. \
174 Enable with: cargo build --features protocol".to_string()
175 ))
176 }
177}
178
179impl Default for KyberKEM {
180 fn default() -> Self {
181 Self::new()
182 }
183}
184
185#[cfg(feature = "pyo3")]
186#[pyo3::prelude::pymethods]
187impl KyberKEM {
188 #[new]
189 pub fn new_py() -> Self {
190 Self::new()
191 }
192
193 pub fn keygen_py(&self) -> Option<(Vec<u8>, Vec<u8>)> {
194 self.keygen().ok()
195 }
196
197 pub fn encapsulate_py(&self, public_key: &[u8]) -> Option<(Vec<u8>, Vec<u8>)> {
198 self.encapsulate(public_key).ok().map(|r| (r.ciphertext, r.shared_secret))
199 }
200
201 pub fn decapsulate_py(&self, ciphertext: &[u8], secret_key: &[u8]) -> Option<Vec<u8>> {
202 self.decapsulate(ciphertext, secret_key).ok()
203 }
204}