[Haskell-beginners] State help request

[Joe]

Hi,

below is an implementation of RC4, which I started after Andrew
Gwozdziewycz from NY Hack & Tell encouraged the group to write it. I
mostly followed the pseudocode on Wikipedia.

I'd appreciate any feedback, but I'm most distressed by the KSA
generation, and would like suggestions for making either version less
nasty.

Thanks.


import Data.Bits
import Data.Vector ((!),(//),Vector,fromList)
import Data.Char
import Control.Monad.State

type PRGA = (Vector Int,Int,Int)
type Key = [Int]

-- identity permutation
permId :: Vector Int
permId = fromList [0..255]

-- generate initial PRGA
ksa :: Key -> PRGA
ksa key = ksaStep permId key 0 0

-- I really don't like passing the counter every time,
ksaStep :: Vector Int -> Key -> Int -> Int -> PRGA
ksaStep s _ 255 _ = (s,0,0)
ksaStep s key i j = let j' = (j + (s!i) + (key !! (i `mod`
keylength))) `mod` 256 in
  ksaStep (s // [(i, s!j'), (j',s!i)]) key (i+1) j'
  where keylength = length key

-- but I tried wedging it into a State, and it's not any clearer
ksa' :: Key -> PRGA
ksa' key = genPRGA
  where genPRGA = snd $ foldl (\s a -> snd $ runState (ksaStep' a) s)
(key,(permId,0,0)) [0..255]

ksaStep' :: Int -> State (Key,PRGA) ()
ksaStep' i = do
  (key, (s,_,j)) <- get
  let j' = (j + (s!i) + (key !! (i `mod` length(key)))) `mod` 256
      s' = s // [(i, s!j'), (j',s!i)]
  put (key,(s',i,j'))

-- a round of the PRGA
prgaStep :: State PRGA Int
prgaStep = do
  (s,i,j) <- get
  let i' = (i + 1) `mod` 256
      j' = (j + (s!i')) `mod` 256
      s' = s // [(i',s!j'), (j',s!i')]
  put (s',i',j')
  return (s!((s'!i' + s'!j') `mod` 256))

keyStream :: PRGA -> [Int]
keyStream p = let (i,p') = runState prgaStep p
               in i : keyStream p'


crypt :: Key -> [Int] -> [Int]
crypt k m = zipWith xor m $ keyStream $ ksa k


pwCrypt :: String -> String -> [Int]
pwCrypt ks ms = crypt key msg
  where key = map ord ks
        msg = map ord ms

pwDecrypt :: String -> [Int] -> String
pwDecrypt k c = map chr $ crypt key c
  where key = map ord k