wireguard-apple/WireGuard/Coordinators/AppCoordinator.swift

//
//  Copyright © 2018 WireGuard LLC. All rights reserved.
//

import Foundation
import NetworkExtension
import os.log
import ZIPFoundation
import PromiseKit

import CoreData
import BNRCoreDataStack

import MobileCoreServices

enum AppCoordinatorError: Error {
    case configImportError(msg: String)
}

extension UINavigationController: Identifyable {}

let APPGROUP = "group.com.wireguard.ios"
let VPNBUNDLE = "com.wireguard.ios.network-extension"

class AppCoordinator: RootViewCoordinator { // swiftlint:disable:this type_body_length

    let persistentContainer = NSPersistentContainer(name: "WireGuard")
    let storyboard = UIStoryboard(name: "Main", bundle: nil)
    var providerManagers: [NETunnelProviderManager]?
    let documentPickerDelegateObject: AppDocumentPickerDelegate

    // MARK: - Properties

    var childCoordinators: [Coordinator] = []

    var rootViewController: UIViewController {
        return self.tunnelsTableViewController
    }

    var tunnelsTableViewController: TunnelsTableViewController!

    /// Window to manage
    let window: UIWindow

    let navigationController: UINavigationController = {
        let navController = UIStoryboard(name: "Main", bundle: nil).instantiateViewController(type: UINavigationController.self)
        return navController
    }()

    // MARK: - Init
    public init(window: UIWindow) {
        self.window = window

        self.window.rootViewController = self.navigationController
        self.window.makeKeyAndVisible()

        documentPickerDelegateObject = AppDocumentPickerDelegate()
        documentPickerDelegateObject.appCoordinator = self

        NotificationCenter.default.addObserver(self,
                                               selector: #selector(VPNStatusDidChange(notification:)),
                                               name: .NEVPNStatusDidChange,
                                               object: nil)
    }

    // MARK: - Functions

    /// Starts the coordinator
    public func start() {
        _ = refreshProviderManagers().then { () -> Promise<Void> in
            self.persistentContainer.viewContext.automaticallyMergesChangesFromParent = true
            self.persistentContainer.loadPersistentStores { [weak self] (_, error) in
                if let error = error {
                    print("Unable to Load Persistent Store. \(error), \(error.localizedDescription)")

                } else {
                    DispatchQueue.main.async {
                        //start
                        if let tunnelsTableViewController = self?.storyboard.instantiateViewController(type: TunnelsTableViewController.self) {
                            self?.tunnelsTableViewController = tunnelsTableViewController
                            self?.tunnelsTableViewController.viewContext = self?.persistentContainer.viewContext
                            self?.tunnelsTableViewController.delegate = self
                            self?.navigationController.viewControllers = [tunnelsTableViewController]
                            do {
                                if let context = self?.persistentContainer.viewContext, try Tunnel.countInContext(context) == 0 {
                                    print("No tunnels ... yet")
                                }
                            } catch {
                                self?.showError(error)
                            }
                        }
                    }
                }
            }
            return Promise.value(())
        }
    }

    func refreshProviderManagers() -> Promise<Void> {
        return Promise { (resolver) in
            NETunnelProviderManager.loadAllFromPreferences { [weak self] (managers, error) in
                if let error = error {
                    os_log("Unable to load provider managers: %{public}@", log: Log.general, type: .error, error.localizedDescription)
                }
                self?.providerManagers = managers
                resolver.fulfill(())
            }
        }
    }

    func importConfig(config: URL) throws {
        do {
            try importConfig(configString: try String(contentsOf: config), title: config.deletingPathExtension().lastPathComponent)
        } catch {
            throw AppCoordinatorError.configImportError(msg: "Failed")
        }
    }

    func importConfig(configString: String, title: String) throws {
        do {
            let addContext = persistentContainer.newBackgroundContext()
            let tunnel = try Tunnel.fromConfig(configString, context: addContext)
            tunnel.title = title
            addContext.saveContext()
            self.saveTunnel(tunnel)
        } catch {
            throw AppCoordinatorError.configImportError(msg: "Failed")
        }
    }

    func importConfigs(configZip: URL) throws {
        if let archive = Archive(url: configZip, accessMode: .read) {
            for entry in archive {
                var entryData = Data(capacity: 0)
                _ = try archive.extract(entry) { (data) in
                    entryData.append(data)
                }
                if let config = String(data: entryData, encoding: .utf8) {
                    try importConfig(configString: config, title: entry.path)
                }
            }
        }
    }

    func checkAndCleanConfigs() {
        _ = refreshProviderManagers().then { () -> Promise<Void> in
            guard let providerManagers = self.providerManagers else {
                return Promise.value(())
            }
            let tunnels = try Tunnel.allInContext(self.persistentContainer.viewContext)
            let tunnelIdentifiers = tunnels.compactMap {$0.tunnelIdentifier}

            let unknownManagers = providerManagers.filter {
                guard let prot = $0.protocolConfiguration as? NETunnelProviderProtocol else {
                    return false
                }
                guard let candidateTunnelIdentifier = prot.providerConfiguration?[PCKeys.tunnelIdentifier.rawValue] as? String else {
                    return false
                }

                return !tunnelIdentifiers.contains(candidateTunnelIdentifier)
            }

            let deletionPromises = unknownManagers.map({ (manager) -> Promise<NETunnelProviderManager> in
                return Promise(resolver: { resolver in
                    return manager.removeFromPreferences(completionHandler: { (error) in
                        if let error = error {
                            resolver.reject(error)
                        } else {
                            resolver.fulfill(manager)
                        }
                    })
                })
            })

            return when(resolved: deletionPromises).asVoid()

        }
    }

    // swiftlint:disable next function_body_length
    func exportConfigs(sourceView: UIView) {
        guard let path = FileManager.default
            .urls(for: .documentDirectory, in: .userDomainMask).first else {
                return
        }
        let saveFileURL = path.appendingPathComponent("wireguard-export.zip")
        do {
            try FileManager.default.removeItem(at: saveFileURL)
        } catch {
            os_log("Failed to delete file: %{public}@ : %{public}@", log: Log.general, type: .error, saveFileURL.absoluteString, error.localizedDescription)
        }

        guard let archive = Archive(url: saveFileURL, accessMode: .create) else {
            return
        }

        do {
            var tunnelsByTitle = [String: [Tunnel]]()
            let tunnels = try Tunnel.allInContext(persistentContainer.viewContext)
            tunnels.forEach {
                guard let title = $0.title ?? $0.tunnelIdentifier else {
                    // there is always a tunnelidentifier.
                    return
                }
                if let tunnels = tunnelsByTitle[title] {
                    tunnelsByTitle[title] = tunnels + [$0]
                } else {
                    tunnelsByTitle[title] = [$0]
                }
            }

            func addEntry(title: String, tunnel: Tunnel) throws {
                let data = tunnel.export().data(using: .utf8)!
                let byteCount: UInt32 = UInt32(data.count)
                try archive.addEntry(with: "\(title).conf", type: .file, uncompressedSize: byteCount, provider: { (position, size) -> Data in
                    return data.subdata(in: position ..< size)
                })
            }

            try tunnelsByTitle.keys.forEach {
                if let tunnels = tunnelsByTitle[$0] {
                    if tunnels.count == 1 {
                        try addEntry(title: $0, tunnel: tunnels[0])
                    } else {
                        for (index, tunnel) in tunnels.enumerated() {
                            try addEntry(title: $0 + "-\(index + 1)", tunnel: tunnel)
                        }
                    }
                }
            }
        } catch {
            os_log("Failed to create archive file: %{public}@ : %{public}@", log: Log.general, type: .error, saveFileURL.absoluteString, error.localizedDescription)
            return
        }

        let activityViewController = UIActivityViewController(
            activityItems: [saveFileURL],
            applicationActivities: nil)
        if let popoverPresentationController = activityViewController.popoverPresentationController {
            popoverPresentationController.sourceView = sourceView
        }
        navigationController.present(activityViewController, animated: true) {
        }
    }

    func exportConfig(tunnel: Tunnel, barButtonItem: UIBarButtonItem) {
        let exportString = tunnel.export()

        guard let path = FileManager.default
            .urls(for: .documentDirectory, in: .userDomainMask).first else {
                return
        }
        let saveFileURL = path.appendingPathComponent("/\(tunnel.title ?? "wireguard").conf")
        do {
            try exportString.write(to: saveFileURL, atomically: true, encoding: .utf8)
        } catch {
            os_log("Failed to export tunnel to: %{public}@", log: Log.general, type: .error, saveFileURL.absoluteString)
            return
        }

        let activityViewController = UIActivityViewController(
            activityItems: [saveFileURL],
            applicationActivities: nil)
        if let popoverPresentationController = activityViewController.popoverPresentationController {
            popoverPresentationController.barButtonItem = barButtonItem
        }
        self.navigationController.present(activityViewController, animated: true) {
        }
    }

    // MARK: - NEVPNManager handling

    @objc private func VPNStatusDidChange(notification: NSNotification) {
        guard let session = notification.object as? NETunnelProviderSession else {
            return
        }

        guard let prot = session.manager.protocolConfiguration as? NETunnelProviderProtocol else {
            return
        }

        guard let changedTunnelIdentifier = prot.providerConfiguration?[PCKeys.tunnelIdentifier.rawValue] as? String else {
            return
        }

        providerManagers?.first(where: { (manager) -> Bool in
            guard let prot = manager.protocolConfiguration as? NETunnelProviderProtocol else {
                return false
            }
            guard let candidateTunnelIdentifier = prot.providerConfiguration?[PCKeys.tunnelIdentifier.rawValue] as? String else {
                return false
            }

            return changedTunnelIdentifier == candidateTunnelIdentifier

        })?.loadFromPreferences(completionHandler: { [weak self] (_) in
            self?.tunnelsTableViewController.updateStatus(for: changedTunnelIdentifier)
        })
    }

    func showTunnelInfoViewController(tunnel: Tunnel, context: NSManagedObjectContext) {
        let tunnelInfoViewController = storyboard.instantiateViewController(type: TunnelInfoTableViewController.self)

        tunnelInfoViewController.configure(context: context, delegate: self, tunnel: tunnel)

        self.navigationController.pushViewController(tunnelInfoViewController, animated: true)
    }

    func showTunnelConfigurationViewController(tunnel: Tunnel?, context: NSManagedObjectContext) {
        let tunnelConfigurationViewController = storyboard.instantiateViewController(type: TunnelConfigurationTableViewController.self)

        tunnelConfigurationViewController.configure(context: context, delegate: self, tunnel: tunnel)

        self.navigationController.pushViewController(tunnelConfigurationViewController, animated: true)
    }

    func showSettings() {
        let settingsTableViewController = storyboard.instantiateViewController(type: SettingsTableViewController.self)

        settingsTableViewController.delegate = self

        self.navigationController.pushViewController(settingsTableViewController, animated: true)
    }

    public func showError(_ error: Error) {
        showAlert(title: NSLocalizedString("Error", comment: "Error alert title"), message: error.localizedDescription)
    }

    func connect(tunnel: Tunnel) {
        _ = refreshProviderManagers().then { () -> Promise<Void> in
            guard let manager = self.providerManager(for: tunnel) else {
                return Promise.value(())
            }
            let block = {
                switch manager.connection.status {
                case .invalid, .disconnected:
                    os_log("connect tunnel: %{public}@", log: Log.general, type: .info, tunnel.description)
                    // Should the manager be enabled?

                    let manager = self.providerManager(for: tunnel)
                    manager?.isEnabled = true
                    manager?.saveToPreferences { (error) in
                        if let error = error {
                            os_log("error saving preferences: %{public}@", log: Log.general, type: .error, error.localizedDescription)
                            return
                        }
                        os_log("saved preferences", log: Log.general, type: .info)

                        let session = manager?.connection as! NETunnelProviderSession //swiftlint:disable:this force_cast
                        do {
                            try session.startTunnel()
                        } catch let error {
                            os_log("error starting tunnel: %{public}@", log: Log.general, type: .error, error.localizedDescription)
                        }
                    }

                default:
                    break
                }
            }

            if manager.connection.status == .invalid {
                manager.loadFromPreferences { (_) in
                    block()
                }
            } else {
                block()
            }

            return Promise.value(())
        }
    }

    func disconnect(tunnel: Tunnel) {
        _ = refreshProviderManagers().then { () -> Promise<Void> in
            let manager = self.providerManager(for: tunnel)!
            let block = {
                switch manager.connection.status {
                case .connected, .connecting:
                    let manager = self.providerManager(for: tunnel)
                    manager?.connection.stopVPNTunnel()
                default:
                    break
                }
            }

            if manager.connection.status == .invalid {
                manager.loadFromPreferences { (_) in
                    block()
                }
            } else {
                block()
            }
            return Promise.value(())
        }
    }

    private func showAlert(title: String, message: String) {
        let alert = UIAlertController(title: title, message: message, preferredStyle: .alert)
        alert.addAction(UIAlertAction(title: NSLocalizedString("OK", comment: "OK button"), style: .default))
        self.navigationController.present(alert, animated: true)
    }

    func providerManager(for tunnel: Tunnel) -> NETunnelProviderManager? {
        return self.providerManagers?.first {
            guard let prot = $0.protocolConfiguration as? NETunnelProviderProtocol else {
                return false
            }
            guard let tunnelIdentifier = prot.providerConfiguration?[PCKeys.tunnelIdentifier.rawValue] as? String else {
                return false
            }
            return tunnelIdentifier == tunnel.tunnelIdentifier
        }
    }

    func extensionGoVersionInformation() -> Promise<String> {
        return Promise(resolver: { (resolver) in
            guard let session = self.providerManagers?.first(where: { $0.isEnabled })?.connection as? NETunnelProviderSession else {
                resolver.reject(GoVersionCoordinatorError.noEnabledSession)
                return
            }
            do {
                try session.sendProviderMessage(ExtensionMessage.requestVersion.data, responseHandler: { (data) in
                    guard let data = data, let responseString = String(data: data, encoding: .utf8) else {
                        resolver.reject(GoVersionCoordinatorError.noResponse)
                        return
                    }
                    resolver.fulfill(responseString)
                })
            } catch {
                resolver.reject(error)
            }
        })
    }
}

class AppDocumentPickerDelegate: NSObject, UIDocumentPickerDelegate {
    weak var appCoordinator: AppCoordinator?

    func documentPicker(_ controller: UIDocumentPickerViewController, didPickDocumentAt url: URL) {
        if url.pathExtension == "conf" {
            do {
                try appCoordinator?.importConfig(config: url)
            } catch {
                os_log("Unable to import config: %{public}@", log: Log.general, type: .error, url.absoluteString)
            }
        } else if url.pathExtension == "zip" {
            do {
                try appCoordinator?.importConfigs(configZip: url)
            } catch {
                os_log("Unable to import config: %{public}@", log: Log.general, type: .error, url.absoluteString)
            }
        }

    }
}

extension NEVPNStatus {
    var statusDescription: String {
        switch self {
        case .connected:
            return "Connected"
        case .connecting:
            return "Connecting"
        case .disconnected:
            return "Disconnected"
        case .disconnecting:
            return "Disconnecting"
        case .invalid:
            return "Invalid"
        case .reasserting:
            return "Reasserting"
        }
    }
}