Hide keyboard shortcuts

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. 

4 

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 

13 

14 

15class EncryptionError(Exception): 

16 """ 

17 raised if an issue happen during encryption 

18 """ 

19 pass 

20 

21 

22def open_input_output(filename, out_filename=None): 

23 """ 

24 Converts *filename* and *out_filename* as streams. 

25 

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) 

46 

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 

60 

61 return filesize, close, st, out_close, ret, sto 

62 

63 

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. 

68 

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() 

79 

80 if out_close: 

81 if out_return: 

82 raise EncryptionError("incompability") 

83 out_stream.close() 

84 

85 if out_return: 

86 return out_stream.getvalue() 

87 else: 

88 return None 

89 

90 

91def get_encryptor(key, algo="AES", chunksize=2 ** 24, **params): 

92 """ 

93 Returns an encryptor with method encrypt and decrypt. 

94 

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 

134 

135 

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/>`_. 

142 

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. 

147 

148 @param filename bytes or Name of the input file 

149 @param out_filename if None, the returns bytes 

150 

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. 

155 

156 @param algo AES (PyCryptodomex) of or fernet (cryptography) 

157 

158 @return filename or bytes 

159 """ 

160 

161 in_size, in_close, in_stream, out_close, out_return, out_stream = open_input_output( 

162 filename, out_filename) 

163 

164 encryptor, origsize, chunksize = get_encryptor( 

165 key, algo, out_stream=out_stream, in_size=in_size, chunksize=chunksize) 

166 

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) 

173 

174 out_stream.write(encryptor.encrypt(chunk)) 

175 

176 return close_input_output(in_size, in_close, in_stream, out_close, out_return, out_stream) 

177 

178 

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/>`_. 

185 

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. 

190 

191 @param filename bytes or Name of the input file 

192 @param out_filename if None, the returns bytes 

193 

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. 

198 

199 @param algo AES (:epkg:`pycryptodomex`) of or fernet (cryptography) 

200 

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) 

205 

206 decryptor, origsize, chunksize = get_encryptor( 

207 key, algo, in_stream=in_stream, chunksize=chunksize) 

208 

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) 

215 

216 return close_input_output(in_size, in_close, in_stream, out_close, out_return, out_stream)