nucleic/Sources/NucleicMDM/MDMProtocol.swift

import Foundation

/// The MDM **check-in** messages a device PUTs to `/checkin` (docs/MACOS_VM_MDM.md §4). Bodies are
/// XML plists carrying a `MessageType`. We model only the three the guest actually sends in our flow.
public enum MDMCheckInMessage: Sendable, Equatable {
    /// First contact: the device offers its UDID (and the profile's Topic). The server accepts.
    case authenticate(udid: String, topic: String?)
    /// The device reports its APNs push token + magic. This is the **pivot**: after a successful
    /// `TokenUpdate` the device immediately polls `/server`, which is where the queued
    /// `InstallProfile` commands drain with no APNs (the load-bearing §2.4 assumption).
    case tokenUpdate(udid: String, pushMagic: String?, tokenHex: String?)
    /// The device is unenrolling (profile removed / `CheckOutWhenRemoved`).
    case checkOut(udid: String)
    /// Any other check-in `MessageType` we don't model (`SetBootstrapToken`, `GetBootstrapToken`,
    /// `DeclarativeManagement`, …). These are acked with an empty 200 rather than rejected: a 400 on
    /// the check-in channel aborts enrollment, and modern `mdmclient` sends several of these during a
    /// normal handshake (finding 8 / Workstream 2.2).
    case other(type: String, udid: String)

    /// Parse a check-in plist body. Only a body that is not a property list with a `MessageType` is a
    /// hard error; an unrecognized type is carried as `.other` so the server can ack it.
    public static func parse(_ body: Data) throws -> MDMCheckInMessage {
        let dict = try MDMPlist.parse(body)
        guard let type = dict["MessageType"] as? String else {
            throw MDMProtocolError.malformed("check-in body has no MessageType")
        }
        let udid = (dict["UDID"] as? String) ?? (dict["EnrollmentID"] as? String) ?? ""
        switch type {
        case "Authenticate":
            return .authenticate(udid: udid, topic: dict["Topic"] as? String)
        case "TokenUpdate":
            let tokenHex = (dict["Token"] as? Data).map { $0.map { String(format: "%02x", $0) }.joined() }
            return .tokenUpdate(
                udid: udid, pushMagic: dict["PushMagic"] as? String, tokenHex: tokenHex)
        case "CheckOut":
            return .checkOut(udid: udid)
        default:
            return .other(type: type, udid: udid)
        }
    }

    public var udid: String {
        switch self {
        case .authenticate(let u, _), .tokenUpdate(let u, _, _), .checkOut(let u),
             .other(_, let u):
            return u
        }
    }
}

/// A device's status report on the command channel (`/server`). The device PUTs one of these; the
/// server answers with the next command (or an empty 200 to end the poll loop).
public struct MDMCommandStatusReport: Sendable, Equatable {
    /// `Idle` (ready for the next command), `Acknowledged`, `Error`, `CommandFormatError`, or
    /// `NotNow` (busy — retry later).
    public var status: String
    /// The command being reported on (absent for the initial `Idle`).
    public var commandUUID: String?
    public var udid: String
    /// A one-line summary of the `ErrorChain` when `status == "Error"`, for logging. Flattened to a
    /// string so the report stays `Sendable` (the raw plist is `[[String: Any]]`).
    public var errorSummary: String?

    public var isIdle: Bool { status == "Idle" }
    public var isTerminal: Bool { status == "Acknowledged" || status == "Error" }

    public static func parse(_ body: Data) throws -> MDMCommandStatusReport {
        let dict = try MDMPlist.parse(body)
        guard let status = dict["Status"] as? String else {
            throw MDMProtocolError.malformed("command report has no Status")
        }
        let udid = (dict["UDID"] as? String) ?? (dict["EnrollmentID"] as? String) ?? ""
        var summary: String?
        if let chain = dict["ErrorChain"] as? [[String: Any]], !chain.isEmpty {
            summary = chain.compactMap { entry in
                (entry["USEnglishDescription"] as? String) ?? (entry["LocalizedDescription"] as? String)
            }.joined(separator: "; ")
        }
        return MDMCommandStatusReport(
            status: status, commandUUID: dict["CommandUUID"] as? String, udid: udid,
            errorSummary: summary)
    }

    public init(status: String, commandUUID: String?, udid: String, errorSummary: String? = nil) {
        self.status = status
        self.commandUUID = commandUUID
        self.udid = udid
        self.errorSummary = errorSummary
    }

    public static func == (l: MDMCommandStatusReport, r: MDMCommandStatusReport) -> Bool {
        l.status == r.status && l.commandUUID == r.commandUUID && l.udid == r.udid
    }
}

/// One MDM command the server delivers on the `/server` channel. We implement the two the
/// notification fix needs — `InstallProfile` and `RemoveProfile` — plus a generic escape hatch for
/// later payload types (P3).
public struct MDMCommand: Sendable, Equatable {
    public var uuid: String
    public var requestType: String
    /// Extra keys merged into the `Command` dict (e.g. `Payload`, `Identifier`). Kept as an ordered
    /// list of (key, value) so serialization is deterministic for tests; values are plist-safe.
    private var fields: [(String, PlistValue)]

    public init(uuid: String = UUID().uuidString, requestType: String, fields: [(String, PlistValue)] = []) {
        self.uuid = uuid
        self.requestType = requestType
        self.fields = fields
    }

    /// `InstallProfile` — installs the given `.mobileconfig` bytes (the PPPC / notifications payloads).
    public static func installProfile(_ profile: Data, uuid: String = UUID().uuidString) -> MDMCommand {
        MDMCommand(uuid: uuid, requestType: "InstallProfile", fields: [("Payload", .data(profile))])
    }

    /// `RemoveProfile` — removes an installed profile by its top-level `PayloadIdentifier`.
    public static func removeProfile(identifier: String, uuid: String = UUID().uuidString) -> MDMCommand {
        MDMCommand(
            uuid: uuid, requestType: "RemoveProfile", fields: [("Identifier", .string(identifier))])
    }

    /// Serialize to the wire plist: `{ CommandUUID, Command: { RequestType, …fields } }`.
    public func plist() throws -> Data {
        var command: [String: Any] = ["RequestType": requestType]
        for (k, v) in fields { command[k] = v.any }
        let top: [String: Any] = ["CommandUUID": uuid, "Command": command]
        return try MDMPlist.serialize(top)
    }

    public static func == (l: MDMCommand, r: MDMCommand) -> Bool {
        l.uuid == r.uuid && l.requestType == r.requestType
    }
}

/// A minimal plist-value enum so `MDMCommand` fields carry only encodable, deterministic values.
public enum PlistValue: Sendable {
    case string(String)
    case data(Data)
    case bool(Bool)
    case int(Int)

    var any: Any {
        switch self {
        case .string(let s): return s
        case .data(let d): return d
        case .bool(let b): return b
        case .int(let i): return i
        }
    }
}

public enum MDMProtocolError: Error, CustomStringConvertible {
    case malformed(String)
    case unsupported(String)

    public var description: String {
        switch self {
        case .malformed(let s): return "malformed MDM message: \(s)"
        case .unsupported(let s): return s
        }
    }
}