#!/usr/bin/env python

"""Unit tests for M2Crypto.RC4.

Copyright (c) 2009 Heikki Toivonen. All rights reserved."""

from M2Crypto import RC4, Rand
from binascii import hexlify

from tests import unittest
from tests.fips import fips_mode


class RC4TestCase(unittest.TestCase):

    @unittest.skipIf(fips_mode, "Can't be run in FIPS mode")
    def test_vectors(self):
        """
        Test with test vectors from Wikipedia: http://en.wikipedia.org/wiki/Rc4
        """
        if fips_mode:
            return
        vectors = (
            (b"Key", b"Plaintext", b"BBF316E8D940AF0AD3"),
            (b"Wiki", b"pedia", b"1021BF0420"),
            (
                b"Secret",
                b"Attack at dawn",
                b"45A01F645FC35B383552544B9BF5",
            ),
        )

        rc4 = RC4.RC4()
        for key, plaintext, ciphertext in vectors:
            rc4.set_key(key)
            self.assertEqual(hexlify(rc4.update(plaintext)).upper(), ciphertext)

        self.assertEqual(rc4.final(), "")

    @unittest.skipIf(fips_mode, "Can't be run in FIPS mode")
    def test_bad(self):
        if fips_mode:
            return
        rc4 = RC4.RC4(b"foo")
        self.assertNotEqual(hexlify(rc4.update(b"bar")).upper(), b"45678")


def suite():
    return unittest.TestLoader().loadTestsFromTestCase(RC4TestCase)


if __name__ == "__main__":
    Rand.load_file("randpool.dat", -1)
    unittest.TextTestRunner().run(suite())
    Rand.save_file("randpool.dat")
