Merge remote-tracking branch 'kali/bug/allow-absolute-paths2' into develop
[leap_pycommon.git] / src / leap / common / crypto.py
1 # -*- coding: utf-8 -*-
2 # crypto.py
3 # Copyright (C) 2013 LEAP
4 #
5 # This program is free software: you can redistribute it and/or modify
6 # it under the terms of the GNU General Public License as published by
7 # the Free Software Foundation, either version 3 of the License, or
8 # (at your option) any later version.
9 #
10 # This program is distributed in the hope that it will be useful,
11 # but WITHOUT ANY WARRANTY; without even the implied warranty of
12 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
13 # GNU General Public License for more details.
14 #
15 # You should have received a copy of the GNU General Public License
16 # along with this program. If not, see <http://www.gnu.org/licenses/>.
17
18 import os
19 import binascii
20
21 from Crypto.Cipher import AES
22 from Crypto.Util import Counter
23 from leap.common.check import leap_assert, leap_assert_type
24
25 #
26 # encryption methods
27 #
28
29 class EncryptionMethods(object):
30     """
31     Representation of encryption methods that can be used.
32     """
33
34     AES_256_CTR = 'aes-256-ctr'
35
36
37 class UnknownEncryptionMethod(Exception):
38     """
39     Raised when trying to encrypt/decrypt with unknown method.
40     """
41     pass
42
43
44 #
45 # encrypt/decrypt functions
46 #
47
48 # In the future, we might want to implement other encryption schemes and
49 # possibly factor out the actual encryption/decryption routines of the
50 # following functions to specific classes, while not changing the API.
51
52 def encrypt_sym(data, key, method=EncryptionMethods.AES_256_CTR):
53     """
54     Encrypt C{data} with C{key}, using C{method} encryption method.
55
56     :param data: The data to be encrypted.
57     :type data: str
58     :param key: The key used to encrypt C{data} (must be 256 bits long).
59     :type key: str
60     :param method: The encryption method to use.
61     :type method: str
62
63     :return: A tuple with the initial value and the encrypted data.
64     :rtype: (long, str)
65     """
66     leap_assert_type(key, str)
67
68     # AES-256 in CTR mode
69     if method == EncryptionMethods.AES_256_CTR:
70         leap_assert(
71             len(key) == 32,  # 32 x 8 = 256 bits.
72             'Wrong key size: %s bits (must be 256 bits long).' % (len(key)*8))
73         iv = os.urandom(8)
74         ctr = Counter.new(64, prefix=iv)
75         cipher = AES.new(key=key, mode=AES.MODE_CTR, counter=ctr)
76         return binascii.b2a_base64(iv), cipher.encrypt(data)
77
78     # raise if method is unknown
79     raise UnknownEncryptionMethod('Unkwnown method: %s' % method)
80
81
82 def decrypt_sym(data, key, method=EncryptionMethods.AES_256_CTR, **kwargs):
83     """
84     Decrypt C{data} with C{key} using C{method} encryption method.
85
86     :param data: The data to be decrypted.
87     :type data: str
88     :param key: The key used to decrypt C{data} (must be 256 bits long).
89     :type key: str
90     :param method: The encryption method to use.
91     :type method: str
92     :param kwargs: Other parameters specific to each encryption method.
93     :type kwargs: dict
94
95     :return: The decrypted data.
96     :rtype: str
97     """
98     leap_assert_type(key, str)
99
100     # AES-256 in CTR mode
101     if method == EncryptionMethods.AES_256_CTR:
102         # assert params
103         leap_assert(
104             len(key) == 32,  # 32 x 8 = 256 bits.
105             'Wrong key size: %s (must be 256 bits long).' % len(key))
106         leap_assert(
107             'iv' in kwargs,
108             'AES-256-CTR needs an initial value given as.')
109         ctr = Counter.new(64, prefix=binascii.a2b_base64(kwargs['iv']))
110         cipher = AES.new(key=key, mode=AES.MODE_CTR, counter=ctr)
111         return cipher.decrypt(data)
112
113     # raise if method is unknown
114     raise UnknownEncryptionMethod('Unkwnown method: %s' % method)