/*
 DnsFilter — Swift-Port von `android/.../filter/DnsFilter.kt`.

 IPv4 + UDP + DNS Packet-Parsing und Response-Builder.

 Verarbeitungsmodell (identisch zu Android):
  - Block-Query  → synthetische NXDOMAIN-Response (RCODE=3), synchron gebaut.
  - Allow-Query  → an Upstream-Resolver (1.1.1.1:53) forwarden, Response wird
                   wieder als IPv4/UDP-Paket verpackt zurückgegeben.
  - Upstream-Fehler → synthetische SERVFAIL-Response (RCODE=2), damit der
                   Client schnell failt statt in seinen Timeout zu laufen.
  - Drop (kein DNS / IPv6 / Nicht-UDP) → nil, das Paket wird verworfen.

 Unterschied zu Android: Android schreibt selbst ins TUN (FileOutputStream).
 Auf iOS schreibt der `NEPacketTunnelProvider` über `packetFlow.writePackets`.
 Daher liefert dieser Filter das fertige Response-Paket nur ZURÜCK; der
 PacketTunnelProvider übernimmt das Schreiben (synchron für Block, via
 Completion-Handler für Forward).

 ── UNGETESTETE ANNAHME ──────────────────────────────────────────────────────
 Das exakte Byte-Layout der NXDOMAIN-Response und der forwardeten Antwort ist
 1:1 aus dem produktiv erprobten Android-`DnsFilter.kt` übernommen. Es konnte
 NICHT auf einem echten iPhone verifiziert werden (NE-Extensions laufen nicht
 im Simulator). Die Annahme: `packetFlow` liefert/erwartet rohe IPv4-Pakete
 inkl. IP-Header (wie Androids TUN-FD). Apple-Doku bestätigt das für
 `NEPacketTunnelFlow` — `protocols` enthält `AF_INET`/`AF_INET6`.
*/

import Foundation
import Network
import os

/// Ergebnis der Paket-Verarbeitung.
enum DnsFilterDecision {
  /// Domain ist geblockt — synthetische NXDOMAIN-Response, sofort schreiben.
  case block(response: Data, domain: String)
  /// Domain ist erlaubt — Query muss an Upstream forwardet werden.
  /// `domain` nur fürs Logging.
  case forward(domain: String)
  /// Paket ist irrelevant (kein IPv4/UDP/DNS) → verwerfen.
  case drop
}

enum DnsFilter {

  private static let log = Logger(
    subsystem: "org.rebreak.app.packettunnel", category: "dnsfilter")

  private static let dnsPort: UInt16 = 53

  // ─── Entry: ein IPv4-Paket vom TUN klassifizieren ──────────────────────────

  /// Parst ein IPv4-Paket und entscheidet block / forward / drop.
  ///
  /// - Parameter packet: rohes IPv4-Paket inkl. IP-Header (wie vom TUN gelesen).
  /// - Parameter hashList: die geladene Blocklist.
  static func classify(packet: Data, hashList: HashList) -> DnsFilterDecision {
    let length = packet.count
    // min IPv4(20) + UDP(8)
    guard length >= 28 else { return .drop }

    return packet.withUnsafeBytes { (raw: UnsafeRawBufferPointer) -> DnsFilterDecision in
      let bytes = raw.bindMemory(to: UInt8.self)

      // IPv4-Check (höchste 4 Bit der Version/IHL-Byte).
      let version = (Int(bytes[0]) >> 4) & 0xF
      guard version == 4 else { return .drop }

      // IHL = Internet Header Length in 32-bit-Words → in Bytes.
      let ihl = (Int(bytes[0]) & 0xF) * 4
      guard ihl >= 20, length >= ihl + 8 else { return .drop }

      // UDP only (IP-Protocol 17).
      let proto = Int(bytes[9]) & 0xFF
      guard proto == 17 else { return .drop }

      // UDP-Header.
      let udpStart = ihl
      let dstPort = (UInt16(bytes[udpStart + 2]) << 8) | UInt16(bytes[udpStart + 3])
      guard dstPort == dnsPort else { return .drop }

      let dnsStart = udpStart + 8
      let dnsLen = length - dnsStart
      guard dnsLen >= 12 else { return .drop }

      // QNAME parsen — beginnt nach dem 12-Byte-DNS-Header.
      // H3-Fix: `parseQname` liefert `nil`, wenn die QNAME nicht parsebar ist
      // (Compression-Pointer in der Frage oder Out-of-Bounds). Früher führte
      // das zu `.drop` → die Auflösung schlug für den Client stillschweigend
      // fehl. Korrekt ist FAIL-OPEN: solche Pakete an den Upstream forwarden,
      // damit reguläre Namensauflösung nicht heimlich kaputtgeht. Ein
      // unparsebares QNAME bedeutet nur, dass WIR es nicht matchen können —
      // nicht, dass die Query verworfen werden darf.
      guard let domain = parseQname(bytes, start: dnsStart + 12, end: length) else {
        return .forward(domain: "<unparsebar>")
      }

      if hashList.matchesAnySuffix(domain) {
        let response = buildErrorResponse(
          bytes, length: length, ihl: ihl, rcode: 3 /* NXDOMAIN */)
        return .block(response: response, domain: domain)
      }
      return .forward(domain: domain)
    }
  }

  // ─── QNAME-Parser ──────────────────────────────────────────────────────────

  /// Parst den QNAME ab `start`. Gibt `nil` zurück bei DNS-Compression-Pointer
  /// (0xC0-Bit gesetzt) oder Out-of-Bounds.
  ///
  /// `nil` heißt NICHT „verwerfen": der Aufrufer (`classify`) behandelt `nil`
  /// als FAIL-OPEN und forwardet die Query an den Upstream (H3-Fix). Ein
  /// nicht-parsebares QNAME darf nie zum stillen Verlust der Auflösung führen.
  private static func parseQname(
    _ bytes: UnsafeBufferPointer<UInt8>, start: Int, end: Int
  ) -> String? {
    var labels: [String] = []
    var pos = start
    while pos < end {
      let len = Int(bytes[pos]) & 0xFF
      if len == 0 { break }
      // Compression-Pointer (zwei höchste Bits) — in einer Query unüblich.
      // Nicht parsebar → nil → Aufrufer forwardet fail-open (H3-Fix).
      if (len & 0xC0) != 0 { return nil }
      pos += 1
      guard pos + len <= end else { return nil }
      var label = [UInt8]()
      label.reserveCapacity(len)
      for i in 0..<len {
        label.append(bytes[pos + i])
      }
      labels.append(String(decoding: label, as: UTF8.self))
      pos += len
    }
    let host = labels.joined(separator: ".")
    return host.isEmpty ? nil : host
  }

  // ─── Question-Section-Ende ─────────────────────────────────────────────────

  /// Ermittelt den Byte-Offset DIREKT HINTER der Question-Section, gemessen ab
  /// dem Paket-Anfang (nicht ab `dnsStart`).
  ///
  /// Layout DNS-Question: QNAME (variabel) + QTYPE (2 B) + QCLASS (2 B).
  /// QNAME ist eine Folge von Längen-prefixed Labels, terminiert durch ein
  /// 0x00-Byte. Ein Compression-Pointer (0xC0-Bit) kommt in einer Query
  /// normalerweise nicht vor — taucht er auf, geben wir `nil` zurück.
  ///
  /// Gibt `nil` zurück, wenn die Question-Section out-of-bounds läuft oder
  /// einen Pointer enthält. Der Aufrufer darf das Paket dann NICHT abschneiden
  /// und sollte fail-open forwarden.
  private static func dnsQuestionEnd(
    _ bytes: UnsafeBufferPointer<UInt8>, dnsStart: Int, length: Int
  ) -> Int? {
    // QNAME beginnt nach dem 12-Byte-DNS-Header.
    var pos = dnsStart + 12
    while pos < length {
      let len = Int(bytes[pos]) & 0xFF
      if len == 0 {
        // 0x00 terminiert die QNAME. Danach QTYPE(2) + QCLASS(2).
        let end = pos + 1 + 4
        guard end <= length else { return nil }
        return end
      }
      // Compression-Pointer in einer Query → nicht abschneidbar.
      if (len & 0xC0) != 0 { return nil }
      pos += 1 + len
    }
    return nil
  }

  // ─── DNS-Error-Response-Builder (NXDOMAIN / SERVFAIL) ──────────────────────

  /// Baut eine synthetische DNS-Error-Response aus der Request-Kopie.
  /// Generischer Builder für NXDOMAIN (RCODE=3) und SERVFAIL (RCODE=2).
  ///
  /// Logik (Basis aus Androids `DnsFilter.kt#buildNxDomainResponse`):
  ///  - Src-/Dst-IP tauschen, Src-/Dst-Port tauschen.
  ///  - DNS-Flags: QR=1, RD=1, RA=1, OPCODE/AA/TC=0, RCODE = Parameter.
  ///  - Answer/Authority/Additional-Counts auf 0.
  ///  - IP-Checksum neu berechnen, UDP-Checksum auf 0 (optional in IPv4).
  ///
  /// K1-Fix: Das Response-Paket wird NACH dem Ende der Question-Section
  /// abgeschnitten. Früher kopierte der Builder das gesamte Request-Paket und
  /// nullte nur ARCOUNT — ein evtl. vorhandener EDNS-OPT-Record blieb dann
  /// physisch als Müll-Bytes hinter der Question-Section stehen, obwohl
  /// ARCOUNT=0 sagt „keine Additional-Records". Manche Resolver/Clients
  /// verwerfen so eine Response als korrupt. Jetzt: Paket auf
  /// IP-Header + UDP-Header + DNS-Header + Question-Section gekürzt,
  /// IP-Total-Length und UDP-Length entsprechend neu berechnet.
  ///
  /// Kann die Question-Section nicht sicher vermessen werden (Pointer /
  /// Out-of-Bounds), fällt der Builder auf das ungeschnittene Paket zurück
  /// (besser eine evtl. „dreckige" Response als gar keine).
  private static func buildErrorResponse(
    _ bytes: UnsafeBufferPointer<UInt8>, length: Int, ihl: Int, rcode: UInt8
  ) -> Data {
    let udpStart = ihl
    let dnsStart = udpStart + 8

    // K1: Ende der Question-Section bestimmen → Paket dahinter abschneiden.
    let respLen: Int
    if let qEnd = dnsQuestionEnd(bytes, dnsStart: dnsStart, length: length) {
      respLen = qEnd
    } else {
      // Fallback: ungeschnitten (Pointer/OOB) — markiert als ungeprüfter Pfad.
      respLen = length
    }

    var resp = [UInt8](repeating: 0, count: respLen)
    for i in 0..<respLen { resp[i] = bytes[i] }

    // IP total-length (Bytes 2..3) auf die neue, evtl. gekürzte Länge setzen.
    resp[2] = UInt8((respLen >> 8) & 0xFF)
    resp[3] = UInt8(respLen & 0xFF)

    // Src-IP (12..15) und Dst-IP (16..19) tauschen.
    for i in 0..<4 {
      let srcByte = bytes[12 + i]
      let dstByte = bytes[16 + i]
      resp[12 + i] = dstByte
      resp[16 + i] = srcByte
    }
    // IP-Header-Checksum-Feld nullen (wird unten neu berechnet).
    resp[10] = 0
    resp[11] = 0

    // UDP-Ports tauschen.
    let srcPort0 = bytes[udpStart]
    let srcPort1 = bytes[udpStart + 1]
    let dstPort0 = bytes[udpStart + 2]
    let dstPort1 = bytes[udpStart + 3]
    resp[udpStart] = dstPort0
    resp[udpStart + 1] = dstPort1
    resp[udpStart + 2] = srcPort0
    resp[udpStart + 3] = srcPort1
    // UDP-Length (Bytes 4..5) = 8 (Header) + DNS-Payload-Länge.
    let udpLen = respLen - udpStart
    resp[udpStart + 4] = UInt8((udpLen >> 8) & 0xFF)
    resp[udpStart + 5] = UInt8(udpLen & 0xFF)
    // UDP-Checksum auf 0 (in IPv4 erlaubt, Empfänger überspringt die Prüfung).
    resp[udpStart + 6] = 0
    resp[udpStart + 7] = 0

    // DNS-Flags: QR=1 OPCODE=0 AA=0 TC=0 RD=1 | RA=1 Z=0 RCODE=<rcode>.
    //   High-Byte 0x81 = QR=1, RD=1.
    //   Low-Byte  0x80 | rcode = RA=1, RCODE = 3 (NXDOMAIN) bzw. 2 (SERVFAIL).
    resp[dnsStart + 2] = 0x81
    resp[dnsStart + 3] = 0x80 | (rcode & 0x0F)
    // ANCOUNT / NSCOUNT / ARCOUNT (Bytes 6..11 im DNS-Header) auf 0 —
    // QDCOUNT (Bytes 4..5) bleibt unangetastet, die Frage spiegeln wir zurück.
    for i in 6...11 {
      resp[dnsStart + i] = 0
    }

    recomputeIpChecksum(&resp, ihl: ihl)
    return Data(resp)
  }

  // ─── SERVFAIL-Response aus einem rohen Request-Paket (K2) ──────────────────

  /// Baut eine SERVFAIL-Response (RCODE=2) direkt aus einem rohen
  /// IPv4/UDP/DNS-Request-Paket — für den Fall, dass das Forwarding an den
  /// Upstream-Resolver fehlschlägt (Timeout, recvError, leere Antwort,
  /// Connection `.failed`).
  ///
  /// K2-Fix: Ohne diese Response bekäme der DNS-Client des iPhones bei jedem
  /// kurzen Upstream-Hänger GAR NICHTS zurück und liefe in seinen eigenen,
  /// mehrsekündigen Timeout → die komplette Namensauflösung des Geräts steht.
  /// Mit einer sofortigen SERVFAIL-Response failt der Client schnell und kann
  /// retryen / den nächsten Resolver fragen.
  ///
  /// Nutzt exakt dieselbe Paket-Bau-Logik wie die NXDOMAIN-Response inkl. des
  /// K1-Abschneidens hinter der Question-Section.
  ///
  /// Gibt `nil` zurück, wenn das Request-Paket kein valides IPv4/UDP-Paket ist.
  static func buildServFailResponse(request: Data) -> Data? {
    let length = request.count
    guard length >= 28 else { return nil }
    return request.withUnsafeBytes { (raw: UnsafeRawBufferPointer) -> Data? in
      let bytes = raw.bindMemory(to: UInt8.self)
      let version = (Int(bytes[0]) >> 4) & 0xF
      guard version == 4 else { return nil }
      let ihl = (Int(bytes[0]) & 0xF) * 4
      guard ihl >= 20, length >= ihl + 8 + 12 else { return nil }
      return buildErrorResponse(bytes, length: length, ihl: ihl, rcode: 2 /* SERVFAIL */)
    }
  }

  // ─── Forward-Response-Builder ──────────────────────────────────────────────

  /// Verpackt die rohe DNS-Antwort des Upstream-Resolvers wieder als IPv4/UDP-
  /// Paket, das ins TUN zurückgeschrieben werden kann. Logik aus
  /// `DnsFilter.kt#forwardAndWrite` (der Paket-Bau-Teil).
  ///
  /// - Parameters:
  ///   - request: die ursprüngliche Query (IPv4-Paket).
  ///   - dnsAnswer: die rohe DNS-Antwort (ab DNS-Header) vom Upstream.
  /// - Returns: das fertige IPv4/UDP/DNS-Response-Paket, oder nil bei Fehler.
  static func buildForwardResponse(request: Data, dnsAnswer: Data) -> Data? {
    let reqLen = request.count
    guard reqLen >= 28 else { return nil }

    return request.withUnsafeBytes { (raw: UnsafeRawBufferPointer) -> Data? in
      let req = raw.bindMemory(to: UInt8.self)
      let ihl = (Int(req[0]) & 0xF) * 4
      guard ihl >= 20, reqLen >= ihl + 8 else { return nil }

      let udpStart = ihl
      let totalLen = ihl + 8 + dnsAnswer.count
      // IPv4 total-length ist ein 16-bit-Feld.
      guard totalLen <= 0xFFFF else { return nil }

      var resp = [UInt8](repeating: 0, count: totalLen)
      // IP-Header aus der Request übernehmen.
      for i in 0..<ihl { resp[i] = req[i] }

      // IP total-length aktualisieren (Bytes 2..3).
      resp[2] = UInt8((totalLen >> 8) & 0xFF)
      resp[3] = UInt8(totalLen & 0xFF)

      // Src-/Dst-IP tauschen.
      for i in 0..<4 {
        let srcByte = req[12 + i]
        let dstByte = req[16 + i]
        resp[12 + i] = dstByte
        resp[16 + i] = srcByte
      }
      resp[10] = 0
      resp[11] = 0

      // UDP-Ports tauschen.
      resp[udpStart] = req[udpStart + 2]
      resp[udpStart + 1] = req[udpStart + 3]
      resp[udpStart + 2] = req[udpStart]
      resp[udpStart + 3] = req[udpStart + 1]
      // UDP-Length = 8 (Header) + Payload.
      let udpLen = 8 + dnsAnswer.count
      resp[udpStart + 4] = UInt8((udpLen >> 8) & 0xFF)
      resp[udpStart + 5] = UInt8(udpLen & 0xFF)
      resp[udpStart + 6] = 0
      resp[udpStart + 7] = 0

      // DNS-Antwort hinter den UDP-Header kopieren.
      let dnsStart = udpStart + 8
      for (i, byte) in dnsAnswer.enumerated() {
        resp[dnsStart + i] = byte
      }

      recomputeIpChecksum(&resp, ihl: ihl)
      return Data(resp)
    }
  }

  // ─── IP-Header-Checksum ────────────────────────────────────────────────────

  /// Berechnet die IPv4-Header-Checksum neu (RFC 791). Identisch zu
  /// `DnsFilter.kt#recomputeIpChecksum` — 16-bit-Einerkomplement-Summe über
  /// den IP-Header, Checksum-Feld zuvor auf 0.
  private static func recomputeIpChecksum(_ packet: inout [UInt8], ihl: Int) {
    packet[10] = 0
    packet[11] = 0
    var sum: UInt32 = 0
    var i = 0
    while i < ihl {
      let word = (UInt32(packet[i]) << 8) | UInt32(packet[i + 1])
      sum += word
      i += 2
    }
    while (sum >> 16) != 0 {
      sum = (sum & 0xFFFF) + (sum >> 16)
    }
    let checksum = ~sum & 0xFFFF
    packet[10] = UInt8((checksum >> 8) & 0xFF)
    packet[11] = UInt8(checksum & 0xFF)
  }

  // ─── DNS-Payload-Extraktion (für Forward) ──────────────────────────────────

  /// Extrahiert die rohe DNS-Query (ab DNS-Header) aus einem IPv4/UDP-Paket —
  /// das ist exakt das, was an den Upstream-Resolver gesendet werden muss.
  static func extractDnsPayload(packet: Data) -> Data? {
    let length = packet.count
    guard length >= 28 else { return nil }
    return packet.withUnsafeBytes { (raw: UnsafeRawBufferPointer) -> Data? in
      let bytes = raw.bindMemory(to: UInt8.self)
      let ihl = (Int(bytes[0]) & 0xF) * 4
      guard ihl >= 20, length >= ihl + 8 else { return nil }
      let dnsStart = ihl + 8
      guard dnsStart < length else { return nil }
      return packet.subdata(in: dnsStart..<length)
    }
  }
}
