Merge pull request #2 from CovenantSQL/feature/e2ee

Add end to end encryption

Author: laodouya

.gitignore

@@ -1,2 +1,4 @@
 .vscode
+.idea
+venv
 *.swp

pycovenantsql/__init__.py

@@ -1,4 +1,3 @@
-import sys
 from ._compat import PY2
 from .converters import escape_dict, escape_sequence, escape_string
 from .constants import FIELD_TYPE
@@ -38,18 +37,18 @@ def __hash__(self):
 
 # TODO it's in pep249 find out meaning and usage of this
 # https://www.python.org/dev/peps/pep-0249/#string
-STRING    = DBAPISet([FIELD_TYPE.ENUM, FIELD_TYPE.STRING,
-                     FIELD_TYPE.VAR_STRING])
-BINARY    = DBAPISet([FIELD_TYPE.BLOB, FIELD_TYPE.LONG_BLOB,
-                     FIELD_TYPE.MEDIUM_BLOB, FIELD_TYPE.TINY_BLOB])
-NUMBER    = DBAPISet([FIELD_TYPE.DECIMAL, FIELD_TYPE.DOUBLE, FIELD_TYPE.FLOAT,
-                     FIELD_TYPE.INT24, FIELD_TYPE.LONG, FIELD_TYPE.LONGLONG,
-                     FIELD_TYPE.TINY, FIELD_TYPE.YEAR])
-DATE      = DBAPISet([FIELD_TYPE.DATE, FIELD_TYPE.NEWDATE])
-TIME      = DBAPISet([FIELD_TYPE.TIME])
+STRING = DBAPISet([FIELD_TYPE.ENUM, FIELD_TYPE.STRING,
+                   FIELD_TYPE.VAR_STRING])
+BINARY = DBAPISet([FIELD_TYPE.BLOB, FIELD_TYPE.LONG_BLOB,
+                   FIELD_TYPE.MEDIUM_BLOB, FIELD_TYPE.TINY_BLOB])
+NUMBER = DBAPISet([FIELD_TYPE.DECIMAL, FIELD_TYPE.DOUBLE, FIELD_TYPE.FLOAT,
+                   FIELD_TYPE.INT24, FIELD_TYPE.LONG, FIELD_TYPE.LONGLONG,
+                   FIELD_TYPE.TINY, FIELD_TYPE.YEAR])
+DATE = DBAPISet([FIELD_TYPE.DATE, FIELD_TYPE.NEWDATE])
+TIME = DBAPISet([FIELD_TYPE.TIME])
 TIMESTAMP = DBAPISet([FIELD_TYPE.TIMESTAMP, FIELD_TYPE.DATETIME])
-DATETIME  = TIMESTAMP
-ROWID     = DBAPISet()
+DATETIME = TIMESTAMP
+ROWID = DBAPISet()
 
 
 def Binary(x):
@@ -59,6 +58,7 @@ def Binary(x):
     else:
         return bytes(x)
 
+
 def Connect(*args, **kwargs):
     """
     Connect to the database; see connections.Connection.__init__() for
@@ -69,6 +69,7 @@ def Connect(*args, **kwargs):
 
 
 from . import connections as _orig_conn
+
 if _orig_conn.Connection.__init__.__doc__ is not None:
     Connect.__doc__ = _orig_conn.Connection.__init__.__doc__
 del _orig_conn
@@ -80,6 +81,7 @@ def get_client_info():  # for MySQLdb compatibility
         version = VERSION[:3]
     return '.'.join(map(str, version))
 
+
 connect = Connection = Connect
 
 NULL = "NULL"

pycovenantsql/e2ee.py

@@ -0,0 +1,74 @@
+from Crypto.Cipher import AES
+from Crypto import Random
+import hashlib
+from binascii import hexlify, unhexlify
+
+BLOCK_SIZE = AES.block_size  # Bytes
+
+salt = unhexlify("3fb8877d37fdc04e4a4765EFb8ab7d36")
+
+
+class PaddingError(Exception):
+    """Exception raised for errors in the padding.
+
+    Attributes:
+        message -- explanation of the error
+    """
+
+    def __init__(self, message):
+        self.message = message
+
+
+pad = lambda s: s + ((BLOCK_SIZE - len(s) % BLOCK_SIZE) *
+                     chr(BLOCK_SIZE - len(s) % BLOCK_SIZE)).encode('ascii')
+
+
+def unpad(s):
+    in_len = len(s)
+    if in_len == 0:
+        raise PaddingError("empty input")
+    pad_char = s[-1]
+    if pad_char > BLOCK_SIZE:
+        raise PaddingError("padding length > 16")
+    for i in s[in_len - pad_char:]:
+        if i != pad_char:
+            raise PaddingError("unexpected padding char")
+    return s[:-pad_char]
+
+
+# kdf does 2 times sha256 and takes the first 16 bytes
+def kdf(raw_key):
+    """
+    kdf does 2 times sha256 and takes the first 16 bytes
+    :param raw_key:
+    :return:
+    """
+    return hashlib.sha256(hashlib.sha256(raw_key + salt).digest()).digest()[:16]
+
+
+def encrypt(raw, password):
+    """
+    encrypt encrypts data with given password by AES-128-CBC PKCS#7, iv will be placed
+    at head of cipher data.
+
+    :param raw: input raw byte array
+    :param password: password byte array
+    :return: encrypted byte array
+    """
+    iv = Random.new().read(AES.block_size)
+    cipher = AES.new(kdf(password), AES.MODE_CBC, iv)
+    return iv + cipher.encrypt(pad(raw))
+
+
+def decrypt(enc, password):
+    """
+    decrypt decrypts data with given password by AES-128-CBC PKCS#7. iv will be read from
+    the head of raw.
+
+    :param enc: input encrypted byte array
+    :param password: password byte array
+    :return: decrypted byte array
+    """
+    iv = enc[:16]
+    cipher = AES.new(kdf(password), AES.MODE_CBC, iv)
+    return unpad(cipher.decrypt(enc[16:]))

pycovenantsql/tests/test_connection.py

@@ -1,7 +1,6 @@
 import datetime
 import sys
 import time
-import unittest2
 import pycovenantsql
 from pycovenantsql.tests import base
 from pycovenantsql._compat import text_type
@@ -41,6 +40,7 @@ def __exit__(self, exc_type, exc_value, traceback):
         if self._created:
             self._c.execute("DROP USER %s" % self._user)
 
+
 class TestConnection(base.PyCovenantSQLTestCase):
 
     def test_largedata(self):
@@ -70,6 +70,7 @@ def test_context(self):
             self.assertEqual(1,cur.fetchone()[0])
             cur.execute('drop table test')
 
+
 # A custom type and function to escape it
 class Foo(object):
     value = "bar"

pycovenantsql/tests/test_e2ee.py

@@ -0,0 +1,80 @@
+# coding: utf-8
+
+import unittest
+from pycovenantsql.e2ee import encrypt, decrypt, unpad, PaddingError
+from binascii import hexlify, unhexlify
+
+# Test cases for all implementations.
+# Because iv is random, so Encrypted data is not always the same,
+# 	but Decrypt(possibleEncrypted) will get raw.
+cases = [
+    {
+        "raw": "11",
+        "pass": ";#K]As9C*6L",
+        "possibleEncrypted": "a372ea2c158a2f99d386e309db4355a659a7a8dd3986fd1d94f7604256061609",
+    },
+    {
+        "raw": "111282C128421286712857128C2128EF" +
+               "128B7671283C128571287512830128EC" +
+               "128391281A1312849128381281E1286A" +
+               "12871128621287A9D12857128C412886" +
+               "128FD12834128DA128F5",
+        "pass": "",
+        "possibleEncrypted": "1bfb6a7fda3e3eb1e14c9afd0baefe86" +
+                             "c90979101f179db7e48a0fa7617881e8" +
+                             "f752c59fb512bb86b8ed69c5644bf2dc" +
+                             "30fbcd3bf79fb20342595c84fad00e46" +
+                             "2fab3e51266492a3d5d085e650c1e619" +
+                             "6278d7f5185c263440ec6fd940ffbb85",
+    },
+    {
+        "raw": "11",
+        "pass": "'K]\"#'pi/1/JD2",
+        "possibleEncrypted": "a83d152777ce3a1c0710b03676ae867c86ab0a47b3ca080f825683ac1079eb41",
+    },
+    {
+        "raw": "11111111111111111111111111111111",
+        "pass": "",
+        "possibleEncrypted": "7dda438c4256a63c62d6816617fcbf9c" +
+                             "7773b9b4f87902b7253848ba2b0ed0ba" +
+                             "f70a3ac976a835b7bc3008e9ba43da74",
+    },
+    {
+        "raw": "11111111111111111111111111111111",
+        "pass": "youofdas1312",
+        "possibleEncrypted": "cab07967cf377dbc010fbf5f84d12bcb" +
+                             "6f8b188e6965738cf9007a671b4bfeb9" +
+                             "f52257aac3808048c341dcaa1c125ca7",
+    },
+    {
+        "raw": "11111111111111111111111111",
+        "pass": "空のBottle😄",
+        "possibleEncrypted": "4384874473945c5b70519ad5ace6305ef6b78c60c3c694add08a8b81899c4171",
+    },
+]
+
+
+class TestE2ee(unittest.TestCase):
+    def test_enc_dec(self):
+        i = 0
+        for case in cases:
+            print("Case: #" + str(i))
+            i += 1
+            enc = encrypt(unhexlify(case["raw"]), case["pass"].encode())
+            dec = decrypt(enc, case["pass"].encode())
+            self.assertEqual(unhexlify(case["raw"]), dec)
+            dec2 = decrypt(unhexlify(case["possibleEncrypted"]), case["pass"].encode())
+            self.assertEqual(unhexlify(case["raw"]), dec2)
+
+    def test_unpad_error(self):
+        self.assertEqual(
+            unpad(unhexlify("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaa01")),
+            unhexlify("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaa")
+        )
+        self.assertEqual(
+            unpad(unhexlify("aaaaaaaaaaaaaaaaaaaaaaaaaaaa0202")),
+            unhexlify("aaaaaaaaaaaaaaaaaaaaaaaaaaaa")
+        )
+        self.assertRaisesRegex(PaddingError, "unexpected padding char", unpad, unhexlify("aaaaaaaaaaaaaaaaaaaaaaaaaaaa0102"))
+        self.assertRaisesRegex(PaddingError, "padding length > 16", unpad, unhexlify("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"))
+        self.assertRaisesRegex(PaddingError, "empty input", unpad, unhexlify(""))

setup.py

@@ -18,6 +18,7 @@
     install_requires=[
         "requests",
         "arrow",
+        "pycrypto",
     ],
     classifiers=[
         'Development Status :: 2 - Pre-Alpha',
@@ -32,10 +33,9 @@
         'Intended Audience :: Developers',
         'Topic :: Database',
     ],
-    keywords=("CovenantSQL","driver","database"),
+    keywords=("CovenantSQL", "driver", "database"),
 
-    author = "laodouya",
-    author_email = "jin.xu@CovenantSQL.io",
-    license = "Apache 2.0 Licence",
+    author="laodouya",
+    author_email="jin.xu@CovenantSQL.io",
+    license="Apache 2.0 Licence",
 )
-