//  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
import OSLog

struct FindMyKeyExtractor {
    // swiftlint:disable identifier_name

    /// This function reads the private keys of the Offline Finding Location system. They will
    /// - Throws: Error when accessing files fails
    /// - Returns: Devices and their respective keys
    static func readPrivateKeys() throws -> [FindMyDevice] {
        var devices = [FindMyDevice]()
        os_log(.debug, "Looking for keys")

        do {

            // The key files have moved with macOS 10.15.4
            let macOS10_15_3Devices = try self.readFromOldLocation()
            devices.append(contentsOf: macOS10_15_3Devices)
        } catch {
            os_log(.error, "Did not find keys for 10.15.3\n%@", String(describing: error))
        }

        do {
            // Tries to discover the new location of the keys
            let macOS10_15_4Devices = try self.findKeyFilesInNewLocation()
            devices.append(contentsOf: macOS10_15_4Devices)
        } catch {
            os_log(.error, "Did not find keys for 10.15.4\n%@", String(describing: error))
        }

        return devices
    }

    // MARK: - macOS 10.15.0 - 10.15.3

    /// Reads the find my keys from the location used until macOS 10.15.3
    /// - Throws: An error if the location is no longer available (e.g. in macOS 10.15.4)
    /// - Returns: An array of find my devices including their keys
    static func readFromOldLocation() throws -> [FindMyDevice] {
        // Access the find my directory where the private advertisement keys are stored unencrypted
        let directoryPath = "com.apple.icloud.searchpartyd/PrivateAdvertisementKeys/"

        let fm = FileManager.default
        let privateKeysPath = fm.urls(for: .libraryDirectory, in: .userDomainMask)
            .first?.appendingPathComponent(directoryPath)
        let folders = try fm.contentsOfDirectory(at: privateKeysPath!,
                                                 includingPropertiesForKeys: nil, options: .skipsHiddenFiles)
        guard folders.isEmpty == false else {throw FindMyError.noFoldersFound}

        print("Found \(folders.count) folders")
        var devices = [FindMyDevice]()

        for folderURL in folders {
            let keyFiles = try fm.contentsOfDirectory(at: folderURL,
                                                      includingPropertiesForKeys: nil, options: .skipsHiddenFiles)
            // Check if keys are available
            print("Found \(keyFiles.count) in folder \(folderURL.lastPathComponent)")
            guard keyFiles.isEmpty == false else {continue}
            var device = FindMyDevice(deviceId: folderURL.lastPathComponent)

            for url in keyFiles {
                do {
                    if url.pathExtension == "keys" {
                        let keyPlist = try Data(contentsOf: url)
                        let keyInfo = try self.parseKeyFile(keyFile: keyPlist)
                        device.keys.append(keyInfo)
                    }
                } catch {
                    print("Could not load key file ", error)
                }

            }

            devices.append(device)
        }

        return devices
    }

    /// Parses the key plist file used until macOS 10.15.3
    /// - Parameter keyFile: Propery list data
    /// - Returns: Find My private Key
    static func parseKeyFile(keyFile: Data) throws -> FindMyKey {
       guard let keyDict = try PropertyListSerialization.propertyList(from: keyFile,
                                                                      options: .init(), format: nil) as? [String: Any],
             let advertisedKey = keyDict["A"] as? Data,
             let privateKey = keyDict["PR"] as?  Data,
             let timeValues = keyDict["D"] as? [Double],
             let pu = keyDict["PU"] as? Data
       else {
        throw FindMyError.parsingFailed
       }

        let hashedKeyDigest = SHA256.hash(data: advertisedKey)
        let hashedKey = Data(hashedKeyDigest)
        let time = Date(timeIntervalSinceReferenceDate: timeValues[0])
        let duration = timeValues[1]

        return FindMyKey(advertisedKey: advertisedKey,
                         hashedKey: hashedKey,
                         privateKey: privateKey,
                         startTime: time,
                         duration: duration,
                         pu: pu,
                         yCoordinate: nil,
                         fullKey: nil)
    }

    // MARK: - macOS 10.15.4 - 10.15.6 (+ Big Sur 11.0 Betas)

    /// Find the randomized key folder which is used since macOS 10.15.4
    /// - Returns: Returns an array of urls that contain keys. Multiple folders are found if the mac has multiple users
    static func findRamdomKeyFolder() -> [URL] {
        os_log(.debug, "Searching for cached keys folder")
        var folderURLs = [URL]()
        let foldersPath = "/private/var/folders/"
        let fm = FileManager.default

        func recursiveSearch(from url: URL, urlArray: inout [URL]) {
            do {
                let randomSubfolders = try fm.contentsOfDirectory(at: url,
                                                                  includingPropertiesForKeys: nil,
                                                                  options: .includesDirectoriesPostOrder)

                for folder in randomSubfolders {
                    if folder.lastPathComponent == "com.apple.icloud.searchpartyd" {
                        urlArray.append(folder.appendingPathComponent("Keys"))
                        os_log(.debug, "Found folder at: %@", folder.path)
                        break
                    } else {
                        recursiveSearch(from: folder, urlArray: &urlArray)
                    }
                }

            } catch {

            }

        }

        recursiveSearch(from: URL(fileURLWithPath: foldersPath), urlArray: &folderURLs)

        return folderURLs

    }

    /// Find the key files in macOS 10.15.4 and newer (not working with fixed version 10.15.6)
    /// - Throws: An error if the key folder cannot be fould
    /// - Returns: An array of devices including their keys
    static func findKeyFilesInNewLocation() throws -> [FindMyDevice] {
        let keysFolders = self.findRamdomKeyFolder()
        guard keysFolders.isEmpty == false else {
            throw NSError(domain: "error", code: NSNotFound, userInfo: nil)
        }

        var devices = [FindMyDevice]()
        for folder in keysFolders {
            if let deviceKeys = try? self.loadNewKeyFilesIn(directory: folder) {
                devices.append(contentsOf: deviceKeys)
            }
        }

        return devices
    }

    /// Load the keys fils in the passed directory
    /// - Parameter directory: Pass a directory url to a location with key files
    /// - Throws: An error if the keys could not be found
    /// - Returns: An array of devices including their keys
    static func loadNewKeyFilesIn(directory: URL) throws -> [FindMyDevice] {
        os_log(.debug, "Loading key files from %@", directory.path)
        let fm = FileManager.default
        let subDirectories = try fm.contentsOfDirectory(at: directory,
                                                        includingPropertiesForKeys: nil, options: .skipsHiddenFiles)

        var devices = [FindMyDevice]()

        for deviceDirectory in subDirectories {
            do {
                var keyFiles = [Data]()
                let keyDirectory = deviceDirectory.appendingPathComponent("Primary")
                let keyFileURLs = try fm.contentsOfDirectory(at: keyDirectory,
                                                             includingPropertiesForKeys: nil,
                                                             options: .skipsHiddenFiles)
                for keyfileURL in keyFileURLs {
                    // Read the key files
                    let keyFile = try Data(contentsOf: keyfileURL)
                    if keyFile.isEmpty == false {
                        keyFiles.append(keyFile)
                    }
                }

                // Decode keys for file
                let decoder = FindMyKeyDecoder()
                var decodedKeys = [FindMyKey]()
                for file in keyFiles {
                    do {
                        let fmKeys = try decoder.parse(keyFile: file)
                        decodedKeys.append(contentsOf: fmKeys)
                    } catch {
                        os_log(.error, "Decoding keys failed %@", error.localizedDescription)
                    }
                }

                let device = FindMyDevice(deviceId: deviceDirectory.lastPathComponent, keys: decodedKeys)
                devices.append(device)
            } catch {
                os_log(.error, "Key directory not found %@", error.localizedDescription)
            }
        }

        return devices
    }

}
