//
//  PacketTunnelProvider.swift
//  TunnelProv
//
//  Created by Stossy11 on 28/03/2025.
//

import NetworkExtension

class PacketTunnelProvider: NEPacketTunnelProvider {
    var tunnelDeviceIp: String = "10.7.0.0"
    var tunnelFakeIp: String = "10.7.0.1"
    var tunnelSubnetMask: String = "255.255.255.0"
    
    private var deviceIpValue: UInt32 = 0
    private var fakeIpValue: UInt32 = 0
    
    override func startTunnel(options: [String : NSObject]?, completionHandler: @escaping (Error?) -> Void) {
        if let deviceIp = options?["TunnelDeviceIP"] as? String {
            tunnelDeviceIp = deviceIp
        }
        if let fakeIp = options?["TunnelFakeIP"] as? String {
            tunnelFakeIp = fakeIp
        }
        
        deviceIpValue = ipToUInt32(tunnelDeviceIp)
        fakeIpValue = ipToUInt32(tunnelFakeIp)
        
        let settings = NEPacketTunnelNetworkSettings(tunnelRemoteAddress: tunnelDeviceIp)
        let ipv4 = NEIPv4Settings(addresses: [tunnelDeviceIp], subnetMasks: [tunnelSubnetMask])
        ipv4.includedRoutes = [NEIPv4Route(destinationAddress: tunnelDeviceIp, subnetMask: tunnelSubnetMask)]
        ipv4.excludedRoutes = [.default()]
        settings.ipv4Settings = ipv4
        
        setTunnelNetworkSettings(settings) { error in
            if error == nil {
                self.readPackets()
            }
            completionHandler(error)
        }
    }
    
    override func stopTunnel(with reason: NEProviderStopReason, completionHandler: @escaping () -> Void) {
        completionHandler()
    }
    
    override func handleAppMessage(_ messageData: Data, completionHandler: ((Data?) -> Void)?) {
        completionHandler?(messageData)
    }
    
    override func sleep(completionHandler: @escaping () -> Void) {
        completionHandler()
    }
    
    override func wake() {}
    
    private func readPackets() {
        packetFlow.readPackets { packets, protocols in
            var output = [Data](repeating: Data(), count: packets.count)
            
            for (i, packet) in packets.enumerated() {
                guard protocols[i].int32Value == AF_INET, packet.count >= 20 else {
                    output[i] = packet
                    continue
                }
                
                output[i] = self.processPacket(packet)
            }
            
            self.packetFlow.writePackets(output, withProtocols: protocols)
            self.readPackets()
        }
    }
    
    private func processPacket(_ packet: Data) -> Data {
        var bytes = [UInt8](packet)
        
        let srcIP = UInt32(bigEndian: bytes.withUnsafeBytes { $0.load(fromByteOffset: 12, as: UInt32.self) })
        let dstIP = UInt32(bigEndian: bytes.withUnsafeBytes { $0.load(fromByteOffset: 16, as: UInt32.self) })
        
        if srcIP == deviceIpValue {
            let replacement = fakeIpValue.bigEndian
            withUnsafeBytes(of: replacement) { bytes.replaceSubrange(12..<16, with: $0) }
        }
        if dstIP == fakeIpValue {
            let replacement = deviceIpValue.bigEndian
            withUnsafeBytes(of: replacement) { bytes.replaceSubrange(16..<20, with: $0) }
        }
        
        bytes.swapAt(12, 16)
        bytes.swapAt(13, 17)
        bytes.swapAt(14, 18)
        bytes.swapAt(15, 19)
        
        return Data(bytes)
    }
    
    private func ipToUInt32(_ ipString: String) -> UInt32 {
        let components = ipString.split(separator: ".")
        guard components.count == 4,
              let b1 = UInt32(components[0]),
              let b2 = UInt32(components[1]),
              let b3 = UInt32(components[2]),
              let b4 = UInt32(components[3]) else {
            return 0
        }
        return (b1 << 24) | (b2 << 16) | (b3 << 8) | b4
    }
}