playtron-sdk 1.0.0

use base64::{engine::general_purpose, Engine as _};
use reqwest::Client;
use serde::Serialize;
use serde_json::Value;
use std::env;
use std::ffi::{CStr, CString};
use std::os::raw::c_char;
use std::ptr;
use zbus::zvariant::Type;

use crate::internal::libpact;

#[derive(Debug, Clone, Serialize, Type)]
pub struct SessionInfo {
    pub nonce: Vec<u8>,
    pub session_id: String,
}

#[derive(Debug, Clone)]
pub struct AttestationClient {
    initialized: bool,
    client: Client,
    base_url: String,
}

#[derive(Debug, Clone)]
pub enum AttestationError {
    SessionCreationFailed(String),
    QuoteGenerationFailed(String),
    ThreadError(String),
}

impl Default for AttestationClient {
    fn default() -> Self {
        Self::new()
    }
}

impl AttestationClient {
    pub fn new() -> Self {
        // Initialize libpact
        let err = unsafe { libpact::LIBPACT.attestation_init() };
        if err != 0 {
            panic!("Unable to construct the AttestationClient: failed to initialize libpact.");
        }

        // Read the attestation server URL
        let base_url = env::var("PACT_ATTESTATION_URL").expect("PACT_ATTESTATION_URL not set");

        // Initialize the HTTP client
        let client = Client::new();

        Self {
            initialized: true,
            client,
            base_url,
        }
    }

    pub async fn create_session(&self) -> Result<SessionInfo, AttestationError> {
        if !self.initialized {
            unreachable!(
                "Should never be able to call 'CreateSession' on an uninitialized client."
            );
        }

        let payload = tokio::task::spawn_blocking(|| {
            // Get the payload required to create a remote attestation session
            let mut output: *mut c_char = ptr::null_mut();
            let status =
                unsafe { libpact::LIBPACT.attestation_get_session_parameters(&mut output) };

            if status < 0 {
                return Err(AttestationError::SessionCreationFailed(format!(
                    "Unable to get session parameters. Error code: {}",
                    status
                )));
            }

            let payload = unsafe { CStr::from_ptr(output).to_string_lossy().into_owned() };
            unsafe { libpact::LIBPACT.buffer_free(output as *mut _) };

            Ok(payload)
        })
        .await
        .map_err(|_| {
            AttestationError::ThreadError("Error joining payload sync thread".to_owned())
        })??;

        // Send the request to the attestation server to create the session
        let response = self
            .client
            .post(format!("{}/api/v1/session/create", self.base_url))
            .header("Accept", "application/json")
            .header("Content-Type", "application/json")
            .body(payload)
            .send()
            .await
            .map_err(|e| {
                AttestationError::SessionCreationFailed(format!(
                    "Create Session request failed: {}",
                    e
                ))
            })?;

        if !response.status().is_success() {
            return Err(AttestationError::SessionCreationFailed(format!(
                "Create Session request failed with status code: {}",
                response.status()
            )));
        }

        // Parse the response
        let response_body = response.text().await.map_err(|e| {
            AttestationError::SessionCreationFailed(format!("Failed to read response body: {}", e))
        })?;
        let deserialized_response: Value = serde_json::from_str(&response_body).map_err(|e| {
            AttestationError::SessionCreationFailed(format!("Failed to parse response body: {}", e))
        })?;

        let nonce = deserialized_response["Nonce"].as_str().ok_or(
            AttestationError::SessionCreationFailed("Missing 'Nonce' in response".to_string()),
        )?;
        let credential = deserialized_response["Credential"].to_string();

        let session_id = tokio::task::spawn_blocking(|| {
            // Decrypt the session id using the AK
            let cred_buffer = CString::new(credential).unwrap();
            let mut session_id_addr: *mut c_char = ptr::null_mut();

            let status = unsafe {
                libpact::LIBPACT.attestation_decrypt_session_id(
                    cred_buffer.as_ptr() as *const _,
                    cred_buffer.as_bytes().len() as i32,
                    &mut session_id_addr,
                )
            };

            if status < 0 {
                return Err(AttestationError::SessionCreationFailed(format!(
                    "Unable to decrypt the session ID. Error code: {}",
                    status
                )));
            }

            let session_id = unsafe {
                CStr::from_ptr(session_id_addr)
                    .to_string_lossy()
                    .into_owned()
            };
            unsafe { libpact::LIBPACT.buffer_free(session_id_addr as *mut _) };

            Ok(session_id)
        })
        .await
        .map_err(|_| {
            AttestationError::ThreadError("Error joining session_id sync thread".to_owned())
        })??;

        // Return the session info
        Ok(SessionInfo {
            nonce: general_purpose::STANDARD.decode(nonce).map_err(|e| {
                AttestationError::SessionCreationFailed(format!("Failed to decode nonce: {}", e))
            })?,
            session_id,
        })
    }

    pub fn get_quote(&self, nonce: &[u8]) -> Result<String, AttestationError> {
        if !self.initialized {
            return Err(AttestationError::QuoteGenerationFailed(
                "Should never be able to call 'GetQuote' on an uninitialized client.".to_string(),
            ));
        }

        // Get the quote from the TPM
        let mut quote_addr: *mut c_char = ptr::null_mut();
        let status = unsafe {
            libpact::LIBPACT.attestation_get_quote(
                nonce.as_ptr() as *const _,
                nonce.len() as i32,
                &mut quote_addr,
            )
        };

        if status < 0 {
            return Err(AttestationError::QuoteGenerationFailed(format!(
                "Unable to get quote from TPM. Error code: {}",
                status
            )));
        }

        // Copy the buffer then free it
        let quote = unsafe { CStr::from_ptr(quote_addr).to_string_lossy().into_owned() };
        unsafe { libpact::LIBPACT.buffer_free(quote_addr as *mut _) };

        Ok(quote)
    }
}