gistfile1.hs

{-# LANGUAGE OverloadedStrings #-}
module Main (main) where

import Prelude hiding ((!!), length)

import Data.Text (Text ())
import qualified Data.Text.IO as TextIO

import Data.Text.Encoding (encodeUtf8)
import Data.ByteString (ByteString (), index, length, pack, unpack)

import Control.Monad (forM_, liftM2, (=<<))

import Data.Word (Word8 ())
import Data.Bits (xor)

import Data.Array.ST (STUArray ())
import Data.Array.MArray (newListArray, readArray, writeArray)

import Control.Monad.ST (ST (), runST)
import Control.Monad.Trans (lift)
import Control.Monad.Trans.State.Strict (StateT (), get, evalStateT, withStateT)

import Data.STRef (STRef (), newSTRef, readSTRef, modifySTRef)

import Text.Printf (printf)

--------------------------------------------------------------------------------
-- MISCELLANEOUS SUPPORT
--------------------------------------------------------------------------------

(!!) :: ByteString -> Word8 -> Word8
bs !! i = index bs $ mod (fromEnum i) (length bs)

--------------------------------------------------------------------------------
-- KEY STATE
--------------------------------------------------------------------------------

type KeyState s = STUArray s Word8 Word8

initialPermutation :: ST s (KeyState s)
initialPermutation = newListArray (0x00, 0xff) [0x00 .. 0xff]

--------------------------------------------------------------------------------
-- THE RC4 STATE MONAD
--------------------------------------------------------------------------------

type RC4 s a = StateT (KeyState s) (ST s) a

--
-- State helpers
--

increment :: Num a => STRef s a -> a -> RC4 s ()
increment ref val = lift $ modifySTRef ref (+ val)

(+=) :: Num a => STRef s a -> a -> RC4 s ()
(+=) = increment

getIndex :: Word8 -> RC4 s Word8
getIndex i = get >>= \ary -> lift $ readArray ary i

putIndex :: Word8 -> Word8 -> RC4 s ()
putIndex idx val = do
    ary <- get
    lift $ writeArray ary idx val

swap :: Word8 -> Word8 -> RC4 s ()
swap i j = do
    temp <- getIndex i
    putIndex i =<< getIndex j
    putIndex j temp

newRC4Ref :: a -> RC4 s (STRef s a)
newRC4Ref = lift . newSTRef

readRC4Ref :: STRef s a -> RC4 s a
readRC4Ref = lift . readSTRef

--
-- RC4 Key Scheduling Algorithm (KSA)
--

scheduleKeys :: ByteString -> RC4 s ()
scheduleKeys key = do
    j <- newRC4Ref 0

    forM_ [minBound .. maxBound] $ \i -> do
        increment j =<< fmap (+ (key !! i)) (getIndex i)
        swap i =<< readRC4Ref j

--
-- RC4 Output Algorithm (PRGA)
--

generateOutput :: ByteString -> RC4 s ByteString
generateOutput plaintext = do
    i <- newRC4Ref 0
    j <- newRC4Ref 0

    output <- newRC4Ref []

    let push ch = lift $ modifySTRef output (ch :)

    forM_ (unpack plaintext) $ \ch -> do
        i += 1
        increment j =<< getIndex =<< readRC4Ref i

        do i <- readRC4Ref i
           j <- readRC4Ref j
           swap i j

        k <- getIndex =<< liftM2 (+) (getIndex =<< readRC4Ref i)
                                     (getIndex =<< readRC4Ref j)
        push (xor k ch)

    fmap (pack . reverse) (readRC4Ref output)

--
-- RC4 Monad
--

runRC4 :: ByteString -> ByteString -> ByteString
runRC4 key plaintext = runST rc4
    where
        rc4   = evalStateT algo undefined
        algo  = do
            state <- lift initialPermutation
            withStateT (const state) (scheduleKeys key >> generateOutput plaintext)


--------------------------------------------------------------------------------
-- PRETTY-PRINTING CIPHERTEXT
--------------------------------------------------------------------------------

prettyCipher :: ByteString -> String
prettyCipher = concatMap (printf "%02X") . unpack
 
--------------------------------------------------------------------------------
-- TEST VECTORS
--------------------------------------------------------------------------------

data TestVector = TestVector { key        :: Text
                             , plaintext  :: Text
                             , ciphertext :: String
                             }

testVectors = [ TestVector "Key"    "Plaintext"       "BBF316E8D940AF0AD3"
              , TestVector "Wiki"   "pedia"           "1021BF0420"
              , TestVector "Secret" "Attack at dawn"  "45A01F645FC35B383552544B9BF5"
              ]

main :: IO ()
main = do
    forM_ testVectors $ \(TestVector key plaintext ciphertext) -> do
        let cipher = runRC4 (encodeUtf8 key) (encodeUtf8 plaintext)

        mapM_ TextIO.putStr ["(", key, ", ", plaintext, ")\n"]
        putStr "Expected: "
        putStr ciphertext
        putChar '\n'

        putStr "  Result: "
        putStr (prettyCipher cipher)
        putChar '\n'

        putChar '\n'