package dev.stave.upnp

import org.jupnp.UpnpServiceConfiguration
import org.jupnp.UpnpServiceImpl
import org.jupnp.model.message.StreamRequestMessage
import org.jupnp.model.message.UpnpResponse
import org.jupnp.model.message.gena.OutgoingSubscribeResponseMessage
import org.jupnp.protocol.ProtocolFactory
import org.jupnp.protocol.ProtocolFactoryImpl
import org.jupnp.protocol.sync.ReceivingSubscribe
import java.net.Inet6Address
import java.net.InetAddress
import java.net.URI
import java.net.URISyntaxException
import java.net.UnknownHostException

/**
 * GENA SUBSCRIBE hardening for the published MediaRenderer (CallStranger,
 * CVE-2020-12695 class).
 *
 * A new subscription names CALLBACK URLs that our event NOTIFYs (current track
 * URI and metadata) are then sent to. jUPnP 3.0.3 keeps only `http://`
 * callbacks and requires each host to resolve to a loopback, link-local or
 * site-local address. That still lets one LAN host point our NOTIFYs at
 * another, and it resolves callback host names through DNS, so a rebinding
 * name can pass the check and later resolve elsewhere.
 *
 * Here a new subscription is accepted only when every callback is `http` and
 * its host is a literal address equal to the address of the peer sending the
 * SUBSCRIBE. Host names are refused without any lookup. Anything else gets
 * 412 Precondition Failed, the status GENA uses for a bad CALLBACK. Renewals
 * (SID, no CALLBACK) and everything else go to jUPnP unchanged, and jUPnP's
 * own checks still run after ours.
 */
internal class PeerOnlySubscriptionsUpnpService(
    configuration: UpnpServiceConfiguration,
) : UpnpServiceImpl(configuration) {

    override fun createProtocolFactory(): ProtocolFactory =
        object : ProtocolFactoryImpl(this@PeerOnlySubscriptionsUpnpService) {
            override fun createReceivingSubscribe(message: StreamRequestMessage): ReceivingSubscribe =
                PeerOnlyReceivingSubscribe(this@PeerOnlySubscriptionsUpnpService, message)
        }
}

private class PeerOnlyReceivingSubscribe(
    service: org.jupnp.UpnpService,
    message: StreamRequestMessage,
) : ReceivingSubscribe(service, message) {

    override fun executeSync(): OutgoingSubscribeResponseMessage? {
        val callbacks = inputMessage.headers[CALLBACK_HEADER]
        if (!callbacks.isNullOrEmpty()) {
            val peer = inputMessage.connection?.remoteAddress
            val reason = callbackRejection(callbacks.joinToString(""), peer)
            if (reason != null) {
                println("[EarhornUPnP] Refused event subscription: $reason")
                return OutgoingSubscribeResponseMessage(UpnpResponse.Status.PRECONDITION_FAILED)
            }
        }
        return super.executeSync()
    }

    private companion object {
        const val CALLBACK_HEADER = "CALLBACK"
    }
}

/**
 * Why a raw CALLBACK header value (`<url><url>…`) from [peer] must be refused,
 * or null when every URL is `http` to [peer]'s own address. Never does a DNS
 * lookup.
 */
internal fun callbackRejection(headerValue: String, peer: InetAddress?): String? {
    if (peer == null) return "sender address unknown"
    val urls = CALLBACK_URL.findAll(headerValue).map { it.groupValues[1].trim() }.toList()
    if (urls.isEmpty()) return "no callback address"
    for (url in urls) {
        val uri = try {
            URI(url)
        } catch (_: URISyntaxException) {
            return "malformed callback address"
        }
        if (!uri.scheme.equals("http", ignoreCase = true)) return "callback is not http"
        val host = uri.host ?: return "callback has no host"
        val address = literalAddress(host) ?: return "callback names a host, not an address"
        if (!address.address.contentEquals(peer.normalised().address)) {
            return "callback points at a different machine than the sender"
        }
    }
    return null
}

private val CALLBACK_URL = Regex("<([^>]*)>")
private val IPV4_LITERAL = Regex("""(\d{1,3})\.(\d{1,3})\.(\d{1,3})\.(\d{1,3})""")

/** [host] as an address if it is an IP literal (IPv6 in brackets, as
 *  [URI.getHost] gives it), else null. No name resolution either way. */
private fun literalAddress(host: String): InetAddress? {
    IPV4_LITERAL.matchEntire(host)?.let { m ->
        val octets = m.groupValues.drop(1).map { it.toInt() }
        if (octets.any { it > 255 }) return null
        return InetAddress.getByAddress(ByteArray(4) { octets[it].toByte() })
    }
    // Bracketed, the JDK parses an IPv6 literal or throws; it never looks it up.
    if (host.length > 2 && host.startsWith("[") && host.endsWith("]")) {
        return try {
            InetAddress.getByName(host).normalised()
        } catch (_: UnknownHostException) {
            null
        }
    }
    return null
}

/** An IPv4-mapped IPv6 peer (`::ffff:a.b.c.d`) compares as its IPv4 address. */
private fun InetAddress.normalised(): InetAddress {
    val bytes = address
    if (this is Inet6Address && bytes.size == 16 &&
        (0 until 10).all { bytes[it] == 0.toByte() } &&
        bytes[10] == 0xff.toByte() && bytes[11] == 0xff.toByte()
    ) {
        return InetAddress.getByAddress(bytes.copyOfRange(12, 16))
    }
    return this
}