Added Dilithium AES support + some optimizations

17 files changed, 423 insertions, 293 deletions

diff --git a/crypto/src/pqc/crypto/crystals/dilithium/DilithiumEngine.cs b/crypto/src/pqc/crypto/crystals/dilithium/DilithiumEngine.cs
index 9b6479eaf..9648167d1 100644
--- a/crypto/src/pqc/crypto/crystals/dilithium/DilithiumEngine.cs
+++ b/crypto/src/pqc/crypto/crystals/dilithium/DilithiumEngine.cs
@@ -42,8 +42,9 @@ namespace Org.BouncyCastle.Pqc.Crypto.Crystals.Dilithium
         public int CryptoSecretKeyBytes { get; private set; }
         public int CryptoBytes { get; private set; }
         public int PolyUniformGamma1NBytes { get; private set; }
-
-        public DilithiumEngine(int mode, SecureRandom random)
+        public Symmetric Symmetric { get; private set; }
+        
+        public DilithiumEngine(int mode, SecureRandom random, bool usingAes)
         {
             Mode = mode;
             switch (Mode)
@@ -90,6 +91,14 @@ namespace Org.BouncyCastle.Pqc.Crypto.Crystals.Dilithium
                 default:
                     throw new ArgumentException("The mode " + mode + "is not supported by Crystals Dilithium!");
             }
+            if(usingAes)
+            {
+                Symmetric = new Symmetric.AesSymmetric();
+            }
+            else
+            {
+                Symmetric = new Symmetric.ShakeSymmetric();
+            }
 
             _random = random;
             PolyVecHPackedBytes = Omega + K;
@@ -99,11 +108,11 @@ namespace Org.BouncyCastle.Pqc.Crypto.Crystals.Dilithium
 
             if (Gamma1 == (1 << 17))
             {
-                PolyUniformGamma1NBytes = ((576 + Symmetric.Shake256Rate - 1) / Symmetric.Shake256Rate);
+                PolyUniformGamma1NBytes = ((576 + Symmetric.Stream256BlockBytes - 1) / Symmetric.Stream256BlockBytes);
             }
             else if (Gamma1 == (1 << 19))
             {
-                PolyUniformGamma1NBytes = ((640 + Symmetric.Shake256Rate - 1) / Symmetric.Shake256Rate);
+                PolyUniformGamma1NBytes = ((640 + Symmetric.Stream256BlockBytes - 1) / Symmetric.Stream256BlockBytes);
             }
             else
             {
@@ -111,23 +120,25 @@ namespace Org.BouncyCastle.Pqc.Crypto.Crystals.Dilithium
             }
         }
         
-        public void GenerateKeyPair(byte[] pk, byte[] sk)
+        public void GenerateKeyPair(out byte[] rho, out byte[] key, out byte[] tr, out byte[] s1_, out byte[] s2_, out byte[] t0_, out byte[] encT1)
         {
             byte[] SeedBuf = new byte[SeedBytes];
             byte[] buf = new byte[2 * SeedBytes + CrhBytes];
-            byte[] tr = new byte[SeedBytes];
-
-            byte[] rho, rhoPrime, key;
-
+            byte[] rhoPrime = new byte[CrhBytes];
+
+            tr = new byte[SeedBytes];
+            rho = new byte[SeedBytes];
+            key = new byte[SeedBytes];
+            s1_ = new byte[L * PolyEtaPackedBytes];
+            s2_ = new byte[K * PolyEtaPackedBytes];
+            t0_ = new byte[K * PolyT0PackedBytes];
             PolyVecMatrix Matrix = new PolyVecMatrix(this);
 
             PolyVecL s1 = new PolyVecL(this), s1Hat;
             PolyVecK s2 = new PolyVecK(this), t1 = new PolyVecK(this), t0 = new PolyVecK(this);
 
             _random.NextBytes(SeedBuf);
-         
-
-
+            
             ShakeDigest Shake256Digest = new ShakeDigest(256);
             Shake256Digest.BlockUpdate(SeedBuf, 0, SeedBytes);
             Shake256Digest.DoFinal(buf, 0, 2 * SeedBytes + CrhBytes);
@@ -156,28 +167,27 @@ namespace Org.BouncyCastle.Pqc.Crypto.Crystals.Dilithium
             t1.ConditionalAddQ();
             t1.Power2Round(t0);
 
-      
-            Packing.PackPublicKey(pk, rho, t1, this);
+            encT1 = Packing.PackPublicKey(t1, this);
 
-            Shake256Digest.BlockUpdate(pk, 0, CryptoPublicKeyBytes);
+            Shake256Digest.BlockUpdate(rho, 0, rho.Length);
+            Shake256Digest.BlockUpdate(encT1, 0, encT1.Length);
             Shake256Digest.DoFinal(tr, 0, SeedBytes);
 
-            Packing.PackSecretKey(sk, rho, tr, key, t0, s1, s2, this);
+            Packing.PackSecretKey(t0_, s1_, s2_, t0, s1, s2, this);
         }
 
-        public byte[] SignSignature(DilithiumEngine engine, byte[] msg, int msglen, byte[] sk)
+        public void SignSignature(byte[] sig, int siglen, byte[] msg, int msglen, byte[] rho, byte[] key, byte[] tr, byte[] t0Enc, byte[] s1Enc, byte[] s2Enc)
         {
-            byte[] sig = new byte[engine.CryptoBytes];
             int n;
             byte[] SeedBuf = new byte[3 * SeedBytes + 2 * CrhBytes];
-            byte[] rho = new byte[SeedBytes], tr = new byte[SeedBytes], key = new byte[SeedBytes], mu = new byte[CrhBytes], rhoPrime = new byte[CrhBytes];
+            byte[] mu = new byte[CrhBytes], rhoPrime = new byte[CrhBytes];
             ushort nonce = 0;
             PolyVecMatrix Matrix = new PolyVecMatrix(this);
             PolyVecL s1 = new PolyVecL(this), y = new PolyVecL(this), z = new PolyVecL(this);
             PolyVecK t0 = new PolyVecK(this), s2 = new PolyVecK(this), w1 = new PolyVecK(this), w0 = new PolyVecK(this), h = new PolyVecK(this);
             Poly cp = new Poly(this);
 
-            Packing.UnpackSecretKey(rho, tr, key, t0, s1, s2, sk, this);
+            Packing.UnpackSecretKey(t0, s1, s2, t0Enc, s1Enc, s2Enc, this);
 
             ShakeDigest ShakeDigest256 = new ShakeDigest(256);
             ShakeDigest256.BlockUpdate(tr, 0, SeedBytes);
@@ -196,96 +206,83 @@ namespace Org.BouncyCastle.Pqc.Crypto.Crystals.Dilithium
                 ShakeDigest256.DoFinal(rhoPrime, 0, CrhBytes);
             }
 
-            Matrix.ExpandMatrix(rho);         
-
+            Matrix.ExpandMatrix(rho);
+            
             s1.Ntt();
             s2.Ntt();
             t0.Ntt();
 
-            int count = 0;
-         
-
-            while (count < 1000)
-            {
-                count++;
-
-                y.UniformGamma1(rhoPrime, nonce++);
-                y.CopyPolyVecL(z);
-                z.Ntt();
-
-                Matrix.PointwiseMontgomery(w1, z);
-
-                w1.Reduce();
-                w1.InverseNttToMont();
-
-                w1.ConditionalAddQ();
-                w1.Decompose(w0);
-
-                w1.PackW1(sig);
-
-                ShakeDigest256.BlockUpdate(mu, 0, CrhBytes);
-                ShakeDigest256.BlockUpdate(sig, 0, K * PolyW1PackedBytes);
-                ShakeDigest256.DoFinal(sig, 0, SeedBytes);
-
-                cp.Challenge(sig);
+        rej:
+            y.UniformGamma1(rhoPrime, nonce++);
+            y.CopyPolyVecL(z);
+            z.Ntt();
 
-                cp.PolyNtt();
+            Matrix.PointwiseMontgomery(w1, z);
+            
+            w1.Reduce();
+            w1.InverseNttToMont();
 
-                z.PointwisePolyMontgomery(cp, s1);
-                z.InverseNttToMont();
-                z.AddPolyVecL(y);
-                z.Reduce();
-                if (z.CheckNorm(Gamma1 - Beta))
-                {
-                    continue;
-                }
+            w1.ConditionalAddQ();
+            w1.Decompose(w0);
+            
+            w1.PackW1(sig);
 
+            ShakeDigest256.BlockUpdate(mu, 0, CrhBytes);
+            ShakeDigest256.BlockUpdate(sig, 0, K * PolyW1PackedBytes);
+            ShakeDigest256.DoFinal(sig, 0, SeedBytes);
 
-                h.PointwisePolyMontgomery(cp, s2);
-                h.InverseNttToMont(); 
+            cp.Challenge(sig);
 
-                w0.Subtract(h);
-                w0.Reduce();
-                if (w0.CheckNorm(Gamma2 - Beta))
-                {
-                    continue;
-                }
+            cp.PolyNtt();
 
-                h.PointwisePolyMontgomery(cp, t0);
-                h.InverseNttToMont();
-                h.Reduce();
-                if (h.CheckNorm(Gamma2))
-                {                  
-                    continue;
-                }
+            z.PointwisePolyMontgomery(cp, s1);
+            z.InverseNttToMont();
+            z.AddPolyVecL(y);
+            z.Reduce();
+            if (z.CheckNorm(Gamma1 - Beta))
+            {
+                goto rej;
+            }
+            
+            h.PointwisePolyMontgomery(cp, s2);
+            h.InverseNttToMont();
 
-                w0.AddPolyVecK(h);
+            w0.Subtract(h);
+            w0.Reduce();
+            if (w0.CheckNorm(Gamma2 - Beta))
+            {
+                goto rej;
+            }
 
+            h.PointwisePolyMontgomery(cp, t0);
+            h.InverseNttToMont();
+            h.Reduce();
+            if (h.CheckNorm(Gamma2))
+            {
+                goto rej;
+            }
 
-                w0.ConditionalAddQ();
+            w0.AddPolyVecK(h);
 
-             
-                n = h.MakeHint(w0, w1);
-                if (n > Omega)
-                {
-                    continue;
-                }
-                       
+            w0.ConditionalAddQ();
 
-                sig =  Packing.PackSignature(sig, z, h, this);
-                return sig;
+            n = h.MakeHint(w0, w1);
+            if (n > Omega)
+            {
+                goto rej;
             }
-            return null;
+
+            Packing.PackSignature(sig, sig, z, h, this);
         }
 
-        public byte[] Sign(DilithiumEngine engine, byte[] msg, int msglen, byte[] sk)
+        public void Sign(byte[] sig, int siglen, byte[] msg, int mlen, byte[] rho, byte[] key, byte[] tr, byte[] t0, byte[] s1, byte[] s2)
         {
-           return SignSignature(engine, msg, msglen, sk);          
+            SignSignature(sig, siglen, msg,  mlen, rho, key, tr, t0, s1, s2);
         }
 
-        public bool SignVerify(byte[] sig, int siglen, byte[] msg, int msglen, byte[] pk)
+        public bool SignVerify(byte[] sig, int siglen, byte[] msg, int msglen, byte[] rho, byte[] encT1)
         {
-            byte[] buf = new byte[K * PolyW1PackedBytes], rho = new byte[SeedBytes], mu = new byte[CrhBytes], c = new byte[SeedBytes], c2 = new byte[SeedBytes];
+            byte[] buf = new byte[K * PolyW1PackedBytes], mu = new byte[CrhBytes], c = new byte[SeedBytes], c2 = new byte[SeedBytes];
             Poly cp = new Poly(this);
             PolyVecMatrix Matrix = new PolyVecMatrix(this);
             PolyVecL z = new PolyVecL(this);
@@ -296,20 +293,24 @@ namespace Org.BouncyCastle.Pqc.Crypto.Crystals.Dilithium
                 return false;
             }
 
-            Packing.UnpackPublicKey(rho, t1, pk, this);
-     
-            if (!Packing.UnpackSignature(c, z, h, sig, this))
+            t1 = Packing.UnpackPublicKey(t1, encT1, this);
+            
+
+
+            if (!Packing.UnpackSignature(z, h, sig, this))
             {
                 return false;
             }
+            c = Arrays.CopyOfRange(sig, 0, SeedBytes);
 
             if (z.CheckNorm(Gamma1 - Beta))
             {
                 return false;
             }
-         
+            
             ShakeDigest Shake256Digest = new ShakeDigest(256);
-            Shake256Digest.BlockUpdate(pk, 0, CryptoPublicKeyBytes);
+            Shake256Digest.BlockUpdate(rho, 0, rho.Length);
+            Shake256Digest.BlockUpdate(encT1, 0, encT1.Length);
             Shake256Digest.DoFinal(mu, 0, SeedBytes);
 
             Shake256Digest.BlockUpdate(mu, 0, SeedBytes);
@@ -352,16 +353,9 @@ namespace Org.BouncyCastle.Pqc.Crypto.Crystals.Dilithium
             return true;
         }
         
-        public bool SignOpen(byte[] msg, int msglen, byte[] sig, int siglen, byte[] pk)
+        public bool SignOpen(byte[] msg, byte[] sig, int siglen, byte[] rho, byte[] t1)
         {
-            int i;
-            if (siglen < CryptoBytes)
-            {
-                return false;
-            }
-         
-            return SignVerify(sig, siglen, msg, msglen, pk);
-            
+            return SignVerify(sig, siglen, msg, msg.Length, rho, t1);
         }
     }
 }
diff --git a/crypto/src/pqc/crypto/crystals/dilithium/DilithiumKeyGenerationParameters.cs b/crypto/src/pqc/crypto/crystals/dilithium/DilithiumKeyGenerationParameters.cs
index a5883bf78..e79129c59 100644
--- a/crypto/src/pqc/crypto/crystals/dilithium/DilithiumKeyGenerationParameters.cs
+++ b/crypto/src/pqc/crypto/crystals/dilithium/DilithiumKeyGenerationParameters.cs
@@ -4,7 +4,7 @@ using Org.BouncyCastle.Security;
 namespace Org.BouncyCastle.Pqc.Crypto.Crystals.Dilithium
 {
     public class DilithiumKeyGenerationParameters
-     : KeyGenerationParameters
+        : KeyGenerationParameters
     {
         private DilithiumParameters parameters;
 
@@ -15,4 +15,4 @@ namespace Org.BouncyCastle.Pqc.Crypto.Crystals.Dilithium
 
         public DilithiumParameters Parameters => parameters;
     }
-}
+}
\ No newline at end of file
diff --git a/crypto/src/pqc/crypto/crystals/dilithium/DilithiumKeyPairGenerator.cs b/crypto/src/pqc/crypto/crystals/dilithium/DilithiumKeyPairGenerator.cs
index 7a71f3fd1..98bf32b2b 100644
--- a/crypto/src/pqc/crypto/crystals/dilithium/DilithiumKeyPairGenerator.cs
+++ b/crypto/src/pqc/crypto/crystals/dilithium/DilithiumKeyPairGenerator.cs
@@ -1,4 +1,3 @@
-
 using Org.BouncyCastle.Crypto;
 using Org.BouncyCastle.Security;
 
@@ -20,15 +19,16 @@ namespace Org.BouncyCastle.Pqc.Crypto.Crystals.Dilithium
         public AsymmetricCipherKeyPair GenerateKeyPair()
         {
             DilithiumEngine engine = parameters.GetEngine(random);
-            byte[] sk = new byte[engine.CryptoSecretKeyBytes];
-            byte[] pk = new byte[engine.CryptoPublicKeyBytes];
-            engine.GenerateKeyPair(pk, sk);
+            byte[] rho, key, tr, s1, s2, t0, encT1;
+            engine.GenerateKeyPair(out rho, out key, out tr, out s1, out s2, out t0, out encT1);
+            
+            //unpack sk
 
-            DilithiumPublicKeyParameters pubKey = new DilithiumPublicKeyParameters(parameters, pk);
-            DilithiumPrivateKeyParameters privKey = new DilithiumPrivateKeyParameters(parameters, sk);
+            DilithiumPublicKeyParameters pubKey = new DilithiumPublicKeyParameters(parameters, rho, encT1);
+            DilithiumPrivateKeyParameters privKey = new DilithiumPrivateKeyParameters(parameters, rho, key, tr, s1, s2, t0, encT1);
 
 
             return new AsymmetricCipherKeyPair(pubKey, privKey);
         }
     }
-}
+}
\ No newline at end of file
diff --git a/crypto/src/pqc/crypto/crystals/dilithium/DilithiumKeyParameters.cs b/crypto/src/pqc/crypto/crystals/dilithium/DilithiumKeyParameters.cs
index 7e0054747..789f3f5ef 100644
--- a/crypto/src/pqc/crypto/crystals/dilithium/DilithiumKeyParameters.cs
+++ b/crypto/src/pqc/crypto/crystals/dilithium/DilithiumKeyParameters.cs
@@ -1,4 +1,3 @@
-
 using Org.BouncyCastle.Crypto;
 
 namespace Org.BouncyCastle.Pqc.Crypto.Crystals.Dilithium
@@ -15,4 +14,4 @@ namespace Org.BouncyCastle.Pqc.Crypto.Crystals.Dilithium
 
         public DilithiumParameters Parameters => parameters;
     }
-}
+}
\ No newline at end of file
diff --git a/crypto/src/pqc/crypto/crystals/dilithium/DilithiumParameters.cs b/crypto/src/pqc/crypto/crystals/dilithium/DilithiumParameters.cs
index a5dd2e338..fb726473a 100644
--- a/crypto/src/pqc/crypto/crystals/dilithium/DilithiumParameters.cs
+++ b/crypto/src/pqc/crypto/crystals/dilithium/DilithiumParameters.cs
@@ -1,4 +1,3 @@
-
 using Org.BouncyCastle.Crypto;
 using Org.BouncyCastle.Security;
 
@@ -6,21 +5,28 @@ namespace Org.BouncyCastle.Pqc.Crypto.Crystals.Dilithium
 {
     public class DilithiumParameters
         : ICipherParameters
-        {
-        public static DilithiumParameters Dilithium2 = new DilithiumParameters(2);
-        public static DilithiumParameters Dilithium3 = new DilithiumParameters(3);
-        public static DilithiumParameters Dilithium5 = new DilithiumParameters(5);
+    {
+        public static DilithiumParameters Dilithium2 = new DilithiumParameters(2, false);
+        public static DilithiumParameters Dilithium2Aes = new DilithiumParameters(2, true);
+        
+        public static DilithiumParameters Dilithium3 = new DilithiumParameters(3, false);
+        public static DilithiumParameters Dilithium3Aes = new DilithiumParameters(3, true);
+        
+        public static DilithiumParameters Dilithium5 = new DilithiumParameters(5, false);
+        public static DilithiumParameters Dilithium5Aes = new DilithiumParameters(5, true);
 
         private int k;
+        private bool usingAes;
 
-        private DilithiumParameters(int param)
+        private DilithiumParameters(int param, bool usingAes)
         {
             k = param;
+            this.usingAes = usingAes;
         }
 
         internal DilithiumEngine GetEngine(SecureRandom Random)
         {
-            return new DilithiumEngine(k, Random);
+            return new DilithiumEngine(k, Random, usingAes);
         }
     }
-}
+}
\ No newline at end of file
diff --git a/crypto/src/pqc/crypto/crystals/dilithium/DilithiumPrivateKeyParameters.cs b/crypto/src/pqc/crypto/crystals/dilithium/DilithiumPrivateKeyParameters.cs
index 510e46eea..ced1932bb 100644
--- a/crypto/src/pqc/crypto/crystals/dilithium/DilithiumPrivateKeyParameters.cs
+++ b/crypto/src/pqc/crypto/crystals/dilithium/DilithiumPrivateKeyParameters.cs
@@ -1,4 +1,3 @@
-
 using Org.BouncyCastle.Utilities;
 
 namespace Org.BouncyCastle.Pqc.Crypto.Crystals.Dilithium
@@ -6,17 +5,65 @@ namespace Org.BouncyCastle.Pqc.Crypto.Crystals.Dilithium
     public class DilithiumPrivateKeyParameters
         : DilithiumKeyParameters
     {
-        private byte[] privateKey;
+        internal byte[] rho;
+        internal byte[] k;
+        internal byte[] tr;
+        internal byte[] s1;
+        internal byte[] s2;
+        internal byte[] t0;
+        
+        private byte[] t1;
 
-        public DilithiumPrivateKeyParameters(DilithiumParameters parameters, byte[] skEncoded)
+        public DilithiumPrivateKeyParameters(DilithiumParameters parameters,  byte[] rho, byte[] K, byte[] tr, byte[] s1, byte[] s2, byte[] t0, byte[] t1)
             : base(true, parameters)
         {
-            privateKey = Arrays.Clone(skEncoded);
+            this.rho = Arrays.Clone(rho);
+            this.k = Arrays.Clone(K);
+            this.tr = Arrays.Clone(tr);
+            this.s1 = Arrays.Clone(s1);
+            this.s2 = Arrays.Clone(s2);
+            this.t0 = Arrays.Clone(t0);
+            this.t1 = Arrays.Clone(t1);
+        }
+        
+        public byte[] GetRho()
+        {
+            return Arrays.Clone(rho);
+        }
+
+        public byte[] GetK()
+        {
+            return Arrays.Clone(k);
+        }
+
+        public byte[] GetTr()
+        {
+            return Arrays.Clone(tr);
+        }
+
+        public byte[] GetS1()
+        {
+            return Arrays.Clone(s1);
+        }
+
+        public byte[] GetS2()
+        {
+            return Arrays.Clone(s2);
+        }
+
+        public byte[] GetT0()
+        {
+            return Arrays.Clone(t0);
+        }
+
+        public byte[] GetT1()
+        {
+            return t1;
         }
 
         public byte[] GetEncoded()
         {
-            return Arrays.Clone(privateKey);
+            return Arrays.ConcatenateAll(rho, k, tr, s1, s2, t0);
         }
     }
-}
+}
\ No newline at end of file
diff --git a/crypto/src/pqc/crypto/crystals/dilithium/DilithiumPublicKeyParameters.cs b/crypto/src/pqc/crypto/crystals/dilithium/DilithiumPublicKeyParameters.cs
index c2891ba3d..138fb4983 100644
--- a/crypto/src/pqc/crypto/crystals/dilithium/DilithiumPublicKeyParameters.cs
+++ b/crypto/src/pqc/crypto/crystals/dilithium/DilithiumPublicKeyParameters.cs
@@ -1,4 +1,3 @@
-
 using Org.BouncyCastle.Utilities;
 
 namespace Org.BouncyCastle.Pqc.Crypto.Crystals.Dilithium
@@ -6,19 +5,35 @@ namespace Org.BouncyCastle.Pqc.Crypto.Crystals.Dilithium
     public class DilithiumPublicKeyParameters
         : DilithiumKeyParameters
     {
+        internal byte[] rho;
+        internal byte[] t1;
 
-    private byte[] publicKey;
+        public DilithiumPublicKeyParameters(DilithiumParameters parameters, byte[] pkEncoded)
+            : base(false, parameters)
+        {
+            this.rho = Arrays.CopyOfRange(pkEncoded, 0, DilithiumEngine.SeedBytes);
+            this.t1 = Arrays.CopyOfRange(pkEncoded, DilithiumEngine.SeedBytes, pkEncoded.Length);
+        }
 
-    public DilithiumPublicKeyParameters(DilithiumParameters parameters, byte[] pkEncoded)
-        : base(false, parameters)
-    {
-        publicKey = Arrays.Clone(pkEncoded);
-    }
+        public DilithiumPublicKeyParameters(DilithiumParameters parameters, byte[] rho, byte[] t1)
+            : base(false, parameters)
+        {
+            this.rho = Arrays.Clone(rho);
+            this.t1 = Arrays.Clone(t1);
+        }
 
-    public byte[] GetEncoded()
-    {
-        return Arrays.Clone(publicKey);
-    }
+        public byte[] GetEncoded()
+        {
+            return Arrays.Concatenate(rho, t1);
+        }
+        public byte[] GetRho()
+        {
+            return Arrays.Clone(rho);
+        }
+        public byte[] GetT1()
+        {
+            return Arrays.Clone(t1);
+        }
 
     }
-}
+}
\ No newline at end of file
diff --git a/crypto/src/pqc/crypto/crystals/dilithium/DilithiumSigner.cs b/crypto/src/pqc/crypto/crystals/dilithium/DilithiumSigner.cs
index 58b9d1803..89519f134 100644
--- a/crypto/src/pqc/crypto/crystals/dilithium/DilithiumSigner.cs
+++ b/crypto/src/pqc/crypto/crystals/dilithium/DilithiumSigner.cs
@@ -41,15 +41,15 @@ namespace Org.BouncyCastle.Pqc.Crypto.Crystals.Dilithium
         public byte[] GenerateSignature(byte[] message)
         {
             DilithiumEngine engine = privKey.Parameters.GetEngine(random);
-          
-            return engine.Sign(engine, message, message.Length, privKey.GetEncoded());
-         
+            byte[] sig = new byte[engine.CryptoBytes];
+            engine.Sign(sig, sig.Length, message, message.Length, privKey.rho, privKey.k, privKey.tr, privKey.t0, privKey.s1, privKey.s2);
+            return sig;
         }
 
         public bool VerifySignature(byte[] message, byte[] signature)
         {
             DilithiumEngine engine = pubKey.Parameters.GetEngine(random);
-            return engine.SignOpen(message, message.Length, signature, signature.Length, pubKey.GetEncoded());
+            return engine.SignOpen(message,signature, signature.Length, pubKey.rho, pubKey.t1 );
         }
     }
-}
+}
\ No newline at end of file
diff --git a/crypto/src/pqc/crypto/crystals/dilithium/Ntt.cs b/crypto/src/pqc/crypto/crystals/dilithium/Ntt.cs
index ba19dd86e..292719058 100644
--- a/crypto/src/pqc/crypto/crystals/dilithium/Ntt.cs
+++ b/crypto/src/pqc/crypto/crystals/dilithium/Ntt.cs
@@ -1,7 +1,4 @@
-using Org.BouncyCastle.Utilities;
-
-
-namespace Org.BouncyCastle.Pqc.Crypto.Crystals.Dilithium
+namespace Org.BouncyCastle.Pqc.Crypto.Crystals.Dilithium
 {
     internal class Ntt
     {
diff --git a/crypto/src/pqc/crypto/crystals/dilithium/Packing.cs b/crypto/src/pqc/crypto/crystals/dilithium/Packing.cs
index 8a5862254..0f1fec1d2 100644
--- a/crypto/src/pqc/crypto/crystals/dilithium/Packing.cs
+++ b/crypto/src/pqc/crypto/crystals/dilithium/Packing.cs
@@ -1,96 +1,74 @@
-using System;
-using Org.BouncyCastle.Utilities;
+using Org.BouncyCastle.Utilities;
+using System;
 
 namespace Org.BouncyCastle.Pqc.Crypto.Crystals.Dilithium
 {
     internal class Packing
     {
-        public static void PackPublicKey(byte[] pk, byte[] rho, PolyVecK t1, DilithiumEngine Engine)
+
+        public static byte[] PackPublicKey(PolyVecK t1, DilithiumEngine Engine)
         {
-            Array.Copy(rho, 0, pk, 0, DilithiumEngine.SeedBytes);
+            byte[] output = new byte[Engine.CryptoPublicKeyBytes - DilithiumEngine.SeedBytes];
+
             for (int i = 0; i < Engine.K; i++)
             {
-                t1.Vec[i].PolyT1Pack(pk, DilithiumEngine.SeedBytes + i * DilithiumEngine.PolyT1PackedBytes);
+                Array.Copy(t1.Vec[i].PolyT1Pack(), 0, output, i * DilithiumEngine.PolyT1PackedBytes, DilithiumEngine.PolyT1PackedBytes );
             }
+            return output;
         }
 
-        public static void UnpackPublicKey(byte[] rho, PolyVecK t1, byte[] pk, DilithiumEngine Engine)
+        public static PolyVecK UnpackPublicKey(PolyVecK t1, byte[] pk, DilithiumEngine Engine)
         {
             int i;
-
-            Array.Copy(pk, 0, rho, 0, DilithiumEngine.SeedBytes);
-
             for (i = 0; i < Engine.K; ++i)
             {
-                t1.Vec[i].PolyT1Unpack(pk, DilithiumEngine.SeedBytes + i * DilithiumEngine.PolyT1PackedBytes);
+                t1.Vec[i].PolyT1Unpack(Arrays.CopyOfRange(pk, i * DilithiumEngine.PolyT1PackedBytes, DilithiumEngine.SeedBytes + (i + 1) * DilithiumEngine.PolyT1PackedBytes));
             }
+
+            return t1;
         }
 
-        public static void PackSecretKey(byte[] sk, byte[] rho, byte[] tr, byte[] key, PolyVecK t0, PolyVecL s1, PolyVecK s2, DilithiumEngine Engine)
+        public static void PackSecretKey(byte[] t0_, byte[] s1_, byte[] s2_, PolyVecK t0, PolyVecL s1, PolyVecK s2, DilithiumEngine Engine)
         {
-            int i, end = 0;
-            Array.Copy(rho, sk, DilithiumEngine.SeedBytes);
-            end += DilithiumEngine.SeedBytes;
-
-            Array.Copy(key, 0, sk, end, DilithiumEngine.SeedBytes);
-            end += DilithiumEngine.SeedBytes;
-
-            Array.Copy(tr, 0, sk, end, DilithiumEngine.SeedBytes);
-            end += DilithiumEngine.SeedBytes;
+            int i;
+            
 
             for (i = 0; i < Engine.L; ++i)
             {
-                s1.Vec[i].PolyEtaPack(sk, end + i * Engine.PolyEtaPackedBytes);
+                s1.Vec[i].PolyEtaPack(s1_, i * Engine.PolyEtaPackedBytes);
             }
-            end += Engine.L * Engine.PolyEtaPackedBytes;
 
             for (i = 0; i < Engine.K; ++i)
             {
-                s2.Vec[i].PolyEtaPack(sk, end + i * Engine.PolyEtaPackedBytes);
+                s2.Vec[i].PolyEtaPack(s2_, i * Engine.PolyEtaPackedBytes);
             }
-            end += Engine.K * Engine.PolyEtaPackedBytes;
 
             for (i = 0; i < Engine.K; ++i)
             {
-                t0.Vec[i].PolyT0Pack(sk, end + i * DilithiumEngine.PolyT0PackedBytes);
+                t0.Vec[i].PolyT0Pack(t0_,i * DilithiumEngine.PolyT0PackedBytes);
             }
         }
 
-        public static void UnpackSecretKey(byte[] rho, byte[] tr, byte[] key, PolyVecK t0, PolyVecL s1, PolyVecK s2, byte[] sk, DilithiumEngine Engine)
+        public static void UnpackSecretKey(PolyVecK t0, PolyVecL s1, PolyVecK s2, byte[] t0Enc, byte[] s1Enc, byte[] s2Enc, DilithiumEngine Engine)
         {
-            int i, end = 0;
-            Array.Copy(sk, 0, rho, 0, DilithiumEngine.SeedBytes);
-            end += DilithiumEngine.SeedBytes;
-
-            Array.Copy(sk, end, key, 0, DilithiumEngine.SeedBytes);
-            end += DilithiumEngine.SeedBytes;
-
-            Array.Copy(sk, end, tr, 0, DilithiumEngine.SeedBytes);
-            end += DilithiumEngine.SeedBytes;
-
+            int i;
             for (i = 0; i < Engine.L; ++i)
             {
-                s1.Vec[i].PolyEtaUnpack(sk, end + i * Engine.PolyEtaPackedBytes);
+                s1.Vec[i].PolyEtaUnpack(s1Enc,i * Engine.PolyEtaPackedBytes);
             }
-            end += Engine.L * Engine.PolyEtaPackedBytes;
-
             for (i = 0; i < Engine.K; ++i)
             {
-                s2.Vec[i].PolyEtaUnpack(sk, end + i * Engine.PolyEtaPackedBytes);
+                s2.Vec[i].PolyEtaUnpack(s2Enc,i * Engine.PolyEtaPackedBytes);
             }
-            end += Engine.K * Engine.PolyEtaPackedBytes;
-
             for (i = 0; i < Engine.K; ++i)
             {
-                t0.Vec[i].PolyT0Unpack(sk, end + i * DilithiumEngine.PolyT0PackedBytes);
+                t0.Vec[i].PolyT0Unpack(t0Enc,i * DilithiumEngine.PolyT0PackedBytes);
             }
         }
 
-        public static byte[] PackSignature(byte[] c, PolyVecL z, PolyVecK h, DilithiumEngine engine)
+        public static void PackSignature(byte[] sig, byte[] c, PolyVecL z, PolyVecK h, DilithiumEngine engine)
         {
             int i, j, k, end = 0;
-            byte[] sig = new byte[engine.CryptoBytes];
-
 
             Array.Copy(c, 0, sig, 0, DilithiumEngine.SeedBytes);
             end += DilithiumEngine.SeedBytes;
@@ -119,21 +97,18 @@ namespace Org.BouncyCastle.Pqc.Crypto.Crystals.Dilithium
                 }
                 sig[end + engine.Omega + i] = (byte)k;
             }
-
-            return sig;
+            //Console.WriteLine("sig = " + Convert.ToHexString(sig));
 
         }
 
-        public static bool UnpackSignature(byte[] c, PolyVecL z, PolyVecK h, byte[] sig, DilithiumEngine Engine)
+        public static bool UnpackSignature(PolyVecL z, PolyVecK h, byte[] sig, DilithiumEngine Engine)
         {
             int i, j, k;
-
-            Array.Copy(sig, c, DilithiumEngine.SeedBytes);
-
+            
             int end = DilithiumEngine.SeedBytes;
             for (i = 0; i < Engine.L; ++i)
             {
-                z.Vec[i].UnpackZ(sig, end + i * Engine.PolyZPackedBytes);
+                z.Vec[i].UnpackZ(Arrays.CopyOfRange(sig, end + i * Engine.PolyZPackedBytes, end + (i + 1) * Engine.PolyZPackedBytes));
             }
             end += Engine.L * Engine.PolyZPackedBytes;
 
diff --git a/crypto/src/pqc/crypto/crystals/dilithium/Poly.cs b/crypto/src/pqc/crypto/crystals/dilithium/Poly.cs
index 9d94cce52..6978796c4 100644
--- a/crypto/src/pqc/crypto/crystals/dilithium/Poly.cs
+++ b/crypto/src/pqc/crypto/crystals/dilithium/Poly.cs
@@ -1,5 +1,5 @@
-using System;
-using Org.BouncyCastle.Crypto.Digests;
+using Org.BouncyCastle.Crypto.Digests;
+using System;
 
 namespace Org.BouncyCastle.Pqc.Crypto.Crystals.Dilithium
 {
@@ -9,25 +9,27 @@ namespace Org.BouncyCastle.Pqc.Crypto.Crystals.Dilithium
 
         private int N;
         private DilithiumEngine Engine;
-        private int PolyUniformNBlocks = (768 + Symmetric.Shake128Rate - 1) / Symmetric.Shake128Rate;
+        private int PolyUniformNBlocks;
+        private Symmetric Symmetric;
 
         public Poly(DilithiumEngine engine)
         {
             N = DilithiumEngine.N;
             Coeffs = new int[N];
             Engine = engine;
+            Symmetric = engine.Symmetric;
+            PolyUniformNBlocks = (768 + Symmetric.Stream128BlockBytes - 1) / Symmetric.Stream128BlockBytes;
         }
 
         public void UniformBlocks(byte[] seed, ushort nonce)
         {
             int i, ctr, off,
-            buflen = PolyUniformNBlocks * Symmetric.Shake128Rate;
+            buflen = PolyUniformNBlocks * Symmetric.Stream128BlockBytes;
             byte[] buf = new byte[buflen + 2];
-            ShakeDigest Shake128Digest = new ShakeDigest(128);
-
-            Symmetric.ShakeStreamInit(Shake128Digest, seed, nonce);
+            
+            Symmetric.Stream128Init(seed, nonce);
 
-            Shake128Digest.DoOutput(buf, 0, buflen + 2);
+            Symmetric.Stream128SqueezeBlocks(buf, 0, buflen);
 
             ctr = RejectUniform(Coeffs, 0, N, buf, buflen);
 
@@ -38,9 +40,9 @@ namespace Org.BouncyCastle.Pqc.Crypto.Crystals.Dilithium
                 {
                     buf[i] = buf[buflen - off + i];
                 }
-                Shake128Digest.DoOutput(buf, buflen + off, 1);
-                buflen = Symmetric.Shake128Rate + off;
-                ctr += RejectUniform(Coeffs, ctr, N, buf, buflen);
+                Symmetric.Stream128SqueezeBlocks(buf, off, Symmetric.Stream128BlockBytes);
+                buflen = Symmetric.Stream128BlockBytes + off;
+                ctr += RejectUniform(Coeffs, ctr, N - ctr, buf, buflen);
             }
 
 
@@ -76,30 +78,29 @@ namespace Org.BouncyCastle.Pqc.Crypto.Crystals.Dilithium
 
             if (Engine.Eta == 2)
             {
-                PolyUniformEtaNBlocks = ((136 + Symmetric.Shake128Rate - 1) / Symmetric.Shake256Rate);
+                PolyUniformEtaNBlocks = ((136 + Symmetric.Stream256BlockBytes - 1) / Symmetric.Stream256BlockBytes);
             }
             else if (Engine.Eta == 4)
             {
-                PolyUniformEtaNBlocks = ((227 + Symmetric.Shake128Rate - 1) / Symmetric.Shake256Rate);
+                PolyUniformEtaNBlocks = ((227 + Symmetric.Stream256BlockBytes - 1) / Symmetric.Stream256BlockBytes);
             }
             else
             {
                 throw new ArgumentException("Wrong Dilithium Eta!");
             }
 
-            int buflen = PolyUniformEtaNBlocks * Symmetric.Shake128Rate;
+            int buflen = PolyUniformEtaNBlocks * Symmetric.Stream256BlockBytes;
 
             byte[] buf = new byte[buflen];
-            ShakeDigest Shake256Digest = new ShakeDigest(256);
 
-            Symmetric.ShakeStreamInit(Shake256Digest, seed, nonce);
-            Shake256Digest.DoOutput(buf, 0, buflen);
+            Symmetric.Stream256Init(seed, nonce);
+            Symmetric.Stream256SqueezeBlocks(buf, 0, buflen);
             ctr = RejectEta(Coeffs, 0, N, buf, buflen, eta);
 
             while (ctr < DilithiumEngine.N)
             {
-                Shake256Digest.DoOutput(buf, buflen, Symmetric.Shake128Rate);
-                ctr += RejectEta(Coeffs, ctr, N - ctr, buf, Symmetric.Shake128Rate, eta);
+                Symmetric.Stream256SqueezeBlocks(buf, 0, Symmetric.Stream256BlockBytes);
+                ctr += RejectEta(Coeffs, ctr, N - ctr, buf, Symmetric.Stream256BlockBytes, eta);
             }
         }
 
@@ -157,26 +158,12 @@ namespace Org.BouncyCastle.Pqc.Crypto.Crystals.Dilithium
             Poly t = new Poly(Engine);
 
             PointwiseMontgomery(u.Vec[0], v.Vec[0]);
-            //Console.Write("temp = [");
-            //for (int j = 0; j < N; ++j)
-            //{
-            //    Console.Write("{0}, ", Coeffs[j]);
-            //}
-            //Console.Write("]\n");
-
-            
 
             for (i = 1; i < Engine.L; ++i)
             {
                 t.PointwiseMontgomery(u.Vec[i], v.Vec[i]);
                 AddPoly(t);
             }
-            //Console.Write("temp = [");
-            //for (int j = 0; j < N; ++j)
-            //{
-            //    Console.Write("{0}, ", Coeffs[j]);
-            //}
-            //Console.Write("]\n");
         }
 
         public void AddPoly(Poly a)
@@ -339,28 +326,30 @@ namespace Org.BouncyCastle.Pqc.Crypto.Crystals.Dilithium
             }
         }
 
-        public void PolyT1Pack(byte[] r, int off)
+        public byte[] PolyT1Pack()
         {
+            byte[] output = new byte[DilithiumEngine.PolyT1PackedBytes];
             for (int i = 0; i < N / 4; ++i)
             {
-                r[off + 5 * i + 0] = (byte)(Coeffs[4 * i + 0] >> 0);
-                r[off + 5 * i + 1] = (byte)((Coeffs[4 * i + 0] >> 8) | (Coeffs[4 * i + 1] << 2));
-                r[off + 5 * i + 2] = (byte)((Coeffs[4 * i + 1] >> 6) | (Coeffs[4 * i + 2] << 4));
-                r[off + 5 * i + 3] = (byte)((Coeffs[4 * i + 2] >> 4) | (Coeffs[4 * i + 3] << 6));
-                r[off + 5 * i + 4] = (byte)(Coeffs[4 * i + 3] >> 2);
+                output[5 * i + 0] = (byte)(Coeffs[4 * i + 0] >> 0);
+                output[5 * i + 1] = (byte)((Coeffs[4 * i + 0] >> 8) | (Coeffs[4 * i + 1] << 2));
+                output[5 * i + 2] = (byte)((Coeffs[4 * i + 1] >> 6) | (Coeffs[4 * i + 2] << 4));
+                output[5 * i + 3] = (byte)((Coeffs[4 * i + 2] >> 4) | (Coeffs[4 * i + 3] << 6));
+                output[5 * i + 4] = (byte)(Coeffs[4 * i + 3] >> 2);
             }
+            return output;
         }
 
-        public void PolyT1Unpack(byte[] a, int off)
+        public void PolyT1Unpack(byte[] a)
         {
             int i;
 
             for (i = 0; i < N / 4; ++i)
             {
-                Coeffs[4 * i + 0] = (((a[off + 5 * i + 0] & 0xFF) >> 0) | ((int)(a[off + 5 * i + 1] & 0xFF) << 8)) & 0x3FF;
-                Coeffs[4 * i + 1] = (((a[off + 5 * i + 1] & 0xFF) >> 2) | ((int)(a[off + 5 * i + 2] & 0xFF) << 6)) & 0x3FF;
-                Coeffs[4 * i + 2] = (((a[off + 5 * i + 2] & 0xFF) >> 4) | ((int)(a[off + 5 * i + 3] & 0xFF) << 4)) & 0x3FF;
-                Coeffs[4 * i + 3] = (((a[off + 5 * i + 3] & 0xFF) >> 6) | ((int)(a[off + 5 * i + 4] & 0xFF) << 2)) & 0x3FF;
+                Coeffs[4 * i + 0] = (((a[5 * i + 0] & 0xFF) >> 0) | ((int)(a[5 * i + 1] & 0xFF) << 8)) & 0x3FF;
+                Coeffs[4 * i + 1] = (((a[5 * i + 1] & 0xFF) >> 2) | ((int)(a[5 * i + 2] & 0xFF) << 6)) & 0x3FF;
+                Coeffs[4 * i + 2] = (((a[5 * i + 2] & 0xFF) >> 4) | ((int)(a[5 * i + 3] & 0xFF) << 4)) & 0x3FF;
+                Coeffs[4 * i + 3] = (((a[5 * i + 3] & 0xFF) >> 6) | ((int)(a[5 * i + 4] & 0xFF) << 2)) & 0x3FF;
             }
         }
 
@@ -443,11 +432,10 @@ namespace Org.BouncyCastle.Pqc.Crypto.Crystals.Dilithium
 
         public void UniformGamma1(byte[] seed, ushort nonce)
         {
-            byte[] buf = new byte[Engine.PolyUniformGamma1NBytes * Symmetric.Shake256Rate];
-            ShakeDigest ShakeDigest256 = new ShakeDigest(256);
-            Symmetric.ShakeStreamInit(ShakeDigest256, seed, nonce);
-            ShakeDigest256.DoFinal(buf, 0, buf.Length);
-            UnpackZ(buf, 0);
+            byte[] buf = new byte[Engine.PolyUniformGamma1NBytes * Symmetric.Stream256BlockBytes];
+            Symmetric.Stream256Init(seed, nonce);
+            Symmetric.Stream256SqueezeBlocks(buf, 0, buf.Length);
+            UnpackZ(buf);
 
         }
 
@@ -496,7 +484,7 @@ namespace Org.BouncyCastle.Pqc.Crypto.Crystals.Dilithium
             }
         }
 
-        public void UnpackZ(byte[] a, int off)
+        public void UnpackZ(byte[] a)
         {
             int i;
             if (Engine.Gamma1 == (1 << 17))
@@ -505,27 +493,27 @@ namespace Org.BouncyCastle.Pqc.Crypto.Crystals.Dilithium
                 {
                     Coeffs[4 * i + 0] =
                         (
-                            (((a[off + 9 * i + 0] & 0xFF)) |
-                                ((a[off + 9 * i + 1] & 0xFF) << 8)) |
-                                ((a[off + 9 * i + 2] & 0xFF) << 16)
+                            (((a[9 * i + 0] & 0xFF)) |
+                                ((a[9 * i + 1] & 0xFF) << 8)) |
+                                ((a[9 * i + 2] & 0xFF) << 16)
                         ) & 0x3FFFF;
                     Coeffs[4 * i + 1] =
                         (
-                            (((a[off + 9 * i + 2] & 0xFF) >> 2) |
-                                ((a[off + 9 * i + 3] & 0xFF) << 6)) |
-                                ((a[off + 9 * i + 4] & 0xFF) << 14)
+                            (((a[9 * i + 2] & 0xFF) >> 2) |
+                                ((a[9 * i + 3] & 0xFF) << 6)) |
+                                ((a[9 * i + 4] & 0xFF) << 14)
                         ) & 0x3FFFF;
                     Coeffs[4 * i + 2] =
                         (
-                            (((a[off + 9 * i + 4] & 0xFF) >> 4) |
-                                ((a[off + 9 * i + 5] & 0xFF) << 4)) |
-                                ((a[off + 9 * i + 6] & 0xFF) << 12)
+                            (((a[9 * i + 4] & 0xFF) >> 4) |
+                                ((a[9 * i + 5] & 0xFF) << 4)) |
+                                ((a[9 * i + 6] & 0xFF) << 12)
                         ) & 0x3FFFF;
                     Coeffs[4 * i + 3] =
                         (
-                            (((a[off + 9 * i + 6] & 0xFF) >> 6) |
-                                ((a[off + 9 * i + 7] & 0xFF) << 2)) |
-                                ((a[off + 9 * i + 8] & 0xFF) << 10)
+                            (((a[9 * i + 6] & 0xFF) >> 6) |
+                                ((a[9 * i + 7] & 0xFF) << 2)) |
+                                ((a[9 * i + 8] & 0xFF) << 10)
                         ) & 0x3FFFF;
 
 
@@ -541,15 +529,15 @@ namespace Org.BouncyCastle.Pqc.Crypto.Crystals.Dilithium
                 {
                     Coeffs[2 * i + 0] =
                         (
-                            (((a[off + 5 * i + 0] & 0xFF)) |
-                                ((a[off + 5 * i + 1] & 0xFF) << 8)) |
-                                ((a[off + 5 * i + 2] & 0xFF) << 16)
+                            (((a[5 * i + 0] & 0xFF)) |
+                                ((a[5 * i + 1] & 0xFF) << 8)) |
+                                ((a[5 * i + 2] & 0xFF) << 16)
                         ) & 0xFFFFF;
                     Coeffs[2 * i + 1] =
                         (
-                            (((a[off + 5 * i + 2] & 0xFF) >> 4) |
-                                ((a[off + 5 * i + 3] & 0xFF) << 4)) |
-                                ((a[off + 5 * i + 4] & 0xFF) << 12)
+                            (((a[5 * i + 2] & 0xFF) >> 4) |
+                                ((a[5 * i + 3] & 0xFF) << 4)) |
+                                ((a[5 * i + 4] & 0xFF) << 12)
                         ) & 0xFFFFF;
 
                     Coeffs[2 * i + 0] = Engine.Gamma1 - Coeffs[2 * i + 0];
@@ -599,11 +587,11 @@ namespace Org.BouncyCastle.Pqc.Crypto.Crystals.Dilithium
         {
             int i, b, pos;
             ulong signs;
-            byte[] buf = new byte[Symmetric.Shake256Rate];
+            byte[] buf = new byte[Symmetric.Stream256BlockBytes];
 
             ShakeDigest ShakeDigest256 = new ShakeDigest(256);
             ShakeDigest256.BlockUpdate(seed, 0, DilithiumEngine.SeedBytes);
-            ShakeDigest256.DoOutput(buf, 0, Symmetric.Shake256Rate);
+            ShakeDigest256.DoOutput(buf, 0, Symmetric.Stream256BlockBytes);
 
             signs = 0;
             for (i = 0; i < 8; ++i)
@@ -622,9 +610,9 @@ namespace Org.BouncyCastle.Pqc.Crypto.Crystals.Dilithium
             {
                 do
                 {
-                    if (pos >= Symmetric.Shake256Rate)
+                    if (pos >= Symmetric.Stream256BlockBytes)
                     {
-                        ShakeDigest256.DoOutput(buf, 0, Symmetric.Shake256Rate);
+                        ShakeDigest256.DoOutput(buf, 0, Symmetric.Stream256BlockBytes);
                         pos = 0;
                     }
                     b = (buf[pos++] & 0xFF);
diff --git a/crypto/src/pqc/crypto/crystals/dilithium/PolyVecK.cs b/crypto/src/pqc/crypto/crystals/dilithium/PolyVecK.cs
index 21718c611..facbc1f33 100644
--- a/crypto/src/pqc/crypto/crystals/dilithium/PolyVecK.cs
+++ b/crypto/src/pqc/crypto/crystals/dilithium/PolyVecK.cs
@@ -1,6 +1,4 @@
-using System;
-
-namespace Org.BouncyCastle.Pqc.Crypto.Crystals.Dilithium
+namespace Org.BouncyCastle.Pqc.Crypto.Crystals.Dilithium
 {
     internal class PolyVecK
     {
diff --git a/crypto/src/pqc/crypto/crystals/dilithium/PolyVecL.cs b/crypto/src/pqc/crypto/crystals/dilithium/PolyVecL.cs
index 67cf4e9f2..75309b374 100644
--- a/crypto/src/pqc/crypto/crystals/dilithium/PolyVecL.cs
+++ b/crypto/src/pqc/crypto/crystals/dilithium/PolyVecL.cs
@@ -1,4 +1,8 @@
 using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Text;
+using System.Threading.Tasks;
 
 namespace Org.BouncyCastle.Pqc.Crypto.Crystals.Dilithium
 {
diff --git a/crypto/src/pqc/crypto/crystals/dilithium/PolyVecMatrix.cs b/crypto/src/pqc/crypto/crystals/dilithium/PolyVecMatrix.cs
index c899850f5..cb41889b9 100644
--- a/crypto/src/pqc/crypto/crystals/dilithium/PolyVecMatrix.cs
+++ b/crypto/src/pqc/crypto/crystals/dilithium/PolyVecMatrix.cs
@@ -1,4 +1,8 @@
 using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Text;
+using System.Threading.Tasks;
 
 namespace Org.BouncyCastle.Pqc.Crypto.Crystals.Dilithium
 {
@@ -40,4 +44,4 @@ namespace Org.BouncyCastle.Pqc.Crypto.Crystals.Dilithium
             }
         }
     }
-}
+}
\ No newline at end of file
diff --git a/crypto/src/pqc/crypto/crystals/dilithium/Reduce.cs b/crypto/src/pqc/crypto/crystals/dilithium/Reduce.cs
index 0deea4698..0f8b95ef1 100644
--- a/crypto/src/pqc/crypto/crystals/dilithium/Reduce.cs
+++ b/crypto/src/pqc/crypto/crystals/dilithium/Reduce.cs
@@ -1,6 +1,4 @@
-using System;
-
-namespace Org.BouncyCastle.Pqc.Crypto.Crystals.Dilithium
+namespace Org.BouncyCastle.Pqc.Crypto.Crystals.Dilithium
 {
     internal class Reduce
     {
@@ -29,4 +27,4 @@ namespace Org.BouncyCastle.Pqc.Crypto.Crystals.Dilithium
 
 
     }
-}
+}
\ No newline at end of file
diff --git a/crypto/src/pqc/crypto/crystals/dilithium/Rounding.cs b/crypto/src/pqc/crypto/crystals/dilithium/Rounding.cs
index 18f0fd7e7..f8deb77f4 100644
--- a/crypto/src/pqc/crypto/crystals/dilithium/Rounding.cs
+++ b/crypto/src/pqc/crypto/crystals/dilithium/Rounding.cs
@@ -1,5 +1,4 @@
 using System;
-
 namespace Org.BouncyCastle.Pqc.Crypto.Crystals.Dilithium
 {
     internal class Rounding
diff --git a/crypto/src/pqc/crypto/crystals/dilithium/Symmetric.cs b/crypto/src/pqc/crypto/crystals/dilithium/Symmetric.cs
index c866ccece..a46c12d90 100644
--- a/crypto/src/pqc/crypto/crystals/dilithium/Symmetric.cs
+++ b/crypto/src/pqc/crypto/crystals/dilithium/Symmetric.cs
@@ -1,21 +1,127 @@
 using Org.BouncyCastle.Crypto.Digests;
+using Org.BouncyCastle.Crypto.Engines;
+using Org.BouncyCastle.Crypto.Modes;
+using Org.BouncyCastle.Crypto.Parameters;
+using Org.BouncyCastle.Utilities;
 
 namespace Org.BouncyCastle.Pqc.Crypto.Crystals.Dilithium
 {
-    static class Symmetric
+    public abstract class Symmetric
     {
-        public const int Shake128Rate = 168;
-        public const int Shake256Rate = 136;
-        public const int Sha3Rate256 = 136;
-        public const int Sha3Rate512 = 72;
+        public int Stream128BlockBytes;
+        public int Stream256BlockBytes;
+        
+        Symmetric(int stream128, int stream256)
+        {
+            Stream128BlockBytes = stream128;
+            Stream256BlockBytes = stream256;
+        }
+        
+        internal abstract void Stream128Init(byte[] seed, ushort nonce);
+
+        internal abstract void Stream256Init(byte[] seed, ushort nonce);
+
+        internal abstract void Stream128SqueezeBlocks(byte[] output, int offset, int size);
+
+        internal abstract void Stream256SqueezeBlocks(byte[] output, int offset, int size);
+        
+        internal class AesSymmetric
+            : Symmetric
+        {
+
+            private SicBlockCipher cipher;
+
+            public AesSymmetric()
+                : base(64, 64)
+            {
+                cipher = new SicBlockCipher(new AesEngine());
+            }
+
+            private void Aes128(byte[] output, int offset, int size)
+            {
+                byte[] buf = new byte[size];   // TODO: there might be a more efficient way of doing this...
+                for (int i = 0; i < size; i+= 16)
+                {
+                    cipher.ProcessBlock(buf, i + offset, output, i + offset);
+                }
+            }
+
+            private void StreamInit(byte[] key, ushort nonce)
+            {
+                byte[] expnonce = new byte[12];
+                expnonce[0] = (byte)nonce;
+                expnonce[1] = (byte)(nonce >> 8);
+                
+                ParametersWithIV kp = new ParametersWithIV(new KeyParameter(Arrays.CopyOfRange(key, 0, 32)), expnonce);
+                cipher.Init(true, kp);
+            }
+
+            internal override void Stream128Init(byte[] seed, ushort nonce)
+            {
+                StreamInit(seed, nonce);
+            }
+
+            internal override void Stream256Init(byte[] seed, ushort nonce)
+            {
+                StreamInit(seed, nonce);
+            }
 
-        public static void ShakeStreamInit(ShakeDigest Digest, byte[] seed, ushort nonce)
+            internal override void Stream128SqueezeBlocks(byte[] output, int offset, int size)
+            {
+                Aes128(output, offset, size);
+            }
+
+            internal override void Stream256SqueezeBlocks(byte[] output, int offset, int size)
+            {
+                Aes128(output, offset, size);
+            }
+        }
+
+
+        internal class ShakeSymmetric
+            : Symmetric
         {
-            byte[] temp = new byte[2];
-            temp[0] = (byte)nonce;
-            temp[1] = (byte)(nonce >> 8);
-            Digest.BlockUpdate(seed, 0, seed.Length);
-            Digest.BlockUpdate(temp, 0, temp.Length);
+            private ShakeDigest digest128;
+            private ShakeDigest digest256;
+
+            public ShakeSymmetric()
+                : base(168, 136)
+            {
+                digest128 = new ShakeDigest(128);
+                digest256 = new ShakeDigest(256);
+            }
+
+            private void StreamInit(ShakeDigest digest, byte[] seed, ushort nonce)
+            {
+                digest.Reset();
+                byte[] temp = new byte[2];
+                temp[0] = (byte)nonce;
+                temp[1] = (byte)(nonce >> 8);
+
+                digest.BlockUpdate(seed, 0, seed.Length);
+                digest.BlockUpdate(temp, 0, temp.Length);
+            }
+
+
+            internal override void Stream128Init(byte[] seed, ushort nonce)
+            {
+                StreamInit(digest128, seed, nonce);
+            }
+
+            internal override void Stream256Init(byte[] seed, ushort nonce)
+            {
+                StreamInit(digest256, seed, nonce);
+            }
+
+            internal override void Stream128SqueezeBlocks(byte[] output, int offset, int size)
+            {
+                digest128.DoOutput(output, offset, size);
+            }
+
+            internal override void Stream256SqueezeBlocks(byte[] output, int offset, int size)
+            {
+                digest256.DoOutput(output, offset, size);
+            }
         }
     }
 }