appchk-app/GlassVPN/zhuhaow-NEKit/IPStack/DNS/DNSMessage.swift

import Foundation

open class DNSMessage {
    //    var sourceAddress: IPv4Address?
    //    var sourcePort: Port?
    //    var destinationAddress: IPv4Address?
    //    var destinationPort: Port?
    open var transactionID: UInt16 = 0
    open var messageType: DNSMessageType = .query
    open var authoritative: Bool = false
    open var truncation: Bool = false
    open var recursionDesired: Bool = false
    open var recursionAvailable: Bool = false
    open var status: DNSReturnStatus = .success
    open var queries: [DNSQuery] = []
    open var answers: [DNSResource] = []
    open var nameservers: [DNSResource] = []
    open var addtionals: [DNSResource] = []

    var payload: Data!

    var bytesLength: Int {
        var len = 12 + queries.reduce(0) {
            $0 + $1.bytesLength
        }
        len += answers.reduce(0) {
            $0 + $1.bytesLength
        }
        len += nameservers.reduce(0) {
            $0 + $1.bytesLength
        }
        len += addtionals.reduce(0) {
            $0 + $1.bytesLength
        }
        return len
    }

    var resolvedIPv4Address: IPAddress? {
        for answer in answers {
            if let address = answer.ipv4Address {
                return address
            }
        }
        return nil
    }

    var type: DNSType? {
        return queries.first?.type
    }

    init() {}

    init?(payload: Data) {
        self.payload = payload
        let scanner = BinaryDataScanner(data: payload, littleEndian: false)

        transactionID = scanner.read16()!

        var bytes = scanner.readByte()!
        if bytes & 0x80 > 0 {
            messageType = .response
        } else {
            messageType = .query
        }

        // ignore OP code

        authoritative = bytes & 0x04 > 0
        truncation = bytes & 0x02 > 0
        recursionDesired = bytes & 0x01 > 0

        bytes = scanner.readByte()!
        recursionAvailable = bytes & 0x80 > 0
        if let status = DNSReturnStatus(rawValue: bytes & 0x0F) {
            self.status = status
        } else {
            DDLogError("Received DNS response with unknown status: \(bytes & 0x0F).")
            self.status = .serverFailure
        }

        let queryCount = scanner.read16()!
        let answerCount = scanner.read16()!
        let nameserverCount = scanner.read16()!
        let addtionalCount = scanner.read16()!

        for _ in 0..<queryCount {
            queries.append(DNSQuery(payload: payload, offset: scanner.position, base: 0)!)
            scanner.advance(by: queries.last!.bytesLength)
        }

        for _ in 0..<answerCount {
            answers.append(DNSResource(payload: payload, offset: scanner.position, base: 0)!)
            scanner.advance(by: answers.last!.bytesLength)
        }

        for _ in 0..<nameserverCount {
            nameservers.append(DNSResource(payload: payload, offset: scanner.position, base: 0)!)
            scanner.advance(by: nameservers.last!.bytesLength)
        }

        for _ in 0..<addtionalCount {
            addtionals.append(DNSResource(payload: payload, offset: scanner.position, base: 0)!)
            scanner.advance(by: addtionals.last!.bytesLength)
        }

    }

    func buildMessage() -> Bool {
        payload = Data(count: bytesLength)
        if transactionID == 0 {
            transactionID = UInt16(arc4random_uniform(UInt32(UInt16.max)))
        }
        setPayloadWithUInt16(transactionID, at: 0, swap: true)
        var byte: UInt8 = 0
        byte += messageType.rawValue << 7
        if authoritative {
            byte += 4
        }
        if truncation {
            byte += 2
        }
        if recursionDesired {
            byte += 1
        }
        setPayloadWithUInt8(byte, at: 2)
        byte = 0
        if recursionAvailable {
            byte += 128
        }

        byte += status.rawValue

        setPayloadWithUInt8(byte, at: 3)
        setPayloadWithUInt16(UInt16(queries.count), at: 4, swap: true)
        setPayloadWithUInt16(UInt16(answers.count), at: 6, swap: true)
        setPayloadWithUInt16(UInt16(nameservers.count), at: 8, swap: true)
        setPayloadWithUInt16(UInt16(addtionals.count), at: 10, swap: true)

        return writeAllRecordAt(12)
    }

    // swiftlint:disable variable_name
    func setPayloadWithUInt8(_ value: UInt8, at: Int) {
        var v = value
        withUnsafeBytes(of: &v) {
            payload.replaceSubrange(at..<at+1, with: $0)
        }
    }

    func setPayloadWithUInt16(_ value: UInt16, at: Int, swap: Bool = false) {
        var v: UInt16
        if swap {
            v = NSSwapHostShortToBig(value)
        } else {
            v = value
        }
        withUnsafeBytes(of: &v) {
            payload.replaceSubrange(at..<at+2, with: $0)
        }
    }

    func setPayloadWithUInt32(_ value: UInt32, at: Int, swap: Bool = false) {
        var v: UInt32
        if swap {
            v = NSSwapHostIntToBig(value)
        } else {
            v = value
        }
        withUnsafeBytes(of: &v) {
            payload.replaceSubrange(at..<at+4, with: $0)
        }
    }

    func setPayloadWithData(_ data: Data, at: Int, length: Int? = nil, from: Int = 0) {
        let length = length ?? data.count - from

        payload.withUnsafeMutableBytes { ptr in
            data.copyBytes(to: ptr.baseAddress!.advanced(by: at).assumingMemoryBound(to: UInt8.self), from: from..<from+length)
        }
    }

    func resetPayloadAt(_ at: Int, length: Int) {
        payload.resetBytes(in: at..<at+length)
    }

    fileprivate func writeAllRecordAt(_ at: Int) -> Bool {
        var position = at
        for query in queries {
            guard writeDNSQuery(query, at: position) else {
                return false
            }
            position += query.bytesLength
        }
        for resources in [answers, nameservers, addtionals] {
            for resource in resources {
                guard writeDNSResource(resource, at: position) else {
                    return false
                }
                position += resource.bytesLength
            }
        }
        return true
    }

    fileprivate func writeDNSQuery(_ query: DNSQuery, at: Int) -> Bool {
        guard DNSNameConverter.setName(query.name, toData: &payload!, at: at) else {
            return false
        }
        setPayloadWithUInt16(query.type.rawValue, at: at + query.nameBytesLength, swap: true)
        setPayloadWithUInt16(query.klass.rawValue, at: at + query.nameBytesLength + 2, swap: true)
        return true
    }

    fileprivate func writeDNSResource(_ resource: DNSResource, at: Int) -> Bool {
        guard DNSNameConverter.setName(resource.name, toData: &payload!, at: at) else {
            return false
        }
        setPayloadWithUInt16(resource.type.rawValue, at: at + resource.nameBytesLength, swap: true)
        setPayloadWithUInt16(resource.klass.rawValue, at: at + resource.nameBytesLength + 2, swap: true)
        setPayloadWithUInt32(resource.TTL, at: at + resource.nameBytesLength + 4, swap: true)
        setPayloadWithUInt16(resource.dataLength, at: at + resource.nameBytesLength + 8, swap: true)
        setPayloadWithData(resource.data, at: at + resource.nameBytesLength + 10)
        return true
    }
}

open class DNSQuery {
    public let name: String
    public let type: DNSType
    public let klass: DNSClass
    let nameBytesLength: Int

    init(name: String, type: DNSType = .a, klass: DNSClass = .internet) {
        self.name = name.trimmingCharacters(in: CharacterSet(charactersIn: "."))
        self.type = type
        self.klass = klass
        self.nameBytesLength = name.utf8.count + 2
    }

    init?(payload: Data, offset: Int, base: Int = 0) {
        (self.name, self.nameBytesLength) = DNSNameConverter.getNamefromData(payload, offset: offset, base: base)

        let scanner = BinaryDataScanner(data: payload, littleEndian: false)
        scanner.skip(to: offset + self.nameBytesLength)

        guard let type = DNSType(rawValue: scanner.read16()!) else {
            DDLogError("Received DNS packet with unknown type.")
            return nil
        }
        self.type = type

        guard let klass = DNSClass(rawValue: scanner.read16()!) else {
            DDLogError("Received DNS packet with unknown class.")
            return nil
        }
        self.klass = klass

    }

    var bytesLength: Int {
        return nameBytesLength + 4
    }
}

open class DNSResource {
    public let name: String
    public let type: DNSType
    public let klass: DNSClass
    public let TTL: UInt32
    let dataLength: UInt16
    public let data: Data

    let nameBytesLength: Int

    init(name: String, type: DNSType = .a, klass: DNSClass = .internet, TTL: UInt32 = 300, data: Data) {
        self.name = name
        self.type = type
        self.klass = klass
        self.TTL = TTL
        dataLength = UInt16(data.count)
        self.data = data
        self.nameBytesLength = name.utf8.count + 2
    }

    static func ARecord(_ name: String, TTL: UInt32 = 300, address: IPAddress) -> DNSResource {
        return DNSResource(name: name, type: .a, klass: .internet, TTL: TTL, data: address.dataInNetworkOrder)
    }

    init?(payload: Data, offset: Int, base: Int = 0) {
        (self.name, self.nameBytesLength) = DNSNameConverter.getNamefromData(payload, offset: offset, base: base)

        let scanner = BinaryDataScanner(data: payload, littleEndian: false)
        scanner.skip(to: offset + self.nameBytesLength)

        guard let type = DNSType(rawValue: scanner.read16()!) else {
            DDLogError("Received DNS packet with unknown type.")
            return nil
        }
        self.type = type

        guard let klass = DNSClass(rawValue: scanner.read16()!) else {
            DDLogError("Received DNS packet with unknown class.")
            return nil
        }
        self.klass = klass
        self.TTL = scanner.read32()!
        dataLength = scanner.read16()!
        self.data = payload.subdata(in: scanner.position..<scanner.position+Int(dataLength))
    }

    var bytesLength: Int {
        return nameBytesLength + 10 + Int(dataLength)
    }

    var ipv4Address: IPAddress? {
        guard type == .a else {
            return nil
        }
        return IPAddress(fromBytesInNetworkOrder: (data as NSData).bytes)
    }
}

class DNSNameConverter {
    static func setName(_ name: String, toData data: inout Data, at: Int) -> Bool {
        let labels = name.components(separatedBy: CharacterSet(charactersIn: "."))
        var position = at

        for label in labels {
            let len = label.utf8.count
            guard len != 0 else {
                // invalid domain name
                return false
            }
            data[position] = UInt8(len)
            position += 1

            data.replaceSubrange(position..<position+len, with: label.data(using: .utf8)!)
            position += len
        }
        data[position] = 0
        return true
    }

    static func getNamefromData(_ data: Data, offset: Int, base: Int = 0) -> (String, Int) {
        let scanner = BinaryDataScanner(data: data, littleEndian: false)
        scanner.skip(to: offset)

        var len: UInt8 = 0
        var name = ""
        var currentReadBytes = 0
        var jumped = false
        var nameBytesLength = 0
        repeat {
            let length = scanner.read16()!
            // is this a pointer?
            if length & 0xC000 == 0xC000 {
                if !jumped {
                    // save the length position
                    nameBytesLength = 2 + currentReadBytes
                    jumped = true
                }
                scanner.skip(to: Int(length & 0x3FFF) + base)
            } else {
                scanner.advance(by: -2)
            }

            len = scanner.readByte()!
            currentReadBytes += 1
            if len == 0 {
                break
            }

            currentReadBytes += Int(len)

            guard let label = String(bytes: scanner.data.subdata(in: scanner.position..<scanner.position+Int(len)), encoding: .utf8) else {
                return ("", currentReadBytes)
            }
            // this is not efficient, but won't take much time, so maybe I'll optimize it later
            name = name.appendingFormat(".%@", label)
            scanner.advance(by: Int(len))
        } while true

        if !jumped {
            nameBytesLength = currentReadBytes
        }

        return (name.trimmingCharacters(in: CharacterSet(charactersIn: ".")), nameBytesLength)
    }
}