Gitfed
bastien-mrq/bt-message / btmessage / Mesh / MeshManager.swift
// MeshManager.swift
// Bluetooth/WiFi mesh via MultipeerConnectivity
// Gère DMs pairwise (XWing) et messages de groupe (AES symétrique)

import Foundation
import MultipeerConnectivity
import Combine
import CryptoKit
import os.log

private let log = Logger(subsystem: "com.btmessage.app", category: "Mesh")

public class MeshManager: NSObject, ObservableObject {
    // MARK: - Published State

    @Published public var connectedPeers: [MCPeerID] = []
    @Published public var knownPublicKeys: [String: HybridPublicKey] = [:]
    @Published public var incomingMessages: [ReceivedMessage] = []

    // MARK: - Internal

    private let serviceType = "btmsg-pqc"
    private let localPeerID: MCPeerID
    private let session: MCSession
    private let advertiser: MCNearbyServiceAdvertiser
    private let browser: MCNearbyServiceBrowser
    private let seenCache = SeenMessageCache()
    private let identity = IdentityManager.shared

    public var onMessageReceived: ((ReceivedMessage) -> Void)?
    public var onGroupKeyReceived: ((GroupModel) -> Void)?

    // MARK: - Init

    public init(displayName: String) {
        self.localPeerID = MCPeerID(displayName: displayName)
        self.session = MCSession(peer: localPeerID,
                                 securityIdentity: nil,
                                 encryptionPreference: .required)
        self.advertiser = MCNearbyServiceAdvertiser(peer: localPeerID,
                                                    discoveryInfo: nil,
                                                    serviceType: serviceType)
        self.browser = MCNearbyServiceBrowser(peer: localPeerID,
                                              serviceType: serviceType)
        super.init()
        session.delegate    = self
        advertiser.delegate = self
        browser.delegate    = self
    }

    // MARK: - Lifecycle

    public func start() {
        advertiser.startAdvertisingPeer()
        browser.startBrowsingForPeers()
        DispatchQueue.main.asyncAfter(deadline: .now() + 1) { [weak self] in
            self?.broadcastPublicKey()
        }
    }

    public func stop() {
        advertiser.stopAdvertisingPeer()
        browser.stopBrowsingForPeers()
        session.disconnect()
    }

    // MARK: - Send DM

    public func sendMessage(text: String, to recipientPeerID: String) throws -> String {
        log.info("sendMessage: to=\(recipientPeerID) keyKnown=\(self.knownPublicKeys[recipientPeerID] != nil)")
        guard let recipientKey = knownPublicKeys[recipientPeerID] else {
            log.error("sendMessage: no key for \(recipientPeerID)")
            throw MeshError.unknownRecipient
        }
        let messageID = UUID().uuidString
        let plaintext = Data(text.utf8)
        let encrypted = try identity.encrypt(message: plaintext, for: recipientKey)
        let chatPayload = ChatPayload(messageID: messageID, encryptedMessage: encrypted)
        guard let payloadData = chatPayload.encoded() else { throw MeshError.encodingFailed }

        let envelope = MeshEnvelope.makeChat(
            senderID: identity.localPeerID,
            recipientID: recipientPeerID,
            payload: payloadData
        )

        // Tracker pour ACK
        pendingAcks.insert(messageID)

        let result = route(envelope: envelope)
        if case .failure = result {
            pendingAcks.remove(messageID)
        }

        return messageID
    }

    // MARK: - Send Group Message

    public func sendGroupMessage(text: String, to groupID: String) throws {
        guard let groupKey = GroupKeyManager.shared.groupKey(for: groupID) else {
            throw MeshError.unknownRecipient
        }
        let plaintext = Data(text.utf8)
        let aesCT = try HybridCrypto.encryptGroup(plaintext: plaintext, groupKey: groupKey)
        let groupPayload = GroupChatPayload(
            groupID: groupID,
            messageID: UUID().uuidString,
            aesCiphertext: aesCT
        )
        guard let payloadData = groupPayload.encoded() else { throw MeshError.encodingFailed }
        let envelope = MeshEnvelope.makeGroupChat(senderID: identity.localPeerID, groupID: groupID, payload: payloadData)
        route(envelope: envelope)
    }

    // MARK: - Distribute Group Key (créateur → membres)

    public func sendGroupKeyDistrib(group: GroupModel, groupKeyData: Data, to recipientPeerID: String) throws {
        guard let recipientKey = knownPublicKeys[recipientPeerID] else {
            throw MeshError.unknownRecipient
        }
        let encryptedKey = try identity.encrypt(message: groupKeyData, for: recipientKey)
        let payload = GroupKeyPayload(
            groupID: group.id,
            groupName: group.name,
            memberIDs: group.memberIDs,
            encryptedGroupKey: encryptedKey
        )
        guard let payloadData = payload.encoded() else { throw MeshError.encodingFailed }
        let envelope = MeshEnvelope.makeGroupKeyDistrib(
            senderID: identity.localPeerID,
            recipientID: recipientPeerID,
            payload: payloadData
        )
        route(envelope: envelope)
    }

    // MARK: - Routing

    /// Erreurs de routage
    enum RoutingError: Error, LocalizedError {
        case noConnectedPeers
        case sendFailed(Error)
        case encodingFailed

        var errorDescription: String? {
            switch self {
            case .noConnectedPeers: return "Aucun pair connecté"
            case .sendFailed(let err): return "Échec d'envoi: \(err.localizedDescription)"
            case .encodingFailed: return "Échec d'encodage"
            }
        }
    }

    @discardableResult
    func route(envelope: MeshEnvelope) -> Result<Void, RoutingError> {
        guard !seenCache.contains(envelope.id) else { return .success(()) }
        seenCache.insert(envelope.id)
        guard let data = envelope.encoded() else {
            log.error("route: encoding failed for envelope \(envelope.id)")
            return .failure(.encodingFailed)
        }

        let peers = session.connectedPeers
        guard !peers.isEmpty else {
            log.warning("route: no connected peers, dropping type=\(envelope.type.rawValue)")
            return .failure(.noConnectedPeers)
        }

        let targetPeers: [MCPeerID]
        if let directPeer = peers.first(where: { $0.displayName == envelope.recipientID }) {
            log.info("route: direct → \(directPeer.displayName) \(data.count)b")
            targetPeers = [directPeer]
        } else {
            log.info("route: flood → \(peers.count) peers \(data.count)b")
            targetPeers = peers
        }

        do {
            try session.send(data, toPeers: targetPeers, with: .reliable)
            return .success(())
        } catch {
            log.error("route: send failed: \(error)")
            return .failure(.sendFailed(error))
        }
    }

    private func handleReceived(envelope: MeshEnvelope, from sender: MCPeerID) {
        guard !seenCache.contains(envelope.id) else { return }

        // Vérification d'intégrité basique: senderID doit correspondre au peer connecté
        // (empêche le spoofing trivial d'identité)
        guard envelope.senderID == sender.displayName else {
            log.warning("Identity spoofing detected: declared \(envelope.senderID) but sent from \(sender.displayName)")
            return
        }

        seenCache.insert(envelope.id)

        switch envelope.type {

        case .chat:
            if envelope.recipientID == identity.localPeerID {
                deliverChatMessage(envelope: envelope)
                sendAck(messageID: envelope.id, to: envelope.senderID)
            } else {
                if let fwd = envelope.forwarded() { _ = route(envelope: fwd) }
            }

        case .keyExchange:
            let peerID = envelope.senderID
            log.info("keyExchange: from=\(peerID) payloadBytes=\(envelope.payload.count)")
            if let pub = HybridPublicKey.from(encoded: envelope.payload, peerID: peerID) {
                log.info("keyExchange: stored key for \(peerID)")
                DispatchQueue.main.async { [weak self] in
                    self?.knownPublicKeys[peerID] = pub
                }
            } else {
                log.error("keyExchange: invalid key from \(peerID), payloadBytes=\(envelope.payload.count) expected=1216")
            }
            if let fwd = envelope.forwarded() { _ = route(envelope: fwd) }

        case .ack:
            handleAck(envelope: envelope)

        case .discovery:
            broadcastPublicKey()
            if let fwd = envelope.forwarded() { _ = route(envelope: fwd) }

        case .groupChat:
            deliverGroupMessage(envelope: envelope)
            // Toujours forwarder — d'autres membres peuvent être derrière des hops
            if let fwd = envelope.forwarded() { _ = route(envelope: fwd) }

        case .groupKeyDistrib:
            if envelope.recipientID == identity.localPeerID {
                receiveGroupKey(envelope: envelope)
            } else {
                if let fwd = envelope.forwarded() { _ = route(envelope: fwd) }
            }
        }
    }

    // MARK: - ACK Handling

    private var pendingAcks: Set<String> = []
    public var onMessageDelivered: ((String) -> Void)?  // Callback pour ACK reçus

    private func handleAck(envelope: MeshEnvelope) {
        guard let messageID = String(data: envelope.payload, encoding: .utf8) else { return }
        let wasPending = pendingAcks.remove(messageID) != nil
        if wasPending {
            log.info("ACK received for message \(messageID.prefix(8))...")
            onMessageDelivered?(messageID)
        }
    }

    // MARK: - Delivery

    private func deliverChatMessage(envelope: MeshEnvelope) {
        guard let payload = ChatPayload.decode(envelope.payload) else {
            log.error("deliverChat: decode failed")
            return
        }
        do {
            let plaintext = try identity.decrypt(message: payload.encryptedMessage)
            let text = String(data: plaintext, encoding: .utf8) ?? "<binary>"
            let msg = ReceivedMessage(
                messageID: payload.messageID,
                senderID: envelope.senderID,
                conversationID: envelope.senderID,
                text: text,
                timestamp: Date(timeIntervalSince1970: envelope.timestamp)
            )
            DispatchQueue.main.async { [weak self] in
                self?.incomingMessages.append(msg)
                self?.onMessageReceived?(msg)
            }
        } catch {
            log.error("deliverChat: decryption failed: \(error)")
        }
    }

    private func deliverGroupMessage(envelope: MeshEnvelope) {
        guard let payload = GroupChatPayload.decode(envelope.payload) else {
            log.error("deliverGroup: decode failed")
            return
        }
        // Ignorer si on n'est pas membre (pas de clé pour ce groupe)
        guard let groupKey = GroupKeyManager.shared.groupKey(for: payload.groupID) else { return }
        // Ignorer les messages qu'on a nous-mêmes envoyés (déjà dans la UI)
        guard envelope.senderID != identity.localPeerID else { return }

        do {
            let plaintext = try HybridCrypto.decryptGroup(aesCiphertext: payload.aesCiphertext, groupKey: groupKey)
            let text = String(data: plaintext, encoding: .utf8) ?? "<binary>"
            let msg = ReceivedMessage(
                messageID: payload.messageID,
                senderID: envelope.senderID,
                conversationID: payload.groupID,
                text: text,
                timestamp: Date(timeIntervalSince1970: envelope.timestamp)
            )
            DispatchQueue.main.async { [weak self] in
                self?.incomingMessages.append(msg)
                self?.onMessageReceived?(msg)
            }
        } catch {
            log.error("deliverGroup: decryption failed: \(error)")
        }
    }

    private func receiveGroupKey(envelope: MeshEnvelope) {
        guard let payload = GroupKeyPayload.decode(envelope.payload) else {
            log.error("receiveGroupKey: decode failed")
            return
        }
        do {
            let groupKeyData = try identity.decrypt(message: payload.encryptedGroupKey)
            guard groupKeyData.count == 32 else {
                log.error("receiveGroupKey: invalid key length \(groupKeyData.count)")
                return
            }
            let group = GroupModel(
                id: payload.groupID,
                name: payload.groupName,
                memberIDs: payload.memberIDs,
                creatorID: envelope.senderID
            )
            GroupKeyManager.shared.store(group: group, keyData: groupKeyData)
            log.info("receiveGroupKey: joined group '\(payload.groupName)' id=\(payload.groupID)")
            DispatchQueue.main.async { [weak self] in
                self?.onGroupKeyReceived?(group)
            }
        } catch {
            log.error("receiveGroupKey: decryption failed: \(error)")
        }
    }

    // MARK: - Helpers

    private func broadcastPublicKey() {
        let encoded = identity.publicKey.encoded
        let envelope = MeshEnvelope.makeKeyExchange(
            senderID: identity.localPeerID,
            publicKeyEncoded: encoded
        )
        guard let data = envelope.encoded() else { return }
        let peers = session.connectedPeers
        guard !peers.isEmpty else { return }
        try? session.send(data, toPeers: peers, with: .reliable)
        log.info("broadcastPublicKey: \(data.count) bytes → \(peers.count) peer(s)")
    }

    private func sendAck(messageID: String, to recipientID: String) {
        route(envelope: MeshEnvelope.makeAck(
            senderID: identity.localPeerID,
            recipientID: recipientID,
            messageID: messageID
        ))
    }
}

// MARK: - MCSessionDelegate

extension MeshManager: MCSessionDelegate {
    public func session(_ session: MCSession, peer peerID: MCPeerID, didChange state: MCSessionState) {
        DispatchQueue.main.async { [weak self] in
            self?.connectedPeers = session.connectedPeers
        }
        log.info("peerState: \(peerID.displayName) → \(state.rawValue)")
        if state == .connected { broadcastPublicKey() }
    }

    public func session(_ session: MCSession, didReceive data: Data, fromPeer peerID: MCPeerID) {
        guard let envelope = MeshEnvelope.decode(data) else {
            log.error("didReceive: failed to decode from \(peerID.displayName)")
            return
        }
        handleReceived(envelope: envelope, from: peerID)
    }

    public func session(_ session: MCSession, didReceive stream: InputStream,
                        withName streamName: String, fromPeer peerID: MCPeerID) {}
    public func session(_ session: MCSession, didStartReceivingResourceWithName resourceName: String,
                        fromPeer peerID: MCPeerID, with progress: Progress) {}
    public func session(_ session: MCSession, didFinishReceivingResourceWithName resourceName: String,
                        fromPeer peerID: MCPeerID, at localURL: URL?, withError error: Error?) {}
}

// MARK: - MCNearbyServiceAdvertiserDelegate

extension MeshManager: MCNearbyServiceAdvertiserDelegate {
    public func advertiser(_ advertiser: MCNearbyServiceAdvertiser,
                           didReceiveInvitationFromPeer peerID: MCPeerID,
                           withContext context: Data?,
                           invitationHandler: @escaping (Bool, MCSession?) -> Void) {
        invitationHandler(true, session)
    }
}

// MARK: - MCNearbyServiceBrowserDelegate

extension MeshManager: MCNearbyServiceBrowserDelegate {
    public func browser(_ browser: MCNearbyServiceBrowser, foundPeer peerID: MCPeerID,
                        withDiscoveryInfo info: [String: String]?) {
        guard !session.connectedPeers.contains(peerID) else { return }
        guard localPeerID.displayName < peerID.displayName else {
            log.debug("foundPeer: \(peerID.displayName) — waiting for invite")
            return
        }
        log.info("foundPeer: inviting \(peerID.displayName)")
        browser.invitePeer(peerID, to: session, withContext: nil, timeout: 10)
    }

    public func browser(_ browser: MCNearbyServiceBrowser, lostPeer peerID: MCPeerID) {}
}

// MARK: - Supporting Types

public struct ReceivedMessage: Identifiable {
    public let id = UUID()
    public let messageID: String
    public let senderID: String
    public let conversationID: String   // peerID pour DM, groupID pour groupe
    public let text: String
    public let timestamp: Date
}

public enum MeshError: Error, LocalizedError {
    case unknownRecipient
    case encodingFailed

    public var errorDescription: String? {
        switch self {
        case .unknownRecipient: return "Clé publique inconnue pour ce destinataire"
        case .encodingFailed:   return "Échec d'encodage du message"
        }
    }
}