// 
// Decompiled by Procyon v0.6.0
// 

package org.bouncycastle.pqc.crypto.mldsa;

import org.bouncycastle.util.Arrays;
import org.bouncycastle.crypto.digests.SHAKEDigest;
import java.security.SecureRandom;

class MLDSAEngine
{
    private final SecureRandom random;
    final SHAKEDigest shake256Digest;
    public static final int DilithiumN = 256;
    public static final int DilithiumQ = 8380417;
    public static final int DilithiumQinv = 58728449;
    public static final int DilithiumD = 13;
    public static final int SeedBytes = 32;
    public static final int CrhBytes = 64;
    public static final int RndBytes = 32;
    public static final int TrBytes = 64;
    public static final int DilithiumPolyT1PackedBytes = 320;
    public static final int DilithiumPolyT0PackedBytes = 416;
    private final int DilithiumPolyVecHPackedBytes;
    private final int DilithiumPolyZPackedBytes;
    private final int DilithiumPolyW1PackedBytes;
    private final int DilithiumPolyEtaPackedBytes;
    private final int DilithiumK;
    private final int DilithiumL;
    private final int DilithiumEta;
    private final int DilithiumTau;
    private final int DilithiumBeta;
    private final int DilithiumGamma1;
    private final int DilithiumGamma2;
    private final int DilithiumOmega;
    private final int DilithiumCTilde;
    private final int CryptoPublicKeyBytes;
    private final int CryptoBytes;
    private final int PolyUniformGamma1NBlocks;
    private final Symmetric symmetric;
    
    protected Symmetric GetSymmetric() {
        return this.symmetric;
    }
    
    int getDilithiumPolyZPackedBytes() {
        return this.DilithiumPolyZPackedBytes;
    }
    
    int getDilithiumPolyW1PackedBytes() {
        return this.DilithiumPolyW1PackedBytes;
    }
    
    int getDilithiumPolyEtaPackedBytes() {
        return this.DilithiumPolyEtaPackedBytes;
    }
    
    int getDilithiumK() {
        return this.DilithiumK;
    }
    
    int getDilithiumL() {
        return this.DilithiumL;
    }
    
    int getDilithiumEta() {
        return this.DilithiumEta;
    }
    
    int getDilithiumTau() {
        return this.DilithiumTau;
    }
    
    int getDilithiumBeta() {
        return this.DilithiumBeta;
    }
    
    int getDilithiumGamma1() {
        return this.DilithiumGamma1;
    }
    
    int getDilithiumGamma2() {
        return this.DilithiumGamma2;
    }
    
    int getDilithiumOmega() {
        return this.DilithiumOmega;
    }
    
    int getDilithiumCTilde() {
        return this.DilithiumCTilde;
    }
    
    int getCryptoPublicKeyBytes() {
        return this.CryptoPublicKeyBytes;
    }
    
    int getPolyUniformGamma1NBlocks() {
        return this.PolyUniformGamma1NBlocks;
    }
    
    MLDSAEngine(final int i, final SecureRandom random) {
        this.shake256Digest = new SHAKEDigest(256);
        switch (i) {
            case 2: {
                this.DilithiumK = 4;
                this.DilithiumL = 4;
                this.DilithiumEta = 2;
                this.DilithiumTau = 39;
                this.DilithiumBeta = 78;
                this.DilithiumGamma1 = 131072;
                this.DilithiumGamma2 = 95232;
                this.DilithiumOmega = 80;
                this.DilithiumPolyZPackedBytes = 576;
                this.DilithiumPolyW1PackedBytes = 192;
                this.DilithiumPolyEtaPackedBytes = 96;
                this.DilithiumCTilde = 32;
                break;
            }
            case 3: {
                this.DilithiumK = 6;
                this.DilithiumL = 5;
                this.DilithiumEta = 4;
                this.DilithiumTau = 49;
                this.DilithiumBeta = 196;
                this.DilithiumGamma1 = 524288;
                this.DilithiumGamma2 = 261888;
                this.DilithiumOmega = 55;
                this.DilithiumPolyZPackedBytes = 640;
                this.DilithiumPolyW1PackedBytes = 128;
                this.DilithiumPolyEtaPackedBytes = 128;
                this.DilithiumCTilde = 48;
                break;
            }
            case 5: {
                this.DilithiumK = 8;
                this.DilithiumL = 7;
                this.DilithiumEta = 2;
                this.DilithiumTau = 60;
                this.DilithiumBeta = 120;
                this.DilithiumGamma1 = 524288;
                this.DilithiumGamma2 = 261888;
                this.DilithiumOmega = 75;
                this.DilithiumPolyZPackedBytes = 640;
                this.DilithiumPolyW1PackedBytes = 128;
                this.DilithiumPolyEtaPackedBytes = 96;
                this.DilithiumCTilde = 64;
                break;
            }
            default: {
                throw new IllegalArgumentException("The mode " + i + "is not supported by Crystals Dilithium!");
            }
        }
        this.symmetric = new Symmetric.ShakeSymmetric();
        this.random = random;
        this.DilithiumPolyVecHPackedBytes = this.DilithiumOmega + this.DilithiumK;
        this.CryptoPublicKeyBytes = 32 + this.DilithiumK * 320;
        this.CryptoBytes = this.DilithiumCTilde + this.DilithiumL * this.DilithiumPolyZPackedBytes + this.DilithiumPolyVecHPackedBytes;
        if (this.DilithiumGamma1 == 131072) {
            this.PolyUniformGamma1NBlocks = (576 + this.symmetric.stream256BlockBytes - 1) / this.symmetric.stream256BlockBytes;
        }
        else {
            if (this.DilithiumGamma1 != 524288) {
                throw new RuntimeException("Wrong Dilithium Gamma1!");
            }
            this.PolyUniformGamma1NBlocks = (640 + this.symmetric.stream256BlockBytes - 1) / this.symmetric.stream256BlockBytes;
        }
    }
    
    byte[][] generateKeyPairInternal(final byte[] array) {
        final byte[] array2 = new byte[128];
        final byte[] array3 = new byte[64];
        final byte[] array4 = new byte[32];
        final byte[] array5 = new byte[64];
        final byte[] array6 = new byte[32];
        final PolyVecMatrix polyVecMatrix = new PolyVecMatrix(this);
        final PolyVecL polyVecL = new PolyVecL(this);
        final PolyVecK polyVecK = new PolyVecK(this);
        final PolyVecK polyVecK2 = new PolyVecK(this);
        final PolyVecK polyVecK3 = new PolyVecK(this);
        this.shake256Digest.update(array, 0, 32);
        this.shake256Digest.update((byte)this.DilithiumK);
        this.shake256Digest.update((byte)this.DilithiumL);
        this.shake256Digest.doFinal(array2, 0, 128);
        System.arraycopy(array2, 0, array4, 0, 32);
        System.arraycopy(array2, 32, array5, 0, 64);
        System.arraycopy(array2, 96, array6, 0, 32);
        polyVecMatrix.expandMatrix(array4);
        polyVecL.uniformEta(array5, (short)0);
        polyVecK.uniformEta(array5, (short)this.DilithiumL);
        final PolyVecL polyVecL2 = new PolyVecL(this);
        polyVecL.copyTo(polyVecL2);
        polyVecL2.polyVecNtt();
        polyVecMatrix.pointwiseMontgomery(polyVecK2, polyVecL2);
        polyVecK2.reduce();
        polyVecK2.invNttToMont();
        polyVecK2.addPolyVecK(polyVecK);
        polyVecK2.conditionalAddQ();
        polyVecK2.power2Round(polyVecK3);
        final byte[] packPublicKey = Packing.packPublicKey(polyVecK2, this);
        this.shake256Digest.update(array4, 0, array4.length);
        this.shake256Digest.update(packPublicKey, 0, packPublicKey.length);
        this.shake256Digest.doFinal(array3, 0, 64);
        final byte[][] packSecretKey = Packing.packSecretKey(array4, array3, array6, polyVecK3, polyVecL, polyVecK, this);
        return new byte[][] { packSecretKey[0], packSecretKey[1], packSecretKey[2], packSecretKey[3], packSecretKey[4], packSecretKey[5], packPublicKey, array };
    }
    
    byte[] deriveT1(final byte[] array, final byte[] array2, final byte[] array3, final byte[] array4, final byte[] array5, final byte[] array6) {
        final PolyVecMatrix polyVecMatrix = new PolyVecMatrix(this);
        final PolyVecL polyVecL = new PolyVecL(this);
        final PolyVecK polyVecK = new PolyVecK(this);
        final PolyVecK polyVecK2 = new PolyVecK(this);
        final PolyVecK polyVecK3 = new PolyVecK(this);
        Packing.unpackSecretKey(polyVecK3, polyVecL, polyVecK, array6, array4, array5, this);
        polyVecMatrix.expandMatrix(array);
        final PolyVecL polyVecL2 = new PolyVecL(this);
        polyVecL.copyTo(polyVecL2);
        polyVecL2.polyVecNtt();
        polyVecMatrix.pointwiseMontgomery(polyVecK2, polyVecL2);
        polyVecK2.reduce();
        polyVecK2.invNttToMont();
        polyVecK2.addPolyVecK(polyVecK);
        polyVecK2.conditionalAddQ();
        polyVecK2.power2Round(polyVecK3);
        return Packing.packPublicKey(polyVecK2, this);
    }
    
    SHAKEDigest getShake256Digest() {
        return new SHAKEDigest(this.shake256Digest);
    }
    
    void initSign(final byte[] array, final boolean b, final byte[] array2) {
        this.shake256Digest.update(array, 0, 64);
        this.absorbCtx(b, array2);
    }
    
    void initVerify(final byte[] array, final byte[] array2, final boolean b, final byte[] array3) {
        final byte[] array4 = new byte[64];
        this.shake256Digest.update(array, 0, array.length);
        this.shake256Digest.update(array2, 0, array2.length);
        this.shake256Digest.doFinal(array4, 0, 64);
        this.shake256Digest.update(array4, 0, 64);
        this.absorbCtx(b, array3);
    }
    
    void absorbCtx(final boolean b, final byte[] array) {
        if (array != null) {
            this.shake256Digest.update((byte)(b ? 1 : 0));
            this.shake256Digest.update((byte)array.length);
            this.shake256Digest.update(array, 0, array.length);
        }
    }
    
    byte[] signInternal(final byte[] array, final int n, final byte[] array2, final byte[] array3, final byte[] array4, final byte[] array5, final byte[] array6, final byte[] array7) {
        final SHAKEDigest shakeDigest = new SHAKEDigest(this.shake256Digest);
        shakeDigest.update(array, 0, n);
        return this.generateSignature(this.generateMu(shakeDigest), shakeDigest, array2, array3, array4, array5, array6, array7);
    }
    
    byte[] generateMu(final SHAKEDigest shakeDigest) {
        final byte[] array = new byte[64];
        shakeDigest.doFinal(array, 0, 64);
        return array;
    }
    
    byte[] generateSignature(final byte[] array, final SHAKEDigest shakeDigest, final byte[] array2, final byte[] array3, final byte[] array4, final byte[] array5, final byte[] array6, final byte[] array7) {
        final byte[] array8 = new byte[this.CryptoBytes];
        final byte[] array9 = new byte[64];
        short n = 0;
        final PolyVecL polyVecL = new PolyVecL(this);
        final PolyVecL polyVecL2 = new PolyVecL(this);
        final PolyVecL polyVecL3 = new PolyVecL(this);
        final PolyVecK polyVecK = new PolyVecK(this);
        final PolyVecK polyVecK2 = new PolyVecK(this);
        final PolyVecK polyVecK3 = new PolyVecK(this);
        final PolyVecK polyVecK4 = new PolyVecK(this);
        final PolyVecK polyVecK5 = new PolyVecK(this);
        final Poly poly = new Poly(this);
        final PolyVecMatrix polyVecMatrix = new PolyVecMatrix(this);
        Packing.unpackSecretKey(polyVecK, polyVecL, polyVecK2, array4, array5, array6, this);
        final byte[] copy = Arrays.copyOf(array3, 128);
        System.arraycopy(array7, 0, copy, 32, 32);
        System.arraycopy(array, 0, copy, 64, 64);
        shakeDigest.update(copy, 0, 128);
        shakeDigest.doFinal(array9, 0, 64);
        polyVecMatrix.expandMatrix(array2);
        polyVecL.polyVecNtt();
        polyVecK2.polyVecNtt();
        polyVecK.polyVecNtt();
        int i = 0;
        while (i < 1000) {
            ++i;
            final PolyVecL polyVecL4 = polyVecL2;
            final byte[] array10 = array9;
            final short n2 = n;
            ++n;
            polyVecL4.uniformGamma1(array10, n2);
            polyVecL2.copyTo(polyVecL3);
            polyVecL3.polyVecNtt();
            polyVecMatrix.pointwiseMontgomery(polyVecK3, polyVecL3);
            polyVecK3.reduce();
            polyVecK3.invNttToMont();
            polyVecK3.conditionalAddQ();
            polyVecK3.decompose(polyVecK4);
            polyVecK3.packW1(this, array8, 0);
            shakeDigest.update(array, 0, 64);
            shakeDigest.update(array8, 0, this.DilithiumK * this.DilithiumPolyW1PackedBytes);
            shakeDigest.doFinal(array8, 0, this.DilithiumCTilde);
            poly.challenge(array8, 0, this.DilithiumCTilde);
            poly.polyNtt();
            polyVecL3.pointwisePolyMontgomery(poly, polyVecL);
            polyVecL3.invNttToMont();
            polyVecL3.addPolyVecL(polyVecL2);
            polyVecL3.reduce();
            if (polyVecL3.checkNorm(this.DilithiumGamma1 - this.DilithiumBeta)) {
                continue;
            }
            polyVecK5.pointwisePolyMontgomery(poly, polyVecK2);
            polyVecK5.invNttToMont();
            polyVecK4.subtract(polyVecK5);
            polyVecK4.reduce();
            if (polyVecK4.checkNorm(this.DilithiumGamma2 - this.DilithiumBeta)) {
                continue;
            }
            polyVecK5.pointwisePolyMontgomery(poly, polyVecK);
            polyVecK5.invNttToMont();
            polyVecK5.reduce();
            if (polyVecK5.checkNorm(this.DilithiumGamma2)) {
                continue;
            }
            polyVecK4.addPolyVecK(polyVecK5);
            polyVecK4.conditionalAddQ();
            if (polyVecK5.makeHint(polyVecK4, polyVecK3) > this.DilithiumOmega) {
                continue;
            }
            Packing.packSignature(array8, polyVecL3, polyVecK5, this);
            return array8;
        }
        return null;
    }
    
    boolean verifyInternalMu(final byte[] array) {
        final byte[] array2 = new byte[64];
        this.shake256Digest.doFinal(array2, 0);
        return Arrays.constantTimeAreEqual(array2, array);
    }
    
    boolean verifyInternalMuSignature(final byte[] array, final byte[] array2, final int n, final SHAKEDigest shakeDigest, final byte[] array3, final byte[] array4) {
        final byte[] array5 = new byte[Math.max(64 + this.DilithiumK * this.DilithiumPolyW1PackedBytes, this.DilithiumCTilde)];
        System.arraycopy(array, 0, array5, 0, array.length);
        return this.doVerifyInternal(array5, array2, n, shakeDigest, array3, array4);
    }
    
    boolean verifyInternal(final byte[] array, final int n, final SHAKEDigest shakeDigest, final byte[] array2, final byte[] array3) {
        final byte[] array4 = new byte[Math.max(64 + this.DilithiumK * this.DilithiumPolyW1PackedBytes, this.DilithiumCTilde)];
        shakeDigest.doFinal(array4, 0);
        return this.doVerifyInternal(array4, array, n, shakeDigest, array2, array3);
    }
    
    private boolean doVerifyInternal(final byte[] array, final byte[] array2, final int n, final SHAKEDigest shakeDigest, final byte[] array3, final byte[] array4) {
        if (n != this.CryptoBytes) {
            return false;
        }
        final PolyVecK polyVecK = new PolyVecK(this);
        final PolyVecL polyVecL = new PolyVecL(this);
        if (!Packing.unpackSignature(polyVecL, polyVecK, array2, this)) {
            return false;
        }
        if (polyVecL.checkNorm(this.getDilithiumGamma1() - this.getDilithiumBeta())) {
            return false;
        }
        final Poly poly = new Poly(this);
        final PolyVecMatrix polyVecMatrix = new PolyVecMatrix(this);
        final PolyVecK polyVecK2 = new PolyVecK(this);
        final PolyVecK polyVecK3 = new PolyVecK(this);
        final PolyVecK unpackPublicKey = Packing.unpackPublicKey(polyVecK2, array4, this);
        poly.challenge(array2, 0, this.DilithiumCTilde);
        polyVecMatrix.expandMatrix(array3);
        polyVecL.polyVecNtt();
        polyVecMatrix.pointwiseMontgomery(polyVecK3, polyVecL);
        poly.polyNtt();
        unpackPublicKey.shiftLeft();
        unpackPublicKey.polyVecNtt();
        unpackPublicKey.pointwisePolyMontgomery(poly, unpackPublicKey);
        polyVecK3.subtract(unpackPublicKey);
        polyVecK3.reduce();
        polyVecK3.invNttToMont();
        polyVecK3.conditionalAddQ();
        polyVecK3.useHint(polyVecK3, polyVecK);
        polyVecK3.packW1(this, array, 64);
        shakeDigest.update(array, 0, 64 + this.DilithiumK * this.DilithiumPolyW1PackedBytes);
        shakeDigest.doFinal(array, 0, this.DilithiumCTilde);
        return Arrays.constantTimeAreEqual(this.DilithiumCTilde, array2, 0, array, 0);
    }
    
    byte[][] generateKeyPair() {
        final byte[] bytes = new byte[32];
        this.random.nextBytes(bytes);
        return this.generateKeyPairInternal(bytes);
    }
}
