ProxyIOS/ProxyCore/Sources/ProxyEngine/HTTPCaptureHandler.swift

import Foundation
import NIOCore
import NIOHTTP1

/// Captures HTTP request/response pairs and writes them to the traffic database.
final class HTTPCaptureHandler: ChannelDuplexHandler {
    typealias InboundIn = HTTPClientResponsePart
    typealias InboundOut = HTTPClientResponsePart
    typealias OutboundIn = HTTPClientRequestPart
    typealias OutboundOut = HTTPClientRequestPart

    private let trafficRepo: TrafficRepository
    private let domain: String
    private let scheme: String

    private var currentRequestId: String?
    private var requestHead: HTTPRequestHead?
    private var requestBody = Data()
    private var responseHead: HTTPResponseHead?
    private var responseBody = Data()
    private var requestStartTime: Double = 0

    private let hardcodedDebugDomain = "okcupid"
    private let hardcodedDebugNeedle = "jill"

    init(trafficRepo: TrafficRepository, domain: String, scheme: String = "https") {
        self.trafficRepo = trafficRepo
        self.domain = domain
        self.scheme = scheme
        ProxyLogger.capture.debug("HTTPCaptureHandler created for \(domain) (\(scheme))")
    }

    // MARK: - Outbound (Request)

    func write(context: ChannelHandlerContext, data: NIOAny, promise: EventLoopPromise<Void>?) {
        let part = unwrapOutboundIn(data)

        switch part {
        case .head(var head):
            currentRequestId = UUID().uuidString
            requestBody = Data()
            requestStartTime = Date().timeIntervalSince1970

            if RulesEngine.shouldStripCache() {
                head.headers.remove(name: "If-Modified-Since")
                head.headers.remove(name: "If-None-Match")
                head.headers.replaceOrAdd(name: "Cache-Control", value: "no-cache")
                head.headers.replaceOrAdd(name: "Pragma", value: "no-cache")
            }

            requestHead = head
            ProxyLogger.capture.info("CAPTURE REQ \(head.method.rawValue) \(self.scheme)://\(self.domain)\(head.uri)")
            context.write(self.wrapOutboundOut(.head(head)), promise: promise)
            return
        case .body(.byteBuffer(let buffer)):
            if requestBody.count < ProxyConstants.maxBodySizeBytes {
                requestBody.append(contentsOf: buffer.readableBytesView)
            }
            ProxyLogger.capture.debug("CAPTURE REQ body chunk: \(buffer.readableBytes) bytes (total: \(self.requestBody.count))")
        case .end:
            ProxyLogger.capture.debug("CAPTURE REQ end — saving to DB")
            saveRequest()
        default:
            break
        }

        context.write(data, promise: promise)
    }

    // MARK: - Inbound (Response)

    func channelRead(context: ChannelHandlerContext, data: NIOAny) {
        let part = unwrapInboundIn(data)

        switch part {
        case .head(var head):
            if RulesEngine.shouldStripCache() {
                head.headers.remove(name: "Expires")
                head.headers.remove(name: "Last-Modified")
                head.headers.remove(name: "ETag")
                head.headers.replaceOrAdd(name: "Expires", value: "0")
                head.headers.replaceOrAdd(name: "Cache-Control", value: "no-cache")
            }

            responseHead = head
            responseBody = Data()
            ProxyLogger.capture.info("CAPTURE RESP \(head.status.code) for \(self.domain)")
            context.fireChannelRead(NIOAny(HTTPClientResponsePart.head(head)))
            return
        case .body(let buffer):
            if responseBody.count < ProxyConstants.maxBodySizeBytes {
                responseBody.append(contentsOf: buffer.readableBytesView)
            }
            ProxyLogger.capture.debug("CAPTURE RESP body chunk: \(buffer.readableBytes) bytes (total: \(self.responseBody.count))")
        case .end:
            ProxyLogger.capture.debug("CAPTURE RESP end — saving to DB")
            saveResponse()
        }

        context.fireChannelRead(data)
    }

    // MARK: - Persistence

    private func saveRequest() {
        guard let head = requestHead, let reqId = currentRequestId else {
            ProxyLogger.capture.error("saveRequest: no head or requestId!")
            return
        }

        let url = "\(scheme)://\(domain)\(head.uri)"
        let headersJSON = encodeHeaders(head.headers)
        let queryParams = extractQueryParams(from: head.uri)
        let shouldHide =
            (IPCManager.shared.hideSystemTraffic && SystemTrafficFilter.isSystemDomain(domain)) ||
            RulesEngine.checkBlockList(url: url, method: head.method.rawValue) == .hideOnly

        let headerCount = head.headers.count
        let bodySize = requestBody.count

        var traffic = CapturedTraffic(
            requestId: reqId, domain: domain, url: url,
            method: head.method.rawValue, scheme: scheme,
            requestHeaders: headersJSON,
            requestBody: requestBody.isEmpty ? nil : requestBody,
            requestBodySize: requestBody.count,
            requestContentType: head.headers.first(name: "Content-Type"),
            queryParameters: queryParams,
            startedAt: requestStartTime,
            isSslDecrypted: scheme == "https",
            isHidden: shouldHide
        )

        do {
            try trafficRepo.insert(&traffic)
            ProxyLogger.capture.info("DB INSERT OK: \(head.method.rawValue) \(self.domain) headers=\(headerCount) body=\(bodySize)B id=\(reqId)")
        } catch {
            ProxyLogger.capture.error("DB INSERT FAILED: \(error.localizedDescription)")
        }
    }

    private func saveResponse() {
        guard let reqId = currentRequestId, let head = responseHead else {
            ProxyLogger.capture.error("saveResponse: no requestId or responseHead!")
            return
        }

        let now = Date().timeIntervalSince1970
        let durationMs = Int((now - requestStartTime) * 1000)
        let headerCount = head.headers.count
        let bodySize = responseBody.count

        do {
            try trafficRepo.updateResponse(
                requestId: reqId,
                statusCode: Int(head.status.code),
                statusText: head.status.reasonPhrase,
                responseHeaders: encodeHeaders(head.headers),
                responseBody: responseBody.isEmpty ? nil : responseBody,
                responseBodySize: responseBody.count,
                responseContentType: head.headers.first(name: "Content-Type"),
                completedAt: now,
                durationMs: durationMs
            )
            ProxyLogger.capture.info("DB UPDATE OK: \(head.status.code) \(self.domain) headers=\(headerCount) body=\(bodySize)B duration=\(durationMs)ms id=\(reqId)")
        } catch {
            ProxyLogger.capture.error("DB UPDATE FAILED for \(reqId): \(error.localizedDescription)")
        }

        logHardcodedBodyDebug(responseHead: head, requestId: reqId)

        // Debounce — don't flood with notifications for every single response
        NotificationThrottle.shared.throttle()
    }

    private func encodeHeaders(_ headers: HTTPHeaders) -> String? {
        var dict: [String: String] = [:]
        for (name, value) in headers { dict[name] = value }
        guard let data = try? JSONEncoder().encode(dict) else { return nil }
        return String(data: data, encoding: .utf8)
    }

    private func extractQueryParams(from uri: String) -> String? {
        guard let url = URLComponents(string: uri),
              let items = url.queryItems, !items.isEmpty else { return nil }
        var dict: [String: String] = [:]
        for item in items { dict[item.name] = item.value ?? "" }
        guard let data = try? JSONEncoder().encode(dict) else { return nil }
        return String(data: data, encoding: .utf8)
    }

    private func logHardcodedBodyDebug(responseHead: HTTPResponseHead, requestId: String) {
        let responseHeaders = headerDictionary(from: responseHead.headers)
        let decodedBody = HTTPBodyDecoder.decodedBodyData(from: responseBody, headers: responseHeaders)
        let searchableBody = HTTPBodyDecoder.searchableText(from: responseBody, headers: responseHeaders) ?? ""
        let preview = decodedBodyPreview(headers: responseHeaders)

        guard domain.localizedCaseInsensitiveContains(hardcodedDebugDomain) ||
                requestHead?.uri.localizedCaseInsensitiveContains(hardcodedDebugDomain) == true ||
                preview.localizedCaseInsensitiveContains(hardcodedDebugNeedle) else {
            return
        }

        let contentType = responseHead.headers.first(name: "Content-Type") ?? "nil"
        let contentEncoding = responseHead.headers.first(name: "Content-Encoding") ?? "nil"
        let containsNeedle = searchableBody.localizedCaseInsensitiveContains(hardcodedDebugNeedle)
        let decodingHint = HTTPBodyDecoder.decodingHint(for: responseBody, headers: responseHeaders)

        ProxyLogger.capture.info(
            """
            HARDCODED DEBUG capture domain=\(self.domain) id=\(requestId) status=\(responseHead.status.code) \
            contentType=\(contentType) contentEncoding=\(contentEncoding) bodyBytes=\(self.responseBody.count) \
            decodedBytes=\(decodedBody?.count ?? 0) decoding=\(decodingHint) containsNeedle=\(containsNeedle)
            """
        )

        if containsNeedle {
            ProxyLogger.capture.info("HARDCODED DEBUG MATCH needle=\(self.hardcodedDebugNeedle) preview=\(preview)")
        } else {
            ProxyLogger.capture.info("HARDCODED DEBUG NO_MATCH needle=\(self.hardcodedDebugNeedle) preview=\(preview)")
        }
    }

    private func decodedSearchableBody(headers: [String: String]) -> String {
        HTTPBodyDecoder.searchableText(from: responseBody, headers: headers) ?? ""
    }

    private func decodedBodyPreview(headers: [String: String]) -> String {
        let raw = decodedSearchableBody(headers: headers)
            .replacingOccurrences(of: "\n", with: " ")
            .replacingOccurrences(of: "\r", with: " ")
        return String(raw.prefix(240))
    }

    private func headerDictionary(from headers: HTTPHeaders) -> [String: String] {
        var dictionary: [String: String] = [:]
        for (name, value) in headers {
            dictionary[name] = value
        }
        return dictionary
    }
}