Skip to main content

pqc_binary_format/
python.rs

1//! Python bindings for PQC Binary Format using PyO3
2//!
3//! This module provides Python bindings for the PQC Binary Format library.
4//! All core functionality is exposed through Python classes and methods.
5
6#![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/// Python wrapper for Algorithm enum
19#[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/// Python wrapper for EncParameters
55#[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/// Python wrapper for KemParameters
80#[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/// Python wrapper for SigParameters
101#[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/// Python wrapper for CompressionParameters
122#[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/// Python wrapper for PqcMetadata
146#[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        // Custom parameters stored in internal HashMap
178    }
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/// Python wrapper for FormatFlags
214#[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/// Python wrapper for PqcBinaryFormat
275#[pyclass(name = "PqcBinaryFormat")]
276pub struct PyPqcBinaryFormat {
277    inner: PqcBinaryFormat,
278}
279
280#[pymethods]
281impl PyPqcBinaryFormat {
282    /// Create a new PQC Binary Format structure
283    ///
284    /// Args:
285    ///     algorithm: Algorithm to use
286    ///     metadata: Metadata container
287    ///     data: Encrypted data bytes
288    ///
289    /// Returns:
290    ///     New PqcBinaryFormat instance
291    #[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    /// Create with specific flags
299    #[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    /// Serialize to bytes
312    ///
313    /// Returns:
314    ///     Bytes object containing serialized format
315    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    /// Deserialize from bytes
324    ///
325    /// Args:
326    ///     data: Bytes to deserialize
327    ///
328    /// Returns:
329    ///     PqcBinaryFormat instance
330    #[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    /// Validate the format structure
338    fn validate(&self) -> PyResult<()> {
339        self.inner
340            .validate()
341            .map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))
342    }
343
344    /// Get algorithm
345    #[getter]
346    fn algorithm(&self) -> PyAlgorithm {
347        PyAlgorithm {
348            inner: self.inner.algorithm(),
349        }
350    }
351
352    /// Get encrypted data
353    #[getter]
354    fn data(&self) -> Vec<u8> {
355        self.inner.data().to_vec()
356    }
357
358    /// Get format flags
359    #[getter]
360    fn flags(&self) -> PyFormatFlags {
361        PyFormatFlags {
362            inner: self.inner.flags(),
363        }
364    }
365
366    /// Get total serialized size
367    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/// Python module initialization
381#[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    // Add version constants
393    m.add("__version__", env!("CARGO_PKG_VERSION"))?;
394    m.add("PQC_BINARY_VERSION", crate::PQC_BINARY_VERSION)?;
395
396    Ok(())
397}