as_bytes method for TimestampResponse

CHANGELOG.md

@@ -6,6 +6,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- TimestampResponse now has a `as_bytes` method to retrieve the original 
+  request bytes ([#XX](https://github.com/trailofbits/rfc3161-client/pull/XX))
+
 ## [0.0.3] - 2024-11-06
 
 ### Added

rust/src/lib.rs

@@ -280,6 +280,17 @@ impl TimeStampResp {
         }
     }
 
+    fn as_bytes<'p>(
+        &self,
+        py: pyo3::Python<'p>,
+    ) -> PyResult<pyo3::Bound<'p, pyo3::types::PyBytes>> {
+        let result = asn1::write_single(&self.raw.borrow_dependent());
+        match result {
+            Ok(response_bytes) => Ok(pyo3::types::PyBytes::new_bound(py, &response_bytes)),
+            Err(e) => Err(pyo3::exceptions::PyValueError::new_err(format!("{e}"))),
+        }
+    }
+
     // Timestamp Token (as_bytes)
     fn time_stamp_token<'p>(
         &self,
@@ -771,13 +782,6 @@ mod tests {
             .unwrap();
 
             let py_bytes = pyo3::types::PyBytes::new_bound(py, &data);
-
-            // Does not work
-            // let raw = OwnedTimeStampResp::try_new(py_bytes.into(), |v| {
-            //     RawTimeStampResp::parse_data(v.as_bytes(py))
-            // }).unwrap();
-
-            // Works
             let raw = OwnedTimeStampResp::try_new(py_bytes.into(), |v| {
                 asn1::parse_single::<RawTimeStampResp>(v.as_bytes(py))
             })

src/rfc3161_client/tsp.py

@@ -99,6 +99,10 @@ def signed_data(self) -> SignedData:
     def time_stamp_token(self) -> bytes:
         """Return the bytes of the TimestampToken field."""
 
+    @abc.abstractmethod
+    def as_bytes(self) -> bytes:
+        """Returns the Timestamp Response as bytes."""
+
 
 TimeStampResponse.register(_rust.TimeStampResp)
 

test/test_tsp.py

@@ -36,6 +36,10 @@ def test_equality(self):
             (_FIXTURE / "response.tsr").read_bytes()
         ) == decode_timestamp_response((_FIXTURE / "response.tsr").read_bytes())
 
+    def test_round_trip(self):
+        timestamp_response = decode_timestamp_response((_FIXTURE / "response.tsr").read_bytes())
+        assert timestamp_response == decode_timestamp_response(timestamp_response.as_bytes())
+
 
 class TestTimestampRequest:
     def test_parsing_good(self):
@@ -53,3 +57,7 @@ def test_equality(self):
         other_request = parse_timestamp_request((_FIXTURE / "other_request.der").read_bytes())
 
         assert timestamp_request != other_request
+
+    def test_round_trip(self):
+        timestamp_request = parse_timestamp_request((_FIXTURE / "request.der").read_bytes())
+        assert timestamp_request == parse_timestamp_request(timestamp_request.as_bytes())