make code python3 compatible

.gitignore

@@ -1,9 +1,12 @@
-_chacha20
+
 .old
 *.o
 *.so
-src/test*
+*.pyc
+
+__pycache__
 *.egg-info
 build
-*.pyc
+dist
+
 Makefile

.travis.yml

@@ -1,6 +1,10 @@
 language: python
 python:
   - 2.7
+  - 3.4
+  - 3.5
+  - 3.6
+  - 3.7
 before_install:
   - sudo apt-get install python-dev
 install:

README.rst

@@ -42,4 +42,14 @@ by itself:
     # decryption (which yields plaintext == message)
     plaintext = aead.verify_and_decrypt(key, nonce, ciphertext, mac, additional_data)
 
+Notes on Python 2 vs 3
+----------------------
+
+In python2 encryption and decryption and tagging will return :code:`str` data while in python3 they will return
+:code:`bytes` data. This is consistent with how much of the python library operates between the two versions (e.g.
+see :code:`binascii.unhexlify`). This can lead to some strange behavior if e.g. you encrypt a :code:`str` value in
+python3 and, after decrypting, your decrypted value does not match your original value because you got :code:`bytes`
+back from the decryption. If the returned type is undesirable it is of course always possible to convert between
+:code:`bytes` and :code:`str` as needed.
+
 .. _RFC7539: https://tools.ietf.org/html/rfc7539

rfc7539/aead.py

@@ -1,34 +1,43 @@
 from hmac import compare_digest  # constant time comparison
+from logging import getLogger
+from struct import pack
 
-from cipher import encrypt
-from mac import tag
+from .cipher import encrypt
+from .mac import tag
+from .util import force_bytes
+
+
+class BadTagException(Exception):
+    def __init__(self, message):
+        super(BadTagException, self).__init__(message)
 
 
 def _len_bytes(n):
     # truncate to 32 bits
     n &= 0xffffffff
-    slen = ''
+    slen = b''
 
     while n:
-        slen = slen + chr(n & 0xff)
+        slen = slen + pack('B', n & 0xff)
         n >>= 8
 
     if len(slen) != 8:
-        slen = slen + (8 - len(slen)) * chr(0x00)
+        slen = slen + (8 - len(slen)) * b'\x00'
 
     return slen
 
 
 def _tag_data(aad, ciphertext):
-    tag_data = aad + chr(0x00) * (16 - (len(aad) % 16))
-    tag_data += ciphertext + chr(0x00) * (16 - (len(ciphertext) % 16))
+    tag_data = aad + b'\x00' * (16 - (len(aad) % 16))
+    tag_data += ciphertext + b'\x00' * (16 - (len(ciphertext) % 16))
     tag_data += _len_bytes(len(aad))
     tag_data += _len_bytes(len(ciphertext))
     return tag_data
 
 
 def encrypt_and_tag(key, nonce, plaintext, aad):
-    tag_key = encrypt(key, nonce, chr(0x00) * 64)
+    aad = force_bytes(aad)
+    tag_key = encrypt(key, nonce, b'\x00' * 64)
     tag_key = tag_key[:32]
     ciphertext = encrypt(key, nonce, plaintext, counter=1)
     tag_input = _tag_data(aad, ciphertext)
@@ -37,12 +46,12 @@ def encrypt_and_tag(key, nonce, plaintext, aad):
 
 
 def verify_and_decrypt(key, nonce, ciphertext, mac, aad):
-    tag_key = encrypt(key, nonce, chr(0x00) * 64)
+    aad = force_bytes(aad)
+    tag_key = encrypt(key, nonce, b'\x00' * 64)
     tag_key = tag_key[:32]
     tag_input = _tag_data(aad, ciphertext)
 
     if not compare_digest(tag(tag_key, tag_input), mac):
-        print 'Got a bad tag, aborting decryption process'
-        return None
+        raise BadTagException('Got a bad tag, aborting decryption process')
 
     return encrypt(key, nonce, ciphertext, counter=1)

rfc7539/cipher.py

@@ -1,5 +1,6 @@
+from .util import force_bytes
 from rfc7539 import _chacha20
 
 
 def encrypt(key, nonce, data, counter=0):
-    return _chacha20.cipher(bytearray(key), bytearray(nonce), bytearray(data), len(data), counter)
+    return _chacha20.cipher(force_bytes(key), force_bytes(nonce), force_bytes(data), len(data), counter)

rfc7539/mac.py

@@ -1,5 +1,6 @@
+from .util import force_bytes
 from rfc7539 import _poly1305
 
 
 def tag(key, msg):
-    return _poly1305.tag(bytearray(key), bytearray(msg), len(msg))
+    return _poly1305.tag(force_bytes(key), force_bytes(msg), len(msg))

rfc7539/test.py

@@ -7,9 +7,9 @@
 class TestChaCha20Keystream(unittest.TestCase):
     # test from https://tools.ietf.org/html/rfc7539#appendix-A.1
     def test_chacha20_keystream1(self):
-        key = chr(0x00) * 32
-        nonce = chr(0x00) * 12
-        data = chr(0x00) * 64
+        key = b'\x00' * 32
+        nonce = b'\x00' * 12
+        data = b'\x00' * 64
         expected = unhexlify(
             '76b8e0ada0f13d90405d6ae55386bd28bdd219b8a08ded1aa836efcc'
             '8b770dc7da41597c5157488d7724e03fb8d84a376a43b8f41518a11c'
@@ -20,9 +20,9 @@ def test_chacha20_keystream1(self):
         self.assertEqual(keystream, expected)
 
     def test_chacha20_keystream2(self):
-        key = chr(0x00) * 32
-        nonce = chr(0x00) * 12
-        data = chr(0x00) * 64
+        key = b'\x00' * 32
+        nonce = b'\x00' * 12
+        data = b'\x00' * 64
         expected = unhexlify(
             '9f07e7be5551387a98ba977c732d080d'
             'cb0f29a048e3656912c6533e32ee7aed'
@@ -34,9 +34,9 @@ def test_chacha20_keystream2(self):
         self.assertEqual(keystream, expected)
 
     def test_chacha20_keystream3(self):
-        key = chr(0x00) * 31 + chr(0x01)
-        nonce = chr(0x00) * 12
-        data = chr(0x00) * 64
+        key = b'\x00' * 31 + b'\x01'
+        nonce = b'\x00' * 12
+        data = b'\x00' * 64
         expected = unhexlify(
             '3aeb5224ecf849929b9d828db1ced4dd'
             '832025e8018b8160b82284f3c949aa5a'
@@ -48,9 +48,9 @@ def test_chacha20_keystream3(self):
         self.assertEqual(keystream, expected)
 
     def test_chacha20_keystream4(self):
-        key = chr(0x00) + chr(0xff) + chr(0x00) * 30
-        nonce = chr(0x00) * 12
-        data = chr(0x00) * 64
+        key = b'\x00' + b'\xff' + b'\x00' * 30
+        nonce = b'\x00' * 12
+        data = b'\x00' * 64
         expected = unhexlify(
             '72d54dfbf12ec44b362692df94137f32'
             '8fea8da73990265ec1bbbea1ae9af0ca'
@@ -62,9 +62,9 @@ def test_chacha20_keystream4(self):
         self.assertEqual(keystream, expected)
 
     def test_chacha20_keystream5(self):
-        key = chr(0x00) * 32
-        nonce = chr(0x00) * 11 + chr(0x02)
-        data = chr(0x00) * 64
+        key = b'\x00' * 32
+        nonce = b'\x00' * 11 + b'\x02'
+        data = b'\x00' * 64
         expected = unhexlify(
             'c2c64d378cd536374ae204b9ef933fcd'
             '1a8b2288b3dfa49672ab765b54ee27c7'
@@ -79,9 +79,9 @@ def test_chacha20_keystream5(self):
 class TestChaCha20Encryption(unittest.TestCase):
     # test from https://tools.ietf.org/html/rfc7539#appendix-A.2
     def test_chacha20_encryption1(self):
-        key = chr(0x00) * 32
-        nonce = chr(0x00) * 12
-        data = chr(0x00) * 64
+        key = b'\x00' * 32
+        nonce = b'\x00' * 12
+        data = b'\x00' * 64
         expected = unhexlify(
             '76b8e0ada0f13d90405d6ae55386bd28bdd219b8a08ded1aa836efcc'
             '8b770dc7da41597c5157488d7724e03fb8d84a376a43b8f41518a11c'
@@ -92,8 +92,8 @@ def test_chacha20_encryption1(self):
         self.assertEqual(keystream, expected)
 
     def test_chacha20_keystream2(self):
-        key = chr(0x00) * 31 + chr(0x01)
-        nonce = chr(0x00) * 11 + chr(0x02)
+        key = b'\x00' * 31 + b'\x01'
+        nonce = b'\x00' * 11 + b'\x02'
         data = 'Any submission to the IETF intended by the Contributor for publication as all ' \
                'or part of an IETF Internet-Draft or RFC and any statement made within the ' \
                'context of an IETF activity is considered an "IETF Contribution". Such ' \
@@ -134,7 +134,7 @@ def test_chacha20_encryption3(self):
             '1c9240a5eb55d38af333888604f6b5f0'
             '473917c1402b80099dca5cbc207075c0'
         )
-        nonce = chr(0x00) * 11 + chr(0x02)
+        nonce = b'\x00' * 11 + b'\x02'
         data = '\'Twas brillig, and the slithy toves\x0aDid gyre and gimble in the wabe:\x0a' \
                'All mimsy were the borogoves,\x0aAnd the mome raths outgrabe.'
         expected = unhexlify(
@@ -155,9 +155,9 @@ def test_chacha20_encryption3(self):
 class TestPoly1305(unittest.TestCase):
     # https://tools.ietf.org/html/rfc7539#appendix-A.3
     def test_poly1305_tag1(self):
-        inp = chr(0x00) * 64
-        key = chr(0x00) * 32
-        expected = chr(0x00) * 16
+        inp = b'\x00' * 64
+        key = b'\x00' * 32
+        expected = b'\x00' * 16
 
         tag = mac.tag(key, inp)
         self.assertEqual(tag, expected)
@@ -208,39 +208,39 @@ def test_poly1305_tag4(self):
 class TestPoly1305Keygen(unittest.TestCase):
     # test from https://tools.ietf.org/html/rfc7539#appendix-A.4
     def test_poly1305_keygen1(self):
-        key = chr(0x00) * 32
-        nonce = chr(0x00) * 12
+        key = b'\x00' * 32
+        nonce = b'\x00' * 12
         expected = unhexlify(
             '76b8e0ada0f13d90405d6ae55386bd28'
             'bdd219b8a08ded1aa836efcc8b770dc7'
         )
 
-        tag_key = cipher.encrypt(key, nonce, chr(0x00) * 64)[:32]
+        tag_key = cipher.encrypt(key, nonce, b'\x00' * 64)[:32]
         self.assertEqual(tag_key, expected)
 
     def test_poly1305_keygen2(self):
-        key = chr(0x00) * 31 + chr(0x01)
-        nonce = chr(0x00) * 11 + chr(0x02)
+        key = b'\x00' * 31 + b'\x01'
+        nonce = b'\x00' * 11 + b'\x02'
         expected = unhexlify(
             'ecfa254f845f647473d3cb140da9e876'
             '06cb33066c447b87bc2666dde3fbb739'
         )
 
-        tag_key = cipher.encrypt(key, nonce, chr(0x00) * 64)[:32]
+        tag_key = cipher.encrypt(key, nonce, b'\x00' * 64)[:32]
         self.assertEqual(tag_key, expected)
 
     def test_poly1305_keygen3(self):
         key = unhexlify(
             '1c9240a5eb55d38af333888604f6b5f0'
             '473917c1402b80099dca5cbc207075c0'
         )
-        nonce = chr(0x00) * 11 + chr(0x02)
+        nonce = b'\x00' * 11 + b'\x02'
         expected = unhexlify(
             '965e3bc6f9ec7ed9560808f4d229f94b'
             '137ff275ca9b3fcbdd59deaad23310ae'
         )
 
-        tag_key = cipher.encrypt(key, nonce, chr(0x00) * 64)[:32]
+        tag_key = cipher.encrypt(key, nonce, b'\x00' * 64)[:32]
         self.assertEqual(tag_key, expected)
 
 

rfc7539/util.py

@@ -0,0 +1,7 @@
+def force_bytes(data):
+    if isinstance(data, bytearray):
+        return data
+    elif not isinstance(data, bytes):
+        return bytearray(data, 'utf8')
+    else:
+        return bytearray(data)

setup.py

@@ -32,7 +32,7 @@ def run(self):
 
 setup(
     name='rfc7539',
-    version='1.0.0',
+    version='1.1.0',
     author='Anton Kueltz',
     author_email='[email protected]',
     license='GNU General Public License v3 (GPLv3)',
@@ -44,11 +44,16 @@ def run(self):
     ext_modules=[_chacha20, _poly1305],
     cmdclass={'test': TestCommand},
     classifiers=[
-        'Development Status :: 3 - Alpha',
+        'Development Status :: 4 - Beta',
         'Intended Audience :: Developers',
         'Topic :: Security :: Cryptography',
+        'Operating System :: POSIX :: Linux',
+        'Operating System :: MacOS :: MacOS X',
         'License :: OSI Approved :: GNU General Public License v3 (GPLv3)',
-        'Programming Language :: Python :: 2',
-        # 'Programming Language :: Python :: 3',
+        'Programming Language :: Python :: 2.7',
+        'Programming Language :: Python :: 3.4',
+        'Programming Language :: Python :: 3.5',
+        'Programming Language :: Python :: 3.6',
+        'Programming Language :: Python :: 3.7',
     ],
 )

src/_chacha20.c

@@ -130,14 +130,40 @@ static PyObject * _chacha20_cipher(PyObject *self, PyObject *args) {
     ChaCha20XOR(out, (unsigned char *)in, msgLen, (unsigned char *)key, (unsigned char *)nonce,
         counter);
 
+#if PY_MAJOR_VERSION >= 3
+    return Py_BuildValue("y#", out, msgLen);
+#else
     return Py_BuildValue("s#", out, msgLen);
+#endif
 }
 
 static PyMethodDef _chacha20__methods__[] = {
     {"cipher", _chacha20_cipher, METH_VARARGS, "Encrypt / Decrypt data via"},
     {NULL, NULL, 0, NULL}        /* Sentinel */
 };
 
+
+#if PY_MAJOR_VERSION >= 3
+static struct PyModuleDef moduledef = {
+    PyModuleDef_HEAD_INIT,
+    "_chacha20",
+    NULL,
+    -1,
+    _chacha20__methods__,
+    NULL,
+    NULL,
+    NULL,
+    NULL,
+};
+
+PyMODINIT_FUNC PyInit__chacha20(void) {
+    PyObject * m = PyModule_Create(&moduledef);
+    return m;
+}
+
+
+#else
 PyMODINIT_FUNC init_chacha20(void) {
     Py_InitModule("_chacha20", _chacha20__methods__);
 }
+#endif