Refactor ffi to sgx_pck_extension crates 0.1.1

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "maturin"
 
 [project]
 name = "intel-sgx-ra"
-version = "2.0.1"
+version = "2.1.0"
 description = "Intel SGX Remote Attestation verification library"
 authors = [
     {name = "Cosmian Tech", email = "tech@cosmian.com"},

rust/Cargo.toml

@@ -8,8 +8,5 @@ name = "lib_sgx_dcap_ratls"
 crate-type = ["cdylib"]
 
 [dependencies]
+pyo3 = { version = "0.20", features = ["extension-module"] }
 sgx_pck_extension = "0.1.1"
-
-[dependencies.pyo3]
-version = "0.19"
-features = ["extension-module"]

rust/src/lib.rs

@@ -1,57 +1,94 @@
-use pyo3::exceptions::PyException;
-use pyo3::prelude::*;
-use pyo3::types::{IntoPyDict, PyBytes};
+use pyo3::{exceptions::PyException, prelude::*, types::PyBytes};
 use sgx_pck_extension::extension::SgxPckExtension;
 
+#[derive(Clone)]
+#[pyclass(name = "Configuration", get_all)]
+pub struct PyConfiguration {
+    pub dynamic_platform: bool,
+    pub cached_keys: bool,
+    pub smt_enabled: bool,
+}
+
+#[derive(Clone)]
+#[pyclass(name = "Tcb", get_all)]
+pub struct PyTcb {
+    pub compsvn: Py<PyBytes>,
+    pub pcesvn: u16,
+    pub cpusvn: Py<PyBytes>,
+}
+
+#[pyclass(name = "SgxPckExtension", get_all)]
+pub struct PySgxPckExtension {
+    pub ppid: Py<PyBytes>,
+    pub tcb: PyTcb,
+    pub pceid: Py<PyBytes>,
+    pub fmspc: Py<PyBytes>,
+    pub sgx_type: u8,
+    pub platform_instance_id: Option<Py<PyBytes>>,
+    pub configuration: Option<PyConfiguration>,
+}
+
+impl From<SgxPckExtension> for PySgxPckExtension {
+    fn from(sgx_pck_extension: SgxPckExtension) -> PySgxPckExtension {
+        let ppid: Py<PyBytes> =
+            Python::with_gil(|py| PyBytes::new(py, sgx_pck_extension.ppid.as_slice()).into());
+
+        let compsvn: Py<PyBytes> = Python::with_gil(|py| {
+            PyBytes::new(py, sgx_pck_extension.tcb.compsvn.as_slice()).into()
+        });
+
+        let cpusvn: Py<PyBytes> =
+            Python::with_gil(|py| PyBytes::new(py, sgx_pck_extension.tcb.cpusvn.as_slice()).into());
+
+        let pceid: Py<PyBytes> =
+            Python::with_gil(|py| PyBytes::new(py, sgx_pck_extension.pceid.as_slice()).into());
+
+        let fmspc: Py<PyBytes> =
+            Python::with_gil(|py| PyBytes::new(py, sgx_pck_extension.fmspc.as_slice()).into());
+
+        let platform_instance_id: Option<Py<PyBytes>> =
+            sgx_pck_extension
+                .platform_instance_id
+                .map(|platform_instance_id| {
+                    Python::with_gil(|py| PyBytes::new(py, platform_instance_id.as_slice()).into())
+                });
+
+        PySgxPckExtension {
+            ppid,
+            tcb: PyTcb {
+                compsvn,
+                pcesvn: sgx_pck_extension.tcb.pcesvn,
+                cpusvn,
+            },
+            pceid,
+            fmspc,
+            sgx_type: sgx_pck_extension.sgx_type as u8,
+            platform_instance_id,
+            configuration: sgx_pck_extension
+                .configuration
+                .map(|configuration| PyConfiguration {
+                    dynamic_platform: configuration.dynamic_platform,
+                    cached_keys: configuration.cached_keys,
+                    smt_enabled: configuration.smt_enabled,
+                }),
+        }
+    }
+}
+
 #[pyfunction]
-fn sgx_pck_extension_from_pem(py: Python<'_>, pem: &[u8]) -> PyResult<PyObject> {
+fn sgx_pck_extension_from_pem(_py: Python<'_>, pem: &[u8]) -> PyResult<PySgxPckExtension> {
     let pck_extension = SgxPckExtension::from_pem_certificate(pem)
         .map_err(|e| PyException::new_err(e.to_string()))?;
 
-    let map: Vec<(&str, PyObject)> = vec![
-        ("ppid", PyBytes::new(py, &pck_extension.ppid).to_object(py)),
-        (
-            "compsvn",
-            PyBytes::new(py, &pck_extension.tcb.compsvn).to_object(py),
-        ),
-        ("pcesvn", pck_extension.tcb.pcesvn.to_object(py)),
-        (
-            "cpusvn",
-            PyBytes::new(py, &pck_extension.tcb.cpusvn).to_object(py),
-        ),
-        (
-            "pceid",
-            PyBytes::new(py, &pck_extension.pceid).to_object(py),
-        ),
-        (
-            "fmspc",
-            PyBytes::new(py, &pck_extension.fmspc).to_object(py),
-        ),
-        ("sgx_type", (pck_extension.sgx_type as u32).to_object(py)),
-        (
-            "platform_instance_id",
-            PyBytes::new(py, &pck_extension.platform_instance_id).to_object(py),
-        ),
-        (
-            "dynamic_platform",
-            pck_extension.configuration.dynamic_platform.to_object(py),
-        ),
-        (
-            "cached_keys",
-            pck_extension.configuration.cached_keys.to_object(py),
-        ),
-        (
-            "smt_enabled",
-            pck_extension.configuration.smt_enabled.to_object(py),
-        ),
-    ];
-
-    Ok(map.into_py_dict(py).to_object(py))
+    Ok(pck_extension.into())
 }
 
 #[pymodule]
 #[pyo3(name = "lib_sgx_dcap_ratls")]
 fn sgx_dcap_ratls(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
+    m.add_class::<PyTcb>()?;
+    m.add_class::<PyConfiguration>()?;
+    m.add_class::<PySgxPckExtension>()?;
     m.add_function(wrap_pyfunction!(sgx_pck_extension_from_pem, m)?)?;
     Ok(())
 }

src/intel_sgx_ra/attest.py

@@ -226,7 +226,6 @@ def retrieve_collaterals(
     ]  # type: x509.Certificate, x509.Certificate, x509.Certificate
 
     sgx_pck_ext: SgxPckExtension = sgx_pck_extension_from_cert(pck_cert)
-    fmspc: bytes = sgx_pck_ext["fmspc"]
 
     common_name: x509.NameAttribute
     common_name, *_ = pck_ca_cert.subject.get_attributes_for_oid(
@@ -250,7 +249,7 @@ def retrieve_collaterals(
             "PCCS returned different Intel SGX PCK Platform/Processor CA"
         )
 
-    tcb_info, _root_ca_cert, tcb_cert = get_tcbinfo(pccs_url, fmspc)
+    tcb_info, _root_ca_cert, tcb_cert = get_tcbinfo(pccs_url, sgx_pck_ext.fmspc)
 
     if _root_ca_cert != root_ca_cert:
         raise CertificateError("PCCS returned different Intel SGX Root CA")

src/intel_sgx_ra/lib_sgx_dcap_ratls.pyi

@@ -1,3 +1,36 @@
-from typing import Any, Dict
+from typing import Optional
 
-def sgx_pck_extension_from_pem(pem: bytes) -> Dict[str, Any]: ...
+class Tcb:
+    @property
+    def compsvn(self) -> bytes: ...
+    @property
+    def pcesvn(self) -> int: ...
+    @property
+    def cpusvn(self) -> bytes: ...
+
+class Configuration:
+    @property
+    def dynamic_platform(self) -> bool: ...
+    @property
+    def cached_keys(self) -> bool: ...
+    @property
+    def smt_enabled(self) -> bool: ...
+
+class SgxPckExtension:
+    @property
+    def ppid(self) -> bytes: ...
+    @property
+    def tcb(self) -> Tcb: ...
+    @property
+    def pceid(self) -> bytes: ...
+    @property
+    def fmspc(self) -> bytes: ...
+    @property
+    def sgx_type(self) -> int: ...
+    @property
+    def platform_instance_id(self) -> Optional[bytes]: ...
+    @property
+    def configuration(self):
+        Optional[Configuration]: ...
+
+def sgx_pck_extension_from_pem(pem: bytes) -> SgxPckExtension: ...

src/intel_sgx_ra/pck.py

@@ -1,41 +1,12 @@
 """intel_sgx_ra.pck module."""
 
-from enum import Enum
-from typing import TypedDict, cast
-
 from cryptography import x509
 from cryptography.hazmat.primitives.serialization import Encoding
 
 # pylint: disable=no-name-in-module,import-error
-from intel_sgx_ra.lib_sgx_dcap_ratls import sgx_pck_extension_from_pem
-
-
-class SgxType(Enum):
-    """Enum SgxType used in SgxPckExtension."""
-
-    Standard = 0
-    Scalable = 1
-
-
-class SgxPckExtension(TypedDict):
-    """Struct SgxPckExtension."""
-
-    ppid: bytes
-    compsvn: bytes
-    pcesvn: int
-    cpusvn: bytes
-    pceid: bytes
-    fmspc: bytes
-    sgx_type: SgxType
-    platform_instance_id: bytes
-    dynamic_platform: bool
-    cached_keys: bool
-    smt_enabled: bool
+from intel_sgx_ra.lib_sgx_dcap_ratls import SgxPckExtension, sgx_pck_extension_from_pem
 
 
 def sgx_pck_extension_from_cert(cert: x509.Certificate) -> SgxPckExtension:
     """Parse Intel SGX PCK ASN.1 extension."""
-    return cast(
-        SgxPckExtension,
-        sgx_pck_extension_from_pem(cert.public_bytes(encoding=Encoding.PEM)),
-    )
+    return sgx_pck_extension_from_pem(cert.public_bytes(encoding=Encoding.PEM))

tests/test_quote.py

@@ -35,6 +35,5 @@ def test_pck_extension(data_path):
 
     pck_extension: SgxPckExtension = sgx_pck_extension_from_cert(pck_cert)
 
-    assert "fmspc" in pck_extension
-    assert isinstance(pck_extension["fmspc"], bytes)
-    assert len(pck_extension["fmspc"]) == 6
+    assert isinstance(pck_extension.fmspc, bytes)
+    assert len(pck_extension.fmspc) == 6