block_template.c: Allow MODE_CTR to behave as a stream cipher

1 files changed, 55 insertions, 14 deletions

diff --git a/src/block_template.c b/src/block_template.c
index d1620fe..7548b21 100644
--- a/src/block_template.c
+++ b/src/block_template.c
@@ -220,7 +220,14 @@ ALGnew(PyObject *self, PyObject *args, PyObject *kwdict)
 	memset(new->oldCipher, 0, BLOCK_SIZE);
 	memcpy(new->IV, IV, IVlen);
 	new->mode = mode;
-	new->count=8;
+	switch(mode) {
+	case MODE_PGP:
+	    new->count=8;
+	    break;
+	case MODE_CTR:
+	default:
+	    new->count=BLOCK_SIZE;   /* stores how many bytes in new->oldCipher have been used */
+	}
 	return new;
 }
 
@@ -242,7 +249,8 @@ ALG_Encrypt(ALGobject *self, PyObject *args)
 		return PyString_FromStringAndSize(NULL, 0);
 	}
 	if ( (len % BLOCK_SIZE) !=0 && 
-	     (self->mode!=MODE_CFB) && (self->mode!=MODE_PGP))
+	     (self->mode!=MODE_CFB) && (self->mode!=MODE_PGP) &&
+	     (self->mode!=MODE_CTR))
 	{
 		PyErr_Format(PyExc_ValueError, 
 			     "Input strings must be "
@@ -358,8 +366,43 @@ ALG_Encrypt(ALGobject *self, PyObject *args)
 		break;
 
 	case(MODE_CTR):
-		for(i=0; i<len; i+=BLOCK_SIZE) 
-		{
+		/* CTR mode is a stream cipher whose keystream is generated by encrypting unique counter values.
+		 * - self->counter points to the Counter callable, which is
+		 *   responsible for generating keystream blocks
+		 * - self->count indicates the current offset within the current keystream block
+		 * - self->IV stores the current keystream block
+		 * - str stores the input string
+		 * - buffer stores the output string
+		 * - len indicates the length if the input and output strings
+		 * - i indicates the current offset within the input and output strings
+		 * - (len-i) is the number of bytes remaining to encrypt
+		 * - (BLOCK_SIZE-self->count) is the number of bytes remaining in the current keystream block
+		 */
+		i = 0;
+		while (i < len) {
+			/* If we don't need more than what remains of the current keystream block, then just XOR it in */
+			if (len-i <= BLOCK_SIZE-self->count) { /* remaining_bytes_to_encrypt <= remaining_bytes_in_IV */
+				/* XOR until the input is used up */
+				for(j=0; j<(len-i); j++) {
+					assert(i+j < len);
+					assert(self->count+j < BLOCK_SIZE);
+					buffer[i+j] = (self->IV[self->count+j] ^= str[i+j]);
+				}
+				self->count += len-i;
+				i = len;
+				continue;
+			}
+
+			/* Use up the current keystream block */
+			for(j=0; j<BLOCK_SIZE-self->count; j++) {
+				assert(i+j < len);
+				assert(self->count+j < BLOCK_SIZE);
+				buffer[i+j] = (self->IV[self->count+j] ^= str[i+j]);
+			}
+			i += BLOCK_SIZE-self->count;
+			self->count = BLOCK_SIZE;
+
+			/* Generate a new keystream block */
 			if (self->counter_shortcut) {
 				/* CTR mode shortcut: If we're using Util.Counter,
 				 * bypass the normal Python function call mechanism
@@ -376,9 +419,8 @@ ALG_Encrypt(ALGobject *self, PyObject *args)
 				}
 				block_encrypt(&(self->st),
 					      (unsigned char *)ctr->val,
-					      temp);
+					      self->IV);
 				ctr->inc_func(ctr);
-
 			} else {
 				PyObject *ctr = PyObject_CallObject(self->counter, NULL);
 				if (ctr == NULL) {
@@ -395,21 +437,20 @@ ALG_Encrypt(ALGobject *self, PyObject *args)
 				}
 				if (PyString_Size(ctr) != BLOCK_SIZE) {
 					PyErr_Format(PyExc_TypeError,
-						 "CTR counter function returned "
-						 "string not of length %i",
-						 BLOCK_SIZE);
+						     "CTR counter function returned "
+						     "string not of length %i",
+						     BLOCK_SIZE);
 					Py_DECREF(ctr);
 					free(buffer);
 					return NULL;
 				}
 				block_encrypt(&(self->st), (unsigned char *)PyString_AsString(ctr),
-					  temp);
+					      self->IV);
 				Py_DECREF(ctr);
 			}
-			for(j=0; j<BLOCK_SIZE; j++)
-			{
-				buffer[i+j] = str[i+j]^temp[j];
-			}
+
+			/* Move the pointer to the start of the keystream block */
+			self->count = 0;
 		}
 		break;