From 6816f2e2ae3d484fbbea92675be0332099cd60c1 Mon Sep 17 00:00:00 2001
From: Dan <14043624+delivrance@users.noreply.github.com>
Date: Thu, 13 Jun 2019 15:15:31 +0200
Subject: [PATCH] Implement input validation checks Closes #6

---
 tgcrypto/tgcrypto.c | 74 +++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 71 insertions(+), 3 deletions(-)

diff --git a/tgcrypto/tgcrypto.c b/tgcrypto/tgcrypto.c
index d9a4edb..51092ad 100644
--- a/tgcrypto/tgcrypto.c
+++ b/tgcrypto/tgcrypto.c
@@ -28,7 +28,28 @@ static PyObject *ige(PyObject *args, uint8_t encrypt) {
     uint8_t *buf;
     PyObject *out;
 
-    PyArg_ParseTuple(args, "y*y*y*", &data, &key, &iv);
+    if (!PyArg_ParseTuple(args, "y*y*y*", &data, &key, &iv))
+        return NULL;
+
+    if (data.len == 0) {
+        PyErr_SetString(PyExc_ValueError, "Data must not be empty");
+        return NULL;
+    }
+
+    if (data.len % 16 != 0) {
+        PyErr_SetString(PyExc_ValueError, "Data size must match a multiple of 16 bytes");
+        return NULL;
+    }
+
+    if (key.len != 32) {
+        PyErr_SetString(PyExc_ValueError, "Key size must be exactly 32 bytes");
+        return NULL;
+    }
+
+    if (iv.len != 32) {
+        PyErr_SetString(PyExc_ValueError, "IV size must be exactly 32 bytes");
+        return NULL;
+    }
 
     buf = ige256(data.buf, data.len, key.buf, iv.buf, encrypt);
 
@@ -55,7 +76,33 @@ static PyObject *ctr256_encrypt(PyObject *self, PyObject *args) {
     uint8_t *buf;
     PyObject *out;
 
-    PyArg_ParseTuple(args, "y*y*y*y*", &data, &key, &iv, &state);
+    if (!PyArg_ParseTuple(args, "y*y*y*y*", &data, &key, &iv, &state))
+        return NULL;
+
+    if (data.len == 0) {
+        PyErr_SetString(PyExc_ValueError, "Data must not be empty");
+        return NULL;
+    }
+
+    if (key.len != 32) {
+        PyErr_SetString(PyExc_ValueError, "Key size must be exactly 32 bytes");
+        return NULL;
+    }
+
+    if (iv.len != 16) {
+        PyErr_SetString(PyExc_ValueError, "IV size must be exactly 16 bytes");
+        return NULL;
+    }
+
+    if (state.len != 1) {
+        PyErr_SetString(PyExc_ValueError, "State size must be exactly 1 byte");
+        return NULL;
+    }
+
+    if (*(uint8_t *) state.buf > 15) {
+        PyErr_SetString(PyExc_ValueError, "State value must be in the range [0, 15]");
+        return NULL;
+    }
 
     buf = ctr256(data.buf, data.len, key.buf, iv.buf, state.buf);
 
@@ -74,7 +121,28 @@ static PyObject *cbc(PyObject *args, uint8_t encrypt) {
     uint8_t *buf;
     PyObject *out;
 
-    PyArg_ParseTuple(args, "y*y*y*", &data, &key, &iv);
+    if (!PyArg_ParseTuple(args, "y*y*y*", &data, &key, &iv))
+        return NULL;
+
+    if (data.len == 0) {
+        PyErr_SetString(PyExc_ValueError, "Data must not be empty");
+        return NULL;
+    }
+
+    if (data.len % 16 != 0) {
+        PyErr_SetString(PyExc_ValueError, "Data size must match a multiple of 16 bytes");
+        return NULL;
+    }
+
+    if (key.len != 32) {
+        PyErr_SetString(PyExc_ValueError, "Key size must be exactly 32 bytes");
+        return NULL;
+    }
+
+    if (iv.len != 16) {
+        PyErr_SetString(PyExc_ValueError, "IV size must be exactly 16 bytes");
+        return NULL;
+    }
 
     buf = cbc256(data.buf, data.len, key.buf, iv.buf, encrypt);