The Matasano Crypto Challenges, Set 1 (http://cryptopals.com/sets/1)

Set1.hs

module Set1 where

import qualified Data.Map.Strict as M
import Data.Bits
import Data.Word
import Data.Char (toLower)
import Data.List (transpose, sort, group)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Char8 as C
import Crypto.Cipher.AES

-- Helpers
stringToBytes :: String -> [Word8]
stringToBytes = BS.unpack . C.pack

charToBytes :: Char -> [Word8]
charToBytes = BS.unpack . C.singleton

bytesToString :: [Word8] -> String
bytesToString = C.unpack . BS.pack

allChars = bytesToString [0..255]


-- Challenge 1: Base64 & Hex encoding ------------------------------------------
base64Chars = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"
base64Dict = M.fromList $ zip [0..] base64Chars
getBase64Char b = M.findWithDefault ' ' b base64Dict

base64DictRev = M.fromList $ zip base64Chars [0..]
base64ToByte c = M.findWithDefault 0 c base64DictRev

-- 3 bytes -> 4 base64 chars
encodeTriple :: Word8 -> Word8 -> Word8 -> String
encodeTriple a b c =
  map getBase64Char $ zipWith (.|.)
    (map transform [(a, -2, 0x3f), (a,  4, 0x30), (b,  2, 0x3c), (c, 0, 0x3f)])
    (map transform [(0,  0,    0), (b, -4, 0x0f), (c, -6, 0x03), (0, 0,    0)])
  where
    transform (a, s, m) = a `shift` s .&. m

-- 4 base64 chars -> 3 bytes
decodeQuad :: String -> [Word8]
decodeQuad (a:b:c:d:xs) =
  zipWith (.|.)
    (map transform [(a,  2, 0xfc), (b,  4, 0xf0), (c, 6, 0xc0)])
    (map transform [(b, -4, 0x03), (c, -2, 0x0f), (d, 0, 0x3f)])
  where
    transform (a, s, m) = base64ToByte a `shift` s .&. m

bytesToBase64 :: [Word8] -> String
bytesToBase64 [] = []
bytesToBase64 [a] = take 2 (encodeTriple a 0 0) ++ "=="
bytesToBase64 [a, b] = take 3 (encodeTriple a b 0) ++ "="
bytesToBase64 (a:b:c:xs) = encodeTriple a b c ++ bytesToBase64 xs

stringToBase64 :: String -> String
stringToBase64 = bytesToBase64 . stringToBytes

base64ToBytes :: String -> [Word8]
base64ToBytes [] = []
base64ToBytes s@(a:b:"==") = take 1 (decodeQuad s)
base64ToBytes s@(a:b:c:"=") = take 2 (decodeQuad s)
base64ToBytes s@(a:b:c:d:xs) = decodeQuad s ++ base64ToBytes xs

hexChars = "0123456789abcdef"
hexDict = M.fromList $ zip hexChars [0..]
hexToByte c = M.findWithDefault 0 (toLower c) hexDict

hexRevDict = M.fromList $ zip [0..] hexChars
byteToHex b = M.findWithDefault ' ' b hexRevDict

hexStringToBytes :: String -> [Word8]
hexStringToBytes [] = []
hexStringToBytes [a] = [0]
hexStringToBytes (a:b:xs) = 16 * hexToByte a + hexToByte b : hexStringToBytes xs

bytesToHexString :: [Word8] -> String
bytesToHexString [] = []
bytesToHexString (x:xs) = byteToHex (x `shiftR` 4 .&. 0xf) : byteToHex (x .&. 0xf) : bytesToHexString xs

hexStringToBase64 :: String -> String
hexStringToBase64 = bytesToBase64 . hexStringToBytes


-- Challenge 2: fixed XOR ------------------------------------------------------
fixedXor :: [Word8] -> [Word8] -> [Word8]
fixedXor = zipWith xor

fixedXorHexStr :: String -> String -> String
fixedXorHexStr a b = bytesToHexString $ fixedXor (hexStringToBytes a) (hexStringToBytes b)


-- Challenge 3: Decrypt single-byte XOR with frequency analysis ----------------

-- Decrypt a hex-encoded message with a single-byte XOR
singleXorHexStr :: String -> Char -> String
singleXorHexStr s c = bytesToString $ fixedXor (hexStringToBytes s) $ cycle (charToBytes c)

-- English letter frequencies (https://en.wikipedia.org/wiki/Letter_frequency)
charFreqEnglish = M.fromList
  [ ('a', 0.08167)
  , ('b', 0.01492)
  , ('c', 0.02782)
  , ('d', 0.04253)
  , ('e', 0.1270)
  , ('f', 0.02228)
  , ('g', 0.02015)
  , ('h', 0.06094)
  , ('i', 0.06966)
  , ('j', 0.00153)
  , ('k', 0.00772)
  , ('l', 0.04025)
  , ('m', 0.02406)
  , ('n', 0.06749)
  , ('o', 0.07507)
  , ('p', 0.01929)
  , ('q', 0.00095)
  , ('r', 0.05987)
  , ('s', 0.06327)
  , ('t', 0.09056)
  , ('u', 0.02758)
  , ('v', 0.00978)
  , ('w', 0.02361)
  , ('x', 0.00150)
  , ('y', 0.01974)
  , ('z', 0.00074) ]

freqExpected c = M.findWithDefault 0 c charFreqEnglish

count c = length . filter (==c)

-- frequency of a Char in a String
freq :: String -> Char -> Double
freq s c = fromIntegral (count c s) / fromIntegral (length s)

-- score for a Char = square deviation from expected frequency
getScore :: String -> Char -> Double
getScore s c = (freq s c - freqExpected c) ^^ 2

-- score an entire string (sum scores for 'a'..'z' and unexpected characters '*')
scoreString :: String -> Double
scoreString s = sum $ map (getScore s') ('*':['a'..'z'])
  where
    s' = map (replaceNonAlpha . toLower) s
    replaceNonAlpha c = if c `elem` ['a'..'z'] then c else '*'

-- Decrypt with every letter, score, and return the result with the lowest score
tryDecryptHex :: String -> (Double, String)
tryDecryptHex s = minimum $ map ((\s -> (scoreString s, s)) . singleXorHexStr s) allChars

tryDecryptBytes bytes =
  minimum $ map (\c ->
    let bytes' = fixedXor bytes $ cycle (charToBytes c)
        s' = bytesToString bytes'
    in (scoreString s', c)
    ) allChars

-- Challenge 4: Find the one line in a file that has been single-byte XOR encrypted
challenge4 :: IO String
challenge4 = do
  content <- readFile "s1c4.txt"
  let l = map tryDecryptHex $ lines content
  return $ (snd . minimum) l


-- Challenge 5: Encrypt Repeating-key XOR --------------------------------------
repeatXor :: String -> String -> String
repeatXor s key = bytesToHexString $ fixedXor (stringToBytes s) $ cycle (stringToBytes key)

-- Challenge 6: Break repeating-key XOR ----------------------------------------

-- Hamming distance: Number of differing bits between two Strings of identical length
hamming :: [Word8] -> [Word8] -> Int
hamming a b = sum $ zipWith (\a b -> popCount (a `xor` b)) a b

-- Normalized Hamming distance
hamNormal :: [Word8] -> [Word8] -> Double
hamNormal a b = fromIntegral (hamming a b) / fromIntegral (length a)

-- Split bytes into blocks of "keySize" length
blockSplit :: Int -> [Word8] -> [[Word8]]
blockSplit keySize [] = []
blockSplit keySize bytes = x : blockSplit keySize xs
  where (x, xs) = splitAt keySize bytes

-- Average normalized Hamming distance of the first 4 blocks of "keySize" length
blockDist :: [Word8] -> Int -> Double
blockDist bytes keySize = sum dists / fromIntegral (length dists)
  where
    b' = blockSplit keySize bytes
    blocks = take 1 b'
    blocks2 = take 3 (drop 1 b')
    dists = hamNormal <$> blocks <*> blocks2

-- Try key sizes 2-40 and return the first couple that produce the
-- smallest Hamming distance between blocks
findKeySizes :: [Word8] -> [Int]
findKeySizes bytes =
  take 4 $ map snd $ sort $ map (\n -> (blockDist bytes n, n)) [2..40]

tryDecryptWithKeysize :: [Word8] -> Int -> String
tryDecryptWithKeysize bytes keySize =
  -- Split bytes into blocks of keySize length, transpose them
  let xBlocks = transpose $ blockSplit keySize bytes
    -- and find the key byte for each block by solving single-byte XOR cypher
      key = stringToBytes $ map (snd . tryDecryptBytes) xBlocks
  -- decrypt bytes with the final key
  in bytesToString $ fixedXor bytes $ cycle key

readBase64File :: String -> IO [Word8]
readBase64File name = do
  content <- readFile name
  return $ base64ToBytes (concat $ lines content) -- remove the linebreaks!

challenge6 = do
  bytes <- readBase64File "s1c6.txt"
  let keySizes = findKeySizes bytes
      -- decrypt with the most promising keySizes
      plains = map (tryDecryptWithKeysize bytes) keySizes
      -- and pick the one with the best frequency score
      bestPlain = snd . minimum $ map (\s -> (scoreString s, s)) plains
  -- 0 seems to be ' ' in the plaintext, so replace:
  return $ map (\c -> if c == '\0' then ' ' else c) bestPlain

-- Challenge 7: Decode AES-128 in ECB mode -------------------------------------
challenge7 = do
  bytes <- readBase64File "s1c7.txt"
  let aes = initAES $ C.pack "YELLOW SUBMARINE"
      plain = decryptECB aes (BS.pack bytes)
  return plain

-- Challenge 8: Detect a line encoded with AES-128 ECB -------------------------

-- Split buffer into strings of length 16,
-- any repeating String suggests ECB mode cipher
isECB :: [Word8] -> Bool
isECB bytes =
  let blocks = map bytesToString $ blockSplit 16 bytes
  in any (> 1) $ map (`count` blocks) blocks

-- read file and filter the lines which were probably encrypted in ECB mode
challenge8 = do
  content <- readFile "s1c8.txt"
  return $ filter (isECB . hexStringToBytes) $ lines content