Forking what is left of ZeroNet and hopefully adding an AT Proto Frontend/Proxy
at main 101 lines 3.7 kB view raw
1import os 2import pyaes 3from .._aes import AES 4 5 6__all__ = ["aes"] 7 8class AESBackend: 9 def _get_algo_cipher_type(self, algo): 10 if not algo.startswith("aes-") or algo.count("-") != 2: 11 raise ValueError("Unknown cipher algorithm {}".format(algo)) 12 key_length, cipher_type = algo[4:].split("-") 13 if key_length not in ("128", "192", "256"): 14 raise ValueError("Unknown cipher algorithm {}".format(algo)) 15 if cipher_type not in ("cbc", "ctr", "cfb", "ofb"): 16 raise ValueError("Unknown cipher algorithm {}".format(algo)) 17 return cipher_type 18 19 20 def is_algo_supported(self, algo): 21 try: 22 self._get_algo_cipher_type(algo) 23 return True 24 except ValueError: 25 return False 26 27 28 def random(self, length): 29 return os.urandom(length) 30 31 32 def encrypt(self, data, key, algo="aes-256-cbc"): 33 cipher_type = self._get_algo_cipher_type(algo) 34 35 # Generate random IV 36 iv = os.urandom(16) 37 38 if cipher_type == "cbc": 39 cipher = pyaes.AESModeOfOperationCBC(key, iv=iv) 40 elif cipher_type == "ctr": 41 # The IV is actually a counter, not an IV but it does almost the 42 # same. Notice: pyaes always uses 1 as initial counter! Make sure 43 # not to call pyaes directly. 44 45 # We kinda do two conversions here: from byte array to int here, and 46 # from int to byte array in pyaes internals. It's possible to fix that 47 # but I didn't notice any performance changes so I'm keeping clean code. 48 iv_int = 0 49 for byte in iv: 50 iv_int = (iv_int * 256) + byte 51 counter = pyaes.Counter(iv_int) 52 cipher = pyaes.AESModeOfOperationCTR(key, counter=counter) 53 elif cipher_type == "cfb": 54 # Change segment size from default 8 bytes to 16 bytes for OpenSSL 55 # compatibility 56 cipher = pyaes.AESModeOfOperationCFB(key, iv, segment_size=16) 57 elif cipher_type == "ofb": 58 cipher = pyaes.AESModeOfOperationOFB(key, iv) 59 60 encrypter = pyaes.Encrypter(cipher) 61 ciphertext = encrypter.feed(data) 62 ciphertext += encrypter.feed() 63 return ciphertext, iv 64 65 66 def decrypt(self, ciphertext, iv, key, algo="aes-256-cbc"): 67 cipher_type = self._get_algo_cipher_type(algo) 68 69 if cipher_type == "cbc": 70 cipher = pyaes.AESModeOfOperationCBC(key, iv=iv) 71 elif cipher_type == "ctr": 72 # The IV is actually a counter, not an IV but it does almost the 73 # same. Notice: pyaes always uses 1 as initial counter! Make sure 74 # not to call pyaes directly. 75 76 # We kinda do two conversions here: from byte array to int here, and 77 # from int to byte array in pyaes internals. It's possible to fix that 78 # but I didn't notice any performance changes so I'm keeping clean code. 79 iv_int = 0 80 for byte in iv: 81 iv_int = (iv_int * 256) + byte 82 counter = pyaes.Counter(iv_int) 83 cipher = pyaes.AESModeOfOperationCTR(key, counter=counter) 84 elif cipher_type == "cfb": 85 # Change segment size from default 8 bytes to 16 bytes for OpenSSL 86 # compatibility 87 cipher = pyaes.AESModeOfOperationCFB(key, iv, segment_size=16) 88 elif cipher_type == "ofb": 89 cipher = pyaes.AESModeOfOperationOFB(key, iv) 90 91 decrypter = pyaes.Decrypter(cipher) 92 data = decrypter.feed(ciphertext) 93 data += decrypter.feed() 94 return data 95 96 97 def get_backend(self): 98 return "fallback" 99 100 101aes = AES(AESBackend())