1 files changed, 72 insertions, 7 deletions

diff --git a/src/solvers/XorDecrypt.scala b/src/solvers/XorDecrypt.scala
index 6b82c34..83fa923 100644
--- a/src/solvers/XorDecrypt.scala
+++ b/src/solvers/XorDecrypt.scala
@@ -2,24 +2,89 @@ package ixee.cryptopals.solvers
 
 import ixee.cryptopals.utils._
 import ixee.cryptopals.utils.ByteUtils._
+import ixee.cryptopals.utils.StreamUtils._
 import ixee.cryptopals.utils.FunctionUtils._
+import ixee.cryptopals.utils.ConversionUtils._
 
 object XorDecrypt {
-  implicit val freq = Frequencies.cornell40kSample
+  implicit val freq = Frequencies.cryptologicalMathematics //cornell40kSample
 
-  def decryptToAscii(ciphertext: Seq[Byte])(key: Byte): String = {
-    new String((ciphertext xor Stream.continually(key)).toArray)
+  def inferKeySizes(xs: Seq[Byte]): Seq[Int] = {
+    inferKeySizesWithHammingDistance(xs).map(_._2)
   }
 
-  def tryBestDecrypt(ciphertext: Seq[Byte]): String = {
-    (findBestSingleByteKey _ :| decryptToAscii(ciphertext) _)(ciphertext)
+  def inferKeySize(xs: Seq[Byte]): Int =
+    inferKeySizes(xs).head
+
+  // not so great for small key sizes.
+  // TODO: Why exactly does this work?
+  def inferKeySizesWithHammingDistance(xs: Seq[Byte]): Seq[(Double, Int)] = {
+    val MinKeySize = 1
+    val MaxKeySize = 40
+
+    (MinKeySize to MaxKeySize)
+      .map(normalizedHammingDistanceForKeySize(xs, _))
+      .zipWithIndex
+      .map( { case (distance, keySize) => (distance, keySize + 1) } )
+      .sortBy(_._1)
+  }
+
+  def normalizedHammingDistanceForKeySize(xs: Seq[Byte], keySize: Int) = {
+    val grouped = xs.grouped(keySize).toArray.dropRight(1)
+      // ez hack to drop non-`keySize` subcomponents of xs
+    (grouped.init zip grouped.tail).map(tup(hammingDistance(_, _))).reduce(_ + _) / (grouped.length - 1.0) / keySize
   }
 
-  def findBestSingleByteKey(ciphertext: Seq[Byte]): Byte =
+  def crackForMultiByteKey(ciphertext: Seq[Byte], inspectSet: Seq[Int] = Seq()): String =
+    decryptToAscii(ciphertext)(findBestMultiByteKey(ciphertext, inspectSet))
+
+  def findBestMultiByteKey(xs: Seq[Byte], inspectSet: Seq[Int] = Seq()): Seq[Byte] =
+    continuous(findBestMultiByteKeyOfSize(inferKeySize(xs))(xs, inspectSet))
+/*
+  def crackForKeySize(xs: Seq[Byte], keySize: Int): String = {
+    (findBestMultiByteKeyOfSize(keySize) _ :| continuous _ :| decryptToAscii(xs) _)(xs)
+  }
+*/
+  def decrypt(ciphertext: Seq[Byte])(key: Seq[Byte]): Seq[Byte] =
+    ciphertext xor key
+
+  def decryptToAscii(ciphertext: Seq[Byte])(key: Seq[Byte]): String =
+    decrypt(ciphertext)(key).asAscii
+
+    // TODO: make an xor key a proper type.
+  def decryptToAsciiSingleByteKey(ciphertext: Seq[Byte])(key: Byte): String =
+    decryptToAscii(ciphertext)(Stream.continually(key))
+
+  def tryBestDecrypt(ciphertext: Seq[Byte]): Seq[Byte] = {
+    (findBestSingleByteKey _ :| { x => Stream.continually(x) } :| decrypt(ciphertext) _)(ciphertext)
+  }
+
+  def tryBestDecryptToText: Seq[Byte] => String =
+    tryBestDecrypt _ :| ((_: Seq[Byte]).asAscii)
+
+  def findBestMultiByteKeyOfSize(keySize: Int)(ciphertext: Seq[Byte], inspectSet: Seq[Int] = Seq()): Seq[Byte] = {
+    // inspectSet is the set of key bytes for which we want debugging information
+    ciphertext.grouped(keySize).toSeq.dropRight(1).transpose.zipWithIndex
+      .map { case (byteCiphertext, idx) => {
+        if (inspectSet.contains(idx)) {
+          println("For idx = " + idx)
+          println(findBestSingleByteKeyWithCandidates(byteCiphertext))
+        }
+        findBestSingleByteKey(byteCiphertext)
+      }}
+  }
+
+  def findBestSingleByteKeyWithCandidates(ciphertext: Seq[Byte]): Seq[(Double, String, Byte)] =
+    candidates(ciphertext).sortBy(_._1).map {
+      case (score, key) => (score, decryptToAsciiSingleByteKey(ciphertext)(key.toByte), key.toByte)
+    }
+
+  def findBestSingleByteKey(ciphertext: Seq[Byte]): Byte = {
     candidates(ciphertext)
       .reduceOption( (a: (Double, Int), b: (Double, Int)) =>
         if (a._1 < b._1) a else b
       ).map(_._2.toByte).getOrElse(0.toByte)
+  }
 
   def candidates(ciphertext: Seq[Byte]) =
     (0 until 256)
@@ -30,6 +95,6 @@ object XorDecrypt {
 
   def candidatesAsAscii(ciphertext: Seq[Byte]) =
     candidates(ciphertext).sortBy(_._1).map(x =>
-      new String((ciphertext xor Stream.continually(x._2.toByte)).toArray)
+      decrypt(ciphertext)(Seq(x._2.toByte)).asAscii
     )
 }