Tusker/NotificationExtension/NotificationService.swift

//
//  NotificationService.swift
//  NotificationExtension
//
//  Created by Shadowfacts on 4/9/24.
//  Copyright © 2024 Shadowfacts. All rights reserved.
//

import UserNotifications
import UserAccounts
import PushNotifications
import CryptoKit
import OSLog
import Pachyderm

private let logger = Logger(subsystem: Bundle.main.bundleIdentifier!, category: "NotificationService")

class NotificationService: UNNotificationServiceExtension {

    override func didReceive(_ request: UNNotificationRequest, withContentHandler contentHandler: @escaping (UNNotificationContent) -> Void) {
        guard let mutableContent = request.content.mutableCopy() as? UNMutableNotificationContent else {
            contentHandler(request.content)
            return
        }
        
        guard request.content.userInfo["v"] as? Int == 1,
              let accountID = request.content.userInfo["ctx"] as? String,
              let account = UserAccountsManager.shared.getAccount(id: accountID),
              let subscription = getSubscription(account: account),
              let encryptedBody = (request.content.userInfo["data"] as? String).flatMap({ Data(base64Encoded: $0) }),
              let salt = (request.content.userInfo["salt"] as? String).flatMap(decodeBase64URL(_:)),
              let serverPublicKeyData = (request.content.userInfo["pk"] as? String).flatMap(decodeBase64URL(_:)) else {
            logger.error("Missing info from push notification")
            contentHandler(request.content)
            return
        }
        
        guard let body = decryptNotification(subscription: subscription, serverPublicKeyData: serverPublicKeyData, salt: salt, encryptedBody: encryptedBody) else {
            contentHandler(request.content)
            return
        }

        let withoutPadding = body.dropFirst(2)
        
        let notification: PushNotification
        do {
            notification = try JSONDecoder().decode(PushNotification.self, from: withoutPadding)
        } catch {
            logger.error("Unable to decode push payload: \(String(describing: error))")
            contentHandler(request.content)
            return
        }
        
        mutableContent.title = notification.title
        mutableContent.body = notification.body
        
        contentHandler(mutableContent)
    }
    
    override func serviceExtensionTimeWillExpire() {
    }
    
    private func getSubscription(account: UserAccountInfo) -> PushNotifications.PushSubscription? {
        DispatchQueue.main.sync {
            // this is necessary because of a swift bug: https://github.com/apple/swift/pull/72507
            MainActor.runUnsafely {
                PushManager.shared.pushSubscription(account: account)
            }
        }
    }
    
    private func decryptNotification(subscription: PushNotifications.PushSubscription, serverPublicKeyData: Data, salt: Data, encryptedBody: Data) -> Data? {
        // See https://github.com/ClearlyClaire/webpush/blob/f14a4d52e201128b1b00245d11b6de80d6cfdcd9/lib/webpush/encryption.rb
        
        var context = Data()
        context.append(0)
        let clientPublicKey = subscription.secretKey.publicKey.x963Representation
        let clientPublicKeyLength = UInt16(clientPublicKey.count)
        context.append(UInt8((clientPublicKeyLength >> 8) & 0xFF))
        context.append(UInt8(clientPublicKeyLength & 0xFF))
        context.append(clientPublicKey)
        let serverPublicKeyLength = UInt16(serverPublicKeyData.count)
        context.append(UInt8((serverPublicKeyLength >> 8) & 0xFF))
        context.append(UInt8(serverPublicKeyLength & 0xFF))
        context.append(serverPublicKeyData)
        
        func info(encoding: String) -> Data {
            var info = Data("Content-Encoding: \(encoding)\0P-256".utf8)
            info.append(context)
            return info
        }

        let sharedSecret: SharedSecret
        do {
            let serverPublicKey = try P256.KeyAgreement.PublicKey(x963Representation: serverPublicKeyData)
            sharedSecret = try subscription.secretKey.sharedSecretFromKeyAgreement(with: serverPublicKey)
        } catch {
            logger.error("Error getting shared secret: \(String(describing: error))")
            return nil
        }
        
        let sharedInfo = Data("Content-Encoding: auth\0".utf8)
        let pseudoRandomKey = sharedSecret.hkdfDerivedSymmetricKey(using: SHA256.self, salt: subscription.authSecret, sharedInfo: sharedInfo, outputByteCount: 32)
        let contentEncryptionKeyInfo = info(encoding: "aesgcm")
        let contentEncryptionKey = HKDF<SHA256>.deriveKey(inputKeyMaterial: pseudoRandomKey, salt: salt, info: contentEncryptionKeyInfo, outputByteCount: 16)
        let nonceInfo = info(encoding: "nonce")
        let nonce = HKDF<SHA256>.deriveKey(inputKeyMaterial: pseudoRandomKey, salt: salt, info: nonceInfo, outputByteCount: 12)
        
        let nonceAndEncryptedBody = nonce.withUnsafeBytes { noncePtr in
            var data = Data(buffer: noncePtr.bindMemory(to: UInt8.self))
            data.append(encryptedBody)
            return data
        }
        do {
            let sealedBox = try AES.GCM.SealedBox(combined: nonceAndEncryptedBody)
            let decrypted = try AES.GCM.open(sealedBox, using: contentEncryptionKey)
            return decrypted
        } catch {
            logger.error("Error decrypting push: \(String(describing: error))")
            return nil
        }
    }

}

extension MainActor {
    @_unavailableFromAsync
    @available(macOS, obsoleted: 14.0)
    @available(iOS, obsoleted: 17.0)
    @available(watchOS, obsoleted: 10.0)
    @available(tvOS, obsoleted: 17.0)
    static func runUnsafely<T>(_ body: @MainActor () throws -> T) rethrows -> T {
        if #available(macOS 14.0, iOS 17.0, watchOS 10.0, tvOS 17.0, *) {
            return try MainActor.assumeIsolated(body)
        }
        
        dispatchPrecondition(condition: .onQueue(.main))
        return try withoutActuallyEscaping(body) { fn in
            try unsafeBitCast(fn, to: (() throws -> T).self)()
        }
    }
}

private func decodeBase64URL(_ s: String) -> Data? {
    var str = s.replacingOccurrences(of: "-", with: "+").replacingOccurrences(of: "_", with: "/")
    if str.count % 4 != 0 {
        str.append(String(repeating: "=", count: 4 - str.count % 4))
    }
    return Data(base64Encoded: str)
}