module Acme.Cipher (Cipher, encode, decode, playfair, substitution, addLetters, subtractLetters, viginere, caesar, Plugboard(..), Reflector(..), RotorMapping(..), enigma) where
import Data.List
import Data.Maybe
import Data.Char
import System.Random
import Control.Monad.ST
import Data.STRef
import Control.Monad
import System.Random.Shuffle
data Cipher key = Cipher (key -> String -> String) (key -> String -> String)
encode ~(Cipher f _) = f
decode ~(Cipher _ f) = f
playfair = Cipher start start where
start k = fn (take 25 (nub (map (\ch -> if ch == 'J' then 'I' else ch) k)) ++ (['A'..'Z'] \\ ('J' : k)))
fn k (c1 : c2 : s) = (k !! (5 * n + y)) : (k !! (5 * x + m)) : fn k s
where
(n, m) = fromJust (findIndex (==c1) k) `divMod` 5
(x, y) = fromJust (findIndex (==c2) k) `divMod` 5
fn _ s = s
swap (x, y) = (y, x)
substitution = Cipher (\k s -> map (\ch -> fromJust $ lookup ch k) s) (\k s -> map (\ch -> fromJust $ lookup ch (map swap k)) s)
letter x = ord x ord 'A' + 1
letter2 x = chr (x + ord 'A' 1)
addLetters x y = if letter2 (letter x + letter y) > 'Z' then letter2 $ letter x + letter y 26 else letter2 $ letter x + letter y
subtractLetters x y = if letter2 (letter x letter y) < 'A' then letter2 $ letter x letter y + 26 else letter2 $ letter x letter y
viginere = Cipher (\k s -> zipWith addLetters s (cycle k)) (\k s -> zipWith subtractLetters s (cycle k))
caesar = Cipher (\k s -> map (`addLetters` k) s) (\k s -> map (`subtractLetters` k) s)
data Rotor s = Rotor { mapping :: [(Char, Char)], position :: STRef s Char, nextRotor :: Maybe (Rotor s) }
advance rotor = do
pos <- readSTRef (position rotor)
let plus1 = addLetters pos 'A'
writeSTRef (position rotor) plus1
when (plus1 == 'A') $ maybe (return ()) advance (nextRotor rotor)
newRotor (RotorMapping mapping) init next = do
pos <- newSTRef init
return (Rotor mapping pos next)
newRotors (map1,map2,map3) (init1,init2,init3) = do
rot1 <- newRotor map1 init1 Nothing
rot2 <- newRotor map2 init2 (Just rot1)
rot3 <- newRotor map3 init3 (Just rot2)
return (rot1, rot2, rot3)
look ls x = fromJust (lookup x ls)
rotorFunction rotor = do
pos <- readSTRef (position rotor)
return (look (mapping rotor) . (`subtractLetters` pos))
reverseRotorFunction rotor = do
pos <- readSTRef (position rotor)
return ((`addLetters` pos) . look (map swap (mapping rotor)))
randomLetter stdgen = randomR ('A','Z') stdgen
newtype Plugboard = Plugboard [(Char, Char)]
newtype Reflector = Reflector [(Char, Char)]
newtype RotorMapping = RotorMapping [(Char, Char)]
encodeEnigma (Plugboard plugboard, Reflector reflector0, mappings) (init1,init2,init3) text = do
let reflector = reflector0 ++ map swap reflector0
(rotor1,rotor2,rotor3) <- newRotors mappings (init1,init2,init3)
mapM (\letter -> do
f1 <- rotorFunction rotor1
f2 <- rotorFunction rotor2
f3 <- rotorFunction rotor3
f4 <- reverseRotorFunction rotor1
f5 <- reverseRotorFunction rotor2
f6 <- reverseRotorFunction rotor3
advance rotor3
return $ look (map swap plugboard) $ f6 $ f5 $ f4 $ look reflector $ f1 $ f2 $ f3 $ look plugboard letter)
text
enigma stdgen = Cipher (\key plaintext -> runST $ do
(init1, stdgen) <- return (randomLetter stdgen)
(init2, stdgen) <- return (randomLetter stdgen)
(init3, _) <- return (randomLetter stdgen)
cipher <- encodeEnigma key (init1,init2,init3) plaintext
return (init1 : init2 : init3 : cipher))
(\key (init1 : init2 : init3 : cipher) -> runST $ encodeEnigma key (init1,init2,init3) cipher)
rot13 = map (\x -> (x, addLetters x 'M')) ['A'..'N']
test = do
ls <- shuffleM ['A'..'Z']
let k = zip ['A'..] ls
let key = (Plugboard k, Reflector rot13, (RotorMapping k, RotorMapping k, RotorMapping k))
stdgen <- getStdGen
putStrLn $ decode (enigma stdgen) key (encode (enigma stdgen) key "TESTINGTESTING")