Refactor CBC utils

5 files changed, 72 insertions, 50 deletions

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)
+}
+