-- Copyright 2013 Thomas Szczarkowski

module CTR where

import Data.Bits
import Data.Word
import Data.List
import qualified Data.Binary as Binary
import Data.Array.IArray
import Data.Function
import Text.Printf
import Control.Monad
import Control.Arrow

import OpenSSL.Symmetric
import OpenSSL.Random
import Bytes
import Ciphers
import CryptoRandom



ctr cipher key nonce input =   
  liftM unblocks $ sequence $ zipWith (\c p -> c >>= (return . xorBuffers p) )
  (ctrStream cipher key nonce)
  (blocks input)
  

-- not implemented: advancing a CTR stream for reuse

ctrStream cipher key (nonce :: Word64) = 
  map (blockCipher Encrypt cipher key) $ 
  map (\j -> littleEndian nonce ++ littleEndian (fromIntegral j :: Word64)) $ [0..]
  

-- Data.Binary has wrong endianness

littleEndian = reverse . decode . Binary.encode
  
  
  
p18 = do
  putStrLn "\n\n==Problem 18==\n"
  let ciphertext = decode $ Base64 "L77na/nrFsKvynd6HzOoG7GHTLXsTVu9qvY/2syLXzhPweyyMTJULu/6/kXX0KSvoOLSFQ=="
  plaintext <- ctr aes (decode "YELLOW SUBMARINE") (0 :: Word64) ciphertext
  putStrLn $ encode plaintext


-- same method as repeating XOR (reusing solution from problem #6)


attackReusedNonce ciphertexts = 
  let slices = transpose ciphertexts
      keystr = map (head . decryptKey . bruteForceXor1) slices
  in map (zipWith xor keystr) ciphertexts
     

p20 = do
  putStrLn "\n\n==Problem 20==\n"
  ciphertexts <- liftM (map (decode.Base64) . lines) $ readFile "gistfile-p20.txt"
  let plaintexts = attackReusedNonce ciphertexts
  putStrLn $ intercalate "\n" $ map encode plaintexts



-- from problem #6
-- no new code follows
  

-- brute-force a one-byte key

bruteForceXor1 cipherText = 
  head $ sortBy (compare `on` negate.decryptScore) $
  [try byte | (byte :: Word8) <- [0..255]]
  where try byte =  
          let key = [byte]
              plaintext = xorCrypt cipherText key
              score = sum $ map (logByteFreqs !) plaintext
          in Decryption key plaintext score

data Decryption = Decryption { 
  decryptKey :: [Word8],
  decryptPlaintext :: [Word8],
  decryptScore :: Float
} deriving Show


-- English letter frequencies: (case-insensitive)
-- https://en.wikipedia.org/wiki/Letter_frequency#Relative_frequencies_of_letters_in_the_English_language
        
freqAZ = [0.08167, 0.01492, 0.02782, 0.04253, 0.12702, 0.02228, 0.02015, 0.06094, 0.06966, 0.00153, 0.00772, 0.04025, 0.02406, 0.06749, 0.07507, 0.01929, 0.00095, 0.05987, 0.06327, 0.09056, 0.02758, 0.00978, 0.02360, 0.00150, 0.01974, 0.00074]

-- extension to byte frequencies 
-- (I made up these numbers, nothing empirical here)

byteFreqs :: Array Word8 Float  -- sums to 1.0
byteFreqs = 
  let lc = zip ['a'..'z'] $ map (0.80*) freqAZ
      uc = zip ['A'..'Z'] $ map (0.03*) freqAZ
      spc = [(' ', 0.15)]
      punc = zip "!@#$%^&*()?/,.:;'\"" $ repeat (0.019 / 18.0)
      rest = zip (encode [0..255]) $ repeat (0.001 / 256.0)
  in accumArray (+) 0 (0,255) $ 
     map (\(c,f) -> (fromIntegral $ head $ decode [c], f)) $ 
     concat [lc,uc,spc,punc,rest]
                                                            

-- (easier to work logarithmically)

logByteFreqs = amap log byteFreqs :: Array Word8 Float

xorCrypt text key = xorBuffers text (concat (repeat key))