nest-data-source-api 0.7.1

use crate::api::{ApiError, PseudonymServiceErrorHandler};
use libpep::data::json::{EncryptedPEPJSONValue, PEPJSONBuilder, PEPJSONValue};
use libpep::factors::PseudonymizationDomain;
use paas_client::pseudonym_service::PseudonymService;
use paas_client::sessions::EncryptionContexts;
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use serde_repr::{Deserialize_repr, Serialize_repr};
use std::fmt::Debug;

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize_repr, Deserialize_repr)]
#[repr(u16)]
pub enum StatusCode {
    Ok = 200,
    BadRequest = 400,
    Unauthorized = 401,
    Forbidden = 403,
    NotFound = 404,
    InternalError = 500,
}

#[derive(Debug, Clone)]
pub struct Record {
    pub participant: String,
    pub extra_data: Value,
    pub result_code: StatusCode,
}

impl Record {
    pub fn ok(participant: String, extra_data: Value) -> Self {
        Self::with_code(participant, extra_data, StatusCode::Ok)
    }

    pub fn not_found(participant: String, extra_data: Value) -> Self {
        Self::with_code(participant, extra_data, StatusCode::NotFound)
    }

    pub fn forbidden(participant: String, extra_data: Value) -> Self {
        Self::with_code(participant, extra_data, StatusCode::Forbidden)
    }

    pub fn server_error(participant: String, extra_data: Value) -> Self {
        Self::with_code(participant, extra_data, StatusCode::InternalError)
    }

    pub fn bad_request(participant: String, error_message: &str) -> Self {
        let error_data = json!({
            "message": error_message
        });
        Self::with_code(participant, error_data, StatusCode::BadRequest)
    }

    pub fn data_processing_error(participant: String, error_message: &str) -> Self {
        let error_data = json!({
            "message": error_message
        });
        Self::with_code(participant, error_data, StatusCode::InternalError)
    }

    pub fn with_code(participant: String, extra_data: Value, code: StatusCode) -> Self {
        Self {
            participant,
            extra_data,
            result_code: code,
        }
    }

    /// Convert to a PEPJSONValue with participants marked as pseudonyms.
    fn to_pep_json(&self) -> Result<PEPJSONValue, String> {
        let fields = json!({
            "participants": self.participant,
            "extra_data": self.extra_data,
            "result_code": self.result_code as u16,
        });
        Ok(PEPJSONBuilder::from_json(&fields, &["participants"])
            .ok_or("Failed to build PEPJSONValue".to_string())?
            .build())
    }

    /// Try to create from a decrypted PEPJSONValue.
    fn from_pep_json(pep_value: &PEPJSONValue) -> Result<Self, String> {
        let json_value = pep_value.to_value().map_err(|e| e.to_string())?;

        let participant = json_value
            .get("participants")
            .ok_or("PEPJSONValue missing 'participants' field".to_string())?
            .as_str()
            .ok_or("'participants' field is not a string".to_string())?
            .to_string();

        let extra_data = json_value.get("extra_data").cloned().unwrap_or(Value::Null);

        let result_code = json_value
            .get("result_code")
            .and_then(|v| v.as_u64())
            .map(|code| match code {
                200 => StatusCode::Ok,
                400 => StatusCode::BadRequest,
                401 => StatusCode::Unauthorized,
                403 => StatusCode::Forbidden,
                404 => StatusCode::NotFound,
                _ => StatusCode::InternalError,
            })
            .unwrap_or(StatusCode::Ok);

        Ok(Self {
            participant,
            extra_data,
            result_code,
        })
    }
}

#[derive(Debug, Serialize, Deserialize)]
pub struct PEPActivityBatchRequest {
    pub record: Vec<EncryptedPEPJSONValue>,
    pub activity: String,
    pub sessions: EncryptionContexts,
    pub domain: PseudonymizationDomain,
    pub domain_to: PseudonymizationDomain,
}

#[derive(Debug)]
pub struct ActivityBatchRequest {
    pub record: Vec<Record>,
    pub activity: String,
    pub domain: PseudonymizationDomain,
}

#[derive(Debug)]
pub struct ActivityBatchResponse {
    pub record: Vec<Record>,
    pub activity: String,
    pub domain: PseudonymizationDomain,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct PEPActivityBatchResponse {
    pub record: Vec<EncryptedPEPJSONValue>,
    pub activity: String,
    pub sessions: EncryptionContexts,
    pub domain: PseudonymizationDomain,
    pub domain_to: PseudonymizationDomain,
}

#[allow(dead_code)]
pub trait HasBatchInfo<T> {
    fn participants(&self) -> Vec<String>;
    fn extra_data(&self) -> Vec<Value>;
    fn record(&self) -> Vec<T>;
    fn domain(&self) -> PseudonymizationDomain;
}

#[allow(dead_code)]
pub trait HasPEPBatchInfo<T> {
    fn encrypted_data(&self) -> Vec<EncryptedPEPJSONValue>;
    fn record(&self) -> Vec<T>;

    fn domain(&self) -> PseudonymizationDomain;
    fn sessions(&self) -> EncryptionContexts;
    fn domain_to(&self) -> PseudonymizationDomain;
}

pub trait PEPBatchMessageType<PepT, T>: Debug {
    type PEPBatchMessage: HasPEPBatchInfo<PepT> + Serialize + Debug;
    type BatchMessage: HasBatchInfo<T> + Debug;

    fn pack(
        request: Self::BatchMessage,
        domain_to: PseudonymizationDomain,
        ps: &mut PseudonymService,
    ) -> Result<Self::PEPBatchMessage, ApiError>;

    #[allow(async_fn_in_trait)]
    async fn unpack(
        request: Self::PEPBatchMessage,
        ps: &mut PseudonymService,
    ) -> Result<Self::BatchMessage, ApiError>;
}

impl HasBatchInfo<Record> for ActivityBatchRequest {
    fn participants(&self) -> Vec<String> {
        self.record.iter().map(|r| r.participant.clone()).collect()
    }

    fn extra_data(&self) -> Vec<Value> {
        self.record.iter().map(|r| r.extra_data.clone()).collect()
    }

    fn record(&self) -> Vec<Record> {
        self.record.clone()
    }

    fn domain(&self) -> PseudonymizationDomain {
        self.domain.clone()
    }
}

impl HasBatchInfo<Record> for ActivityBatchResponse {
    fn participants(&self) -> Vec<String> {
        self.record.iter().map(|r| r.participant.clone()).collect()
    }

    fn extra_data(&self) -> Vec<Value> {
        self.record.iter().map(|r| r.extra_data.clone()).collect()
    }

    fn record(&self) -> Vec<Record> {
        self.record.clone()
    }

    fn domain(&self) -> PseudonymizationDomain {
        self.domain.clone()
    }
}

impl HasPEPBatchInfo<EncryptedPEPJSONValue> for PEPActivityBatchRequest {
    fn encrypted_data(&self) -> Vec<EncryptedPEPJSONValue> {
        self.record.clone()
    }

    fn record(&self) -> Vec<EncryptedPEPJSONValue> {
        self.record.clone()
    }

    fn domain(&self) -> PseudonymizationDomain {
        self.domain.clone()
    }

    fn sessions(&self) -> EncryptionContexts {
        self.sessions.clone()
    }

    fn domain_to(&self) -> PseudonymizationDomain {
        self.domain_to.clone()
    }
}

impl HasPEPBatchInfo<EncryptedPEPJSONValue> for PEPActivityBatchResponse {
    fn encrypted_data(&self) -> Vec<EncryptedPEPJSONValue> {
        self.record.clone()
    }

    fn record(&self) -> Vec<EncryptedPEPJSONValue> {
        self.record.clone()
    }

    fn domain(&self) -> PseudonymizationDomain {
        self.domain.clone()
    }

    fn sessions(&self) -> EncryptionContexts {
        self.sessions.clone()
    }

    fn domain_to(&self) -> PseudonymizationDomain {
        self.domain_to.clone()
    }
}

fn encrypt_records(
    records: &[Record],
    ps: &mut PseudonymService,
) -> Result<(Vec<EncryptedPEPJSONValue>, EncryptionContexts), ApiError> {
    if records.is_empty() {
        return Err(ApiError::Internal("No data points to encrypt".to_string()));
    }

    let mut rng = rand::rng();

    let json_records: Vec<PEPJSONValue> = records
        .iter()
        .map(|r| r.to_pep_json())
        .collect::<Result<Vec<_>, _>>()
        .map_err(ApiError::Internal)?;

    ps.encrypt_batch(&json_records, &mut rng)
        .handle_pseudonym_error()
}

async fn decrypt_records(
    pep_message: &impl HasPEPBatchInfo<EncryptedPEPJSONValue>,
    ps: &mut PseudonymService,
) -> Result<Vec<Record>, ApiError> {
    let transcrypted = ps
        .transcrypt_batch(
            pep_message.record(),
            &pep_message.sessions(),
            &pep_message.domain(),
            &pep_message.domain_to(),
        )
        .await
        .handle_pseudonym_error()?;

    let decrypted = ps.decrypt_batch(&transcrypted).handle_pseudonym_error()?;

    decrypted
        .iter()
        .map(|pep_value| Record::from_pep_json(pep_value).map_err(ApiError::Internal))
        .collect::<Result<Vec<_>, _>>()
}

#[derive(Debug)]
pub struct ActivityBatchRequestMessageType;

impl PEPBatchMessageType<EncryptedPEPJSONValue, Record> for ActivityBatchRequestMessageType {
    type PEPBatchMessage = PEPActivityBatchRequest;
    type BatchMessage = ActivityBatchRequest;

    fn pack(
        request: Self::BatchMessage,
        domain_to: PseudonymizationDomain,
        ps: &mut PseudonymService,
    ) -> Result<Self::PEPBatchMessage, ApiError> {
        let (record, sessions) = encrypt_records(&request.record, ps)?;
        Ok(Self::PEPBatchMessage {
            record,
            activity: request.activity,
            domain: request.domain,
            domain_to,
            sessions,
        })
    }

    async fn unpack(
        request: Self::PEPBatchMessage,
        ps: &mut PseudonymService,
    ) -> Result<Self::BatchMessage, ApiError> {
        let domain = request.domain_to.clone();
        let activity = request.activity.clone();
        let record = decrypt_records(&request, ps).await?;
        Ok(Self::BatchMessage {
            record,
            domain,
            activity,
        })
    }
}

#[derive(Debug)]
pub struct ActivityBatchResponseMessageType;

impl PEPBatchMessageType<EncryptedPEPJSONValue, Record> for ActivityBatchResponseMessageType {
    type PEPBatchMessage = PEPActivityBatchResponse;
    type BatchMessage = ActivityBatchResponse;

    fn pack(
        request: Self::BatchMessage,
        domain_to: PseudonymizationDomain,
        ps: &mut PseudonymService,
    ) -> Result<Self::PEPBatchMessage, ApiError> {
        let (record, sessions) = encrypt_records(&request.record, ps)?;
        Ok(Self::PEPBatchMessage {
            record,
            activity: request.activity,
            domain: request.domain,
            domain_to,
            sessions,
        })
    }

    async fn unpack(
        request: Self::PEPBatchMessage,
        ps: &mut PseudonymService,
    ) -> Result<Self::BatchMessage, ApiError> {
        let domain = request.domain_to.clone();
        let activity = request.activity.clone();
        let record = decrypt_records(&request, ps).await?;
        Ok(Self::BatchMessage {
            record,
            domain,
            activity,
        })
    }
}