python: fix python examples

openwsn-berkeley · Nov 3, 2024 · f48d242 · f48d242
1 parent 9b23dea
commit f48d242
Show file tree

Hide file tree

Showing 3 changed files with 60 additions and 5 deletions.
diff --git a/lakers-python/src/initiator.rs b/lakers-python/src/initiator.rs
@@ -10,6 +10,7 @@ pub struct PyEdhocInitiator {
     wait_m2: WaitM2,
     processing_m2: ProcessingM2,
     processed_m2: ProcessedM2,
+    wait_m4: WaitM4,
     completed: Completed,
 }
 
@@ -34,6 +35,7 @@ impl PyEdhocInitiator {
             wait_m2: WaitM2::default(),
             processing_m2: ProcessingM2::default(),
             processed_m2: ProcessedM2::default(),
+            wait_m4: WaitM4::default(),
             completed: Completed::default(),
         }
     }
@@ -122,7 +124,7 @@ impl PyEdhocInitiator {
             &ead_3,
         ) {
             Ok((state, message_3, prk_out)) => {
-                self.completed = state;
+                self.wait_m4 = state;
                 Ok((
                     PyBytes::new_bound(py, message_3.as_slice()),
                     PyBytes::new_bound(py, prk_out.as_slice()),
@@ -132,6 +134,32 @@ impl PyEdhocInitiator {
         }
     }
 
+    pub fn completed_without_message_4<'a>(&mut self, py: Python<'a>) -> PyResult<()> {
+        match i_complete_without_message_4(&self.wait_m4) {
+            Ok(state) => {
+                self.completed = state;
+                Ok(())
+            }
+            Err(error) => Err(error.into()),
+        }
+    }
+
+    pub fn propcess_message_4<'a>(
+        &mut self,
+        py: Python<'a>,
+        message_4: Vec<u8>,
+    ) -> PyResult<Option<EADItem>> {
+        let message_4 = EdhocMessageBuffer::new_from_slice(message_4.as_slice())?;
+
+        match i_process_message_4(&mut self.wait_m4, &mut default_crypto(), &message_4) {
+            Ok((state, ead_4)) => {
+                self.completed = state;
+                Ok(ead_4)
+            }
+            Err(error) => Err(error.into()),
+        }
+    }
+
     pub fn edhoc_exporter<'a>(
         &mut self,
         py: Python<'a>,

diff --git a/lakers-python/src/responder.rs b/lakers-python/src/responder.rs
@@ -11,6 +11,7 @@ pub struct PyEdhocResponder {
     processing_m1: ProcessingM1,
     wait_m3: WaitM3,
     processing_m3: ProcessingM3,
+    processed_m3: ProcessedM3,
     completed: Completed,
 }
 
@@ -34,6 +35,7 @@ impl PyEdhocResponder {
             processing_m1: ProcessingM1::default(),
             wait_m3: WaitM3::default(),
             processing_m3: ProcessingM3::default(),
+            processed_m3: ProcessedM3::default(),
             completed: Completed::default(),
         })
     }
@@ -112,13 +114,38 @@ impl PyEdhocResponder {
         let valid_cred_i = valid_cred_i.to_credential()?;
         match r_verify_message_3(&mut self.processing_m3, &mut default_crypto(), valid_cred_i) {
             Ok((state, prk_out)) => {
-                self.completed = state;
+                self.processed_m3 = state;
                 Ok(PyBytes::new_bound(py, prk_out.as_slice()))
             }
             Err(error) => Err(error.into()),
         }
     }
 
+    #[pyo3(signature = (ead_4=None))]
+    fn prepare_message_4<'a>(
+        &mut self,
+        py: Python<'a>,
+        ead_4: Option<EADItem>,
+    ) -> PyResult<Bound<'a, PyBytes>> {
+        match r_prepare_message_4(&self.processed_m3, &mut default_crypto(), &ead_4) {
+            Ok((state, message_4)) => {
+                self.completed = state;
+                Ok(PyBytes::new_bound(py, message_4.as_slice()))
+            }
+            Err(error) => Err(error.into()),
+        }
+    }
+
+    pub fn completed_without_message_4<'a>(&mut self, py: Python<'a>) -> PyResult<()> {
+        match r_complete_without_message_4(&self.processed_m3) {
+            Ok(state) => {
+                self.completed = state;
+                Ok(())
+            }
+            Err(error) => Err(error.into()),
+        }
+    }
+
     pub fn edhoc_exporter<'a>(
         &mut self,
         py: Python<'a>,

diff --git a/shared/src/lib.rs b/shared/src/lib.rs
@@ -390,23 +390,23 @@ pub struct ProcessingM3 {
     pub ead_3: Option<EADItem>,
 }
 
-#[derive(Debug)]
+#[derive(Default, Debug)]
 pub struct PreparingM3 {
     pub prk_3e2m: BytesHashLen,
     pub prk_4e3m: BytesHashLen,
     pub th_3: BytesHashLen,
     pub mac_3: BytesMac3,
 }
 
-#[derive(Debug)]
+#[derive(Default, Debug)]
 pub struct ProcessedM3 {
     pub prk_4e3m: BytesHashLen,
     pub th_4: BytesHashLen,
     pub prk_out: BytesHashLen,
     pub prk_exporter: BytesHashLen,
 }
 
-#[derive(Debug)]
+#[derive(Default, Debug)]
 pub struct WaitM4 {
     pub prk_4e3m: BytesHashLen,
     pub th_4: BytesHashLen,