/*
 * Copyright (C) 2023 MURENA SAS
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */
package foundation.e.advancedprivacy.trackers.service

import android.os.ParcelFileDescriptor
import foundation.e.advancedprivacy.trackers.service.usecases.ResolveDNSUseCase
import kotlinx.coroutines.CancellationException
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.Job
import kotlinx.coroutines.isActive
import kotlinx.coroutines.launch
import org.pcap4j.packet.DnsPacket
import org.pcap4j.packet.IpPacket
import org.pcap4j.packet.IpSelector
import org.pcap4j.packet.IpV4Packet
import org.pcap4j.packet.IpV6Packet
import org.pcap4j.packet.UdpPacket
import org.pcap4j.packet.namednumber.IpNumber
import org.pcap4j.packet.namednumber.UdpPort
import timber.log.Timber
import java.io.DataOutputStream
import java.io.FileInputStream
import java.io.FileOutputStream
import java.io.IOException
import java.net.Inet6Address
import java.util.Arrays

class TunLooper(
    private val resolveDNSUseCase: ResolveDNSUseCase,
) {
    private var vpnInterface: ParcelFileDescriptor? = null
    private var fileInputStream: FileInputStream? = null
    private var dataOutputStream: DataOutputStream? = null

    private fun closeStreams() {
        fileInputStream?.close()
        fileInputStream = null

        dataOutputStream?.close()
        dataOutputStream = null

        vpnInterface?.close()
        vpnInterface = null
    }

    fun listenJob(
        vpnInterface: ParcelFileDescriptor,
        scope: CoroutineScope
    ): Job = scope.launch(Dispatchers.IO) {
        this@TunLooper.vpnInterface = vpnInterface
        val fis = FileInputStream(vpnInterface.fileDescriptor)
        this@TunLooper.fileInputStream = fis
        dataOutputStream = DataOutputStream(FileOutputStream(vpnInterface.fileDescriptor))

        while (isActive) {
            runCatching {
                val buffer = ByteArray(Config.MTU)
                val pLen = fis.read(buffer)

                if (pLen > 0) {
                    scope.launch { handleIpPacket(buffer, pLen) }
                }
            }.onFailure {
                if (it is CancellationException) {
                    closeStreams()
                    throw it
                } else {
                    Timber.w(it, "while reading from VPN fd")
                }
            }
        }
        closeStreams()
    }

    private suspend fun handleIpPacket(buffer: ByteArray, pLen: Int) {
        val pdata = Arrays.copyOf(buffer, pLen)
        try {
            val packet = IpSelector.newPacket(pdata, 0, pdata.size)
            if (packet is IpPacket) {
                val ipPacket = packet
                if (isPacketDNS(ipPacket)) {
                    handleDnsPacket(ipPacket)
                }
            }
        } catch (e: Exception) {
            Timber.w(e, "Can't parse packet, ignore it.")
        }
    }

    private fun isPacketDNS(p: IpPacket): Boolean {
        if (p.header.protocol === IpNumber.UDP) {
            val up = p.payload as UdpPacket
            return up.header.dstPort === UdpPort.DOMAIN
        }
        return false
    }

    private suspend fun handleDnsPacket(ipPacket: IpPacket) {
        try {
            val udpPacket = ipPacket.payload as UdpPacket
            val dnsRequest = udpPacket.payload as DnsPacket
            val dnsResponse = resolveDNSUseCase.processDNS(dnsRequest)

            if (dnsResponse != null) {
                val dnsBuilder = dnsResponse.builder

                val udpBuilder = UdpPacket.Builder(udpPacket)
                    .srcPort(udpPacket.header.dstPort)
                    .dstPort(udpPacket.header.srcPort)
                    .srcAddr(ipPacket.getHeader().getDstAddr())
                    .dstAddr(ipPacket.getHeader().getSrcAddr())
                    .correctChecksumAtBuild(true)
                    .correctLengthAtBuild(true)
                    .payloadBuilder(dnsBuilder)

                val respPacket: IpPacket? = if (ipPacket is IpV4Packet) {
                    val ipV4Packet = ipPacket
                    val ipv4Builder = IpV4Packet.Builder()
                    ipv4Builder
                        .version(ipV4Packet.header.version)
                        .protocol(ipV4Packet.header.protocol)
                        .tos(ipV4Packet.header.tos)
                        .srcAddr(ipV4Packet.header.dstAddr)
                        .dstAddr(ipV4Packet.header.srcAddr)
                        .correctChecksumAtBuild(true)
                        .correctLengthAtBuild(true)
                        .dontFragmentFlag(ipV4Packet.header.dontFragmentFlag)
                        .reservedFlag(ipV4Packet.header.reservedFlag)
                        .moreFragmentFlag(ipV4Packet.header.moreFragmentFlag)
                        .ttl(Integer.valueOf(64).toByte())
                        .payloadBuilder(udpBuilder)
                    ipv4Builder.build()
                } else if (ipPacket is IpV6Packet) {
                    IpV6Packet.Builder(ipPacket as IpV6Packet?)
                        .srcAddr(ipPacket.getHeader().getDstAddr() as Inet6Address)
                        .dstAddr(ipPacket.getHeader().getSrcAddr() as Inet6Address)
                        .payloadBuilder(udpBuilder)
                        .build()
                } else null

                respPacket?.let {
                    try {
                        dataOutputStream?.write(it.rawData)
                    } catch (e: IOException) {
                        Timber.e(e, "error writing to VPN fd")
                    }
                }
            }
        } catch (ioe: java.lang.Exception) {
            Timber.e(ioe, "could not parse DNS packet")
        }
    }
}