enigma in Haskell

enigma.lhs

Emulating an Enigma machine
===

In this gist I'll go over the development of an Enigma machine encoder
in Haskell.

Here are some useful background info about how the Enigma machine works:

- Enigma Simulator: http://enigmaco.de/enigma/enigma.html
- "How Enigma Machines Work" http://enigma.louisedade.co.uk/howitworks.html

Starting Out
===

We start with some definitions for the set of
letters we will be working with:

> module Enigma where
>
> import Data.Char
> import Data.List
> import Control.Applicative ((<|>))
>
> letters = [ 'A' .. 'Z' ]
> letterCount = length letters
> 
> letterOrd :: Char -> Int
> letterOrd ch = ord ch - ord 'A'
> 
> letterChr :: Int -> Char
> letterChr i = chr ( ord 'A' + (mod i letterCount) )
> 
> rotateLetter :: Int -> Char -> Char
> rotateLetter i ch = letterChr (letterOrd ch + i) 

Note how `leterChr` always returns a letter in the range A-Z.
It's nice the have functions which give you guarantees - such
as always returning a valid letter.

Next we define a type for permutations of the letters A-Z.
Even though it is just an alias for a String, giving it a name tells the reader
that it is a special kind of string and aids in understanding
the code.

We also define a Cipher type which is just a pair of permutations.

> -- A Permutation should have exactly 26 letters and be a permutation of ['A'..'Z'].
> type Permutation = String
> 
> permEncode :: Permutation -> Char -> Char
> permEncode perm ch = perm !! (letterOrd ch)
> 
> inversePerm :: Permutation -> Permutation
> inversePerm perm = map snd $ sort [ (b,a) | a <- letters, let b = permEncode perm a ]
>
> data Cipher = Cipher { _perm :: Permutation, _invperm :: Permutation }
>        deriving (Show,Eq)
>
> makeCipher :: Permutation -> Cipher
> makeCipher perm = Cipher perm (inversePerm perm)
>
> -- Create a cipher from pairs
> makeCipherFromPairs :: [(Char,Char)] -> Cipher
> makeCipherFromPairs pairs = Cipher perm (inversePerm perm)
>   where go a = let Just b = lookup a pairs <|> lookup a invPairs <|> (Just a)
>                in b
>         perm = map go letters
>         invPairs = [ (b,a) | (a,b) <- pairs ]

At this point at lot of peroperties may be tested, i.e.:

- does `permEncode ['A'..'Z']` result in the identity function?
- does `permEncode (inversePerm perm) . permEncode perm` result in the identity function?
- does `makeCipherFromPairs` work correctly on specific cases

Now we define the effect of rotating a permutation. This is just a conjugation
 of `permEncode` with `rotateLetter`:

> permShiftEncode :: Permutation -> Int -> Char -> Char
> permShiftEncode perm shift = rotateLetter (negate shift) . permEncode perm . rotateLetter shift

Now we are ready to define the main encoding function.

> -- encode a letter through a chain of conjugations
> encodeLetter :: [(Cipher,Int)] -> Permutation -> Char -> Char
> encodeLetter [] perm = permEncode perm
> encodeLetter ((c,i):cs) perm = permShiftEncode (_invperm c) i
>                                 . encodeLetter cs perm
>                                 . permShiftEncode (_perm c) i

There are two big advantages to writing `permShiftEncode` and `encodeLetter`
this way:

1. It follows the mathematical definition, so we have more confidence that
it works corectly.
2. We've defined it without reference to any other data structures,
so it's testable at this point in our development.

EnigmaConfig
===

We are now ready to define structures related to the Enigma machine itself,
a Wheel type and the EnigmaConfig type.

> data Wheel = Wheel { _cipher :: Cipher, _turnovers :: String, _ringSetting :: Int }
>        deriving (Show,Eq)
>
> data EnigmaConfig = EnigmaConfig { _plug      :: Cipher
>                                  , _reflector :: Permutation
>                                  , _wheels    :: [Wheel]
>                                  , _positions :: [Position]
>                                  }
>        deriving (Show, Eq)

Here is the function to encode a letter using an EnigmaConfig:

> enigmaEncodeLetter :: EnigmaConfig -> Char -> Char
> enigmaEncodeLetter ec = encodeLetter ciphers (_reflector ec)
>   where ciphers = [ (_plug ec, 0) ]
>                     ++ zip (map _cipher $ _wheels ec) [ p-1 | p <- _positions ec]

Note that separating out the plug board and reflector allows
us to ensure that their rotational shifts are always 0.

Also, because Positions are 1-based we have to subtract 1 from them
before passing them to encodeLetter. This might make a good case
for making positions 0-based.

Stepping
===

The way the wheel positions change is the most complex part
of the Enigma. Basically stepping occurs in a "ripple"
fashion - the rightmost wheel is always stepped, and
the next wheel is stepped only if previous wheel
was stepped and was at a turn over position.

However, the logic for the middle wheel is special - in
addition to stepped if the right wheel was at turnover
position it will also be stepped if itself is at a turnover
position. In this case the left wheel will also be stepped.

Since our EnigmaConfig allows for an aribtrary
list of wheels, we have a two options:

- require that `_wheels` has length 3
- in the case that the number of wheels is not 3, apply the unique
  stepping logic for the middle wheel to the penultimate wheel
  in the list

Below are three possible implementations.

> -- Positions are 1-based and have the range 1..letterCount
> type Position = Int
>
> -- increment a position (which is 1-based, not 0-based)
> incPosition :: Position -> Position
> incPosition p = 1 + mod p letterCount
>
> isTurnOver :: Wheel -> Position -> Bool
> isTurnOver wheel pos = elem window (_turnovers wheel)
>   where window = letterChr (pos + _ringSetting wheel - 2)
>         -- note: -2 here because both pos and ring setting are 1 based
>
> -- ripple stepping - not Enigma compliant
> stepPositions0 :: [(Wheel,Position)] -> [Position]
> stepPositions0 [] = []
> stepPositions0 ((w,p):rest) = p' : if isTurnOver w p then stepPositions0 rest
>                                                      else map snd rest
>   where p' = incPosition p
>
> -- assume there are only three wheels
> stepPositions1 ::[ (Wheel,Position) ] -> [Position]
> stepPositions1 [ (rw,rp), (mw, mp), ( lw, lp) ] = [ rp', mp', lp' ]
>   where rp' = incPosition rp
>         mp' = if (isTurnOver mw mp) || (isTurnOver rw rp)
>                then incPosition mp
>                else mp
>         lp' = if isTurnOver mw mp
>                 then incPosition lp
>                 else lp
> stepPositions1 _ = error "need exactly three wheels"
> 
> -- apply special stepping logic to penultimate wheel
> stepPositions2 :: [(Wheel,Position)] -> [Position]
> stepPositions2 pairs = go True pairs
>   where
>     go b [ (w1,p1),(w2,p2) ] =  [ p1', p2']  -- at last two wheels
>       where p1' = if b || isTurnOver w1 p1 
>                     then incPosition p1
>                     else p1
>             p2' = if isTurnOver w1 p1 then incPosition p2 else p2
>     go b ((w,p):rest)    = p' : go (isTurnOver w p) rest
>       where p' = if b then incPosition p else p
>     go _ []              = []                -- only happens if number of wheels == 1

The use of the Position type here is helpful to remind the reader
that the value is 1-based instead of 0-based.

Encoding a Message
===

> enigmaStep :: EnigmaConfig -> EnigmaConfig
> enigmaStep ec = ec { _positions = stepPositions1 $ zip (_wheels ec) (_positions ec) }
>
> -- encode a message
> enigmaMessage :: EnigmaConfig -> String -> (EnigmaConfig, String)
> enigmaMessage ec [] = (ec, [])
> enigmaMessage ec (a:as) = let b          = enigmaEncodeLetter ec' a
>                               ec'        = enigmaStep ec
>                               (ec'', bs) = enigmaMessage ec' as
>                           in (ec'', (b:bs))

The reason I decided on writing enigmaMessage with this signature is because 
then it may be defined with `mapAccumL`:

> enigmaMessage' :: EnigmaConfig -> String -> (EnigmaConfig, String)
> enigmaMessage' = mapAccumL go 
>   where go ec a = let ec' = enigmaStep ec
>                   in (ec', enigmaEncodeLetter ec' a)

Putting it all together
===

First a utility function to create a Wheel:

> makeWheel :: Position -> Permutation -> String -> Wheel
> makeWheel ringSetting perm turnLetters = Wheel (makeCipher perm) turnLetters ringSetting
>
> makeWheel0 :: Permutation -> String -> Wheel
> makeWheel0 = makeWheel 1
>
> compI    = makeWheel0 "EKMFLGDQVZNTOWYHXUSPAIBRCJ" "Q"
> compII   = makeWheel0 "AJDKSIRUXBLHWTMCQGZNPYFVOE" "E"
> compIII  = makeWheel0 "BDFHJLCPRTXVZNYEIWGAKMUSQO" "V"
> compIV   = makeWheel0 "ESOVPZJAYQUIRHXLNFTGKDCMWB" "J"
> compV    = makeWheel0 "VZBRGITYUPSDNHLXAWMJQOFECK" "Z"
> compVI   = makeWheel0 "JPGVOUMFYQBENHZRDKASXLICTW" "ZM"
> compVII  = makeWheel0 "NZJHGRCXMYSWBOUFAIVLPEKQDT" "ZM"
> compVIII = makeWheel0 "FKQHTLXOCBJSPDZRAMEWNIUYGV" "ZM"
> 
> reflectorBeta  = "LEYJVCNIXWPBQMDRTAKZGFUHOS" 
> reflectorGamma = "FSOKANUERHMBTIYCWLQPZXVGJD" 
> reflectorA     = "EJMZALYXVBWFCRQUONTSPIKHGD"
> reflectorB     = "YRUHQSLDPXNGOKMIEBFZCWVJAT"
> reflectorC     = "FVPJIAOYEDRZXWGCTKUQSBNMHL"
> reflectorb     = "ENKQAUYWJICOPBLMDXZVFTHRGS"
> reflectorc     = "RDOBJNTKVEHMLFCWZAXGYIPSUQ"
> reflectorNone  = "ABCDEFGHIJKLMNOPQRSTUVWXYZ" 
> 
> plug   = makeCipherFromPairs [ ('A','E'), ('B','F'), ('C','M'), ('D','Q'), ('H','U'),
>                                ('J','N'), ('L','X'), ('P','R'), ('S','Z'), ('V','W') ]
> 
> cfg = EnigmaConfig { _plug = plug
>                    , _reflector = reflectorBeta
>                    , _wheels =    [ compVIII { _ringSetting = 12 }
>                                   , compVI   { _ringSetting = 5  }
>                                   , compV    { _ringSetting = 15 }
>                                   ]
>                    , _positions = [ 25, 16, 15 ]
>                    }

Thank you, @erantapaa.
I know this gist is very old but if someone will need it. This gist has couple of mistakes:
1.

permShiftEncode :: Permutation -> Int -> Int -> Char -> Char
permShiftEncode perm shift start = rotateLetter (letterCount + start - shift) . permEncode perm . rotateLetter (letterCount + shift - start)

encodeLetter :: [(Cipher,Int,Int)] -> Permutation -> Char -> Char
encodeLetter [] perm = permEncode perm
encodeLetter ((c,i,j):cs) perm = permShiftEncode (_invperm c) i j
                               . encodeLetter cs perm
                               . permShiftEncode (_perm c) i j

enigmaEncodeLetter :: EnigmaConfig -> Char -> Char
enigmaEncodeLetter ec = encodeLetter ciphers (_reflector ec)
  where ciphers = (_plug ec, 0, 0)
                : [ (_cipher (_wheels ec !! i), (_positions ec !! i) - 1, _ringSetting (_wheels ec !! i) - 1)
                  | i <- [0..((length . _positions) ec - 1)]
                  ]

isTurnOver :: Wheel -> Position -> Bool
isTurnOver wheel pos = window `elem` _turnovers wheel
  where window = letterChr (pos - 1)