ProxyIOS/ProxyCore/Sources/ProxyEngine/HTTPCaptureHandler.swift

import Foundation
import NIOCore
import NIOHTTP1

/// Captures HTTP request/response pairs and writes them to the traffic database.
/// Inserted into the pipeline after TLS termination (MITM) or for plain HTTP.
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

    init(trafficRepo: TrafficRepository, domain: String, scheme: String = "https") {
        self.trafficRepo = trafficRepo
        self.domain = domain
        self.scheme = scheme
    }

    // MARK: - Outbound (Request)

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

        switch part {
        case .head(let head):
            currentRequestId = UUID().uuidString
            requestHead = head
            requestBody = Data()
            requestStartTime = Date().timeIntervalSince1970
        case .body(.byteBuffer(let buffer)):
            if requestBody.count < ProxyConstants.maxBodySizeBytes {
                requestBody.append(contentsOf: buffer.readableBytesView)
            }
        case .end:
            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(let head):
            responseHead = head
            responseBody = Data()
        case .body(let buffer):
            if responseBody.count < ProxyConstants.maxBodySizeBytes {
                responseBody.append(contentsOf: buffer.readableBytesView)
            }
        case .end:
            saveResponse()
        }

        context.fireChannelRead(data)
    }

    // MARK: - Persistence

    private func saveRequest() {
        guard let head = requestHead, let reqId = currentRequestId else { return }

        let url = "\(scheme)://\(domain)\(head.uri)"
        let headersJSON = encodeHeaders(head.headers)
        let queryParams = extractQueryParams(from: head.uri)

        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"
        )

        try? trafficRepo.insert(&traffic)
    }

    private func saveResponse() {
        guard let reqId = currentRequestId, let head = responseHead else { return }

        let now = Date().timeIntervalSince1970
        let durationMs = Int((now - requestStartTime) * 1000)

        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
        )

        IPCManager.shared.post(.newTrafficCaptured)
    }

    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)
    }
}