diff --git a/.gitignore b/.gitignore index b5d8df1..d448e5e 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,4 @@ .vscode +.idea +venv *.swp diff --git a/pycovenantsql/__init__.py b/pycovenantsql/__init__.py index aae2a32..a278d31 100644 --- a/pycovenantsql/__init__.py +++ b/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" diff --git a/pycovenantsql/e2ee.py b/pycovenantsql/e2ee.py new file mode 100644 index 0000000..2626a04 --- /dev/null +++ b/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:])) diff --git a/pycovenantsql/tests/test_connection.py b/pycovenantsql/tests/test_connection.py index 48b71d7..3a3e7b4 100644 --- a/pycovenantsql/tests/test_connection.py +++ b/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" diff --git a/pycovenantsql/tests/test_e2ee.py b/pycovenantsql/tests/test_e2ee.py new file mode 100644 index 0000000..cab782c --- /dev/null +++ b/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("")) diff --git a/setup.py b/setup.py index bded4cd..139145e 100644 --- a/setup.py +++ b/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", ) -