//
//  PacketTunnelProvider.swift
//  punchnet
//
//  Created by 安礼成 on 2025/8/3.
//


//
//  PacketTunnelProvider.swift
//  Tun
//
//  Created by 安礼成 on 2024/1/17.
//

import NetworkExtension

class PacketTunnelProvider: NEPacketTunnelProvider {
    var context: SDLContext?
    private var rootTask: Task<Void, Error>?
    
    override func startTunnel(options: [String : NSObject]?, completionHandler: @escaping (Error?) -> Void) {
        // host: "192.168.0.101", port: 1265
        guard let options else {
            return
        }
        
        // 如果当前在运行状态，不允许重复请求
        guard self.context == nil else {
            return
        }
        
        // let token = options["token"] as! String
        let installed_channel = options["installed_channel"] as! String
        let superIp = options["super_ip"] as! String
        let superPort = options["super_port"] as! Int
        let stunServersStr = options["stun_servers"] as! String
        let noticePort = options["notice_port"] as! Int
        let token = options["token"] as! String
        let networkCode = options["network_code"] as! String
        let clientId = options["client_id"] as! String
        let remoteDnsServer = options["remote_dns_server"] as! String
        let hostname = options["hostname"] as! String

        let stunServers = stunServersStr.split(separator: ";").compactMap { server -> SDLConfiguration.StunServer? in
            let parts = server.split(separator: ":", maxSplits: 2)
            guard parts.count == 2 else {
                return nil
            }
            
            let ports = parts[1].split(separator: ",", maxSplits: 2)
            guard ports.count == 2, let port1 = Int(String(ports[0])), let port2 = Int(String(ports[1])) else {
                return nil
            }
            
            return .init(host: String(parts[0]), ports: [port1, port2])
        }
        
        guard stunServers.count >= 2 else {
            NSLog("stunServers配置错误")
            return
        }
        
        NSLog("[PacketTunnelProvider] client_id: \(clientId), token: \(token), network_code: \(networkCode)")

        let config = SDLConfiguration(version: 1,
                                      installedChannel: installed_channel,
                                      superHost: superIp,
                                      superPort: superPort,
                                      stunServers: stunServers,
                                      clientId: clientId,
                                      noticePort: noticePort,
                                      token: token,
                                      networkCode: networkCode,
                                      remoteDnsServer: remoteDnsServer,
                                      hostname: hostname)
        // 加密算法
        let rsaCipher = try! CCRSACipher(keySize: 1024)
        let aesChiper = CCAESChiper()
        
        self.rootTask = Task {
            do {
                self.context = SDLContext(provider: self, config: config, rsaCipher: rsaCipher, aesCipher: aesChiper, logger: SDLLogger(level: .debug))
                try await self.context?.start()
            } catch let err {
                NSLog("[PacketTunnelProvider] exit with error: \(err)")
                exit(-1)
            }
        }
        completionHandler(nil)
    }
    
    override func stopTunnel(with reason: NEProviderStopReason, completionHandler: @escaping () -> Void) {
        // Add code here to start the process of stopping the tunnel.
        self.rootTask?.cancel()
        Task {
            await self.context?.stop()
        }
        self.context = nil
        self.rootTask = nil
        
        completionHandler()
    }
    
    override func handleAppMessage(_ messageData: Data, completionHandler: ((Data?) -> Void)?) {
        // Add code here to handle the message.
        if let handler = completionHandler {
            handler(messageData)
        }
    }
    
    override func sleep(completionHandler: @escaping () -> Void) {
        // Add code here to get ready to sleep.
        completionHandler()
    }
    
    override func wake() {
        // Add code here to wake up.
    }
    
}

// 获取物理网卡ip地址
extension PacketTunnelProvider {
    
    public static var viaInterface: NetworkInterface? = {
        let interfaces = NetworkInterfaceManager.getInterfaces()
        
        return interfaces.first {$0.name == "en0"}
    }()
    
    struct CCRSACipher: RSACipher {
        var pubKey: String
        let privateKeyDER: Data
        
        init(keySize: Int) throws {
            let (privateKey, publicKey) = try Self.loadKeys(keySize: keySize)
            let privKeyStr = SwKeyConvert.PrivateKey.derToPKCS1PEM(privateKey)
            
            self.pubKey = SwKeyConvert.PublicKey.derToPKCS8PEM(publicKey)
            self.privateKeyDER = try SwKeyConvert.PrivateKey.pemToPKCS1DER(privKeyStr)
        }
        
        public func decode(data: Data) throws -> Data {
            let tag = Data()
            let (decryptedData, _) = try CC.RSA.decrypt(data, derKey: self.privateKeyDER, tag: tag, padding: .pkcs1, digest: .none)
            
            return decryptedData
        }
        
        private static func loadKeys(keySize: Int) throws -> (Data, Data) {
            if let privateKey = UserDefaults.standard.data(forKey: "privateKey"),
               let publicKey = UserDefaults.standard.data(forKey: "publicKey") {
                
                return (privateKey, publicKey)
            } else {
                let (privateKey, publicKey) = try CC.RSA.generateKeyPair(keySize)
                UserDefaults.standard.setValue(privateKey, forKey: "privateKey")
                UserDefaults.standard.setValue(publicKey, forKey: "publicKey")
                
                return (privateKey, publicKey)
            }
        }
    }
    
    struct CCAESChiper: AESCipher {
        func decypt(aesKey: Data, data: Data) throws -> Data {
            let ivData = Data(aesKey.prefix(16))
            return try CC.crypt(.decrypt, blockMode: .cbc, algorithm: .aes, padding: .pkcs7Padding, data: data, key: aesKey, iv: ivData)
        }
        
        func encrypt(aesKey: Data, data: Data) throws -> Data {
            let ivData = Data(aesKey.prefix(16))
            
            return try CC.crypt(.encrypt, blockMode: .cbc, algorithm: .aes, padding: .pkcs7Padding, data: data, key: aesKey, iv: ivData)
        }
    }
    
}