Forking what is left of ZeroNet and hopefully adding an AT Proto Frontend/Proxy
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())