-- RC4.hs (C) 2002 HardCore SoftWare, Doug Hoyte
--
-- This program is distributed under the terms of the GNU GPL.
-- See www.gnu.org for more information. 
--
-- Haskell implementation of the (Alleged) RC4 stream cipher.
--
-- The function streamOfRC4 takes a String as an argument (the key)
-- and returns an infinite list of the corresponding RC4 stream.

module RC4 (streamOfRC4) where
import Array


clip  :: Int -> Int
clip x =  x `mod` 256


swap        :: Array Int Int -> Int -> Int -> Array Int Int
swap arr x y = if x == y then arr
                 else (arr // [(x, arr ! y), (y, arr ! x)])


initialArray :: Array Int Int
initialArray  = array (0,255) [(i, i) | i <- [0..255]]


setKey :: String -> Int -> Int -> Int -> Array Int Int -> Array Int Int
setKey key ki si 256 arr = arr
setKey key ki si i arr   = setKey key
                                  ((ki+1) `mod` (length key))
                                  newsi
                                  (i+1)
                                  (swap arr newsi i)
                         where
                           newsi = clip (si + fromEnum (key!!ki) + arr!i)


getRC4Stream        :: Array Int Int -> Int -> Int -> [Char]
getRC4Stream arr x y = toEnum (newarr ! clip (sx + arr ! nexty))
                         : getRC4Stream newarr nextx nexty
                     where
                       newarr = swap arr nextx nexty
                       nextx  = clip (x + 1)
                       nexty  = clip (sx + y)
                       sx     = arr ! nextx


streamOfRC4    :: String -> [Char]
streamOfRC4 key = getRC4Stream (setKey key 0 0 0 initialArray) 0 0