From f48d2428801a756a961ba1d607311edfb9ea9e05 Mon Sep 17 00:00:00 2001
From: Elsa Lopez Perez <elsa.lopez-perez@inria.fr>
Date: Sun, 3 Nov 2024 12:21:35 +0100
Subject: [PATCH] python: fix python examples

---
 lakers-python/src/initiator.rs | 30 +++++++++++++++++++++++++++++-
 lakers-python/src/responder.rs | 29 ++++++++++++++++++++++++++++-
 shared/src/lib.rs              |  6 +++---
 3 files changed, 60 insertions(+), 5 deletions(-)

diff --git a/lakers-python/src/initiator.rs b/lakers-python/src/initiator.rs
index c9e080b0..78550581 100644
--- 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
index 16d6aaad..78e6cfec 100644
--- 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
index 3ac970f4..62b9c2db 100644
--- a/shared/src/lib.rs
+++ b/shared/src/lib.rs
@@ -390,7 +390,7 @@ pub struct ProcessingM3 {
     pub ead_3: Option<EADItem>,
 }
 
-#[derive(Debug)]
+#[derive(Default, Debug)]
 pub struct PreparingM3 {
     pub prk_3e2m: BytesHashLen,
     pub prk_4e3m: BytesHashLen,
@@ -398,7 +398,7 @@ pub struct PreparingM3 {
     pub mac_3: BytesMac3,
 }
 
-#[derive(Debug)]
+#[derive(Default, Debug)]
 pub struct ProcessedM3 {
     pub prk_4e3m: BytesHashLen,
     pub th_4: BytesHashLen,
@@ -406,7 +406,7 @@ pub struct ProcessedM3 {
     pub prk_exporter: BytesHashLen,
 }
 
-#[derive(Debug)]
+#[derive(Default, Debug)]
 pub struct WaitM4 {
     pub prk_4e3m: BytesHashLen,
     pub th_4: BytesHashLen,