vignere.hs

#!/usr/bin/env stack
-- stack --resolver lts-16.16 script --package unordered-containers --ghc-options -Wall --ghc-options -fbreak-on-exception
{-# LANGUAGE MultiWayIf, LambdaCase #-}

module Main where

import Data.Char
import qualified Data.HashMap.Strict as HM
import Data.List (foldl')

alphabets :: [Char]
alphabets = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"

data ShiftDir
  = ShiftFwd
  | ShiftBwd

shift :: Char -> Char -> ShiftDir -> Char
shift a k shiftDir =
  let aInt = (fromEnum a) - (fromEnum 'A')
      kInt = (fromEnum k) - (fromEnum 'A')
      xInt =
        case shiftDir of
          ShiftFwd -> (aInt + kInt) `mod` 26
          ShiftBwd -> (aInt - kInt) `mod` 26
      x'Int = xInt + fromEnum 'A'
      x = toEnum x'Int
   in x

enc :: [Char] -> [Char] -> [Char]
enc [] _ = []
enc (c:cs) (k:ks) =
  let c' = shift c k ShiftFwd
   in c' : (enc cs ks)
enc _ [] = error "enc ran out of key"

dec :: [Char] -> [Char] -> [Char]
dec [] _ = []
dec (c:cs) (k:ks) =
  let c' = shift c k ShiftBwd
   in c' : (dec cs ks)
dec _ [] = error "dec ran out of key"

sanitize :: [Char] -> [Char]
sanitize cs = filter isAlpha $ map toUpper cs

chunksOf :: Int -> [a] -> [[a]]
chunksOf n [] = take n $ repeat []
chunksOf n cs =
  let (p1, p2) = splitAt n cs
   in p1 : (chunksOf n p2)

newtype FreqMap =
  FreqMap (HM.HashMap Char Int)

fmNew :: FreqMap
fmNew = FreqMap $ HM.empty

fmGet :: FreqMap -> Char -> Int
fmGet (FreqMap hm) c = HM.lookupDefault 0 c hm

fmIncr :: Char -> FreqMap -> FreqMap
fmIncr c (FreqMap hm) =
  FreqMap $
  HM.alter
    (\case
       Nothing -> Just $ 1
       Just x -> Just $ x + 1)
    c
    hm

fmEnglish :: [(Char, Double)]
fmEnglish =
  [ ('A', 0.082)
  , ('B', 0.015)
  , ('C', 0.028)
  , ('D', 0.043)
  , ('E', 0.13)
  , ('F', 0.022)
  , ('G', 0.02)
  , ('H', 0.061)
  , ('I', 0.07)
  , ('J', 0.0015)
  , ('K', 0.0077)
  , ('L', 0.04)
  , ('M', 0.024)
  , ('N', 0.067)
  , ('O', 0.075)
  , ('P', 0.019)
  , ('Q', 0.00095)
  , ('R', 0.06)
  , ('S', 0.063)
  , ('T', 0.091)
  , ('U', 0.028)
  , ('V', 0.0098)
  , ('W', 0.024)
  , ('X', 0.0015)
  , ('Y', 0.02)
  , ('Z', 0.00074)
  ]

buildFreqMap :: [Char] -> [FreqMap] -> [FreqMap]
buildFreqMap [] hms = hms
buildFreqMap _ [] = error "Not enough FreqMaps"
buildFreqMap (c:cs) (h:hms) = (fmIncr c h) : (buildFreqMap cs hms)

buildChunkWiseFreqMap :: Int -> [Char] -> [FreqMap]
buildChunkWiseFreqMap n str =
  let emptyFreqMaps = take n $ repeat fmNew
      chunks = chunksOf n str
   in foldl' (flip buildFreqMap) emptyFreqMaps chunks

fmPrint :: FreqMap -> IO ()
fmPrint fm =
  sequence_ $
  map
    (\c -> do
       putStr [c]
       putStr ": "
       putStrLn $ show $ fmGet fm c)
    alphabets

fmList :: FreqMap -> [(Char, Int)]
fmList fm = map (\c -> (c, fmGet fm c)) alphabets

printFms :: [FreqMap] -> Int -> IO ()
printFms [] _ = return ()
printFms (fm:fms) i = do
  putStr "Chunk "
  putStr $ show i
  putStrLn ":"
  fmPrint fm
  putStrLn ""
  printFms fms (i + 1)

-- rotateBwd :: Int -> [a] -> [a]
-- rotateBwd n lst =
--   let (p1, p2) = splitAt n lst
--    in p2 <> p1
rotateFwd :: Int -> [a] -> [a]
rotateFwd n lst =
  let (p1, p2) = splitAt ((length lst) - n) lst
   in p2 <> p1

maxOn :: Ord a => (b -> a) -> [b] -> b
maxOn _ [] = error "Empty list"
maxOn _ [a] = a
maxOn f (a:as) =
  let v1 = a
      v2 = maxOn f as
   in if f v1 > f v2
        then v1
        else v2

identifyShiftCount :: FreqMap -> Int
identifyShiftCount fm =
  let fmEnglishEnum = zip [0,1 .. 25] (repeat fmEnglish)
      shiftedFms = map (\(n, lst) -> (n, rotateFwd n lst)) fmEnglishEnum
      shiftedFmZipped = zip shiftedFms (repeat (fmList fm))
      scores =
        map
          (\((n, lst1), lst2) ->
             ( n
             , sum $
               map (\(c1, c2) -> c1 * (fromIntegral c2)) $
               zip (map snd lst1) (map snd lst2)))
          shiftedFmZipped
      maxScore = maxOn snd scores
   in fst maxScore

identifyKeyLength :: String -> Int
identifyKeyLength cipherText = 
  let scores = 
        map (\tau -> 
          let fm = steppedFreqMap cipherText tau fmNew
              probSq = fmProbSq fm
              -- score = abs (probSq - 0.065)
              score = probSq
          in (tau, score)) [1,2..10]
      (maxTau, _) = maxOn snd scores
  in maxTau

steppedFreqMap :: String -> Int -> FreqMap -> FreqMap
steppedFreqMap [] _ fm = fm
steppedFreqMap (s:ss) tau fm =
  let fm' = fmIncr s fm
      (_, p2) = splitAt (tau - 1) ss
   in steppedFreqMap p2 tau fm'

fmSum :: FreqMap -> Int
fmSum (FreqMap hm) = sum $ map snd $ HM.toList hm

fmProbSq :: FreqMap -> Double
fmProbSq fm@(FreqMap hm) =
  let sum_ = fromIntegral $ fmSum fm
   in sum $
      map
        (\(_, x) ->
           let xf = fromIntegral x
            in (xf / sum_) ** 2) $
      HM.toList hm

attack :: String -> Int -> String
attack cipherText keyLen =
  let fmChunks = buildChunkWiseFreqMap keyLen cipherText
   in map
        (\fm ->
           let shiftCount = identifyShiftCount fm
            in toEnum (fromEnum 'A' + shiftCount))
        fmChunks

main :: IO ()
main = do
  let keyBase = "ABCDEFG"
  let keyLen = length keyBase
  let key = cycle keyBase
  let plainText =
        sanitize $
        "The zebra puzzle is a well-known logic puzzle. Many versions of the puzzle. If it is important to always be able to interrupt such threads, you should turn this optimization off. Consider also recompiling all libraries with this optimization turned off, if you need to guarantee interruptibility. A total of 1400 individual tests were created, which is comforting. We can increase the depth easily enough, but to find out exactly how well the code is being tested we should turn to the built in code coverage tool."
  putStrLn "Key: "
  putStrLn keyBase
  putStrLn "Key Length: "
  putStrLn $ show keyLen
  putStrLn ""
  putStrLn "Plain text:"
  putStrLn plainText
  putStrLn ""
  let cipherText = enc plainText key
  putStrLn "Cipher text:"
  putStrLn cipherText
  putStrLn ""
  --
  let keyLen' = identifyKeyLength cipherText
  putStrLn "Estimated Key Length:"
  print keyLen'
  let keyBase' = attack cipherText keyLen'
  let key' = cycle keyBase'
  putStrLn "Estimated Key:"
  putStrLn keyBase'
  putStrLn ""
  let plainText' = dec cipherText key'
  putStrLn "Deciphered text:"
  putStrLn plainText'
  putStrLn ""