Refactor CBC utils

8 files changed, 98 insertions, 90 deletions

diff --git a/src/solvers/Challenge10.scala b/src/solvers/Challenge10.scala
index a130106..7880327 100644
--- a/src/solvers/Challenge10.scala
+++ b/src/solvers/Challenge10.scala
@@ -6,9 +6,8 @@ import ixee.cryptopals.utils.StreamUtils._
 import ixee.cryptopals.utils.ConversionUtils._
 import ixee.cryptopals.utils.FunctionUtils._
 import ixee.cryptopals.utils._
+import ixee.cryptopals.utils.crypto._
 import io.github.marklister.base64.Base64._
-import javax.crypto.Cipher
-import javax.crypto.spec.SecretKeySpec
 
 object Challenge10 {
   val path = "./data/10.txt"
@@ -22,30 +21,23 @@ object Challenge10 {
       .toByteArray
 
   def run = {
-    val encInstance = CryptoUtils.cbcEncryptInstance(
-      "AES",
-      "YELLOW SUBMARINE".asBytes.toArray,
-      Stream.continually(0.toByte).take(16)
-    )
-
-    val decInstance = CryptoUtils.cbcDecryptInstance(
-      "AES",
-      "YELLOW SUBMARINE".asBytes.toArray,
-      Stream.continually(0.toByte).take(16)
-    )
-
-    val s = "foo bar frobnicator the quick brown fox jumps over the lazy dog".asBytes
-
-    val enc = encInstance.enc(s) ++ encInstance.end()
-    val dec = decInstance.dec(enc)
-    println(dec.startsWith(enc))
-
-    val decInstance2 = CryptoUtils.cbcDecryptInstance(
-      "AES",
-      "YELLOW SUBMARINE".asBytes.toArray,
-      Stream.continually(0.toByte).take(16)
-    )
-
-    decInstance2.dec(ciphertext)
+    def cbcBuilder = SchemeBuilder("AES", "YELLOW SUBMARINE".asBytes)
+      .cbc(Stream.continually(0.toByte).take(16))
+
+    val encInstance = cbcBuilder.encrypt
+    val decInstance = cbcBuilder.decrypt
+
+    val s = "fooo bar frobnicator the quick brown fox jumps over the lazy dog".asBytes
+
+    println(s)
+    val enc = encInstance.end(s)
+    val dec = decInstance.end(enc)
+    println(enc)
+    println(dec)
+    println(dec.startsWith(s))
+
+    val decInstance2 = cbcBuilder.decrypt
+
+    decInstance2.end(ciphertext)
   }
 }
diff --git a/src/utils/CryptoUtils.scala b/src/utils/CryptoUtils.scala
index 430934b..6c2849e 100644
--- a/src/utils/CryptoUtils.scala
+++ b/src/utils/CryptoUtils.scala
@@ -11,17 +11,4 @@ object CryptoUtils {
     s ++ Stream.continually(padLength.toByte).take(padLength)
   }
 
-  def cbcDecryptInstance(cipherPrim: String, key: Seq[Byte], iv: Seq[Byte]): CBCDecrypter = {
-    val keySpec = new SecretKeySpec(key.toArray, cipherPrim)
-    val cipher = Cipher.getInstance(s"$cipherPrim/ECB/NoPadding")
-    cipher.init(Cipher.DECRYPT_MODE, keySpec)
-    new CBCDecrypter(cipher, iv)
-  }
-
-  def cbcEncryptInstance(cipherPrim: String, key: Seq[Byte], iv: Seq[Byte]): CBCEncrypter = {
-    val keySpec = new SecretKeySpec(key.toArray, cipherPrim)
-    val cipher = Cipher.getInstance(s"$cipherPrim/ECB/NoPadding")
-    cipher.init(Cipher.ENCRYPT_MODE, keySpec)
-    new CBCEncrypter(cipher, iv)
-  }
 }
diff --git a/src/utils/FunctionUtils.scala b/src/utils/FunctionUtils.scala
index 21ad8d0..3d6610c 100644
--- a/src/utils/FunctionUtils.scala
+++ b/src/utils/FunctionUtils.scala
@@ -9,4 +9,11 @@ object FunctionUtils {
   implicit class Compositor[A, B](f: A => B) {
     def :|[C](g: B => C): A => C = f.andThen(g)
   }
+
+  implicit class Tap[A](x: A) {
+    def tap(f: A => Any): A = {
+      f(x)
+      x
+    }
+  }
 }
diff --git a/src/utils/crypto/CBCCipher.scala b/src/utils/crypto/CBCCipher.scala
index d6d99ff..3bd0784 100644
--- a/src/utils/crypto/CBCCipher.scala
+++ b/src/utils/crypto/CBCCipher.scala
@@ -7,28 +7,38 @@ import ixee.cryptopals.utils.CryptoUtils._
 import ixee.cryptopals.utils.TupleUtils._
 import ixee.cryptopals.utils.ByteUtils._
 
-class CBCEncrypter(cipher: Cipher, init: Seq[Byte]) {
-  val blockSize = 16
+class CBCCipher(private[this] val cipher: Cipher, private[this] val iv: Seq[Byte], mode: Int) extends IxeeCipher {
+  val blockSize = cipher.getBlockSize
 
-  var state: Seq[Byte] = init
+  var state: Seq[Byte] = iv
   var leftover: Seq[Byte] = Seq()
 
-  def enc(data: Seq[Byte]): Seq[Byte] = {
-    val (blocks, newLeftover) = blockized(leftover ++ data)(blockSize)
-    leftover = newLeftover
-    blocks.map(encBlock _).foldLeft(Seq[Byte]())(_ ++ _)
-  }
+  lazy val handleBlock: Seq[Byte] => Seq[Byte] =
+    if (mode == Cipher.ENCRYPT_MODE) encBlock _
+    else decBlock _
 
-  def encBlock(data: Seq[Byte]): Seq[Byte] = {
-    state = cipher.update((data xor state).toArray)
-    state
+  def update(data: Seq[Byte]): Seq[Byte] = {
+    val blocks = blockized(leftover ++ data).tap(updateLeftover)._1
+    blocks.foldLeft(Seq[Byte]())(_ ++ handleBlock(_))
   }
 
+  private def decBlock(data: Seq[Byte]): Seq[Byte] =
+    (cipher.update(data.toArray).toSeq xor state).tap(_ => state = data)
+
+  private def encBlock(data: Seq[Byte]): Seq[Byte] =
+    cipher.update((data xor state).toArray).tap(state = _)
+
+  // wouldn't hurt to invalidate this object afterward, but meh
+  // TODO: strip padding!
   def end(): Seq[Byte] =
-    cipher.doFinal((pkcs7pad(leftover, blockSize) xor state).toArray)
+    if (mode == Cipher.DECRYPT_MODE)
+    else cipher.update((pkcs7pad(leftover, blockSize) xor state).toArray)
 
-  def blockized(data: Seq[Byte])(size: Int): (Seq[Seq[Byte]], Seq[Byte]) =
-    groupBlocks <-: data.splitAt(data.length - (data.length % size))
+  def blockized(data: Seq[Byte]): (Seq[Seq[Byte]], Seq[Byte]) =
+    groupBlocks <-: data.splitAt(data.length - (data.length % blockSize))
 
   def groupBlocks: Seq[Byte] => Seq[Seq[Byte]] = _.grouped(blockSize).toSeq
+
+  def updateLeftover(pair: (Seq[Seq[Byte]], Seq[Byte])) =
+    leftover = pair._2
 }
diff --git a/src/utils/crypto/CBCDecrypter.scala b/src/utils/crypto/CBCDecrypter.scala
deleted file mode 100644
index a6cf855..0000000
--- a/src/utils/crypto/CBCDecrypter.scala
+++ /dev/null
@@ -1,36 +0,0 @@
-
-package ixee.cryptopals.utils.crypto
-
-import javax.crypto.Cipher
-import ixee.cryptopals.utils.ConversionUtils._
-import ixee.cryptopals.utils.FunctionUtils._
-import ixee.cryptopals.utils.CryptoUtils._
-import ixee.cryptopals.utils.TupleUtils._
-import ixee.cryptopals.utils.ByteUtils._
-
-class CBCDecrypter(cipher: Cipher, init: Seq[Byte]) {
-  val blockSize = 16
-
-  var state: Seq[Byte] = init
-  var leftover: Seq[Byte] = Seq()
-
-  def dec(data: Seq[Byte]): Seq[Byte] = {
-    val (blocks, newLeftover) = blockized(leftover ++ data)(blockSize)
-    leftover = newLeftover
-    blocks.map(decBlock _).foldLeft(Seq[Byte]())(_ ++ _)
-  }
-
-  def decBlock(data: Seq[Byte]): Seq[Byte] = {
-    val ret = cipher.update(data.toArray).toSeq xor state
-    state = data
-    ret
-  }
-
-  def end(): Seq[Byte] =
-    cipher.doFinal((pkcs7pad(leftover, blockSize) xor state).toArray)
-
-  def blockized(data: Seq[Byte])(size: Int): (Seq[Seq[Byte]], Seq[Byte]) =
-    groupBlocks <-: data.splitAt(data.length - (data.length % size))
-
-  def groupBlocks: Seq[Byte] => Seq[Seq[Byte]] = _.grouped(blockSize).toSeq
-}
diff --git a/src/utils/crypto/CipherGenerator.scala b/src/utils/crypto/CipherGenerator.scala
new file mode 100644
index 0000000..e7651cb
--- /dev/null
+++ b/src/utils/crypto/CipherGenerator.scala
@@ -0,0 +1,27 @@
+package ixee.cryptopals.utils.crypto
+
+import javax.crypto.Cipher
+
+sealed trait CipherGenerator {
+  def encrypt: IxeeCipher // feels kinda gross to type this
+  def decrypt: IxeeCipher // writing crypto feels wrong
+}
+
+case class EcbBuilder(primitives: SchemeBuilder) extends CipherGenerator {
+  // lol
+  def encrypt = ???
+  def decrypt = ???
+}
+
+case class CbcBuilder(primitives: SchemeBuilder, iv: Seq[Byte]) extends CipherGenerator {
+  def encrypt =
+    setup(Cipher.ENCRYPT_MODE)
+  def decrypt =
+    setup(Cipher.DECRYPT_MODE)
+
+  private[this] def setup(cipherMode: Int) = {
+    val cipher = primitives.cipher
+    cipher.init(cipherMode, primitives.key)
+    new CBCCipher(cipher, iv, cipherMode)
+  }
+}
diff --git a/src/utils/crypto/IxeeCipher.scala b/src/utils/crypto/IxeeCipher.scala
new file mode 100644
index 0000000..e4a8617
--- /dev/null
+++ b/src/utils/crypto/IxeeCipher.scala
@@ -0,0 +1,8 @@
+package ixee.cryptopals.utils.crypto
+
+trait IxeeCipher {
+  def update(data: Seq[Byte]): Seq[Byte]
+  def end(): Seq[Byte]
+  def end(data: Seq[Byte]): Seq[Byte] =
+    update(data) ++ end()
+}
diff --git a/src/utils/crypto/SchemeBuilder.scala b/src/utils/crypto/SchemeBuilder.scala
new file mode 100644
index 0000000..ae92062
--- /dev/null
+++ b/src/utils/crypto/SchemeBuilder.scala
@@ -0,0 +1,13 @@
+package ixee.cryptopals.utils.crypto
+
+import javax.crypto.spec.SecretKeySpec
+import javax.crypto.Cipher
+
+case class SchemeBuilder(cipherAlgo: String, private val bareKey: Seq[Byte]) {
+  def cipher = Cipher.getInstance(s"$cipherAlgo/ECB/NoPadding")
+  lazy val key = new SecretKeySpec(bareKey.toArray, cipherAlgo)
+
+  def ecb = EcbBuilder(this)
+  def cbc(iv: Seq[Byte]) = CbcBuilder(this, iv)
+}
+