Create a gist now

Instantly share code, notes, and snippets.

What would you like to do?
RC4 via SBV library
/*
* Oh-so-slightly modified version of the RC4 code posted to cypherpunks that just adds a main() function driver.
*
* http://tinyurl.com/66lbx9d
*
*/
#define DEBUG
#include <stdio.h>
#include <assert.h>
#define countof(a) (sizeof(a)/sizeof(*(a)))
typedef struct rc4_key
{
unsigned char state[256];
unsigned char x;
unsigned char y;
} rc4_key;
static void swap_byte(unsigned char *a, unsigned char *b);
void prepare_key(unsigned char *key_data_ptr, int key_data_len,
rc4_key *key)
{
unsigned char index1;
unsigned char index2;
unsigned char* state;
short counter;
state = &key->state[0];
for(counter = 0; counter < 256; counter++) {
state[counter] = counter;
}
key->x = 0;
key->y = 0;
index1 = 0;
index2 = 0;
for(counter = 0; counter < 256; counter++)
{
index2 = (key_data_ptr[index1] + state[counter] + index2) % 256;
index1 = (index1 + 1) % key_data_len;
swap_byte(&state[counter], &state[index2]);
}
}
void rc4(unsigned char *buffer_ptr, int buffer_len, rc4_key
*key)
{
unsigned char x;
unsigned char y;
unsigned char* state;
unsigned char xorIndex;
short counter;
x = key->x;
y = key->y;
state = &key->state[0];
for(counter = 0; counter < buffer_len; counter ++)
{
x = (x + 1) % 256;
y = (state[x] + y) % 256;
swap_byte(&state[x], &state[y]);
xorIndex = state[x] + (state[y]) % 256;
buffer_ptr[counter] ^= state[xorIndex];
}
key->x = x;
key->y = y;
}
static void swap_byte(unsigned char *a, unsigned char *b)
{
unsigned char swapByte;
swapByte = *a;
*a = *b;
*b = swapByte;
}
int main() {
struct rc4_key k1;
struct rc4_key k2;
unsigned char s[] = {1,2,3,4,5}; /* secret */
unsigned char m1[] = {0x43,0x54,0x02,0x65,0x00,0x76}; /* msg */
unsigned char m2[] = {0x43,0x54,0x02,0x65,0x00,0x76};
prepare_key(s,5,&k1);
prepare_key(s,5,&k2);
rc4(m1, countof(m1), &k1);
rc4(m1, countof(m1), &k2);
for(int i=0; i<countof(m2);i++) {
int b = m1[i]==m2[i];
assert(b); /* Assert truthiness */
}
printf("\n\n");
return 0;
}
-- | An RC4 implementation utilizing SBV.
{-# LANGUAGE ScopedTypeVariables, MultiParamTypeClasses #-}
module RC4
( -- * Types
S -- :: *
, Key -- :: *
, RC4 -- :: *
-- * Swapping array values
, swap -- :: SWord8 -> SWord8 -> S -> S
-- * Key-scheduling algorithm
, initState -- :: S
, ksa -- :: Key -> S -> S
-- * Psuedo-random generator algorithm
, prga -- :: RC4 -> (SWord8, RC4)
-- * Interface
, initCtx -- :: Key -> RC4
, encrypt -- :: [SWord8] -> RC4 -> ([SWord8], RC4)
, decrypt -- :: [SWord8] -> RC4 -> ([SWord8], RC4)
-- * 'ByteString' interface
, initCtx' -- :: [Word8] -> RC4
, encryptBS -- :: RC4 -> ByteString -> (ByteString, RC4)
, decryptBS -- :: RC4 -> ByteString -> RC4 -> (ByteString, RC4)
-- * 'crypto-api' instances
, RC4Key -- :: *
, RC4IV -- :: *
, createRC4IV -- :: RC4Key -> RC4IV
-- * Tests & Verification
, doAllTests -- :: IO ()
-- * C Code generation
, codegenRC4 -- :: IO ()
) where
import Data.SBV
import Data.Maybe (fromJust)
import qualified Data.ByteString as B
import Control.Applicative ((<$>))
import Data.Serialize
import Crypto.Classes
import qualified Crypto.Cipher.RC4 as Crypto
import qualified Data.Vector.Unboxed as V
import Test.QuickCheck (quickCheck)
--------------------------
-- Types
-- | Denotes the permutation of all possibly 256 bytes, used as an RC4 state.
type S = SFunArray Word8 Word8
-- | A key for RC4 is merely a stream of Word8 value
type Key = [SWord8]
-- | Represents the current contex the RC4 cipher is in.
type RC4 = (S, SWord8, SWord8)
--------------------------
-- Swapping array values
-- | Swap two elements in an array, in this case, our RC4 `S` array used
-- in the PRGA/KSA
swap :: SWord8 -> SWord8 -> S -> S
swap i j a
= let tmp = readArray a i
a' = writeArray a i (readArray a j)
in writeArray a' j tmp
--------------------------
-- Key scheduling & PRGA
-- | Initial state table - this is the initial `S` array used in the RC4 implementation,
-- and it is initially populated with 256 values from the range `[0,255]`.
initState :: S
initState
= foldr (\i a -> writeArray a i i) (mkSFunArray 0) [0..255]
-- | This function implements the key-scheduling algorithm (KSA) and initializes
-- the RC4 `S` array with a given key.
ksa :: Key -> S -> S
ksa key = ite (kl ./= 0) go id
where kl :: SWord8
kl = literal $ fromIntegral (length key)
idx :: [a] -> SWord8 -> a
idx arr i | Just x <- unliteral i = arr !! (fromIntegral x)
| otherwise = error "unsafeIdx: non-concrete argument"
-- perform the key scheduling
go :: S -> S
go st = let (v,_,_) = k 0 (st, 0, 0) in v
where k :: SWord16 -> (S, SWord8, SWord8) -> (S, SWord8, SWord8)
k c v@(s, j, i)
= ite (c .== 256) v $
let lo = snd $ split c -- truncate the low bits
j' = ((key `idx` i) + (s .! lo) + j)
i' = snd $ (i+1) `bvQuotRem` kl
s' = swap lo j' s
in k (c+1) (s',j',i')
-- | The PRGA (*psuedo-random-generation-algorithm*) that generates the RC4 keystream
-- after the initial key scheduling is done.
prga :: RC4 -> (SWord8, RC4)
prga (st,i,j)
= let i' = (i + 1)
j' = ((st .! i') + j)
st' = swap i' j' st
k = (st' .! i') + (st' .! j')
in (st' .! k, (st', i', j'))
----------------------------------------
-- Interface
-- | Initialize an RC4 context by supplying the encryption key. This routine performs
-- key exchange.
initCtx :: Key -> RC4
initCtx k
= (ksa k initState, 0, 0)
-- | Encrypt some data using an encryption context. Returns the encrypted data,
-- and the updated key to use for further encryption. Do not use the contexts
-- you have already used to encrypt future data.
encrypt :: RC4 -> [SWord8] -> ([SWord8], RC4)
encrypt c plaintxt = k ([],c) plaintxt
where k :: ([SWord8], RC4) -> [SWord8] -> ([SWord8], RC4)
k v [] = v
k (out,ctx) (pt:pt')
= let (ks,ctx') = prga ctx
ct = pt `xor` ks
in k (out ++ [ct],ctx') pt'
-- | Decrypt some data with the given context, and return the decrypted data
-- and new context. Do not use the old context to decrypt any data in the future.
decrypt :: RC4 -> [SWord8] -> ([SWord8], RC4)
decrypt = encrypt
-- ByteStrings
-- | Initialize a context using a list of 'Word8's instead of a 'SWord8's
initCtx' :: [Word8] -> RC4
initCtx' = initCtx . w2sw
-- | Encrypt a 'ByteString'.
encryptBS :: RC4 -> B.ByteString -> (B.ByteString, RC4)
encryptBS ctx bs
= let (s,ctx') = encrypt ctx $ w2sw $ B.unpack bs
in (toBS s,ctx')
where toBS = B.pack . sw2w
-- | Decrypt a 'ByteString'.
decryptBS :: RC4 -> B.ByteString -> (B.ByteString, RC4)
decryptBS = encryptBS
-- CryptoAPI
data RC4Key = RC4Key [Word8]
data RC4IV = RC4IV RC4
instance Serialize RC4Key where
put (RC4Key v) = putWord8 (fromIntegral $ length v) >> mapM_ putWord8 v
get = RC4Key <$> (getWord8 >>= \x -> sequence $ replicate (fromIntegral x) getWord8)
instance StreamCipher RC4Key RC4IV where
buildStreamKey bs
| B.length bs < 1 = Nothing
| otherwise = Just $ RC4Key (B.unpack bs)
encryptStream _ (RC4IV iv) b = liftIV $ encryptBS iv b
decryptStream _ (RC4IV iv) b = liftIV $ decryptBS iv b
streamKeyLength (RC4Key k) = (length k) * 8
-- | Creates an 'RC4IV' after you've gotten an 'RC4Key' via 'buildStreamKey'
createRC4IV :: RC4Key -> RC4IV
createRC4IV (RC4Key k) = RC4IV $ initCtx' k
liftIV :: (B.ByteString,RC4) -> (B.ByteString,RC4IV)
liftIV (b,v) = (b,RC4IV v)
--------------------------
-- Verification & Tests
-- | Basic quickcheck properties for the crypto-api interface and some checks against
-- `cryptocipher`'s RC4 implementation.
checkRc4Props :: IO ()
checkRc4Props = sequence_ [ test "crypto api interface" $ quickCheck prop_cryptoapi
, test "cryptocipher ksa equivalence" $ quickCheck prop_crypto_ksa_eq
, test "cryptocipher enc equivalence" $ quickCheck prop_crypto_enc_eq
]
where prop_cryptoapi :: [Word8] -> [Word8] -> Bool
prop_cryptoapi key pt = (length key) > 1 ==> enc1 == enc2
where (enc1,_) = encryptBS (initCtx' key) $ B.pack pt
(enc2,_) = let k = fromJust (buildStreamKey $ B.pack key) :: RC4Key
iv = createRC4IV k
in encryptStream k iv $ B.pack pt
prop_crypto_ksa_eq :: [Word8] -> Bool
prop_crypto_ksa_eq key = (length key) > 1 ==> k1 == k2
where k1 = let (v,_,_) = initCtx' key in sw2w $ fst (foldr ra ([],v) [0..255] )
where ra i (o,v) = let x = readArray v (literal i)
in (x:o,v)
k2 = let (v,_,_) = Crypto.initCtx key in V.toList v
prop_crypto_enc_eq :: [Word8] -> [Word8] -> Bool
prop_crypto_enc_eq key pt = (length key) > 1 ==> enc1 == enc2
where (enc1,_) = encryptBS (initCtx' key) $ B.pack pt
(_,enc2) = Crypto.encrypt (Crypto.initCtx key) $ B.pack pt
-- | This gives a correctness theorem for 'swap' and proves it via the
-- SMT solver: if we 'swap' the same elements of an array twice in an
-- imperative fashion, the initial array and the resulting array
-- values are equivalent (i.e. @swap i j (swap i j a) == a@).
proveSwapIsCorrect :: IO ThmResult
proveSwapIsCorrect
= proveWith timingSMTCfg $ forAll_ $ \i j -> swapIsCorrect i j initState
where swapIsCorrect :: SWord8 -> SWord8 -> S -> SBool
swapIsCorrect i j a
= let (x,y) = (a .! i, a .! j)
a1 = swap i j a
a2 = swap i j a1
(x',y') = (a2 .! i, a2 .! j)
in (x .== x' &&& y .== y')
-- | Ensures encryption followed by decryption is the identity function for 40bit keys and 40bit plaintext,
-- and proves it via the SMT solver. So we shouldn't need to QuickCheck it :)
proveRc4IsCorrect :: IO ThmResult
proveRc4IsCorrect
= proveWith timingSMTCfg $ forAll_ $ \plaintext key -> rc4IsCorrect plaintext key
where rc4IsCorrect :: (SWord8,SWord8,SWord8,SWord8,SWord8) -> (SWord8,SWord8,SWord8,SWord8,SWord8) -> SBool
rc4IsCorrect (pt1,pt2,pt3,pt4,pt5) (k1,k2,k3,k4,k5)
= let ctx = initCtx [k1,k2,k3,k4,k5]
(enc,_) = encrypt ctx [pt1,pt2,pt3,pt4,pt5]
(dec,_) = decrypt ctx enc
in dec .== [pt1,pt2,pt3,pt4,pt5]
-- | Run all tests and verification functions.
doAllTests :: IO ()
doAllTests
= sequence_ [ checkRc4Props
, verif "swap is correct" proveSwapIsCorrect
, verif "rc4 is correct" proveRc4IsCorrect
]
test, verif :: String -> IO a -> IO ()
test x f = putStrLn ("testing: "++x) >> f >> return ()
verif x f = putStrLn ("verifying: "++x) >> f >> return ()
--------------------------
-- Code generation
-- | Generate C code for our RC4 implementation.
-- TODO
codegenRC4 :: IO ()
codegenRC4 = return ()
--------------------------
-- Utilities
(.!) :: S -> SWord8 -> SWord8
(.!) = readArray
sw2w :: [SWord8] -> [Word8]
sw2w = map (fromJust . unliteral)
w2sw :: [Word8] -> [SWord8]
w2sw = map literal
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment