IterNW.hs

{-# LANGUAGE
    GADTs
  , FlexibleContexts
  , RankNTypes
  , ScopedTypeVariables #-}


import Data.Word
import Data.List( find, sortBy )
import Data.Maybe( fromJust )
import Data.Function( on )
import Control.Applicative

import System.Random
import Data.Array.IO
import Data.Array
import Data.Array.ST
import Control.Monad
import Control.Monad.ST


type Quartet = (Int, Int, Int, Int)
type Quartets = [Quartet]

type SubMap = [(Quartet, Quartet)]
type PMatrix = [Int]  -- just a list of indices, UNSAFISH
type RoundSpec = ([SubMap], PMatrix)


sbox :: SubMap -> Quartet -> Quartet
sbox m x = snd . fromJust $ find ((== x) . fst) m

invertSubMap :: SubMap -> SubMap
invertSubMap = map flip
    where flip (a, b) = (b, a)


listToQuartet :: [Int] -> Quartet
listToQuartet [a, b, c, d] = (a, b, c, d)
listToQuartet q = error $ "Malformed LIST quartet: " ++ show q

quartetToList :: Quartet -> [Int]
quartetToList (a, b, c, d) = [a, b, c, d]

int16ToBits :: Word16 -> [Int]
int16ToBits x = pad ++ z
    where pad = replicate (16 - unPaddedLen) 0
          z = toBase 2 x
          unPaddedLen = length z

bitsToWord16 :: [Int] -> Word16
bitsToWord16 = fromOctets 2

fromQuartets :: Quartets -> Word16
fromQuartets = bitsToWord16 . concatMap quartetToList

toQuartets :: Word16 -> Quartets
toQuartets = map listToQuartet . takeFours . int16ToBits

takeFours :: [a] -> [[a]]
takeFours [] = []
takeFours lst = fst : takeFours rst
    where (fst, rst) = (take 4 lst, drop 4 lst)


toBase :: (Integral b1, Num b) => b1 -> b1 -> [b]
toBase x = 
   map fromIntegral .
   reverse .
   map (`mod` x) .
   takeWhile (/=0) .
   iterate (`div` x)

powersOf :: (Integral a) => a -> [a]
powersOf n = 1 : map (*n) (powersOf n)

fromOctets :: (Integral a, Integral b, Integral c) => a -> [c] -> b
fromOctets n x = 
   fromIntegral $ 
   sum $ 
   zipWith (*) (powersOf n) (reverse (map fromIntegral x))

propIsoQuartets :: Word16 -> Bool
propIsoQuartets = isomorphic fromQuartets toQuartets

isomorphic :: Eq a => (b -> a) -> (a -> b) -> a -> Bool
isomorphic f g x = (f . g) x == x

testQuartets :: IO ()
testQuartets = print $ all id $ map propIsoQuartets [0 ..]

randomSubMap :: IO SubMap
randomSubMap = zip allQuarters `fmap` shuffle allQuarters

randomPMatrix :: IO [Int]
randomPMatrix = shuffle [0 .. 15]

randomRoundSpec :: IO RoundSpec
randomRoundSpec = (,) <$> replicateM 4 randomSubMap <*> randomPMatrix

allQuarters :: Quartets
allQuarters = map (last . toQuartets) $ take 16 [0 ..]

shuffle :: [a] -> IO [a]
shuffle xs = do
    ar <- newArray n xs
    forM [1..n] $ \i -> do
        j <- randomRIO (i,n)
        vi <- readArray ar i
        vj <- readArray ar j
        writeArray ar j vi
        return vj
        where
            n = length xs
            newArray :: Int -> [a] -> IO (IOArray Int a)
            newArray n =  newListArray (1,n)

permute :: forall a. PMatrix -> [a] -> [a]
permute is lst = runST $ do
    ar <- fromList lst :: ST s (STArray s Int a)
    let zs = zip is lst
    forM_ zs $ \(i,x) -> do
        writeArray ar i x
    getElems ar
  where fromList lst = newListArray (0, length lst - 1) lst

pbox :: PMatrix -> Quartets -> Quartets
pbox pmatrix = map listToQuartet
             . takeFours
             . permute pmatrix
             . concatMap quartetToList

invertPMatrix :: PMatrix -> PMatrix
invertPMatrix = map fst . sortBySecond . zip [0..]
    where sortBySecond = sortBy (compare `on` snd)

encryptRound :: RoundSpec -> Quartets -> Quartets
encryptRound (smaps, pmat) =
        pbox pmat
      . zipWith sbox smaps

decryptRound :: RoundSpec -> Quartets -> Quartets
decryptRound (smaps, pmat) =
        zipWith sbox (map invertSubMap smaps)
      . pbox (invertPMatrix pmat)

encrypt :: [RoundSpec] -> Word16 -> Word16
encrypt spec x = fromQuartets $
    foldl (|>) x' $ map encryptRound spec
    where x' = toQuartets x

decrypt :: [RoundSpec] -> Word16 -> Word16
decrypt spec x = fromQuartets $
    foldr ($) x' $ map decryptRound spec
    where x' = toQuartets x

onQuartets :: (Quartets -> Quartets) -> Word16 -> Word16
onQuartets f = fromQuartets . f . toQuartets

(|>) :: a -> (a -> b) -> b
(|>) = flip ($)

main :: IO ()
main = do
    roundsSpec <- replicateM 4 randomRoundSpec
    let plain = toQuartets 1
        roundSpec = head roundsSpec
        encrypted = encryptRound roundSpec plain
        decrypted = decryptRound roundSpec encrypted
    print plain
    print encrypted
    print decrypted

    -- NB: this is great for property-based testing
    print $ all id $ map (isomorphic (encrypt roundsSpec) (decrypt roundsSpec)) [0 ..]