Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1"""
2@file
3@brief Encryption functionalities.
5Inspired from `AES encryption of files in Python with PyCrypto
6<http://eli.thegreenplace.net/2010/06/25/aes-encryption-of-files-in-python-with-pycrypto>`_
7"""
8import random
9import os
10import struct
11import base64
12from io import BytesIO as StreamIO
15class EncryptionError(Exception):
16 """
17 raised if an issue happen during encryption
18 """
19 pass
22def open_input_output(filename, out_filename=None):
23 """
24 Converts *filename* and *out_filename* as streams.
26 @param filename bytes or filename or BytesIO
27 @param out_filename BytesIO or filename or None
28 @return in_size, in_close, in_stream, out_close, out_return, out_stream
29 """
30 # input
31 typstr = str # unicode #
32 if isinstance(filename, typstr):
33 if not os.path.exists(filename):
34 raise FileNotFoundError(filename)
35 st = open(filename, "rb")
36 close = True
37 filesize = os.path.getsize(filename)
38 elif isinstance(filename, StreamIO):
39 st = filename
40 close = False
41 filesize = len(st.getvalue())
42 else:
43 st = StreamIO(filename)
44 close = False
45 filesize = len(filename)
47 # output
48 if out_filename is None:
49 sto = StreamIO()
50 ret = True
51 out_close = False
52 elif isinstance(out_filename, StreamIO):
53 sto = out_filename
54 ret = False
55 out_close = False
56 else:
57 sto = open(out_filename, "wb")
58 ret = False
59 out_close = True
61 return filesize, close, st, out_close, ret, sto
64def close_input_output(in_size, in_close, in_stream, out_close, out_return, out_stream):
65 """
66 Takes the output of @see fn open_input_output and closes streams
67 and return expected values.
69 @param in_size size of input
70 @param in_close should it close the input stream
71 @param in_stream input stream
72 @param out_close should it closes the output stream
73 @param out_return should it returns something
74 @param out_stream output stream
75 @return None or content of output stream
76 """
77 if in_close:
78 in_stream.close()
80 if out_close:
81 if out_return:
82 raise EncryptionError("incompability")
83 out_stream.close()
85 if out_return:
86 return out_stream.getvalue()
87 else:
88 return None
91def get_encryptor(key, algo="AES", chunksize=2 ** 24, **params):
92 """
93 Returns an encryptor with method encrypt and decrypt.
95 @param key key
96 @param algo AES or fernet
97 @param chunksize Fernet does not allow streaming
98 @param params additional parameters
99 @return encryptor, origsize
100 """
101 if algo == "fernet":
102 from cryptography.fernet import Fernet
103 if hasattr(key, "encode"):
104 # it a string
105 bkey = key.encode()
106 else:
107 bkey = key
108 bkey = base64.b64encode(bkey)
109 encryptor = Fernet(bkey)
110 origsize = None
111 chunksize = None
112 elif algo == "AES":
113 from Cryptodome.Cipher import AES
114 ksize = {16, 32, 64, 128, 256}
115 chunksize = chunksize # pylint: disable=W0127
116 if len(key) not in ksize:
117 raise EncryptionError(
118 "len(key)=={0} should be of length {1}".format(len(key), str(ksize)))
119 if "out_stream" in params:
120 iv = bytes([random.randint(0, 0xFF) for i in range(16)])
121 params["out_stream"].write(struct.pack('<Q', params["in_size"]))
122 params["out_stream"].write(iv)
123 encryptor = AES.new(key, AES.MODE_CBC, iv)
124 origsize = params["in_size"]
125 else:
126 origsize = struct.unpack(
127 '<Q', params["in_stream"].read(struct.calcsize('Q')))[0]
128 iv = params["in_stream"].read(16)
129 encryptor = AES.new(key, AES.MODE_CBC, iv) # decryptor
130 else:
131 raise ValueError("unknown algorithm: {0}, should be in {1}".format(
132 algo, ["fernet", "AES"]))
133 return encryptor, origsize, chunksize
136def encrypt_stream(key, filename, out_filename=None, chunksize=2 ** 18, algo="AES"):
137 """
138 Encrypts a file using AES (CBC mode) with the given key.
139 The function relies on module :epkg:`pycrypto`, :epkg:`cryptography`,
140 algoritm `AES <https://fr.wikipedia.org/wiki/Advanced_Encryption_Standard>`_,
141 `Fernet <https://cryptography.io/en/latest/fernet/>`_.
143 @param key The encryption key - a string that must be
144 either 16, 24 or 32 bytes long. Longer keys
145 are more secure. If the data to encrypt is in bytes,
146 the key must be given in bytes too.
148 @param filename bytes or Name of the input file
149 @param out_filename if None, the returns bytes
151 @param chunksize Sets the size of the chunk which the function
152 uses to read and encrypt the file. Larger chunk
153 sizes can be faster for some files and machines.
154 chunksize must be divisible by 16.
156 @param algo AES (PyCryptodomex) of or fernet (cryptography)
158 @return filename or bytes
159 """
161 in_size, in_close, in_stream, out_close, out_return, out_stream = open_input_output(
162 filename, out_filename)
164 encryptor, origsize, chunksize = get_encryptor(
165 key, algo, out_stream=out_stream, in_size=in_size, chunksize=chunksize)
167 while True:
168 chunk = in_stream.read(chunksize)
169 if len(chunk) == 0:
170 break
171 if len(chunk) % 16 != 0 and origsize is not None:
172 chunk += b' ' * (16 - len(chunk) % 16)
174 out_stream.write(encryptor.encrypt(chunk))
176 return close_input_output(in_size, in_close, in_stream, out_close, out_return, out_stream)
179def decrypt_stream(key, filename, out_filename=None, chunksize=3 * 2 ** 13, algo="AES"):
180 """
181 Decrypts a file using AES (CBC mode) with the given key.
182 The function relies on module :epkg:`pycrypto`, :epkg:`cryptography`,
183 algoritm `AES <https://fr.wikipedia.org/wiki/Advanced_Encryption_Standard>`_,
184 `Fernet <https://cryptography.io/en/latest/fernet/>`_.
186 @param key The encryption key - a string that must be
187 either 16, 24 or 32 bytes long. Longer keys
188 are more secure. If the data to encrypt is in bytes,
189 the key must be given in bytes too.
191 @param filename bytes or Name of the input file
192 @param out_filename if None, the returns bytes
194 @param chunksize Sets the size of the chunk which the function
195 uses to read and encrypt the file. Larger chunk
196 sizes can be faster for some files and machines.
197 chunksize must be divisible by 16.
199 @param algo AES (:epkg:`pycryptodomex`) of or fernet (cryptography)
201 @return filename or bytes
202 """
203 in_size, in_close, in_stream, out_close, out_return, out_stream = open_input_output(
204 filename, out_filename)
206 decryptor, origsize, chunksize = get_encryptor(
207 key, algo, in_stream=in_stream, chunksize=chunksize)
209 while True:
210 chunk = in_stream.read(chunksize)
211 if len(chunk) == 0:
212 break
213 out_stream.write(decryptor.decrypt(chunk))
214 out_stream.truncate(origsize)
216 return close_input_output(in_size, in_close, in_stream, out_close, out_return, out_stream)