-
-
Save ownclo/b347ce3ca381b82be1a2 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
{-# 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 ..] |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment