diff --git a/telethon/crypto/aes_ctr.py b/telethon/crypto/aes_ctr.py index 7bd7b79a..711f093e 100644 --- a/telethon/crypto/aes_ctr.py +++ b/telethon/crypto/aes_ctr.py @@ -1,21 +1,91 @@ -import pyaes +from ctypes.util import find_library +lib = find_library('ssl') + +if not lib: + # libssl.so not found; use pyaes + import pyaes + + class AESModeCTR: + """Wrapper around pyaes.AESModeOfOperationCTR mode with custom IV""" + # TODO Maybe make a pull request to pyaes to support iv on CTR + + def __init__(self, key, iv): + # TODO Use libssl if available + assert isinstance(key, bytes) + self._aes = pyaes.AESModeOfOperationCTR(key) + + assert isinstance(iv, bytes) + assert len(iv) == 16 + self._aes._counter._counter = list(iv) + + def encrypt(self, data): + return self._aes.encrypt(data) + + def decrypt(self, data): + return self._aes.decrypt(data) + +else: + + _libssl = ctypes.cdll.LoadLibrary(lib) + + AES_MAXNR = 14 + class AES_KEY(ctypes.Structure): + _fields_ = [ + #uint32_t rd_key[4 * (AES_MAXNR + 1)]; + ('rd_key', ctypes.c_uint32 * (4*(AES_MAXNR + 1))), + ('rounds', ctypes.c_uint), + ] -class AESModeCTR: - """Wrapper around pyaes.AESModeOfOperationCTR mode with custom IV""" - # TODO Maybe make a pull request to pyaes to support iv on CTR + class AESModeCTR: + """Wrapper around pyaes.AESModeOfOperationCTR mode with custom IV""" - def __init__(self, key, iv): - # TODO Use libssl if available - assert isinstance(key, bytes) - self._aes = pyaes.AESModeOfOperationCTR(key) + def __init__(self, key, iv): + assert isinstance(key, bytes) + self.aeskey = AES_KEY() + self.liv = len(iv) + self.civ = (ctypes.c_ubyte * self.liv)(*iv) + ckey = (ctypes.c_ubyte * len(key))(*key) + cklen = ctypes.c_int(len(key)*8) + _libssl.AES_set_encrypt_key(ckey, cklen, ctypes.byref(self.aeskey)) - assert isinstance(iv, bytes) - assert len(iv) == 16 - self._aes._counter._counter = list(iv) + def encrypt(self, data): - def encrypt(self, data): - return self._aes.encrypt(data) + cin = (ctypes.c_ubyte * len(plain))(*data) + ctlen = ctypes.c_size_t(len(data)) + cout = (ctypes.c_ubyte * len(data))() + ccount = (ctypes.c_ubyte * self.liv)(*(b'\x00'*self.liv)) + cnum = ctypes.c_uint(0) - def decrypt(self, data): - return self._aes.decrypt(data) + _libssl.AES_ctr128_encrypt( + ctypes.byref(cin), ctypes.byref(cout), + ctlen, ctypes.byref(aeskey), + self.civ, + ccount, + ctypes.byref(cnum), + ) + + return bytes(cout) + + def decrypt(self, data): + # AES-CTR is symetric + return self.encrypt(data) + +""" +# define AES_MAXNR 14 +struct aes_key_st { +# ifdef AES_LONG + unsigned long rd_key[4 * (AES_MAXNR + 1)]; +# else + unsigned int rd_key[4 * (AES_MAXNR + 1)]; +# endif + int rounds; +}; +typedef struct aes_key_st AES_KEY; + +void AES_ctr128_encrypt(const unsigned char *in, unsigned char *out, + size_t length, const AES_KEY *key, + unsigned char ivec[AES_BLOCK_SIZE], + unsigned char ecount_buf[AES_BLOCK_SIZE], + unsigned int *num); +"""