//  OpenHaystack – Tracking personal Bluetooth devices via Apple's Find My network
//
//  Copyright © 2021 Secure Mobile Networking Lab (SEEMOO)
//  Copyright © 2021 The Open Wireless Link Project
//
//  SPDX-License-Identifier: AGPL-3.0-only

import Foundation
import CryptoKit

struct DecryptReports {

    /// Decrypt a find my report with the according key
    /// - Parameters:
    ///   - report: An encrypted FindMy Report
    ///   - key: A FindMyKey
    /// - Throws: Errors if the decryption fails
    /// - Returns: An decrypted location report
    static func decrypt(report: FindMyReport, with key: FindMyKey) throws -> FindMyLocationReport {
        let payloadData = report.payload
        let keyData = key.privateKey

        let privateKey = keyData
        let ephemeralKey = payloadData.subdata(in: 5..<62)

        guard let sharedKey = BoringSSL.deriveSharedKey(
                fromPrivateKey: privateKey,
                andEphemeralKey: ephemeralKey) else {
            throw FindMyError.decryptionError(description: "Failed generating shared key")
        }

        let derivedKey = self.kdf(fromSharedSecret: sharedKey, andEphemeralKey: ephemeralKey)

        print("Derived key \(derivedKey.base64EncodedString())")

        let encData = payloadData.subdata(in: 62..<72)
        let tag = payloadData.subdata(in: 72..<payloadData.endIndex)

        let decryptedContent = try self.decryptPayload(payload: encData, symmetricKey: derivedKey, tag: tag)
        let locationReport = self.decode(content: decryptedContent, report: report)
        print(locationReport)
        return locationReport
    }

    /// Decrypt the payload
    /// - Parameters:
    ///   - payload: Encrypted payload part
    ///   - symmetricKey: Symmetric key
    ///   - tag: AES GCM tag
    /// - Throws: AES GCM error
    /// - Returns: Decrypted error
    static func decryptPayload(payload: Data, symmetricKey: Data, tag: Data) throws -> Data {
        let decryptionKey = symmetricKey.subdata(in: 0..<16)
        let iv = symmetricKey.subdata(in: 16..<symmetricKey.endIndex)

        print("Decryption Key \(decryptionKey.base64EncodedString())")
        print("IV \(iv.base64EncodedString())")

        let sealedBox = try AES.GCM.SealedBox(nonce: AES.GCM.Nonce(data: iv), ciphertext: payload, tag: tag)
        let symKey = SymmetricKey(data: decryptionKey)
        let decrypted = try AES.GCM.open(sealedBox, using: symKey)

        return decrypted
    }

    static func decode(content: Data, report: FindMyReport) -> FindMyLocationReport {
        var longitude: Int32 = 0
        _ = withUnsafeMutableBytes(of: &longitude, {content.subdata(in: 4..<8).copyBytes(to: $0)})
        longitude = Int32(bigEndian: longitude)

        var latitude: Int32 = 0
        _ = withUnsafeMutableBytes(of: &latitude, {content.subdata(in: 0..<4).copyBytes(to: $0)})
        latitude = Int32(bigEndian: latitude)

        var accuracy: UInt8 = 0
        _ = withUnsafeMutableBytes(of: &accuracy, {content.subdata(in: 8..<9).copyBytes(to: $0)})

        let latitudeDec = Double(latitude)/10000000.0
        let longitudeDec = Double(longitude)/10000000.0

        return FindMyLocationReport(lat: latitudeDec, lng: longitudeDec, acc: accuracy, dP: report.datePublished, t: report.timestamp, c: report.confidence)
    }

    static func kdf(fromSharedSecret secret: Data, andEphemeralKey ephKey: Data) -> Data {

        var shaDigest = SHA256()
        shaDigest.update(data: secret)
        var counter: Int32 = 1
        let counterData = Data(Data(bytes: &counter, count: MemoryLayout.size(ofValue: counter)).reversed())
        shaDigest.update(data: counterData)
        shaDigest.update(data: ephKey)

        let derivedKey = shaDigest.finalize()

        return Data(derivedKey)
    }
}
