1#![allow(non_local_definitions)]
7#![allow(missing_docs)]
8
9use pyo3::prelude::*;
10use pyo3::types::PyBytes;
11use std::collections::HashMap;
12
13use crate::{
14 Algorithm, CompressionParameters, EncParameters, FormatFlags, KemParameters, PqcBinaryFormat,
15 PqcMetadata, SigParameters,
16};
17
18#[pyclass(name = "Algorithm")]
20#[derive(Clone)]
21pub struct PyAlgorithm {
22 inner: Algorithm,
23}
24
25#[pymethods]
26impl PyAlgorithm {
27 #[new]
28 fn new(name: &str) -> PyResult<Self> {
29 let inner = Algorithm::from_name(name).ok_or_else(|| {
30 PyErr::new::<pyo3::exceptions::PyValueError, _>(format!("Unknown algorithm: {}", name))
31 })?;
32 Ok(Self { inner })
33 }
34
35 #[getter]
36 fn name(&self) -> String {
37 self.inner.name().to_string()
38 }
39
40 #[getter]
41 fn id(&self) -> u16 {
42 self.inner.as_id()
43 }
44
45 fn __str__(&self) -> String {
46 self.name()
47 }
48
49 fn __repr__(&self) -> String {
50 format!("Algorithm('{}')", self.name())
51 }
52}
53
54#[pyclass(name = "EncParameters")]
56#[derive(Clone)]
57pub struct PyEncParameters {
58 #[pyo3(get, set)]
59 pub iv: Vec<u8>,
60 #[pyo3(get, set)]
61 pub tag: Vec<u8>,
62}
63
64#[pymethods]
65impl PyEncParameters {
66 #[new]
67 fn new(iv: Vec<u8>, tag: Vec<u8>) -> Self {
68 Self { iv, tag }
69 }
70
71 fn to_dict(&self) -> HashMap<String, Vec<u8>> {
72 let mut map = HashMap::new();
73 map.insert("iv".to_string(), self.iv.clone());
74 map.insert("tag".to_string(), self.tag.clone());
75 map
76 }
77}
78
79#[pyclass(name = "KemParameters")]
81#[derive(Clone)]
82pub struct PyKemParameters {
83 #[pyo3(get, set)]
84 pub public_key: Vec<u8>,
85 #[pyo3(get, set)]
86 pub ciphertext: Vec<u8>,
87}
88
89#[pymethods]
90impl PyKemParameters {
91 #[new]
92 fn new(public_key: Vec<u8>, ciphertext: Vec<u8>) -> Self {
93 Self {
94 public_key,
95 ciphertext,
96 }
97 }
98}
99
100#[pyclass(name = "SigParameters")]
102#[derive(Clone)]
103pub struct PySigParameters {
104 #[pyo3(get, set)]
105 pub public_key: Vec<u8>,
106 #[pyo3(get, set)]
107 pub signature: Vec<u8>,
108}
109
110#[pymethods]
111impl PySigParameters {
112 #[new]
113 fn new(public_key: Vec<u8>, signature: Vec<u8>) -> Self {
114 Self {
115 public_key,
116 signature,
117 }
118 }
119}
120
121#[pyclass(name = "CompressionParameters")]
123#[derive(Clone)]
124pub struct PyCompressionParameters {
125 #[pyo3(get, set)]
126 pub algorithm: String,
127 #[pyo3(get, set)]
128 pub level: u8,
129 #[pyo3(get, set)]
130 pub original_size: u64,
131}
132
133#[pymethods]
134impl PyCompressionParameters {
135 #[new]
136 fn new(algorithm: String, level: u8, original_size: u64) -> Self {
137 Self {
138 algorithm,
139 level,
140 original_size,
141 }
142 }
143}
144
145#[pyclass(name = "PqcMetadata")]
147#[derive(Clone)]
148pub struct PyPqcMetadata {
149 #[pyo3(get, set)]
150 pub enc_params: PyEncParameters,
151 #[pyo3(get, set)]
152 pub kem_params: Option<PyKemParameters>,
153 #[pyo3(get, set)]
154 pub sig_params: Option<PySigParameters>,
155 #[pyo3(get, set)]
156 pub compression_params: Option<PyCompressionParameters>,
157}
158
159#[pymethods]
160impl PyPqcMetadata {
161 #[new]
162 fn new(
163 enc_params: PyEncParameters,
164 kem_params: Option<PyKemParameters>,
165 sig_params: Option<PySigParameters>,
166 compression_params: Option<PyCompressionParameters>,
167 ) -> Self {
168 Self {
169 enc_params,
170 kem_params,
171 sig_params,
172 compression_params,
173 }
174 }
175
176 fn add_custom(&mut self, _key: String, _value: Vec<u8>) {
177 }
179}
180
181impl PyPqcMetadata {
182 fn to_rust(&self) -> PqcMetadata {
183 PqcMetadata {
184 kem_params: self.kem_params.as_ref().map(|k| KemParameters {
185 public_key: k.public_key.clone(),
186 ciphertext: k.ciphertext.clone(),
187 params: HashMap::new(),
188 }),
189 sig_params: self.sig_params.as_ref().map(|s| SigParameters {
190 public_key: s.public_key.clone(),
191 signature: s.signature.clone(),
192 params: HashMap::new(),
193 }),
194 enc_params: EncParameters {
195 iv: self.enc_params.iv.clone(),
196 tag: self.enc_params.tag.clone(),
197 params: HashMap::new(),
198 },
199 compression_params: self
200 .compression_params
201 .as_ref()
202 .map(|c| CompressionParameters {
203 algorithm: c.algorithm.clone(),
204 level: c.level,
205 original_size: c.original_size,
206 params: HashMap::new(),
207 }),
208 custom: HashMap::new(),
209 }
210 }
211}
212
213#[pyclass(name = "FormatFlags")]
215#[derive(Clone)]
216pub struct PyFormatFlags {
217 inner: FormatFlags,
218}
219
220#[pymethods]
221impl PyFormatFlags {
222 #[new]
223 fn new() -> Self {
224 Self {
225 inner: FormatFlags::new(),
226 }
227 }
228
229 fn with_compression(&mut self) -> Self {
230 Self {
231 inner: self.inner.with_compression(),
232 }
233 }
234
235 fn with_streaming(&mut self) -> Self {
236 Self {
237 inner: self.inner.with_streaming(),
238 }
239 }
240
241 fn with_additional_auth(&mut self) -> Self {
242 Self {
243 inner: self.inner.with_additional_auth(),
244 }
245 }
246
247 fn with_experimental(&mut self) -> Self {
248 Self {
249 inner: self.inner.with_experimental(),
250 }
251 }
252
253 #[getter]
254 fn has_compression(&self) -> bool {
255 self.inner.has_compression()
256 }
257
258 #[getter]
259 fn has_streaming(&self) -> bool {
260 self.inner.has_streaming()
261 }
262
263 #[getter]
264 fn has_additional_auth(&self) -> bool {
265 self.inner.has_additional_auth()
266 }
267
268 #[getter]
269 fn has_experimental(&self) -> bool {
270 self.inner.has_experimental()
271 }
272}
273
274#[pyclass(name = "PqcBinaryFormat")]
276pub struct PyPqcBinaryFormat {
277 inner: PqcBinaryFormat,
278}
279
280#[pymethods]
281impl PyPqcBinaryFormat {
282 #[new]
292 fn new(algorithm: PyAlgorithm, metadata: PyPqcMetadata, data: Vec<u8>) -> Self {
293 let rust_metadata = metadata.to_rust();
294 let inner = PqcBinaryFormat::new(algorithm.inner, rust_metadata, data);
295 Self { inner }
296 }
297
298 #[staticmethod]
300 fn with_flags(
301 algorithm: PyAlgorithm,
302 flags: PyFormatFlags,
303 metadata: PyPqcMetadata,
304 data: Vec<u8>,
305 ) -> Self {
306 let rust_metadata = metadata.to_rust();
307 let inner = PqcBinaryFormat::with_flags(algorithm.inner, flags.inner, rust_metadata, data);
308 Self { inner }
309 }
310
311 fn to_bytes<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
316 let bytes = self
317 .inner
318 .to_bytes()
319 .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
320 Ok(PyBytes::new(py, &bytes))
321 }
322
323 #[staticmethod]
331 fn from_bytes(data: &[u8]) -> PyResult<Self> {
332 let inner = PqcBinaryFormat::from_bytes(data)
333 .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
334 Ok(Self { inner })
335 }
336
337 fn validate(&self) -> PyResult<()> {
339 self.inner
340 .validate()
341 .map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))
342 }
343
344 #[getter]
346 fn algorithm(&self) -> PyAlgorithm {
347 PyAlgorithm {
348 inner: self.inner.algorithm(),
349 }
350 }
351
352 #[getter]
354 fn data(&self) -> Vec<u8> {
355 self.inner.data().to_vec()
356 }
357
358 #[getter]
360 fn flags(&self) -> PyFormatFlags {
361 PyFormatFlags {
362 inner: self.inner.flags(),
363 }
364 }
365
366 fn total_size(&self) -> usize {
368 self.inner.total_size()
369 }
370
371 fn __repr__(&self) -> String {
372 format!(
373 "PqcBinaryFormat(algorithm='{}', data_len={})",
374 self.inner.algorithm().name(),
375 self.inner.data().len()
376 )
377 }
378}
379
380#[pymodule]
382fn pqc_binary_format(m: &Bound<'_, PyModule>) -> PyResult<()> {
383 m.add_class::<PyAlgorithm>()?;
384 m.add_class::<PyEncParameters>()?;
385 m.add_class::<PyKemParameters>()?;
386 m.add_class::<PySigParameters>()?;
387 m.add_class::<PyCompressionParameters>()?;
388 m.add_class::<PyPqcMetadata>()?;
389 m.add_class::<PyFormatFlags>()?;
390 m.add_class::<PyPqcBinaryFormat>()?;
391
392 m.add("__version__", env!("CARGO_PKG_VERSION"))?;
394 m.add("PQC_BINARY_VERSION", crate::PQC_BINARY_VERSION)?;
395
396 Ok(())
397}