blob: f95a3a54554048b4522dc5f7a99880d27340fd23 [file] [log] [blame]
Roman Okhrimenko977b3752022-03-31 14:40:48 +03001"""
2Copyright (c) 2021 Cypress Semiconductor Corporation
3
4Licensed under the Apache License, Version 2.0 (the "License");
5you may not use this file except in compliance with the License.
6You may obtain a copy of the License at
7
8 http://www.apache.org/licenses/LICENSE-2.0
9
10Unless required by applicable law or agreed to in writing, software
11distributed under the License is distributed on an "AS IS" BASIS,
12WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13See the License for the specific language governing permissions and
14limitations under the License.
15"""
16import os
17import struct
18from cryptography.hazmat.primitives.ciphers import (
19 Cipher, algorithms, modes
20)
21
22NONCE_SIZE = 12
23
24class EncryptorMXS40Sv2:
25 def __init__(self, key, nonce, initial_counter=0):
26 # with open(key_path, 'rb') as f:
27 # self.key = f.read()
28 self.key = key
29 self.nonce = nonce
30 self.counter = 0
31 self.initial_counter = initial_counter
32 cipher = Cipher(algorithms.AES(key), modes.ECB())
33 self.encryptor = cipher.encryptor()
34
35 def _load(self, image_path):
36 with open(image_path, 'rb') as f:
37 image = f.read()
38 return image
39
40 def _save(self, data, output_path):
41 with open(output_path, 'wb') as f:
42 f.write(data)
43
44 def update(self, data):
45 """
46 Encrypts a byte array using a customized AES-CTR mode
47 where a counter is incremented by 16 per block.
48 A nonce format is (128 bit):
49 bits 0...31 - counter + initial values
50 bits 32...127 - random nonce
51 """
52 chunk_size = 16
53 counter = self.counter
54 ciphertext = bytes()
55 for i in range(0, len(image), chunk_size):
56 indata = struct.pack('<I', initial_counter + counter) + nonce
57 counter += chunk_size
58 cipher_block = self.encryptor.update(indata)
59 chunk = image[i:i + chunk_size]
60 ciphertext += bytes(a ^ b for a, b in zip(chunk, cipher_block))
61 self.counter = counter
62 return ciphertext
63
64 def encrypt(self, image, initial_counter=0):
65 """
66 Encrypts a byte array using a customized AES-CTR mode
67 where a counter is incremented by 16 per block.
68 A nonce format is (128 bit):
69 bits 0...31 - counter + initial values
70 bits 32...127 - random nonce
71 """
72 chunk_size = 16
73 counter = 0
74 # self.initial_counter = initial_counter
75 # counter = 0 if initial_counter is None else int(initial_counter, 0)
76 # counter = initial_counter
77 ciphertext = bytes()
78 for i in range(0, len(image), chunk_size):
79 indata = struct.pack('<I', initial_counter + counter) + self.nonce[:12]
80 counter += chunk_size
81 cipher_block = self.encryptor.update(indata)
82 chunk = image[i:i + chunk_size]
83 ciphertext += bytes(a ^ b for a, b in zip(chunk, cipher_block))
84 self.encryptor.finalize()
85 # return ciphertext, nonce
86 return ciphertext
87
88 def encrypt_image(self, input_path, initial_counter=None, output_path=None, nonce_path=None):
89 """
90 Encrypts an image each time using new random nonce.
91 Saves the nonce and the encrypted image to specified locations.
92 If the output locations are not given the output files are saved
93 in the same location as the input image with predefined names.
94 """
95 image = self._load(input_path)
96
97 nonce = os.urandom(NONCE_SIZE)
98 init = 0 if initial_counter is None else int(initial_counter, 0)
99 ciphertext, nonce = self._encrypt(image, nonce, init)
100
101 if output_path is None:
102 output_path = '{0}_{2}{1}'.format(*os.path.splitext(input_path) + ('encrypted',))
103
104 if nonce_path is None:
105 nonce_path = '{0}_{2}{1}'.format(*os.path.splitext(input_path) + ('nonce',))
106
107 self._save(ciphertext, output_path)
108 self._save(nonce, nonce_path)
109
110 return output_path, nonce_path