claudetools/projects/msp-tools/guru-connect/server/src/support_codes.rs

//! Support session codes management
//!
//! Handles generation and validation of 6-digit support codes
//! for one-time remote support sessions.

use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use chrono::{DateTime, Utc};
use rand::Rng;
use serde::{Deserialize, Serialize};
use uuid::Uuid;

/// A support session code
#[derive(Debug, Clone, Serialize)]
pub struct SupportCode {
    pub code: String,
    pub session_id: Uuid,
    pub created_by: String,
    pub created_at: DateTime<Utc>,
    pub status: CodeStatus,
    pub client_name: Option<String>,
    pub client_machine: Option<String>,
    pub connected_at: Option<DateTime<Utc>>,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(rename_all = "lowercase")]
pub enum CodeStatus {
    Pending,    // Waiting for client to connect
    Connected,  // Client connected, session active
    Completed,  // Session ended normally
    Cancelled,  // Code cancelled by tech
}

/// Request to create a new support code
#[derive(Debug, Deserialize)]
pub struct CreateCodeRequest {
    pub technician_id: Option<String>,
    pub technician_name: Option<String>,
}

/// Response when a code is validated
#[derive(Debug, Serialize)]
pub struct CodeValidation {
    pub valid: bool,
    pub session_id: Option<String>,
    pub server_url: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub error: Option<String>,
}

/// Manages support codes
#[derive(Clone)]
pub struct SupportCodeManager {
    codes: Arc<RwLock<HashMap<String, SupportCode>>>,
    session_to_code: Arc<RwLock<HashMap<Uuid, String>>>,
}

impl SupportCodeManager {
    pub fn new() -> Self {
        Self {
            codes: Arc::new(RwLock::new(HashMap::new())),
            session_to_code: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    /// Generate a unique 6-digit code
    async fn generate_unique_code(&self) -> String {
        let codes = self.codes.read().await;
        let mut rng = rand::thread_rng();
        
        loop {
            let code: u32 = rng.gen_range(100000..999999);
            let code_str = code.to_string();
            
            if !codes.contains_key(&code_str) {
                return code_str;
            }
        }
    }

    /// Create a new support code
    pub async fn create_code(&self, request: CreateCodeRequest) -> SupportCode {
        let code = self.generate_unique_code().await;
        let session_id = Uuid::new_v4();
        
        let support_code = SupportCode {
            code: code.clone(),
            session_id,
            created_by: request.technician_name.unwrap_or_else(|| "Unknown".to_string()),
            created_at: Utc::now(),
            status: CodeStatus::Pending,
            client_name: None,
            client_machine: None,
            connected_at: None,
        };

        let mut codes = self.codes.write().await;
        codes.insert(code.clone(), support_code.clone());

        let mut session_to_code = self.session_to_code.write().await;
        session_to_code.insert(session_id, code);

        support_code
    }

    /// Validate a code and return session info
    pub async fn validate_code(&self, code: &str) -> CodeValidation {
        let codes = self.codes.read().await;
        
        match codes.get(code) {
            Some(support_code) => {
                if support_code.status == CodeStatus::Pending || support_code.status == CodeStatus::Connected {
                    CodeValidation {
                        valid: true,
                        session_id: Some(support_code.session_id.to_string()),
                        server_url: Some("wss://connect.azcomputerguru.com/ws/support".to_string()),
                        error: None,
                    }
                } else {
                    CodeValidation {
                        valid: false,
                        session_id: None,
                        server_url: None,
                        error: Some("This code has expired or been used".to_string()),
                    }
                }
            }
            None => CodeValidation {
                valid: false,
                session_id: None,
                server_url: None,
                error: Some("Invalid code".to_string()),
            },
        }
    }

    /// Mark a code as connected
    pub async fn mark_connected(&self, code: &str, client_name: Option<String>, client_machine: Option<String>) {
        let mut codes = self.codes.write().await;
        if let Some(support_code) = codes.get_mut(code) {
            support_code.status = CodeStatus::Connected;
            support_code.client_name = client_name;
            support_code.client_machine = client_machine;
            support_code.connected_at = Some(Utc::now());
        }
    }

    /// Link a support code to an actual WebSocket session
    pub async fn link_session(&self, code: &str, real_session_id: Uuid) {
        let mut codes = self.codes.write().await;
        if let Some(support_code) = codes.get_mut(code) {
            // Update session_to_code mapping with real session ID
            let old_session_id = support_code.session_id;
            support_code.session_id = real_session_id;

            // Update the reverse mapping
            let mut session_to_code = self.session_to_code.write().await;
            session_to_code.remove(&old_session_id);
            session_to_code.insert(real_session_id, code.to_string());
        }
    }

    /// Get code by its code string
    pub async fn get_code(&self, code: &str) -> Option<SupportCode> {
        let codes = self.codes.read().await;
        codes.get(code).cloned()
    }

    /// Mark a code as completed
    pub async fn mark_completed(&self, code: &str) {
        let mut codes = self.codes.write().await;
        if let Some(support_code) = codes.get_mut(code) {
            support_code.status = CodeStatus::Completed;
        }
    }

    /// Cancel a code (works for both pending and connected)
    pub async fn cancel_code(&self, code: &str) -> bool {
        let mut codes = self.codes.write().await;
        if let Some(support_code) = codes.get_mut(code) {
            if support_code.status == CodeStatus::Pending || support_code.status == CodeStatus::Connected {
                support_code.status = CodeStatus::Cancelled;
                return true;
            }
        }
        false
    }

    /// Check if a code is cancelled
    pub async fn is_cancelled(&self, code: &str) -> bool {
        let codes = self.codes.read().await;
        codes.get(code).map(|c| c.status == CodeStatus::Cancelled).unwrap_or(false)
    }

    /// Check if a code is valid for connection (exists and is pending)
    pub async fn is_valid_for_connection(&self, code: &str) -> bool {
        let codes = self.codes.read().await;
        codes.get(code).map(|c| c.status == CodeStatus::Pending).unwrap_or(false)
    }

    /// List all codes (for dashboard)
    pub async fn list_codes(&self) -> Vec<SupportCode> {
        let codes = self.codes.read().await;
        codes.values().cloned().collect()
    }

    /// List active codes only
    pub async fn list_active_codes(&self) -> Vec<SupportCode> {
        let codes = self.codes.read().await;
        codes.values()
            .filter(|c| c.status == CodeStatus::Pending || c.status == CodeStatus::Connected)
            .cloned()
            .collect()
    }

    /// Get code by session ID
    pub async fn get_by_session(&self, session_id: Uuid) -> Option<SupportCode> {
        let session_to_code = self.session_to_code.read().await;
        let code = session_to_code.get(&session_id)?;

        let codes = self.codes.read().await;
        codes.get(code).cloned()
    }

    /// Get the status of a code as a string (for auth checks)
    pub async fn get_status(&self, code: &str) -> Option<String> {
        let codes = self.codes.read().await;
        codes.get(code).map(|c| match c.status {
            CodeStatus::Pending => "pending".to_string(),
            CodeStatus::Connected => "connected".to_string(),
            CodeStatus::Completed => "completed".to_string(),
            CodeStatus::Cancelled => "cancelled".to_string(),
        })
    }
}

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