File: aead.py

package info (click to toggle)
python-libnacl 2.1.0-3
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid
  • size: 488 kB
  • sloc: python: 2,634; makefile: 149; sh: 3
file content (110 lines) | stat: -rw-r--r-- 3,955 bytes parent folder | download | duplicates (2)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
# -*- coding: utf-8 -*-
'''
Utilities to make secret box encryption simple
'''
# Import libnacl
import libnacl
import libnacl.utils
import libnacl.base


class AEAD(libnacl.base.BaseKey):
    '''
    Manage AEAD encryption using the IETF ChaCha20-Poly1305(default) or AES-GCM algorithm
    '''

    def __init__(self, key=None):
        if key is None:
            key = libnacl.utils.aead_key()
        if len(key) != libnacl.crypto_aead_chacha20poly1305_ietf_KEYBYTES:  # same size for both
            raise ValueError('Invalid key')
        self.sk = key
        self.usingAES = False
        self.usingXCHACHA = False
        super().__init__()

    def useAESGCM(self):
        self.usingAES = True
        return self

    def useXCHACHA(self):
        self.usingXCHACHA = True
        return self

    def encrypt(self, msg, aad, nonce=None, pack_nonce_aad=True):
        '''
        Encrypt the given message. If a nonce is not given it will be
        generated via the rand_nonce function
        '''
        if nonce is None:
            if self.usingXCHACHA:
                nonce = libnacl.utils.rand_aead_xchacha_nonce()
            else:
                nonce = libnacl.utils.rand_aead_nonce()
        if self.usingXCHACHA:
            if len(nonce) != libnacl.crypto_aead_xchacha20poly1305_ietf_NPUBBYTES:
                raise ValueError('Invalid nonce')
        else:
            if len(nonce) != libnacl.crypto_aead_aes256gcm_NPUBBYTES:
                raise ValueError('Invalid nonce')
        if self.usingXCHACHA:
            ctxt = libnacl.crypto_aead_xchacha20poly1305_ietf_encrypt(msg, aad, nonce, self.sk)
        elif self.usingAES:
            ctxt = libnacl.crypto_aead_aes256gcm_encrypt(msg, aad, nonce, self.sk)
        else:
            ctxt = libnacl.crypto_aead_chacha20poly1305_ietf_encrypt(msg, aad, nonce, self.sk)

        if pack_nonce_aad:
            return aad + nonce + ctxt
        else:
            return aad, nonce, ctxt

    def decrypt(self, ctxt, aadLen):
        '''
        Decrypt the given message, if no nonce or aad are given they will be
        extracted from the message
        '''
        aad = ctxt[:aadLen]
        if self.usingXCHACHA:
            nonce = ctxt[aadLen:aadLen+libnacl.crypto_aead_xchacha20poly1305_ietf_NPUBBYTES]
            ctxt = ctxt[aadLen+libnacl.crypto_aead_xchacha20poly1305_ietf_NPUBBYTES:]
            if len(nonce) != libnacl.crypto_aead_xchacha20poly1305_ietf_NPUBBYTES:
                raise ValueError('Invalid nonce')
        else:
            nonce = ctxt[aadLen:aadLen+libnacl.crypto_aead_aes256gcm_NPUBBYTES]
            ctxt = ctxt[aadLen+libnacl.crypto_aead_aes256gcm_NPUBBYTES:]
            if len(nonce) != libnacl.crypto_aead_aes256gcm_NPUBBYTES:
                raise ValueError('Invalid nonce')
        if self.usingXCHACHA:
            return libnacl.crypto_aead_xchacha20poly1305_ietf_decrypt(ctxt, aad, nonce, self.sk)
        if self.usingAES:
            return libnacl.crypto_aead_aes256gcm_decrypt(ctxt, aad, nonce, self.sk)
        return libnacl.crypto_aead_chacha20poly1305_ietf_decrypt(ctxt, aad, nonce, self.sk)

    def decrypt_unpacked(self, aad, nonce, ctxt):
        '''
        Decrypt the given message, if no nonce or aad are given they will be
        extracted from the message
        '''
        if len(nonce) != libnacl.crypto_aead_aes256gcm_NPUBBYTES:
            raise ValueError('Invalid nonce')
        if self.usingAES:
            return libnacl.crypto_aead_aes256gcm_decrypt(ctxt, aad, nonce, self.sk)
        return libnacl.crypto_aead_chacha20poly1305_ietf_decrypt(ctxt, aad, nonce, self.sk)


class AEAD_AESGCM(AEAD):
    def __init__(self, key=None):
        super().__init__(key)
        self.useAESGCM()


class AEAD_XCHACHA(AEAD):
    def __init__(self, key=None):
        super().__init__(key)
        self.useXCHACHA()


class AEAD_CHACHA(AEAD):
    def __init__(self, key=None):
        super().__init__(key)