1 files changed, 30 insertions, 21 deletions
diff --git a/src/utils/CryptoUtils.scala b/src/utils/CryptoUtils.scala
index 2678895..658786a 100644
--- a/src/utils/CryptoUtils.scala
+++ b/src/utils/CryptoUtils.scala
@@ -58,36 +58,45 @@ object CryptoUtils {
   }
 
   def extractUnknownViaEcbOracle(encrypt: Seq[Byte] => Seq[Byte]) = {
-    val blockSize = detectEcbBlockSize(encrypt)
-    val baseCiphertext = encrypt(Seq())
+    class Ciphertext(encrypt: Seq[Byte] => Seq[Byte], prefix: Seq[Byte] = Seq(), private val overrideBlockSize: Option[Int] = None) {
+      private lazy val cipherBytes = encrypt(prefix)
+      lazy val blockSize = overrideBlockSize.getOrElse(detectEcbBlockSize(encrypt))
+      lazy val blocks = cipherBytes.grouped(blockSize).toSeq
+      lazy val blockCount = blocks.length
+      lazy val paddingSize = prefix.length
+    }
+
+    val baseCiphert = new Ciphertext(encrypt)
+
+    val cipherts = Seq(baseCiphert) ++
+      (1 to 15)
+        .map( x => ("@" * x).asBytes )
+        .map(new Ciphertext(encrypt, _, overrideBlockSize = Some(baseCiphert.blockSize)))
+
+    val blockSize = baseCiphert.blockSize //detectEcbBlockSize(encrypt)
+    //val baseCiphertext = baseCiphert.cipherBytes //encrypt(Seq())
     def blocksIn(xs: Seq[Byte]) = xs.length / blockSize
-    val baseBlockCount = blocksIn(baseCiphertext)
+    val baseBlockCount = baseCiphert.blockCount //blocksIn(baseCiphertext)
 
-    def rainbow(prefix: Seq[Byte]): Map[Seq[Byte], Byte] = {
+    def rainbow(const: Seq[Byte], generator: (Byte, Seq[Byte]) => Seq[Byte]) =
       (0 to 255)
         .map(_.toByte)
-        .map(prefix :+ _)
+        .map(generator(_, const))
         .map(encrypt)
-        .map(_.take(16).toSeq)
+        .map(_.take(blockSize).toSeq)
         .zipWithIndex
         .map(_ :-> { _.toByte } )
         .toMap
-    }
 
-    def suffixRainbow(prefix: Seq[Byte]): Map[Seq[Byte], Byte] = {
-      (0 to 255)
-        .map(_.toByte)
-        .map(_ +: prefix)
-        .map(encrypt)
-        .map(_.take(16).toSeq)
-        .zipWithIndex
-        .map(_ :-> { _.toByte } )
-        .toMap
-    }
+    def rainbowSuffix(prefix: Seq[Byte]): Map[Seq[Byte], Byte] =
+      rainbow(prefix, (byte, seq) => seq :+ byte)
+
+    def rainbowPrefix(suffix: Seq[Byte]): Map[Seq[Byte], Byte] =
+      rainbow(suffix, (byte, seq) => byte +: seq)
 
     def probeFirstBlockAndPaddings: (Seq[Byte], Map[Int, Seq[Byte]]) = {
       def prefix(known: Seq[Byte]) = ("@" * (blockSize - 1 - known.length)).asBytes
-      def genRainbow(known: Seq[Byte]) = rainbow(prefix(known) ++ known)
+      def genRainbow(known: Seq[Byte]) = rainbowSuffix(prefix(known) ++ known)
       def firstCryptedBlock(known: Seq[Byte]) = encrypt(prefix(known)).take(blockSize).toSeq
       def nextByte(known: Seq[Byte]) = genRainbow(known)(firstCryptedBlock(known))
       (0 until 16).foldLeft((Seq[Byte](), Map[Int, Seq[Byte]]())) { (ac, idx) =>
@@ -126,7 +135,7 @@ object CryptoUtils {
         val middleBlocks = curr._2.take(baseBlockCount).tail
 
         (ac._1 :+ middleBlocks, maybeLastBlock +: ac._2)
-      }) :-> { _.flatten :+ baseCiphertext.takeRight(16)}
+      }) :-> { _.flatten :+ baseCiphert.blocks.last } //baseCiphertext.takeRight(16)}
 
     def breakLastBlock(blocks: Seq[Seq[Byte]]) = {
       /*
@@ -138,14 +147,14 @@ object CryptoUtils {
       // drop the first block because it will be 16 16 16 16 16 16 ... 16
       blocks.tail.foldLeft(Seq[Byte]()) { (bytes, block) => {
         val postfix = pkcs7pad("?".asBytes ++ bytes, blockSize).tail
-        suffixRainbow(postfix)(block) +: bytes
+        rainbowPrefix(postfix)(block) +: bytes
       }}
     }
 
     def breakBlock(plaintext: Seq[Byte], padded: Seq[Seq[Byte]]): Seq[Byte] =
       padded.foldLeft(Seq[Byte]()) { (blockText: Seq[Byte], block: Seq[Byte]) => {
         val prefix = (plaintext ++ blockText).takeRight(blockSize - 1)
-        val b = rainbow(prefix)(block)
+        val b = rainbowSuffix(prefix)(block)
         blockText :+ b
       }}