cg_api_secure-webshare/crates/cgcx-moderation/src/lib.rs

use cgcx_config::{Config, ShareMode};
use cgcx_core::Result;
use serde::{Deserialize, Serialize};
use std::collections::HashSet;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use tracing::{info, warn};

#[derive(Debug, Clone, Default, Deserialize, Serialize)]
pub struct ModerationLists {
    pub blacklisted: HashSet<i64>,
    pub whitelisted: HashSet<i64>,
}

pub struct ModerationEngine {
    lists: Arc<std::sync::RwLock<ModerationLists>>,
    share_mode: ShareMode,
    blacklist_path: PathBuf,
    whitelist_path: PathBuf,
}

impl ModerationEngine {
    pub fn new(config: &Config, base_data_dir: PathBuf) -> Self {
        Self {
            lists: Arc::new(std::sync::RwLock::new(ModerationLists::default())),
            share_mode: config.content.share_mode.clone(),
            blacklist_path: base_data_dir.join("blacklisted_ids.json"),
            whitelist_path: base_data_dir.join("whitelisted_ids.json"),
        }
    }

    pub async fn load(&self) -> Result<()> {
        let blacklisted = load_id_set(&self.blacklist_path).await?;
        let whitelisted = load_id_set(&self.whitelist_path).await?;
        let mut lists = self.lists.write().unwrap();
        *lists = ModerationLists { blacklisted, whitelisted };
        info!(
            "Moderation lists loaded: {} blacklisted, {} whitelisted",
            lists.blacklisted.len(),
            lists.whitelisted.len()
        );
        Ok(())
    }

    pub async fn is_allowed(&self, user_id: i64) -> bool {
        let lists = self.lists.read().unwrap();
        match self.share_mode {
            ShareMode::B => !lists.blacklisted.contains(&user_id),
            ShareMode::W => lists.whitelisted.contains(&user_id),
        }
    }

    pub async fn blacklist(&self, user_id: i64) -> Result<()> {
        {
            let mut lists = self.lists.write().unwrap();
            lists.blacklisted.insert(user_id);
        }
        self.save_blacklist().await?;
        Ok(())
    }

    pub async fn whitelist(&self, user_id: i64) -> Result<()> {
        {
            let mut lists = self.lists.write().unwrap();
            lists.whitelisted.insert(user_id);
        }
        self.save_whitelist().await?;
        Ok(())
    }

    pub async fn remove_blacklist(&self, user_id: i64) -> Result<()> {
        {
            let mut lists = self.lists.write().unwrap();
            lists.blacklisted.remove(&user_id);
        }
        self.save_blacklist().await?;
        Ok(())
    }

    pub async fn remove_whitelist(&self, user_id: i64) -> Result<()> {
        {
            let mut lists = self.lists.write().unwrap();
            lists.whitelisted.remove(&user_id);
        }
        self.save_whitelist().await?;
        Ok(())
    }

    async fn save_blacklist(&self) -> Result<()> {
        let ids = {
            let lists = self.lists.read().unwrap();
            lists.blacklisted.iter().copied().collect()
        };
        let data = IdListFile {
            ids,
            updated_at: chrono::Utc::now().to_rfc3339(),
        };
        let json = serde_json::to_string_pretty(&data)
            .map_err(|e| cgcx_core::CgcxError::Moderation(e.to_string()))?;
        tokio::fs::write(&self.blacklist_path, json)
            .await
            .map_err(|e| cgcx_core::CgcxError::Moderation(e.to_string()))?;
        Ok(())
    }

    async fn save_whitelist(&self) -> Result<()> {
        let ids = {
            let lists = self.lists.read().unwrap();
            lists.whitelisted.iter().copied().collect()
        };
        let data = IdListFile {
            ids,
            updated_at: chrono::Utc::now().to_rfc3339(),
        };
        let json = serde_json::to_string_pretty(&data)
            .map_err(|e| cgcx_core::CgcxError::Moderation(e.to_string()))?;
        tokio::fs::write(&self.whitelist_path, json)
            .await
            .map_err(|e| cgcx_core::CgcxError::Moderation(e.to_string()))?;
        Ok(())
    }

    pub fn spawn_reload_task(self: Arc<Self>) {
        tokio::spawn(async move {
            let mut interval = tokio::time::interval(std::time::Duration::from_secs(30));
            loop {
                interval.tick().await;
                if let Err(e) = self.load().await {
                    warn!("Moderation list reload failed: {}", e);
                }
            }
        });
    }
}

#[derive(Debug, Clone, Deserialize, Serialize)]
struct IdListFile {
    ids: Vec<i64>,
    updated_at: String,
}

async fn load_id_set(path: &Path) -> Result<HashSet<i64>> {
    if !path.exists() {
        return Ok(HashSet::new());
    }
    let json = tokio::fs::read_to_string(path)
        .await
        .map_err(|e| cgcx_core::CgcxError::Moderation(e.to_string()))?;
    let file: IdListFile = serde_json::from_str(&json)
        .map_err(|e| cgcx_core::CgcxError::Moderation(e.to_string()))?;
    Ok(file.ids.into_iter().collect())
}