another.im-ios/AnotherXMPP/modules/auth/AuthorizationMechanisms.swift

import Foundation

enum SaslType {
    case sasl1
    case sasl2

    var xmlns: String {
        switch self {
        case .sasl1: "urn:ietf:params:xml:ns:xmpp-sasl"
        case .sasl2: "urn:xmpp:sasl:2"
        }
    }
}

// TODO: Implement SHA-256/SHA-512, the difference only in hmac/pbkdf2 lenght
enum AuthorizationMechanismType: String {
    case plain = "PLAIN"
    case scramSha1 = "SCRAM-SHA-1"
    // case scramSha1Plus = "SCRAM-SHA-1-PLUS" // TODO: check whats wrong with cahnnel binding

    var priority: Int { // less - better
        switch self {
        case .plain: 1000
        // case .scramSha1Plus: 10
        case .scramSha1: 20
        }
    }

    var isPlus: Bool {
        switch self {
        // case .scramSha1Plus: return true
        default: return false
        }
    }
}

protocol AuthorizationMechanism {
    var initRequest: XMLElement? { get }
    func challenge(xml: XMLElement) async -> Event
}

final class AuthorizationMechanismImpl: AuthorizationMechanism {
    private let type: AuthorizationMechanismType
    private let jid: JID
    private let credentials: Credentials?
    private let userAgent: UserAgent
    private let saslType: SaslType
    private let channelBind: String?
    private let inlines: XMLElement? // TODO: Implement inlines

    private let alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789"
    private var clientNonce = ""
    private var serverSignature = ""

    init(type: AuthorizationMechanismType, jid: JID, credentials: Credentials?, saslType: SaslType, userAgent: UserAgent, channelBind: String?, inlines: XMLElement?) {
        self.type = type
        self.jid = jid
        self.credentials = credentials
        self.userAgent = userAgent
        self.saslType = saslType
        self.channelBind = channelBind
        self.inlines = inlines
    }

    var initRequest: XMLElement? {
        switch type {
        case .plain:
            return plainRequest

        case .scramSha1: // .scramSha1Plus:
            return scramSha1Request
        }
    }

    func challenge(xml: XMLElement) async -> Event {
        switch type {
        case .plain:
            return await plainChallenge(xml: xml)

        case .scramSha1: // .scramSha1Plus:
            return await scramSha1Challenge(xml: xml)
        }
    }
}

// MARK: PLAIN
private extension AuthorizationMechanismImpl {
    var plainRequest: XMLElement? {
        guard let pass = credentials?["password"], !pass.isEmpty else {
            print("no credentials...")
            return nil
        }
        let lreq = "\0\(jid.localPart)\0\(pass)"
        let utf8str = lreq.data(using: .utf8)
        let base64 = utf8str?.base64EncodedString(options: NSData.Base64EncodingOptions(rawValue: 0))
        let str = base64 ?? ""
        switch saslType {
        case .sasl1:
            return XMLElement(
                name: "auth",
                xmlns: saslType.xmlns,
                attributes: ["mechanism": "PLAIN"],
                content: str,
                nodes: []
            )

        case .sasl2:
            return XMLElement(
                name: "authenticate",
                xmlns: saslType.xmlns,
                attributes: ["mechanism": "PLAIN"],
                content: nil,
                nodes: [
                    XMLElement(name: "initial-response", xmlns: nil, attributes: [:], content: str, nodes: []),
                    userAgentXml
                ]
            )
        }
    }

    func plainChallenge(xml: XMLElement) async -> Event {
        if xml.name == "success" {
            return succesEvent(xml: xml)
        } else {
            let error = xml.nodes.first?.name ?? "unknown"
            let text = xml.nodes.first(where: { $0.name == "text" })?.content ?? ""
            return .gotAuthError(.mechanismError("\(error) \(text)"))
        }
    }
}

// MARK: SCRAM-SHA-1
private extension AuthorizationMechanismImpl {
    var scramSha1Request: XMLElement? {
        guard let pass = credentials?["password"], !pass.isEmpty else {
            print("no credentials...")
            return nil
        }

        let (requestStr, clientNonce) = scramSha1InitRequestString()
        self.clientNonce = clientNonce
        switch saslType {
        case .sasl1:
            return XMLElement(
                name: "auth",
                xmlns: saslType.xmlns,
                attributes: ["mechanism": type.rawValue],
                content: requestStr,
                nodes: []
            )

        case .sasl2:
            return XMLElement(
                name: "authenticate",
                xmlns: saslType.xmlns,
                attributes: ["mechanism": type.rawValue],
                content: nil,
                nodes: [
                    XMLElement(name: "initial-response", xmlns: nil, attributes: [:], content: requestStr, nodes: []),
                    userAgentXml
                ]
            )
        }
    }

    func scramSha1Challenge(xml: XMLElement) async -> Event {
        switch xml.name {
        case "challenge":
            guard let pass = credentials?["password"], !pass.isEmpty else {
                return .gotAuthError(.mechanismError("No password"))
            }
            // decode base64 challenge with options
            guard let content = xml.content else { return .gotAuthError(.mechanismError("Wrong challenge string")) }
            guard let msgBytes = Foundation.Data(base64Encoded: content, options: NSData.Base64DecodingOptions(rawValue: 0)) else {
                return .gotAuthError(.mechanismError("Wrong challenge string"))
            }
            let msg = String(decoding: msgBytes, as: UTF8.self)

            // get payload
            let dict = msg.split(separator: ",")
                .compactMap { part -> (String, String)? in
                    guard part.count > 2 else { return nil }
                    let line1 = part.prefix(1)
                    let line2 = part.dropFirst(2)
                    return (String(line1), String(line2))
                }
                .reduce(into: [String: String]()) {
                    $0[$1.0] = $1.1
                }

            guard
                let serverNonce = dict["r"],
                let salt = dict["s"],
                let itrs = dict["i"],
                let iterations = Int(itrs)
            else {
                return .gotAuthError(.mechanismError("Wrong challenge string"))
            }

            guard serverNonce.hasPrefix(clientNonce) else {
                return .gotAuthError(.mechanismError("Server nonce incorrect"))
            }

            guard let saltData = Data(base64Encoded: salt, options: Data.Base64DecodingOptions(rawValue: 0)) else {
                return .gotAuthError(.mechanismError("Error forming challenge response"))
            }

            let saltedPass = Data(pass.utf8).pbkdf2(salt: saltData, rounds: iterations)
            let clientFinalMessageBare = "c=biws,r=\(serverNonce)"
            let clientKeyData = Data("Client Key".utf8)
            let serverKeyData = Data("Server Key".utf8)
            let clientKey = clientKeyData.hmac(key: saltedPass)
            let storedKey = clientKey.sha1()
            let authMessage = "n=\(jid.localPart),r=\(clientNonce),\(msg),\(clientFinalMessageBare)"
            let clientSignature = authMessage.data(using: .utf8)?.hmac(key: storedKey) ?? Data()
            let clientProof = clientKey.xor(other: clientSignature)
            let serverKey = serverKeyData.hmac(key: saltedPass)
            serverSignature = authMessage.data(using: .utf8)?.hmac(key: serverKey).base64EncodedString() ?? ""
            let clientFinalMessage = "\(clientFinalMessageBare),p=\(clientProof.base64EncodedString())"

            // make challenge response
            let req = XMLElement(
                name: "response",
                xmlns: saslType.xmlns,
                attributes: [:],
                content: clientFinalMessage.base64Encoded,
                nodes: []
            )
            return .xmlOutbound(req)

        case "success":
            // get server signature
            let signature: String?
            switch saslType {
            case .sasl1:
                signature = xml.content?.base64Decoded

            case .sasl2:
                signature = xml.nodes.first(where: { $0.name == "additional-data" })?.content?.base64Decoded
            }

            // check signature
            if let signature, signature == "v=\(serverSignature)" {
                return succesEvent(xml: xml)
            } else {
                return .gotAuthError(.mechanismError("Wrong server signature"))
            }

        case "failure":
            let error = xml.nodes.first?.name ?? "unknown"
            let text = xml.nodes.first(where: { $0.name == "text" })?.content ?? ""
            return .gotAuthError(.mechanismError("\(error) \(text)"))

        default:
            return .gotAuthError(.mechanismError("Unknown server challenge \(xml.name) \(xml.content ?? "")"))
        }
    }

    func scramSha1InitRequestString() -> (String, String) {
        let gssHeader = "n,,"
        let randString = randomString(length: 20)
        let lreq = "\(gssHeader)n=\(jid.localPart),r=\(randString)"
        let utf8str = lreq.data(using: .utf8)
        let base64 = utf8str?.base64EncodedString(options: NSData.Base64EncodingOptions(rawValue: 0))
        return (base64 ?? "", randString)
    }

    func randomString(length: Int) -> String {
        String((0 ..< length).map { _ in alphabet.randomElement() ?? "A" })
    }

    var userAgentXml: XMLElement {
        XMLElement(
            name: "user-agent",
            xmlns: nil,
            attributes: ["id": userAgent.uuid],
            content: nil,
            nodes: [
                XMLElement(name: "device", xmlns: nil, attributes: [:], content: userAgent.device, nodes: []),
                XMLElement(name: "software", xmlns: nil, attributes: [:], content: userAgent.software, nodes: [])
            ]
        )
    }

    func succesEvent(xml: XMLElement) -> Event {
        var args: [String: String] = [:]
        switch saslType {
        case .sasl1:
            break

        case .sasl2:
            if let authId = xml.nodes.first(where: { $0.name == "authorization-identifier" })?.content {
                args["authorization-identifier"] = authId
            }
        }
        return .authDone(sasl: saslType, args: args)
    }
}