1# Copyright 2013-2018 Donald Stufft and individual contributors
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15from __future__ import absolute_import, division, print_function
16
17from nacl import exceptions as exc
18from nacl._sodium import ffi, lib
19from nacl.exceptions import ensure
20
21
22crypto_secretstream_xchacha20poly1305_ABYTES = \
23    lib.crypto_secretstream_xchacha20poly1305_abytes()
24crypto_secretstream_xchacha20poly1305_HEADERBYTES = \
25    lib.crypto_secretstream_xchacha20poly1305_headerbytes()
26crypto_secretstream_xchacha20poly1305_KEYBYTES = \
27    lib.crypto_secretstream_xchacha20poly1305_keybytes()
28crypto_secretstream_xchacha20poly1305_MESSAGEBYTES_MAX = \
29    lib.crypto_secretstream_xchacha20poly1305_messagebytes_max()
30crypto_secretstream_xchacha20poly1305_STATEBYTES = \
31    lib.crypto_secretstream_xchacha20poly1305_statebytes()
32
33
34crypto_secretstream_xchacha20poly1305_TAG_MESSAGE = \
35    lib.crypto_secretstream_xchacha20poly1305_tag_message()
36crypto_secretstream_xchacha20poly1305_TAG_PUSH = \
37    lib.crypto_secretstream_xchacha20poly1305_tag_push()
38crypto_secretstream_xchacha20poly1305_TAG_REKEY = \
39    lib.crypto_secretstream_xchacha20poly1305_tag_rekey()
40crypto_secretstream_xchacha20poly1305_TAG_FINAL = \
41    lib.crypto_secretstream_xchacha20poly1305_tag_final()
42
43
44def crypto_secretstream_xchacha20poly1305_keygen():
45    """
46    Generate a key for use with
47    :func:`.crypto_secretstream_xchacha20poly1305_init_push`.
48
49    """
50    keybuf = ffi.new(
51        "unsigned char[]",
52        crypto_secretstream_xchacha20poly1305_KEYBYTES,
53    )
54    lib.crypto_secretstream_xchacha20poly1305_keygen(keybuf)
55    return ffi.buffer(keybuf)[:]
56
57
58class crypto_secretstream_xchacha20poly1305_state(object):
59    """
60    An object wrapping the crypto_secretstream_xchacha20poly1305 state.
61
62    """
63    __slots__ = ['statebuf', 'rawbuf', 'tagbuf']
64
65    def __init__(self):
66        """ Initialize a clean state object."""
67        self.statebuf = ffi.new(
68            "unsigned char[]",
69            crypto_secretstream_xchacha20poly1305_STATEBYTES,
70        )
71
72        self.rawbuf = None
73        self.tagbuf = None
74
75
76def crypto_secretstream_xchacha20poly1305_init_push(state, key):
77    """
78    Initialize a crypto_secretstream_xchacha20poly1305 encryption buffer.
79
80    :param state: a secretstream state object
81    :type state: crypto_secretstream_xchacha20poly1305_state
82    :param key: must be
83                :data:`.crypto_secretstream_xchacha20poly1305_KEYBYTES` long
84    :type key: bytes
85    :return: header
86    :rtype: bytes
87
88    """
89    ensure(
90        isinstance(state, crypto_secretstream_xchacha20poly1305_state),
91        'State must be a crypto_secretstream_xchacha20poly1305_state object',
92        raising=exc.TypeError,
93    )
94    ensure(
95        isinstance(key, bytes),
96        'Key must be a bytes sequence',
97        raising=exc.TypeError,
98    )
99    ensure(
100        len(key) == crypto_secretstream_xchacha20poly1305_KEYBYTES,
101        'Invalid key length',
102        raising=exc.ValueError,
103    )
104
105    headerbuf = ffi.new(
106        "unsigned char []",
107        crypto_secretstream_xchacha20poly1305_HEADERBYTES,
108    )
109
110    rc = lib.crypto_secretstream_xchacha20poly1305_init_push(
111        state.statebuf, headerbuf, key)
112    ensure(rc == 0, 'Unexpected failure', raising=exc.RuntimeError)
113
114    return ffi.buffer(headerbuf)[:]
115
116
117def crypto_secretstream_xchacha20poly1305_push(
118    state,
119    m,
120    ad=None,
121    tag=crypto_secretstream_xchacha20poly1305_TAG_MESSAGE,
122):
123    """
124    Add an encrypted message to the secret stream.
125
126    :param state: a secretstream state object
127    :type state: crypto_secretstream_xchacha20poly1305_state
128    :param m: the message to encrypt, the maximum length of an individual
129              message is
130              :data:`.crypto_secretstream_xchacha20poly1305_MESSAGEBYTES_MAX`.
131    :type m: bytes
132    :param ad: additional data to include in the authentication tag
133    :type ad: bytes or None
134    :param tag: the message tag, usually
135                :data:`.crypto_secretstream_xchacha20poly1305_TAG_MESSAGE` or
136                :data:`.crypto_secretstream_xchacha20poly1305_TAG_FINAL`.
137    :type tag: int
138    :return: ciphertext
139    :rtype: bytes
140
141    """
142    ensure(
143        isinstance(state, crypto_secretstream_xchacha20poly1305_state),
144        'State must be a crypto_secretstream_xchacha20poly1305_state object',
145        raising=exc.TypeError,
146    )
147    ensure(isinstance(m, bytes), 'Message is not bytes', raising=exc.TypeError)
148    ensure(
149        len(m) <= crypto_secretstream_xchacha20poly1305_MESSAGEBYTES_MAX,
150        'Message is too long',
151        raising=exc.ValueError,
152    )
153    ensure(
154        ad is None or isinstance(ad, bytes),
155        'Additional data must be bytes or None',
156        raising=exc.TypeError,
157    )
158
159    clen = len(m) + crypto_secretstream_xchacha20poly1305_ABYTES
160    if state.rawbuf is None or len(state.rawbuf) < clen:
161        state.rawbuf = ffi.new('unsigned char[]', clen)
162
163    if ad is None:
164        ad = ffi.NULL
165        adlen = 0
166    else:
167        adlen = len(ad)
168
169    rc = lib.crypto_secretstream_xchacha20poly1305_push(
170        state.statebuf,
171        state.rawbuf, ffi.NULL,
172        m, len(m),
173        ad, adlen,
174        tag,
175    )
176    ensure(rc == 0, 'Unexpected failure', raising=exc.RuntimeError)
177
178    return ffi.buffer(state.rawbuf, clen)[:]
179
180
181def crypto_secretstream_xchacha20poly1305_init_pull(state, header, key):
182    """
183    Initialize a crypto_secretstream_xchacha20poly1305 decryption buffer.
184
185    :param state: a secretstream state object
186    :type state: crypto_secretstream_xchacha20poly1305_state
187    :param header: must be
188                :data:`.crypto_secretstream_xchacha20poly1305_HEADERBYTES` long
189    :type header: bytes
190    :param key: must be
191                :data:`.crypto_secretstream_xchacha20poly1305_KEYBYTES` long
192    :type key: bytes
193
194    """
195    ensure(
196        isinstance(state, crypto_secretstream_xchacha20poly1305_state),
197        'State must be a crypto_secretstream_xchacha20poly1305_state object',
198        raising=exc.TypeError,
199    )
200    ensure(
201        isinstance(header, bytes),
202        'Header must be a bytes sequence',
203        raising=exc.TypeError,
204    )
205    ensure(
206        len(header) == crypto_secretstream_xchacha20poly1305_HEADERBYTES,
207        'Invalid header length',
208        raising=exc.ValueError,
209    )
210    ensure(
211        isinstance(key, bytes),
212        'Key must be a bytes sequence',
213        raising=exc.TypeError,
214    )
215    ensure(
216        len(key) == crypto_secretstream_xchacha20poly1305_KEYBYTES,
217        'Invalid key length',
218        raising=exc.ValueError,
219    )
220
221    if state.tagbuf is None:
222        state.tagbuf = ffi.new('unsigned char *')
223
224    rc = lib.crypto_secretstream_xchacha20poly1305_init_pull(
225        state.statebuf, header, key)
226    ensure(rc == 0, 'Unexpected failure', raising=exc.RuntimeError)
227
228
229def crypto_secretstream_xchacha20poly1305_pull(state, c, ad=None):
230    """
231    Read a decrypted message from the secret stream.
232
233    :param state: a secretstream state object
234    :type state: crypto_secretstream_xchacha20poly1305_state
235    :param c: the ciphertext to decrypt, the maximum length of an individual
236              ciphertext is
237              :data:`.crypto_secretstream_xchacha20poly1305_MESSAGEBYTES_MAX` +
238              :data:`.crypto_secretstream_xchacha20poly1305_ABYTES`.
239    :type c: bytes
240    :param ad: additional data to include in the authentication tag
241    :type ad: bytes or None
242    :return: (message, tag)
243    :rtype: (bytes, int)
244
245    """
246    ensure(
247        isinstance(state, crypto_secretstream_xchacha20poly1305_state),
248        'State must be a crypto_secretstream_xchacha20poly1305_state object',
249        raising=exc.TypeError,
250    )
251    ensure(
252        state.tagbuf is not None,
253        (
254            'State must be initialized using '
255            'crypto_secretstream_xchacha20poly1305_init_pull'
256        ),
257        raising=exc.ValueError,
258    )
259    ensure(
260        isinstance(c, bytes),
261        'Ciphertext is not bytes',
262        raising=exc.TypeError,
263    )
264    ensure(
265        len(c) >= crypto_secretstream_xchacha20poly1305_ABYTES,
266        'Ciphertext is too short',
267        raising=exc.ValueError,
268    )
269    ensure(
270        len(c) <= (
271            crypto_secretstream_xchacha20poly1305_MESSAGEBYTES_MAX +
272            crypto_secretstream_xchacha20poly1305_ABYTES
273        ),
274        'Ciphertext is too long',
275        raising=exc.ValueError,
276    )
277    ensure(
278        ad is None or isinstance(ad, bytes),
279        'Additional data must be bytes or None',
280        raising=exc.TypeError,
281    )
282
283    mlen = len(c) - crypto_secretstream_xchacha20poly1305_ABYTES
284    if state.rawbuf is None or len(state.rawbuf) < mlen:
285        state.rawbuf = ffi.new('unsigned char[]', mlen)
286
287    if ad is None:
288        ad = ffi.NULL
289        adlen = 0
290    else:
291        adlen = len(ad)
292
293    rc = lib.crypto_secretstream_xchacha20poly1305_pull(
294        state.statebuf,
295        state.rawbuf, ffi.NULL,
296        state.tagbuf,
297        c, len(c),
298        ad, adlen,
299    )
300    ensure(rc == 0, 'Unexpected failure', raising=exc.RuntimeError)
301
302    return (ffi.buffer(state.rawbuf, mlen)[:], int(state.tagbuf[0]))
303
304
305def crypto_secretstream_xchacha20poly1305_rekey(state):
306    """
307    Explicitly change the encryption key in the stream.
308
309    Normally the stream is re-keyed as needed or an explicit ``tag`` of
310    :data:`.crypto_secretstream_xchacha20poly1305_TAG_REKEY` is added to a
311    message to ensure forward secrecy, but this method can be used instead
312    if the re-keying is controlled without adding the tag.
313
314    :param state: a secretstream state object
315    :type state: crypto_secretstream_xchacha20poly1305_state
316
317    """
318    ensure(
319        isinstance(state, crypto_secretstream_xchacha20poly1305_state),
320        'State must be a crypto_secretstream_xchacha20poly1305_state object',
321        raising=exc.TypeError,
322    )
323    lib.crypto_secretstream_xchacha20poly1305_rekey(state.statebuf)
324