Merge pull request #4 from ooni/python-bindings

mmaker · web-flow · commit ca0a9063bd8c · 2025-09-29T15:52:58.000+02:00
Python bindings
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -1,9 +1,14 @@
 [workspace]
-members = [
-    "ooniauth-core"
-]
+members = ["ooniauth-core", "ooniauth-py"]
 resolver = "2"
 
+[workspace.dependencies]
+rand = "0.8.5"
+bincode = "1"
+cmz = {  version = "0.1.0-rc1" }
+serde = "1.0.219"
+thiserror = "2.0.12"
+
 [patch.crates-io]
 cmz = { path = "include/cmz" }
 sigma-compiler-core = { path = "include/sigma-compiler/sigma-compiler-core" }
diff --git a/ooniauth-core/Cargo.toml b/ooniauth-core/Cargo.toml
@@ -8,16 +8,16 @@ cmz = "0.1.0-rc1"
 curve25519-dalek = { version = "4", features = ["digest", "group", "rand_core", "serde"] }
 ff = "0.13.1"
 group = "0.13.0"
-rand = "0.8.5"
-serde = "1.0.219"
+rand = {workspace = true}
+serde = {workspace = true}
 serde_with = "3.12.0"
-bincode = "1"
+bincode = {workspace = true}
 serde_bytes = "0.11.17"
 sha2 = "0.10.9"
 chrono = "0.4.41"
 time = "0.3.41"
 subtle = "2.6.1"
-thiserror = "2.0.12"
+thiserror = {workspace = true}
 syn = "2.0.103"
 hex = "0.4"
 
@@ -28,7 +28,6 @@ criterion = { version = "0.5"}
 name = "basic_usage"
 test = true
 
-
 [[bench]]
 name = "bench_server"
 harness = false
diff --git a/ooniauth-core/examples/basic_usage.rs b/ooniauth-core/examples/basic_usage.rs
@@ -11,7 +11,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
     // 1. Server initialization
     let now = Instant::now();
     println!("1. Initializing server...");
-    let mut server = ServerState::new(&mut rng);
+    let server = ServerState::new(&mut rng);
     let public_params = server.public_parameters();
     println!("   Key generation completed in {} ms", now.elapsed().as_millis());
 
@@ -69,7 +69,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
     let probe_asn = "AS1234".to_string();
 
     // Set valid age range (credential valid for 30 days)
-    let today = server.today();
+    let today = ServerState::today();
     let age_range = (today - 30)..(today + 1);
     let measurement_count_range = 0..100;
 
diff --git a/ooniauth-core/src/lib.rs b/ooniauth-core/src/lib.rs
@@ -18,13 +18,16 @@ pub mod submit;
 #[derive(Debug, Serialize, Deserialize)]
 pub struct ServerState {
     /// The private key for the main User Auth credential
-    sk: CMZPrivkey<G>,
-    pp: CMZPubkey<G>,
+    sk: SecretKey,
+    pp: PublicParameters,
 }
 
+pub type PublicParameters = CMZPubkey<G>;
+pub type SecretKey = CMZPrivkey<G>;
+
 pub struct UserState {
     /// The public parameters for the client
-    pub pp: CMZPubkey<G>,
+    pub pp: PublicParameters,
     pub(crate) credential: Option<UserAuthCredential>,
 }
 
@@ -38,13 +41,25 @@ impl ServerState {
         Self { sk, pp }
     }
 
+    pub fn secret_key_ref(&self) -> &SecretKey {
+        &self.sk
+    }
+
+    pub fn public_parameters_ref(&self) -> &PublicParameters {
+        &self.pp
+    }
+
+    pub fn from_creds(sk : SecretKey, pp : PublicParameters) -> Self {
+        Self {sk, pp}
+    }
+
     /// Get the public parameters for credential operations
-    pub fn public_parameters(&self) -> CMZPubkey<G> {
+    pub fn public_parameters(&self) -> PublicParameters {
         self.pp.clone()
     }
 
     /// Get today's (real or simulated) date as u32
-    pub fn today(&self) -> u32 {
+    pub fn today() -> u32 {
         // We will not encounter negative Julian dates (~6700 years ago)
         // or ones larger than 32 bits
         (time::OffsetDateTime::now_utc().date())
diff --git a/ooniauth-core/src/registration.rs b/ooniauth-core/src/registration.rs
@@ -64,7 +64,7 @@ impl UserState {
 
 impl ServerState {
     pub fn open_registration(
-        &mut self,
+        &self,
         req: open_registration::Request,
     ) -> Result<open_registration::Reply, CMZError> {
         let mut rng = rand::thread_rng();
@@ -78,7 +78,7 @@ impl ServerState {
             |UAC: &mut UserAuthCredential| {
                 UAC.set_privkey(&self.sk);
                 UAC.measurement_count = Some(Scalar::ZERO);
-                UAC.age = Some(self.today().into());
+                UAC.age = Some(ServerState::today().into());
                 Ok(())
             },
             |_UAC: &UserAuthCredential| Ok(()),
@@ -97,7 +97,7 @@ mod tests {
     fn test_registration() {
         let rng = &mut rand::thread_rng();
         // Initialize group first for gen_keys
-        let mut server_state = ServerState::new(rng);
+        let server_state = ServerState::new(rng);
         // Note: request() will call cmz_group_init again, but that's okay
         let mut user_state = UserState::new(server_state.public_parameters());
 
diff --git a/ooniauth-core/src/submit.rs b/ooniauth-core/src/submit.rs
@@ -140,7 +140,7 @@ impl UserState {
 
 impl ServerState {
     pub fn handle_submit(
-        &mut self,
+        &self,
         rng: &mut (impl RngCore + CryptoRng),
         req: submit::Request,
         nym: &[u8; 32],
@@ -259,7 +259,7 @@ mod tests {
         // Test submit request with valid parameters
         let probe_cc = "US".to_string();
         let probe_asn = "AS1234".to_string();
-        let today = server_state.today();
+        let today = ServerState::today();
         let age_range = (today - 30)..(today + 1); // Credential valid for 30 days
         let measurement_count_range = 0..100;
 
diff --git a/ooniauth-py/Cargo.toml b/ooniauth-py/Cargo.toml
@@ -0,0 +1,19 @@
+[package]
+name = "ooniauth_py"
+version = "0.1.0"
+edition = "2024"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+[lib]
+name = "ooniauth_py"
+crate-type = ["cdylib", "rlib"]
+
+[dependencies]
+cmz = { workspace = true }
+ooniauth-core = {path = "../ooniauth-core"}
+pyo3 = "0.26.0"
+pyo3-stub-gen = "0.13.1"
+rand = {workspace = true}
+bincode = {workspace = true}
+serde = {workspace = true}
+thiserror = {workspace = true}
diff --git a/ooniauth-py/Makefile b/ooniauth-py/Makefile
@@ -0,0 +1,11 @@
+all: stubs
+	maturin build --release
+
+dev: stubs
+	maturin develop
+
+dev-release: stubs
+	maturin develop --release
+
+stubs: 
+	cargo run --bin gen_stubs
diff --git a/ooniauth-py/ooniauth_py.pyi b/ooniauth-py/ooniauth_py.pyi
@@ -0,0 +1,66 @@
+# This file is automatically generated by pyo3_stub_gen
+# ruff: noqa: E501, F401
+
+import builtins
+import typing
+
+class CredentialError(builtins.Exception):
+    r"""
+    An authentication error
+    """
+    ...
+
+class DeserializationFailed(builtins.Exception):
+    r"""
+    An error trying to deserialize a binary buffer
+    """
+    ...
+
+class ProtocolError(builtins.Exception):
+    r"""
+    An error performing the protocol
+    """
+    ...
+
+class ServerState:
+    def __new__(cls) -> ServerState: ...
+    @staticmethod
+    def from_creds(public_parameters:bytes, secret_key:bytes) -> ServerState:
+        r"""
+        Create a new server state from binary-serialized public and private keys
+        This is meant to be used by the server, so it can store the keys somewhere and recreate the
+        state when needed
+        """
+    def get_secret_key(self) -> bytes: ...
+    def get_public_parameters(self) -> bytes: ...
+    def handle_registration_request(self, registration_request:bytes) -> bytes: ...
+    @staticmethod
+    def today() -> builtins.int: ...
+    def handle_submit_request(self, nym:bytes, request:bytes, probe_cc:str, probe_asn:str, age_range:list, measurement_count_range:list) -> bytes: ...
+
+class SubmitRequest:
+    @property
+    def nym(self) -> bytes: ...
+    @property
+    def request(self) -> bytes: ...
+
+class UserState:
+    def __new__(cls, public_params:bytes) -> UserState: ...
+    def get_credential(self) -> typing.Optional[bytes]: ...
+    def make_registration_request(self) -> bytes: ...
+    def handle_registration_response(self, resp:bytes) -> None:
+        r"""
+        Handle a registration response sent by the server, updating your credentials
+        
+        Note that this function will only work if you previously called
+        `make_registration_request`
+        """
+    def make_submit_request(self, probe_cc:str, probe_asn:str, emission_date:builtins.int) -> SubmitRequest: ...
+    def handle_submit_response(self, response:bytes) -> None:
+        r"""
+        Handle a submit response sent by the server, updating your credentials
+        
+        Note that this function will only work if you previously called
+        `make_submit_request`
+        """
+
diff --git a/ooniauth-py/pyproject.toml b/ooniauth-py/pyproject.toml
@@ -0,0 +1,16 @@
+[build-system]
+requires = ["maturin>=1.9,<2.0"]
+build-backend = "maturin"
+
+[project]
+name = "ooniauth-py"
+requires-python = ">=3.8"
+classifiers = [
+    "Programming Language :: Rust",
+    "Programming Language :: Python :: Implementation :: CPython",
+    "Programming Language :: Python :: Implementation :: PyPy",
+]
+dynamic = ["version"]
+[tool.maturin]
+features = ["pyo3/extension-module"]
+
diff --git a/ooniauth-py/src/bin/gen_stubs.rs b/ooniauth-py/src/bin/gen_stubs.rs
@@ -0,0 +1,9 @@
+use ooniauth_py::stub_info;
+use pyo3_stub_gen::Result;
+
+fn main() -> Result<()> {
+    // `stub_info` is a function defined by `define_stub_info_gatherer!` macro.
+    let stub = stub_info()?;
+    stub.generate()?;
+    Ok(())
+}
diff --git a/ooniauth-py/src/exceptions.rs b/ooniauth-py/src/exceptions.rs
@@ -0,0 +1,79 @@
+use cmz::CMZError;
+use ooniauth_core::errors;
+use pyo3::PyErr;
+use pyo3::exceptions::PyException;
+use pyo3_stub_gen::create_exception;
+use thiserror::Error;
+
+// Python excepions: This is what the user sees from the python side when running into an error
+
+// Note that the module name must be the name of the pyproject module name to prevent errors,
+// we should create an issue in the pyo3_stub_gen repo about this. The thing is that
+// if the value of the module argument here is different of the pyproject one except for
+// _ and -, it confuses itself and thinks that its a different module while writing in
+// the same resulting .pyi, deleting parts of the content
+
+create_exception! {
+    ooniauth-py,
+    ProtocolError,
+    PyException,
+    "An error performing the protocol"
+}
+
+create_exception! {
+    ooniauth-py,
+    DeserializationFailed,
+    PyException,
+    "An error trying to deserialize a binary buffer"
+}
+
+// TODO Q: should I create a new python exception per variant of errors::CredentialError?
+create_exception! {
+    ooniauth-py,
+    CredentialError,
+    PyException,
+    "An authentication error"
+}
+
+// The following errors are useful to map rust errors to their corresponding python exceptions
+// defined above
+
+#[derive(Debug, Error)]
+pub enum OoniErr {
+    #[error("Protocol Error: {reason}")]
+    ProtocolError { reason: CMZError },
+
+    #[error("Credential Error: {reason}")]
+    CredentialError { reason: errors::CredentialError },
+
+    #[error("Deserialization Error: {reason}")]
+    DeserializationFailed { reason: String },
+}
+
+pub type OoniResult<T> = Result<T, OoniErr>;
+
+impl From<OoniErr> for PyErr {
+    fn from(value: OoniErr) -> Self {
+        // This function maps from rust error enums to their corresponding python exception
+        match value {
+            OoniErr::ProtocolError { reason } => ProtocolError::new_err(format!("{reason}")),
+            OoniErr::DeserializationFailed { reason } => DeserializationFailed::new_err(reason),
+            OoniErr::CredentialError {
+                reason: errors::CredentialError::CMZError(e),
+            } => ProtocolError::new_err(format!("{e}")),
+            OoniErr::CredentialError { reason } => CredentialError::new_err(format!("{reason}")),
+        }
+    }
+}
+
+impl From<CMZError> for OoniErr {
+    fn from(value: CMZError) -> Self {
+        OoniErr::ProtocolError { reason: value }
+    }
+}
+
+impl From<errors::CredentialError> for OoniErr {
+    fn from(value: errors::CredentialError) -> Self {
+        OoniErr::CredentialError { reason: value }
+    }
+}
diff --git a/ooniauth-py/src/lib.rs b/ooniauth-py/src/lib.rs
@@ -0,0 +1,27 @@
+use pyo3::prelude::*;
+use pyo3_stub_gen::define_stub_info_gatherer;
+
+pub mod exceptions;
+pub mod protocol;
+mod utils;
+
+pub use exceptions::*;
+pub use protocol::*;
+
+/// Here we define the python module itself and its members
+#[pymodule]
+fn ooniauth_py(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    m.add_class::<ServerState>()?;
+    m.add_class::<UserState>()?;
+    m.add_class::<SubmitRequest>()?;
+    m.add("ProtocolError", m.py().get_type::<ProtocolError>())?;
+    m.add("CredentialError", m.py().get_type::<CredentialError>())?;
+    m.add(
+        "DeserializationFailed",
+        m.py().get_type::<DeserializationFailed>(),
+    )?;
+    Ok(())
+}
+
+// Define a function to gather stub information.
+define_stub_info_gatherer!(stub_info);
diff --git a/ooniauth-py/src/protocol.rs b/ooniauth-py/src/protocol.rs
diff --git a/ooniauth-py/src/utils.rs b/ooniauth-py/src/utils.rs
