Update to v1.1.1

README.rst

@@ -47,28 +47,51 @@ what you should do next:
 API
 ===
 
-TgCrypto API consists of these four functions:
+TgCrypto API consists of these four methods:
 
 .. code-block:: python
 
-    def ige_encrypt(data: bytes, key: bytes, iv: bytes) -> bytes:
+    def ige256_encrypt(data: bytes, key: bytes, iv: bytes) -> bytes:
 
 .. code-block:: python
 
-    def ige_decrypt(data: bytes, key: bytes, iv: bytes) -> bytes:
+    def ige256_decrypt(data: bytes, key: bytes, iv: bytes) -> bytes:
 
 .. code-block:: python
 
-    def ctr_encrypt(data: bytes, key: bytes, iv: bytes) -> bytes:
+    def ctr256_encrypt(data: bytes, key: bytes, iv: bytes, state: bytes) -> bytes:
 
 .. code-block:: python
 
-    def ctr_decrypt(data: bytes, key: bytes, iv: bytes) -> bytes:
+    def ctr256_decrypt(data: bytes, key: bytes, iv: bytes, state: bytes) -> bytes:
 
 Usage
 =====
 
-TgCrypto is as simple as this example:
+IGE Mode:
+---------
+
+**Note**: Data must be padded to match a multiple of the block size (16 bytes).
+
+.. code-block:: python
+
+    import os
+    import tgcrypto
+
+    data = os.urandom(10 * 1024 * 1024 + 7)  # 10 MB of random data + 7 bytes to show padding
+    key = os.urandom(32)  # Random Key
+    iv = os.urandom(32)  # Random IV
+
+    # Pad with zeroes: -7 % 16 = 9
+    data += bytes(-len(data) % 16)
+
+    ige_encrypted = tgcrypto.ige256_encrypt(data, key, iv)
+    ige_decrypted = tgcrypto.ige256_decrypt(ige_encrypted, key, iv)
+
+    print(data == ige_decrypted)  # True
+    
+CTR Mode (single chunk):
+------------------------
 
 .. code-block:: python
 
@@ -76,13 +99,61 @@ TgCrypto is as simple as this example:
     import tgcrypto
 
     data = os.urandom(10 * 1024 * 1024)  # 10 MB of random data
+    
     key = os.urandom(32)  # Random Key
-    iv = os.urandom(32)  # Random IV
 
-    ige_encrypted = tgcrypto.ige_encrypt(data, key, iv)
+    enc_iv = bytearray(os.urandom(16))  # Random IV
-    ige_decrypted = tgcrypto.ige_decrypt(ige_encrypted, key, iv)
+    dec_iv = enc_iv.copy()  # Keep a copy for decryption
 
-    assert data == ige_decrypted
+    ctr_encrypted = tgcrypto.ctr256_encrypt(data, key, enc_iv, bytes(1))
+    ctr_decrypted = tgcrypto.ctr256_decrypt(ctr_encrypted, key, dec_iv, bytes(1))
+
+    print(data == ctr_decrypted)  # True
+
+CTR Mode (stream):
+------------------
+
+.. code-block:: python
+
+    import os
+    import tgcrypto
+    from io import BytesIO
+
+    data = BytesIO(os.urandom(10 * 1024 * 1024))  # 10 MB of random data
+
+    key = os.urandom(32)  # Random Key
+
+    enc_iv = bytearray(os.urandom(16))  # Random IV
+    dec_iv = enc_iv.copy()  # Keep a copy for decryption
+
+    enc_state = bytes(1)  # Encryption state, starts from 0
+    dec_state = bytes(1)  # Decryption state, starts from 0
+
+    encrypted_data = BytesIO()  # Encrypted data buffer
+    decrypted_data = BytesIO()  # Decrypted data buffer
+
+    while True:
+        chunk = data.read(1024)
+
+        if not chunk:
+            break
+
+        # Write 1K encrypted bytes into the encrypted data buffer
+        encrypted_data.write(tgcrypto.ctr256_encrypt(chunk, key, enc_iv, enc_state))
+
+    # Reset position. We need to read it now
+    encrypted_data.seek(0)
+
+    while True:
+        chunk = encrypted_data.read(1024)
+
+        if not chunk:
+            break
+
+        # Write 1K decrypted bytes into the decrypted data buffer
+        decrypted_data.write(tgcrypto.ctr256_decrypt(chunk, key, dec_iv, dec_state))
+
+    print(data.getvalue() == decrypted_data.getvalue())  # True
 
 Contribution
 ============
@@ -130,8 +201,7 @@ License
 
     <h1 align="center">
         <a href="https://github.com/pyrogram/tgcrypto">
-            <div><img src="https://pyrogram.ml/images/icon.png" alt="Pyrogram Icon"></div>
+            <div><img src="https://raw.githubusercontent.com/pyrogram/logos/master/logos/tgcrypto_logo2.png" alt="TgCrypto Logo"></div>
-            <div><img src="https://pyrogram.ml/images/tgcrypto.png" alt="TgCrypto Logo"></div>
         </a>
     </h1>
 
@@ -151,13 +221,18 @@ License
             Community
         </a>
         <br><br>
+        <a href="https://github.com/pyrogram/pyrogram">
+            <img src="https://img.shields.io/badge/PYROGRAM-V0.8.0-eda738.svg?longCache=true&style=for-the-badge&colorA=262b30"
+                alt="TgCrypto">
+        </a>
+        
         <a href="https://github.com/pyrogram/tgcrypto">
-            <img src="https://www.pyrogram.ml/images/tgcrypto_version.svg"
+            <img src="https://img.shields.io/badge/TGCRYPTO-V1.1.1-eda738.svg?longCache=true&style=for-the-badge&colorA=262b30"
                 alt="TgCrypto">
         </a>
     </p>
 
-.. |logo| image:: https://pyrogram.ml/images/tgcrypto_logo.png
+.. |logo| image:: https://raw.githubusercontent.com/pyrogram/logos/master/logos/tgcrypto_logo2.png
     :target: https://github.com/pyrogram/tgcrypto
     :alt: TgCrypto
 

setup.py

@@ -27,7 +27,7 @@ with open("README.rst", encoding="utf-8") as f:
 
 setup(
     name="TgCrypto",
-    version="1.0.4",
+    version="1.1.1",
     description="Fast Telegram Crypto Library for Python",
     long_description=readme,
     url="https://github.com/pyrogram",
@@ -65,7 +65,7 @@ setup(
         "Source": "https://github.com/pyrogram/tgcrypto",
         "Documentation": "https://docs.pyrogram.ml",
     },
-    python_requires="~=3.3",
+    python_requires="~=3.4",
     packages=find_packages(),
     zip_safe=False,
     ext_modules=[

tgcrypto/aes256.h

@@ -18,17 +18,26 @@
 
 #include <stdint.h>
 #include <stdlib.h>
+#include <string.h>
 
 #ifndef AES256_H
 #define AES256_H
 
 #define AES_BLOCK_SIZE 16
-#define KEY_SCHEDULE_SIZE 60
+#define EXPANDED_KEY_SIZE 60
 
-void aes256_key_expansion(const uint8_t key[32], uint32_t w[60]);
+#define LROTL(x) (((x) << 8) | ((x) >> 24))
+#define LROTR(x) ((x) >> 8) | ((x) << 24))
+#define SWAP(x) ((LROTL((x)) & 0x00ff00ff) | ((LROTR((x)) & 0xff00ff00))
+#define GET(p) SWAP(*((uint32_t *)(p)))
+#define PUT(ct, st) {*((uint32_t *)(ct)) = SWAP((st));}
 
-void aes256_encrypt(const uint8_t in[16], uint8_t out[16], const uint32_t w[60]);
+void aes256_set_encryption_key(const uint8_t key[32], uint32_t expandedKey[60]);
 
-void aes256_decrypt(const uint8_t in[16], uint8_t out[16], const uint32_t w[60]);
+void aes256_set_decryption_key(const uint8_t key[32], uint32_t expandedKey[60]);
 
-#endif   // AES256_H
+void aes256_encrypt(const uint8_t in[16], uint8_t out[16], const uint32_t expandedKey[60]);
+
+void aes256_decrypt(const uint8_t in[16], uint8_t out[16], const uint32_t expandedKey[60]);
+
+#endif  // AES256_H

tgcrypto/ctr256.c

@@ -18,40 +18,34 @@
 
 #include "aes256.h"
 
-uint8_t *ctr256(const uint8_t in[], uint32_t length, const uint8_t key[32], const uint8_t iv[16]) {
+#define MIN(a, b) (((a) < (b)) ? (a) : (b))
+
+uint8_t *ctr256(const uint8_t in[], uint32_t length, const uint8_t key[32], uint8_t iv[16], uint8_t *state) {
     uint8_t *out = (uint8_t *) malloc(length * sizeof(uint8_t));
-    uint8_t iv_buf[AES_BLOCK_SIZE], out_buf[AES_BLOCK_SIZE];
+    uint8_t chunk[AES_BLOCK_SIZE];
-    uint32_t key_schedule[KEY_SCHEDULE_SIZE];
+    uint32_t expandedKey[EXPANDED_KEY_SIZE];
-    uint32_t i, j;
+    uint32_t i, j, k;
 
     memcpy(out, in, length);
-    memcpy(iv_buf, iv, AES_BLOCK_SIZE);
+    aes256_set_encryption_key(key, expandedKey);
-    aes256_key_expansion(key, key_schedule);
 
-    if (length > AES_BLOCK_SIZE)
+    aes256_encrypt(iv, chunk, expandedKey);
-        for (i = 0; i < length - AES_BLOCK_SIZE; i += AES_BLOCK_SIZE) {
-            aes256_encrypt(iv_buf, out_buf, key_schedule);
 
-            for (j = 0; j < AES_BLOCK_SIZE; ++j)
+    for (i = 0; i < length; i += AES_BLOCK_SIZE)
-                out[i + j] ^= out_buf[j];
+        for (j = 0; j < MIN(length - i, AES_BLOCK_SIZE); ++j) {
+            out[i + j] ^= chunk[(*state)++];
 
-            for (j = AES_BLOCK_SIZE - 1; j >= 0; --j)
+            if (*state >= AES_BLOCK_SIZE)
-                if (++iv_buf[j])
+                *state = 0;
-                    break;
+
+            if (*state == 0) {
+                for (k = AES_BLOCK_SIZE - 1; k >= 0; --k)
+                    if (++iv[k])
+                        break;
+
+                aes256_encrypt(iv, chunk, expandedKey);
+            }
         }
 
-    aes256_encrypt(iv_buf, out_buf, key_schedule);
-
-    for (j = 0; j < length - i; ++j)
-        out[i + j] ^= out_buf[j];
-
     return out;
 }
-
-uint8_t *ctr256_encrypt(uint8_t *in, uint32_t length, uint8_t *key, uint8_t *iv) {
-    return ctr256(in, length, key, iv);
-}
-
-uint8_t *ctr256_decrypt(uint8_t *in, uint32_t length, uint8_t *key, uint8_t *iv) {
-    return ctr256(in, length, key, iv);
-}

tgcrypto/ctr256.h

@@ -19,8 +19,6 @@
 #ifndef CTR256_H
 #define CTR256_H
 
-uint8_t *ctr256_encrypt(const uint8_t in[], uint32_t length, const uint8_t key[32], const uint8_t iv[16]);
+uint8_t *ctr256(const uint8_t in[], uint32_t length, const uint8_t key[32], const uint8_t iv[16], uint8_t *state);
 
-uint8_t *ctr256_decrypt(const uint8_t in[], uint32_t length, const uint8_t key[32], const uint8_t iv[16]);
+#endif  // CTR256_H
-
-#endif //IGE256_H

tgcrypto/ige256.c

@@ -20,37 +20,29 @@
 
 uint8_t *ige256(const uint8_t in[], uint32_t length, const uint8_t key[32], const uint8_t iv[32], uint8_t encrypt) {
     uint8_t *out = (uint8_t *) malloc(length * sizeof(uint8_t));
-    uint8_t iv_1[AES_BLOCK_SIZE], iv_2[AES_BLOCK_SIZE];
+    uint8_t iv1[AES_BLOCK_SIZE], iv2[AES_BLOCK_SIZE];
     uint8_t chunk[AES_BLOCK_SIZE], buffer[AES_BLOCK_SIZE];
-    uint32_t key_schedule[KEY_SCHEDULE_SIZE];
+    uint32_t expandedKey[EXPANDED_KEY_SIZE];
     uint32_t i, j;
 
-    memcpy(encrypt ? iv_1 : iv_2, (uint8_t *) iv, AES_BLOCK_SIZE);
+    memcpy(encrypt ? iv1 : iv2, (uint8_t *) iv, AES_BLOCK_SIZE);
-    memcpy(encrypt ? iv_2 : iv_1, (uint8_t *) iv + AES_BLOCK_SIZE, AES_BLOCK_SIZE);
+    memcpy(encrypt ? iv2 : iv1, (uint8_t *) iv + AES_BLOCK_SIZE, AES_BLOCK_SIZE);
-    aes256_key_expansion(key, key_schedule);
+    (encrypt ? aes256_set_encryption_key : aes256_set_decryption_key)(key, expandedKey);
 
     for (i = 0; i < length; i += AES_BLOCK_SIZE) {
         memcpy(chunk, &in[i], AES_BLOCK_SIZE);
 
         for (j = 0; j < AES_BLOCK_SIZE; ++j)
-            buffer[j] = in[i + j] ^ iv_1[j];
+            buffer[j] = in[i + j] ^ iv1[j];
 
-        (encrypt ? aes256_encrypt : aes256_decrypt)((uint8_t *) &buffer, &out[i], key_schedule);
+        (encrypt ? aes256_encrypt : aes256_decrypt)((uint8_t *) &buffer, &out[i], expandedKey);
 
         for (j = 0; j < AES_BLOCK_SIZE; ++j)
-            out[i + j] ^= iv_2[j];
+            out[i + j] ^= iv2[j];
 
-        memcpy(iv_1, &out[i], AES_BLOCK_SIZE);
+        memcpy(iv1, &out[i], AES_BLOCK_SIZE);
-        memcpy(iv_2, chunk, AES_BLOCK_SIZE);
+        memcpy(iv2, chunk, AES_BLOCK_SIZE);
     }
 
     return out;
 }
-
-uint8_t *ige256_encrypt(const uint8_t in[], uint32_t length, const uint8_t key[32], const uint8_t iv[32]) {
-    return ige256(in, length, key, iv, 1);
-}
-
-uint8_t *ige256_decrypt(const uint8_t in[], uint32_t length, const uint8_t key[32], const uint8_t iv[32]) {
-    return ige256(in, length, key, iv, 0);
-}

tgcrypto/ige256.h

@@ -19,8 +19,6 @@
 #ifndef IGE256_H
 #define IGE256_H
 
-uint8_t *ige256_encrypt(const uint8_t in[], uint32_t length, const uint8_t key[32], const uint8_t iv[32]);
+uint8_t *ige256(const uint8_t in[], uint32_t length, const uint8_t key[32], const uint8_t iv[32], uint8_t encrypt);
 
-uint8_t *ige256_decrypt(const uint8_t in[], uint32_t length, const uint8_t key[32], const uint8_t iv[32]);
+#endif  // IGE256_H
-
-#endif //IGE256_H

tgcrypto/tgcrypto.c

@@ -17,18 +17,19 @@
 // along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
 
 #include <Python.h>
+
 #include "aes256.h"
 #include "ige256.h"
 #include "ctr256.h"
 
-static PyObject* tgcrypto(PyObject *args, uint8_t mode, uint8_t encrypt) {
+static PyObject* ige(PyObject *args, uint8_t encrypt) {
     Py_buffer data, key, iv;
-    uint8_t *(*fn)(), *buf;
+    uint8_t *buf;
-    PyObject* out;
+    PyObject *out;
 
     PyArg_ParseTuple(args, "y*y*y*", &data, &key, &iv);
-    fn = encrypt? mode? ctr256_encrypt: ige256_encrypt: mode? ctr256_decrypt: ige256_decrypt;
+
-    buf = fn(data.buf, data.len, key.buf, iv.buf);
+    buf = ige256(data.buf, data.len, key.buf, iv.buf, encrypt);
 
     PyBuffer_Release(&data);
     PyBuffer_Release(&key);
@@ -40,30 +41,39 @@ static PyObject* tgcrypto(PyObject *args, uint8_t mode, uint8_t encrypt) {
     return out;
 }
 
-static PyObject* ige_encrypt(PyObject *self, PyObject *args) {
+static PyObject *ige256_encrypt(PyObject *self, PyObject *args) {
-    return tgcrypto(args, 0, 1);
+    return ige(args, 1);
 }
 
-
+static PyObject *ige256_decrypt(PyObject *self, PyObject *args) {
-static PyObject* ige_decrypt(PyObject *self, PyObject *args) {
+    return ige(args, 0);
-    return tgcrypto(args, 0, 0);
 }
 
-static PyObject* ctr_encrypt(PyObject *self, PyObject *args) {
+static PyObject *ctr256_encrypt(PyObject *self, PyObject *args) {
-    return tgcrypto(args, 1, 1);
+    Py_buffer data, key, iv, state;
-}
+    uint8_t *buf;
+    PyObject *out;
 
+    PyArg_ParseTuple(args, "y*y*y*y*", &data, &key, &iv, &state);
 
-static PyObject* ctr_decrypt(PyObject *self, PyObject *args) {
+    buf = ctr256(data.buf, data.len, key.buf, iv.buf, state.buf);
-    return tgcrypto(args, 1, 0);
+
+    PyBuffer_Release(&data);
+    PyBuffer_Release(&key);
+    PyBuffer_Release(&iv);
+
+    out = Py_BuildValue("y#", buf, data.len);
+    free(buf);
+
+    return out;
 }
 
 static PyMethodDef methods[] = {
-    {"ige_encrypt", (PyCFunction) ige_encrypt, METH_VARARGS, "AES-256-IGE Encryption"},
+    {"ige256_encrypt", (PyCFunction) ige256_encrypt, METH_VARARGS, "AES-256-IGE Encryption"},
-    {"ige_decrypt", (PyCFunction) ige_decrypt, METH_VARARGS, "AES-256-IGE Decryption"},
+    {"ige256_decrypt", (PyCFunction) ige256_decrypt, METH_VARARGS, "AES-256-IGE Decryption"},
-    {"ctr_encrypt", (PyCFunction) ctr_encrypt, METH_VARARGS, "AES-256-CTR Encryption"},
+    {"ctr256_encrypt", (PyCFunction) ctr256_encrypt, METH_VARARGS, "AES-256-CTR Encryption"},
-    {"ctr_decrypt", (PyCFunction) ctr_decrypt, METH_VARARGS, "AES-256-CTR Decryption"},
+    {"ctr256_decrypt", (PyCFunction) ctr256_encrypt, METH_VARARGS, "AES-256-CTR Decryption"},
-    {NULL,      NULL,                      0,            NULL}
+    {NULL, NULL, 0, NULL}
 };
 
 static struct PyModuleDef module = {