-- Copyright 2013 Thomas Szczarkowski

import qualified Data.ByteString as BS
import qualified Data.ByteString.Char8 as Char8
import qualified Data.ByteString.Base64 as Base64
import Data.Char
import Data.Function
import Data.Array.IArray
import Data.List
import Data.Bits
import Data.Word
import Control.Monad
import Text.Printf

import SymmetricCiphers 


main = p1 >> p2 >> p3 >> p4 >> p5 >> p6 >> p7 >> p8


p1 = do
  putStrLn "\n==Problem 1==\n"
  let x = decodeHex "49276d206b696c6c696e6720796f757220627261696e206c696b65206120706f69736f6e6f7573206d757368726f6f6d"
  let y = encodeB64 x
  let x' = decodeB64 y
  [encodeHex x, y, encodeHex x'] `each` putStrLn


each xs f = sequence $ map f xs
infixl 1 `each`

decodeHex (x:y:rest) = (read ("0x" ++ [x,y])) : decodeHex rest
decodeHex _ = []
encodeHex = concat . (map (\x -> printf "%02x" x))

decodeB64 = BS.unpack . Base64.decodeLenient . Char8.pack
encodeB64 = Char8.unpack . Base64.encode . BS.pack

fromAscii = BS.unpack . Char8.pack
toAscii = Char8.unpack . BS.pack



p2 = do
  putStrLn "\n==Problem 2==\n"
  let x = decodeHex "1c0111001f010100061a024b53535009181c"
  let y = decodeHex "686974207468652062756c6c277320657965"
  let z = xorBuffers x y
  [x,y,z] `each` putStrLn . encodeHex
  
    
xorBuffers = zipWith xor



p3 = do    
  putStrLn "\n==Problem 3==\n"  
  let cipherText = "1b37373331363f78151b7f2b783431333d78397828372d363c78373e783a393b3736"
  let (Decryption key plainText score) = 
        bruteForceXor1 $ decodeHex cipherText
  putStrLn $ encodeHex key
  putStrLn $ toAscii plainText
  putStrLn $ printf "log10(likelihood ratio): %0.3f\n" 
    (logLikelihoodRatio plainText)


-- XOR cipher with repeating key
  
xorCrypt text key = xorBuffers text (concat (repeat key))


-- struct representing a a possible decryption

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


-- 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


-- 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 (map chr [0..255]) $ repeat (0.001 / 256.0)
  in accumArray (+) 0 (0,255) $ 
     map (\(c,f) -> (fromIntegral $ ord c, f)) $ 
     concat [lc,uc,spc,punc,rest]
                                                            

-- (easier to work logarithmically)

logByteFreqs = amap log byteFreqs :: Array Word8 Float


-- scoring function

logLikelihoodRatio text = 
  let language = sum $ map (logByteFreqs !) text
      randomNoise = -(fromIntegral $ length text) * log 256
  in  (1/(log 10)) * (language - randomNoise)
     


p4 = do 
  putStrLn "\n==Problem 4==\n"
  ciphers <- liftM ((map decodeHex) . words) $ 
             readFile "gistfile-p4.txt"

  let candidates = filter ((> 0.0) . logLikelihoodRatio . decryptPlaintext) $
                   map bruteForceXor1 ciphers
                   
  putStrLn "probable decryptions:" 
  putStrLn "log10(likelihood ratio) | plaintext"
  sortBy (compare `on` negate.decryptScore) candidates 
    `each` display
    where display x = putStrLn $ printf "%.3f\t\t%s" 
                      (logLikelihoodRatio $ decryptPlaintext x) 
                      (toAscii $ decryptPlaintext x)


bruteForceXor cipher = 
  head $ sortBy (compare `on` negate.decryptScore) $
  [_knownKeySize cipher sz | sz <- [2..40]]
  
  
_knownKeySize cipher sz = 
  let slices = slice cipher sz    -- one slice per key byte; 
                                   -- contains the ciphertext bytes that key byte encrypts
      sub = map bruteForceXor1 slices
      key = concat $ map decryptKey sub
      plaintext = unslice $ map decryptPlaintext sub
      score = sum $ map decryptScore sub
  in Decryption key plaintext score


slice list stride = 
  map (map (list !!)) $
  [[j, j+stride .. length list - 1] | j <- [0..stride-1]]
        

unslice = concat . transpose



p5 = do              
  putStrLn "\n==Problem 5==\n"  
  let plainText = fromAscii "Burning 'em, if you ain't quick and nimble\nI go crazy when I hear a cymbal"
  let key = fromAscii "ICE"
  let cipherText = xorCrypt plainText key
  putStrLn $ encodeHex cipherText
  
  
p6 = do
  putStrLn "\n==Problem 6==\n"    
  cipherText <- liftM (decodeB64 . concat . words) $ 
                readFile "gistfile-p6.txt"
  let decrypt = bruteForceXor cipherText
  ["key: ",
   encodeHex $ decryptKey decrypt,
   toAscii $ decryptKey decrypt,
   "\nplaintext: ",
   toAscii $ decryptPlaintext decrypt] `each` putStrLn
                             
  
p7 = do
  putStrLn "\n==Problem 7==\n"
    -- FFI wrapping to OpenSSL
  cipher <- SymmetricCiphers.findCipher "aes-128-ecb"
  let key = Char8.pack "YELLOW SUBMARINE"
  let iv = Char8.pack ['\0' | _ <- [1..16]]
  cipherText <- liftM (BS.pack . decodeB64 . concat . words) $ 
                readFile "gistfile-p7.txt"
  plainText <- SymmetricCiphers.decrypt cipher key iv cipherText
  putStrLn $ Char8.unpack plainText


p8 = do
  putStrLn "\n==Problem 8==\n"
  cipherTexts <- liftM (map decodeHex . words) $ 
                 readFile "gistfile-p8.txt"  
                 
  -- secure modes won't have repeating blocks (expected rate 2^-128)
  
  let maybeECB = filter (hasDuplicates . blocks 16) cipherTexts
  maybeECB `each` putStrLn . encodeHex
  

hasDuplicates x = (x /= nub x)

blocks sz text
  | null text = []
  | otherwise = let (b,bs) = splitAt sz text
                in b : blocks sz bs