1# -*- coding: utf-8 -*-
2#
3#  Cipher/ARC4.py : ARC4
4#
5# ===================================================================
6# The contents of this file are dedicated to the public domain.  To
7# the extent that dedication to the public domain is not available,
8# everyone is granted a worldwide, perpetual, royalty-free,
9# non-exclusive license to exercise all rights associated with the
10# contents of this file for any purpose whatsoever.
11# No rights are reserved.
12#
13# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
14# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
15# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
16# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS
17# BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
18# ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
19# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
20# SOFTWARE.
21# ===================================================================
22
23from Crypto.Util.py3compat import b
24
25from Crypto.Util._raw_api import (load_pycryptodome_raw_lib, VoidPointer,
26                                  create_string_buffer, get_raw_buffer,
27                                  SmartPointer, c_size_t, c_uint8_ptr)
28
29
30_raw_arc4_lib = load_pycryptodome_raw_lib("Crypto.Cipher._ARC4", """
31                    int ARC4_stream_encrypt(void *rc4State, const uint8_t in[],
32                                            uint8_t out[], size_t len);
33                    int ARC4_stream_init(uint8_t *key, size_t keylen,
34                                         void **pRc4State);
35                    int ARC4_stream_destroy(void *rc4State);
36                    """)
37
38
39class ARC4Cipher:
40    """ARC4 cipher object. Do not create it directly. Use
41    :func:`Crypto.Cipher.ARC4.new` instead.
42    """
43
44    def __init__(self, key, *args, **kwargs):
45        """Initialize an ARC4 cipher object
46
47        See also `new()` at the module level."""
48
49        if len(args) > 0:
50            ndrop = args[0]
51            args = args[1:]
52        else:
53            ndrop = kwargs.pop('drop', 0)
54
55        if len(key) not in key_size:
56            raise ValueError("Incorrect ARC4 key length (%d bytes)" %
57                             len(key))
58
59        self._state = VoidPointer()
60        result = _raw_arc4_lib.ARC4_stream_init(c_uint8_ptr(key),
61                                                c_size_t(len(key)),
62                                                self._state.address_of())
63        if result != 0:
64            raise ValueError("Error %d while creating the ARC4 cipher"
65                             % result)
66        self._state = SmartPointer(self._state.get(),
67                                   _raw_arc4_lib.ARC4_stream_destroy)
68
69        if ndrop > 0:
70            # This is OK even if the cipher is used for decryption,
71            # since encrypt and decrypt are actually the same thing
72            # with ARC4.
73            self.encrypt(b'\x00' * ndrop)
74
75        self.block_size = 1
76        self.key_size = len(key)
77
78    def encrypt(self, plaintext):
79        """Encrypt a piece of data.
80
81        :param plaintext: The data to encrypt, of any size.
82        :type plaintext: bytes, bytearray, memoryview
83        :returns: the encrypted byte string, of equal length as the
84          plaintext.
85        """
86
87        ciphertext = create_string_buffer(len(plaintext))
88        result = _raw_arc4_lib.ARC4_stream_encrypt(self._state.get(),
89                                                   c_uint8_ptr(plaintext),
90                                                   ciphertext,
91                                                   c_size_t(len(plaintext)))
92        if result:
93            raise ValueError("Error %d while encrypting with RC4" % result)
94        return get_raw_buffer(ciphertext)
95
96    def decrypt(self, ciphertext):
97        """Decrypt a piece of data.
98
99        :param ciphertext: The data to decrypt, of any size.
100        :type ciphertext: bytes, bytearray, memoryview
101        :returns: the decrypted byte string, of equal length as the
102          ciphertext.
103        """
104
105        try:
106            return self.encrypt(ciphertext)
107        except ValueError as e:
108            raise ValueError(str(e).replace("enc", "dec"))
109
110
111def new(key, *args, **kwargs):
112    """Create a new ARC4 cipher.
113
114    :param key:
115        The secret key to use in the symmetric cipher.
116        Its length must be in the range ``[5..256]``.
117        The recommended length is 16 bytes.
118    :type key: bytes, bytearray, memoryview
119
120    :Keyword Arguments:
121        *   *drop* (``integer``) --
122            The amount of bytes to discard from the initial part of the keystream.
123            In fact, such part has been found to be distinguishable from random
124            data (while it shouldn't) and also correlated to key.
125
126            The recommended value is 3072_ bytes. The default value is 0.
127
128    :Return: an `ARC4Cipher` object
129
130    .. _3072: http://eprint.iacr.org/2002/067.pdf
131    """
132    return ARC4Cipher(key, *args, **kwargs)
133
134# Size of a data block (in bytes)
135block_size = 1
136# Size of a key (in bytes)
137key_size = range(5, 256+1)
138