package de.upapp

import org.bouncycastle.asn1.x500.X500Name
import org.bouncycastle.asn1.x509.BasicConstraints
import org.bouncycastle.asn1.x509.Extension
import org.bouncycastle.cert.jcajce.JcaX509CertificateConverter
import org.bouncycastle.cert.jcajce.JcaX509v3CertificateBuilder
import org.bouncycastle.crypto.util.PrivateKeyFactory
import org.bouncycastle.operator.jcajce.JcaContentSignerBuilder
import org.bouncycastle.tls.Certificate
import org.bouncycastle.tls.CipherSuite
import org.bouncycastle.tls.DefaultTlsServer
import org.bouncycastle.tls.HashAlgorithm
import org.bouncycastle.tls.SignatureAlgorithm
import org.bouncycastle.tls.SignatureAndHashAlgorithm
import org.bouncycastle.tls.TlsContext
import org.bouncycastle.tls.crypto.TlsCryptoParameters
import org.bouncycastle.tls.crypto.impl.bc.BcDefaultTlsCredentialedSigner
import org.bouncycastle.tls.PSKTlsServer
import org.bouncycastle.tls.TlsPSKIdentityManager
import org.bouncycastle.tls.TlsServerProtocol
import org.bouncycastle.tls.crypto.impl.bc.BcTlsCrypto
import org.junit.Assert.assertArrayEquals
import org.junit.Assert.assertEquals
import org.junit.Assert.assertNull
import org.junit.Assert.assertThrows
import org.junit.Test
import java.io.IOException
import java.io.InputStream
import java.io.OutputStream
import java.math.BigInteger
import java.security.KeyPairGenerator
import java.security.SecureRandom
import java.security.cert.X509Certificate
import java.util.Date

class ProtocolTest {
    @Test
    fun decodesSegments() {
        assertNull(decode(frame(7, TRANSMIT, u16(3) + byteArrayOf(1, 2))))
        val (s, size) = decode(frame(0xC0E1, TRANSMIT, u16(2) + byteArrayOf(1, 2)) + frame(0, HEARTBEAT))!!
        assertEquals(0xC0E1, s.port)
        assertArrayEquals(byteArrayOf(1, 2), s.data)
        assertEquals(7, size)
        // Example from the MIBBridge protocol notes: OPEN from port 0xC0E1 to port 80.
        assertArrayEquals(byteArrayOf(0xE1.toByte(), 0xC0.toByte(), 1, 0x50, 0), frame(0xC0E1, OPEN, u16(80)))
    }

    @Test
    fun plainHttpWithChunkedBody() {
        val tunnel = fakeCar { _, input, output ->
            input.readHead()
            output.write("HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n3\r\nWVW\r\n2\r\nZZ\r\n0\r\n\r\n".toByteArray())
        }
        val r = tunnel.open(80).use { exchange(it.input, it.output, "GET", "h", "/car/info/vin") }
        assertEquals(200, r.status)
        assertEquals("WVWZZ", r.text)
    }

    @Test
    fun pskSessionMatchesCarCipher() {
        val tunnel = fakeCar { _, input, output ->
            val tls = TlsServerProtocol(input, output)
            tls.accept(object : PSKTlsServer(BcTlsCrypto(SecureRandom()), object : TlsPSKIdentityManager {
                override fun getHint(): ByteArray? = null
                override fun getPSK(identity: ByteArray) = if (identity.decodeToString() == "user") "secret".toByteArray() else null
            }) {
                override fun getSupportedCipherSuites() = intArrayOf(CipherSuite.TLS_PSK_WITH_AES_128_CBC_SHA256)
            })
            val head = tls.inputStream.readHead()!!
            val body = """{"data": [], "request": "${head.lines().first()}"}"""
            tls.outputStream.write("HTTP/1.1 200 OK\r\nContent-Length: ${body.length}\r\n\r\n$body".toByteArray())
            tls.close()
        }
        val r = tunnel.open(443).use { ch ->
            val tls = sessionTls(ch.input, ch.output, "user", "secret")
            exchange(tls.inputStream, tls.outputStream, "GET", "h", "/car/batteries").also { runCatching { tls.close() } }
        }
        assertEquals("""{"data": [], "request": "GET /car/batteries HTTP/1.1"}""", r.text)
    }

    @Test
    fun pinMatchesMapsAndMore() {
        // Computed with the maps+more algorithm in Python.
        assertEquals("304163", pin(ByteArray(16) { it.toByte() }, ByteArray(16) { (it + 16).toByte() }))
    }

    @Test
    fun registrationChecksCertificateAndAgreesOnPin() {
        val ca = TestCa()
        var serverPin = ""
        val tunnel = fakeCar { _, input, output ->
            val tls = TlsServerProtocol(input, output)
            tls.accept(object : DefaultTlsServer(BcTlsCrypto(SecureRandom())) {
                override fun getSupportedCipherSuites() = intArrayOf(CipherSuite.TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256)
                override fun getRSASignerCredentials() = ca.signer(context, crypto as BcTlsCrypto)
                override fun notifyHandshakeComplete() {
                    super.notifyHandshakeComplete()
                    serverPin = registrationPin(context)
                }
            })
            tls.inputStream.readHead()
            tls.outputStream.write("HTTP/1.1 200 OK\r\nContent-Length: 9\r\n\r\nuser,pass".toByteArray())
            tls.close()
        }
        val r = tunnel.open(443).use { ch ->
            val (tls, pin) = registrationTls(ch.input, ch.output, ca.root)
            assertEquals(serverPin, pin)
            assertEquals(6, pin.length)
            exchange(tls.inputStream, tls.outputStream, "GET", "h", "/auth/x").also { runCatching { tls.close() } }
        }
        assertEquals("user,pass", r.text)

        val stranger = TestCa()
        val ch = tunnel.open(443)
        assertThrows(IOException::class.java) { registrationTls(ch.input, ch.output, stranger.root) }
    }

    /** A root CA and a server certificate signed by it. The server certificate expired an hour ago. */
    private class TestCa {
        private val gen = KeyPairGenerator.getInstance("RSA").apply { initialize(2048) }
        private val rootKey = gen.generateKeyPair()
        private val leafKey = gen.generateKeyPair()
        private val now = System.currentTimeMillis()
        val root = cert("CN=Test CA", rootKey.public, true, now - 86_400_000, now + 86_400_000)
        private val leaf = cert("CN=Car", leafKey.public, false, now - 7_200_000, now - 3_600_000)

        private fun cert(subject: String, key: java.security.PublicKey, ca: Boolean, from: Long, to: Long): X509Certificate {
            val builder = JcaX509v3CertificateBuilder(X500Name("CN=Test CA"), BigInteger.valueOf(from), Date(from), Date(to), X500Name(subject), key)
                .addExtension(Extension.basicConstraints, true, BasicConstraints(ca))
            return JcaX509CertificateConverter().getCertificate(builder.build(JcaContentSignerBuilder("SHA256withRSA").build(rootKey.private)))
        }

        fun signer(context: TlsContext, crypto: BcTlsCrypto) = BcDefaultTlsCredentialedSigner(
            TlsCryptoParameters(context), crypto, PrivateKeyFactory.createKey(leafKey.private.encoded),
            Certificate(arrayOf(crypto.createCertificate(leaf.encoded), crypto.createCertificate(root.encoded))),
            SignatureAndHashAlgorithm(HashAlgorithm.sha256, SignatureAlgorithm.rsa),
        )
    }

    @Test
    fun reusesSessionAndRetriesAfterCarCloses() {
        var connections = 0
        val tunnel = fakeCar { port, input, output ->
            if (port == 80) {
                input.readHead()
                output.write("HTTP/1.1 200 OK\r\nContent-Length: 3\r\n\r\nVIN".toByteArray())
                return@fakeCar
            }
            val n = ++connections
            val tls = TlsServerProtocol(input, output)
            tls.accept(object : PSKTlsServer(BcTlsCrypto(SecureRandom()), object : TlsPSKIdentityManager {
                override fun getHint(): ByteArray? = null
                override fun getPSK(identity: ByteArray) = "secret".toByteArray()
            }) {
                override fun getSupportedCipherSuites() = intArrayOf(CipherSuite.TLS_PSK_WITH_AES_128_CBC_SHA256)
            })
            // The car closes after three requests without notice.
            repeat(3) {
                val path = tls.inputStream.readHead()?.lines()?.first()?.split(' ')?.get(1) ?: return@repeat
                val body = "$n $path"
                tls.outputStream.write("HTTP/1.1 200 OK\r\nContent-Length: ${body.length}\r\n\r\n$body".toByteArray())
            }
            tls.close()
        }
        val car = Car(tunnel, storedCredentials("user,secret"))
        car.connect {}
        assertEquals("1 /a", car.get("a"))
        assertEquals("1 /b", car.get("b"))
        assertEquals("2 /c", car.get("c"))
        assertEquals(2, connections)
    }

    private fun InputStream.readHead(): String? {
        val sb = StringBuilder()
        while (!sb.endsWith("\r\n\r\n")) sb.append(read().takeIf { it >= 0 }?.toChar() ?: return null)
        return sb.toString()
    }

    private fun fakeCar(serve: (port: Int, InputStream, OutputStream) -> Unit) = CarSide(serve).tunnel
}