package model.pki.cert;

import model.TestConstants;
import model.asn1.ASN1Object;
import model.asn1.Null;
import model.asn1.ObjectIdentifier;
import model.asn1.exceptions.ParseException;
import model.asn1.parsing.BytesReader;
import org.junit.jupiter.api.Test;

import java.math.BigInteger;
import java.time.ZoneId;
import java.time.ZonedDateTime;
import java.util.Arrays;

import static model.TestConstants.mutate;
import static org.junit.jupiter.api.Assertions.*;

public class TbsCertificateTest {
    @Test
    void testConstructor() {
        assertEquals(TbsCertificate.VERSION_V3, TestConstants.CERT_GENERATED.getVersion().getLong());
        assertEquals(100, TestConstants.CERT_GENERATED.getSerialNumber().getLong());
        assertArrayEquals(ObjectIdentifier.OID_RSA_ENCRYPTION,
                TestConstants.CERT_GENERATED.getSignature().getType().getInts());
        assertEquals("CN=Test CA,C=CA", TestConstants.CERT_GENERATED.getIssuer().toString());
        assertEquals(TestConstants.NOW, TestConstants.CERT_GENERATED.getValidity().getNotBefore().getTimestamp());
        assertEquals(TestConstants.NOW.plusYears(1),
                TestConstants.CERT_GENERATED.getValidity().getNotAfter().getTimestamp());
        assertEquals("CN=Yuuta Liang,C=CA", TestConstants.CERT_GENERATED.getSubject().toString());
        assertArrayEquals(ObjectIdentifier.OID_EC_PUBLIC_KEY,
                TestConstants.CERT_GENERATED.getSubjectPublicKeyInfo().getAlgorithm().getType().getInts());
        assertEquals(2, TestConstants.CERT_GENERATED.getExtensions().getExtensions().length);
        assertArrayEquals(ObjectIdentifier.OID_BASIC_CONSTRAINTS,
                TestConstants.CERT_GENERATED.getExtensions().getExtensions()[0].getExtnId().getInts());
        assertArrayEquals(ObjectIdentifier.OID_KEY_USAGE,
                TestConstants.CERT_GENERATED.getExtensions().getExtensions()[1].getExtnId().getInts());
    }

    @Test
    void testGetExtension() throws ParseException {
        TbsCertificate parsed = new TbsCertificate(new BytesReader(trimToTbs(TestConstants.CERT_L1_ECC)),
                false);
        assertNull(parsed.getExtension(ObjectIdentifier.OID_EXTENSION_REQUEST));
        assertNotNull(parsed.getExtension(ObjectIdentifier.OID_BASIC_CONSTRAINTS));
        parsed = new TbsCertificate(new BytesReader(trimToTbs(TestConstants.CERT_V1)),
                false);
        assertNull(parsed.getExtension(ObjectIdentifier.OID_EXTENSION_REQUEST));
        assertNull(parsed.getExtension(ObjectIdentifier.OID_BASIC_CONSTRAINTS));
    }

    @Test
    void testParse() throws ParseException {
        TbsCertificate parsed = new TbsCertificate(new BytesReader(trimToTbs(TestConstants.CERT_L1_ECC)),
                false);
        assertEquals(TbsCertificate.VERSION_V3, parsed.getVersion().getLong());
        assertEquals(0, parsed.getSerialNumber().getValue()
                .compareTo(new BigInteger("644983544608556543477205958886697401602227090424")));
        assertArrayEquals(ObjectIdentifier.OID_ECDSA_WITH_SHA256,
                parsed.getSignature().getType().getInts());
        assertNull(parsed.getSignature().getParameters());
        assertEquals("CN=Yuuta Root CA,C=CA", parsed.getIssuer().toString());
        assertEquals(ZonedDateTime.of(2023, 6, 23,
                        2, 50, 46, 0, ZoneId.of("UTC")),
                parsed.getValidity().getNotBefore().getTimestamp());
        assertEquals(ZonedDateTime.of(2048, 6, 23,
                        2, 50, 46, 0, ZoneId.of("UTC")),
                parsed.getValidity().getNotAfter().getTimestamp());
        assertEquals("CN=Yuuta Root CA,C=CA", parsed.getSubject().toString());
        assertArrayEquals(ObjectIdentifier.OID_EC_PUBLIC_KEY,
                parsed.getSubjectPublicKeyInfo().getAlgorithm().getType().getInts());
        assertArrayEquals(ObjectIdentifier.OID_PRIME256_V1,
                ((ObjectIdentifier) parsed.getSubjectPublicKeyInfo().getAlgorithm().getParameters()).getInts());
        assertEquals(4, parsed.getExtensions().getExtensions().length);
        assertArrayEquals(ObjectIdentifier.OID_SUBJECT_KEY_IDENTIFIER,
                parsed.getExtensions().getExtensions()[0].getExtnId().getInts());
        assertArrayEquals(ObjectIdentifier.OID_AUTHORITY_KEY_IDENTIFIER,
                parsed.getExtensions().getExtensions()[1].getExtnId().getInts());
        assertArrayEquals(ObjectIdentifier.OID_BASIC_CONSTRAINTS,
                parsed.getExtensions().getExtensions()[2].getExtnId().getInts());
        assertArrayEquals(ObjectIdentifier.OID_KEY_USAGE,
                parsed.getExtensions().getExtensions()[3].getExtnId().getInts());

        parsed = new TbsCertificate(
                new BytesReader(trimToTbs(TestConstants.CERT_L2_RSA)),
                false);
        assertEquals(TbsCertificate.VERSION_V3, parsed.getVersion().getLong());
        assertEquals(0, parsed.getSerialNumber().getValue()
                .compareTo(new BigInteger("354327098948136693059815576591331472151989570311")));
        assertArrayEquals(ObjectIdentifier.OID_ECDSA_WITH_SHA512,
                parsed.getSignature().getType().getInts());
        assertNull(parsed.getSignature().getParameters());
        assertEquals("CN=Yuuta Root CA,C=CA", parsed.getIssuer().toString());
        assertEquals(ZonedDateTime.of(2023, 6, 24,
                        0, 15, 22, 0, ZoneId.of("UTC")),
                parsed.getValidity().getNotBefore().getTimestamp());
        assertEquals(ZonedDateTime.of(2033, 6, 21,
                        0, 15, 22, 0, ZoneId.of("UTC")),
                parsed.getValidity().getNotAfter().getTimestamp());
        assertEquals("DC=MOE,DC=YUUTA,DC=AD,CN=Yuuta Home Issuing CA", parsed.getSubject().toString());
        assertArrayEquals(ObjectIdentifier.OID_RSA_ENCRYPTION,
                parsed.getSubjectPublicKeyInfo().getAlgorithm().getType().getInts());
        assertEquals(Null.TAG.getNumber(),
                parsed.getSubjectPublicKeyInfo().getAlgorithm().getParameters().getTag().getNumber());
        assertEquals(526, parsed.getSubjectPublicKeyInfo().getSubjectPublicKey().getVal().length);
        assertEquals(6, parsed.getExtensions().getExtensions().length);
        assertArrayEquals(ObjectIdentifier.OID_SUBJECT_KEY_IDENTIFIER,
                parsed.getExtensions().getExtensions()[0].getExtnId().getInts());
        assertArrayEquals(ObjectIdentifier.OID_AUTHORITY_KEY_IDENTIFIER,
                parsed.getExtensions().getExtensions()[1].getExtnId().getInts());
        assertArrayEquals(ObjectIdentifier.OID_BASIC_CONSTRAINTS,
                parsed.getExtensions().getExtensions()[2].getExtnId().getInts());
        assertArrayEquals(ObjectIdentifier.OID_KEY_USAGE,
                parsed.getExtensions().getExtensions()[3].getExtnId().getInts());
        assertArrayEquals(ObjectIdentifier.OID_CRL_DISTRIBUTION_POINTS,
                parsed.getExtensions().getExtensions()[4].getExtnId().getInts());
        assertArrayEquals(ObjectIdentifier.OID_AUTHORITY_INFO_ACCESS,
                parsed.getExtensions().getExtensions()[5].getExtnId().getInts());

        parsed = new TbsCertificate(
                new BytesReader(Arrays.stream(TestConstants.CERT_V1).skip(4).toArray(Byte[]::new)),
                false);
        assertNull(parsed.getVersion());
        assertNull(parsed.getExtensions());
    }

    @Test
    void testParseFail() throws ParseException {
        final Byte[] in = trimToTbs(TestConstants.CERT_L2_RSA);
        // Wrong version parent tag
        assertThrows(ParseException.class, () ->
                new TbsCertificate(new BytesReader(mutate(in, 4, -96, 2)), false));
        // Wrong version inner tag
        assertThrows(ParseException.class, () ->
                new TbsCertificate(new BytesReader(mutate(in, 6, 0x2, 3)), false));
        // Wrong serial number tag
        assertThrows(ParseException.class, () ->
                new TbsCertificate(new BytesReader(mutate(in, 9, 0x2, 3)), false));
        // Wrong signature tag
        assertThrows(ParseException.class, () ->
                new TbsCertificate(new BytesReader(mutate(in, 31, 0x30, 3)), false));
        // Wrong issuer tag
        assertThrows(ParseException.class, () ->
                new TbsCertificate(new BytesReader(mutate(in, 43, 0x30, 0x31)), false));
        // Wrong validity tag
        assertThrows(ParseException.class, () ->
                new TbsCertificate(new BytesReader(mutate(in, 82, 0x30, 0x31)), false));
        // Wrong subject tag
        assertThrows(ParseException.class, () ->
                new TbsCertificate(new BytesReader(mutate(in, 114, 0x30, 0x31)), false));
        // Wrong subject public key info tag
        assertThrows(ParseException.class, () ->
                new TbsCertificate(new BytesReader(mutate(in, 212, 0x30, 0x31)), false));
        // Wrong extensions parent tag
        assertThrows(ParseException.class, () ->
                new TbsCertificate(new BytesReader(mutate(in, 762, -93, 0x31)), false));
        // Wrong extensions inner tag
        assertThrows(ParseException.class, () ->
                new TbsCertificate(new BytesReader(mutate(in, 765, 0x30, 0x31)), false));
        // Extensions exist, but wrong version
        assertThrows(ParseException.class, () ->
                new TbsCertificate(new BytesReader(mutate(in, 8, 0x2, TbsCertificate.VERSION_V2)),
                        false));
        // Totally wrong version
        assertThrows(ParseException.class, () ->
                new TbsCertificate(new BytesReader(mutate(in, 8, 0x2, TbsCertificate.VERSION_V3 + 1)),
                        false));
        assertThrows(ParseException.class, () ->
                new TbsCertificate(new BytesReader(mutate(in, 8, 0x2, -10)),
                        false));

        // Extensions exist, but no version
        final TbsCertificate certV1 = new TbsCertificate(new BytesReader(trimToTbs(TestConstants.CERT_V1)), false);
        final TbsCertificate certV3 = new TbsCertificate(new BytesReader(trimToTbs(TestConstants.CERT_L2_RSA)), false);
        Byte[] wrongCert = new TbsCertificate(ASN1Object.TAG_SEQUENCE, null,
                null,
                certV1.getSerialNumber(),
                certV1.getSignature(),
                certV1.getIssuer(),
                certV1.getValidity(),
                certV1.getSubject(),
                certV1.getSubjectPublicKeyInfo(),
                certV3.getExtensions())
                .encodeDER();
        assertThrows(ParseException.class, () ->
                new TbsCertificate(new BytesReader(wrongCert), false));
    }

    @Test
    void testEncode() throws ParseException {
        Byte[] in = trimToTbs(TestConstants.CERT_L1_ECC);
        assertArrayEquals(Arrays.copyOfRange(in, 0, 345), new TbsCertificate(new BytesReader(in), false).encodeDER());
        in = trimToTbs(TestConstants.CERT_L2_RSA);
        assertArrayEquals(Arrays.copyOfRange(in, 0, 989), new TbsCertificate(new BytesReader(in), false).encodeDER());
        in = trimToTbs(TestConstants.CERT_V1);
        assertArrayEquals(Arrays.copyOfRange(in, 0, 583), new TbsCertificate(new BytesReader(in), false).encodeDER());
    }

    private static Byte[] trimToTbs(Byte[] in) {
        return Arrays.stream(in).skip(4).toArray(Byte[]::new);
    }
}