package ixee.cryptopals.utils

import ixee.cryptopals.utils.crypto._
import ixee.cryptopals.utils.TupleUtils._
import ixee.cryptopals.utils.StreamUtils._
import ixee.cryptopals.utils.FunctionUtils._
import ixee.cryptopals.utils.ConversionUtils._
import javax.crypto.Cipher
import javax.crypto.spec.SecretKeySpec

object CryptoUtils {

  def pkcs7pad(s: Seq[Byte], blockSize: Int): Seq[Byte] = {
    val padLength = blockSize - (s.length % blockSize)
    s ++ Stream.continually(padLength.toByte).take(padLength)
  }

  def stripPkcs7Pad(s: Seq[Byte]): Seq[Byte] =
    s.dropRight(s.last)

  def cbcEncrypt(builder: CbcBuilder)(data: Seq[Byte]) =
    builder.encrypt.end(data)

  def cbcDecrypt(builder: CbcBuilder)(data: Seq[Byte]) =
    stripPkcs7Pad(builder.decrypt.end(data))

  def ecbEncrypt(builder: EcbBuilder)(data: Seq[Byte]) =
    builder.encrypt.end(data)

  def ecbDecrypt(builder: EcbBuilder)(data: Seq[Byte]) =
    stripPkcs7Pad(builder.decrypt.end(data))

  def detectMode(xs: Seq[Byte]): String = {
    def dupBlocks(xs: Seq[Seq[Byte]]) =
      pairsOf(xs).map(tup(_ == _)).count(_ == true)

    def countDupBlocks(xs: Seq[Byte]): Int =
      dupBlocks(xs.grouped(16).toSeq.init.toStream)

                              //.... well very probably.
    if (countDupBlocks(xs) > 0) "ECB"
    else "CBC"
  }

  def isEcb(xs: Seq[Byte]): Boolean = detectMode(xs) == "ECB"

  def detectEcbBlockSize(encryptor: Seq[Byte] => Seq[Byte]): Int = {
    val pad = (2 to 512 by 2).map("@" * _).zipWithIndex.map {
      _.mapAll(_1 = _.asBytes, _2 = _ + 1)
    }

    val minPadSize = pad
      .map( { encryptor(_) } <-: _ )
      .find( { x => isEcb(x._1) })
      .map(_._2)

    minPadSize.get // if this was a None we have serious possibility of this not being ECB
  }

  def extractUnknownViaEcbOracle(encrypt: Seq[Byte] => Seq[Byte]) = {
    val blockSize = detectEcbBlockSize(encrypt)

    def rainbow(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 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 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) =>
        (ac._1 :+ nextByte(ac._1), ac._2 + (idx -> encrypt(prefix(ac._1))))
      }
    }

    def probeLastBlockSize: Int = {
      val baseBlockCount = encrypt(Seq[Byte]()).length
      val firstLargerCiphertext = (0 until blockSize)
        .map(" " * _).map(_.asBytes).map(encrypt)
        .zipWithIndex
        .find(_._1.length != baseBlockCount)

        // this will always be Some(_) because
        // somewhere between 0..blockSize WILL grow the text.
      firstLargerCiphertext.get._2
    }

    println("Last block is " + probeLastBlockSize + " bytes")

    val (firstBlock, ciphertexts) = probeFirstBlockAndPaddings
    /*
     * zip together cipher blocks so that they look like
     * rot0b0, rot1b0, rot2b0, rot3b0, rot4b0, rot5b0, rot6b0, rot7b0
     * rot0b1, rot1b1, rot2b1, ...
     *
     */
    val transposed = ciphertexts.toSeq.sortBy(_._1).map(_ :-> { x =>
      val y = x.grouped(16).toSeq
      println(y.length)
      y
    })

//    println("Also...")
//    println(rainbow(firstBlock.tail)(transposed(1)(0)))

    val prefix = firstBlock.tail
    val currRainbow = rainbow(prefix)
    println(ciphertexts(0).drop(16).take(16))
    val next = currRainbow(ciphertexts(0).drop(16).take(16).toSeq)
    println("Next: " + new String(Array(next.toByte)))
    val nowPrefix = (prefix.tail :+ next)
    val r2 = rainbow(nowPrefix)
    val next2 = r2(ciphertexts(1).drop(16).take(16).toSeq)
    println("Next: " + new String(Array(next2.toByte)))
    val pref3 = (nowPrefix.tail :+ next2)
    val r3 = rainbow(pref3)
    val next3 = r3(ciphertexts(2).drop(16).take(16).toSeq)
    println("Next: " + new String(Array(next3.toByte)))

    firstBlock
  }

}