module Crypto.Classical.Cipher.Enigma where
import Control.Applicative
import Control.Monad.Trans.State.Lazy
import Crypto.Classical.Types
import Crypto.Classical.Util
import qualified Data.ByteString.Lazy.Char8 as B
import Data.Char
import Data.Map.Lazy (Map)
import qualified Data.Map.Lazy as M
import Data.Maybe (fromJust)
import Data.Modular
import Lens.Micro
import Lens.Micro.TH
newtype Enigma a = Enigma { _enigma :: a } deriving (Eq,Show,Functor)
makeLenses ''Enigma
instance Applicative Enigma where
pure = Enigma
Enigma f <*> Enigma a = Enigma $ f a
instance Monad Enigma where
return = pure
Enigma a >>= f = f a
instance Cipher EnigmaKey Enigma where
decrypt = encrypt
encrypt k m = pure . B.pack $ evalState (traverse f $ B.unpack m) k'
where k' = withInitPositions k
f c | not $ isLetter c = return c
| isLower c = f $ toUpper c
| otherwise = do
modify (& rotors %~ turn)
(EnigmaKey rots _ rl pl) <- get
let rs = rots ^.. traverse . circuit
rs' = reverse $ map mapInverse rs
pl' = mapInverse pl
cmp = foldl1 compose
e = pl |.| cmp rs |.| rl |.| cmp rs' |.| pl'
return . letter . fromJust . flip M.lookup e $ int c
withInitPositions :: EnigmaKey -> EnigmaKey
withInitPositions k = k & rotors .~ zipWith f (k ^. rotors) (k ^. settings)
where f r s = (r & circuit %~ rotate (int s)
& turnover %~ (\n -> n int s))
turn :: [Rotor] -> [Rotor]
turn [] = []
turn (r:rs) = if (r' ^. turnover) == 25 then r' : turn rs else r' : rs
where r' = r & circuit %~ rotate 1 & turnover %~ (\n -> n 1)
rotate :: ℤ/26 -> Map (ℤ/26) (ℤ/26) -> Map (ℤ/26) (ℤ/26)
rotate n r = M.fromList (M.toList r & traverse . both %~ (\n' -> n' n))