Merge pull request #5 from mind-network/dev

improve performance and add interface

mindlakesdk/__init__.py

@@ -2,6 +2,7 @@
 
 from eth_account.messages import encode_defunct
 from web3 import Web3
+import requests
 
 import mindlakesdk.settings as settings
 import mindlakesdk.utils
@@ -22,6 +23,7 @@ def __init__(self, walletPrivateKey: str, appKey: str, gateway: str = None):
         logging.debug(__name__)
         self.__session = mindlakesdk.utils.Session()
         session = self.__session
+        session.requstSession = requests.Session()
         self.datalake = DataLake(session)
         self.cryptor = Cryptor(session)
         self.permission = Permission(session)

mindlakesdk/cryptor.py

@@ -21,36 +21,15 @@ class EncType(Enum):
 
     def __init__(self, session: Session):
         self.__session = session
+        self.cryptParamsByColumn = {}
+        self.cryptParamsByID = {}
 
-    def encrypt(self, data, columnOrType: str|DataType) -> ResultType:
-        if isinstance(columnOrType, DataType):
-            dataType = columnOrType
-            result = mindlakesdk.message.getDKbyName(self.__session)
-            if not result:
-                return result
-        else:
-            tableName, columnName = columnOrType.split('.')
-            result = mindlakesdk.message.getDataTypeByName(self.__session, tableName, columnName)
-            if not result:
-                return result
-            dataType = DataType(result.data)
-            result = mindlakesdk.message.getDKbyName(self.__session, self.__session.walletAddress, tableName, columnName)
-            # Temporary solution for MS not returning Error Code
-            if result.code == 40010:
-                # DK not found, create one
-                result = mindlakesdk.keyhelper.genDK(self.__session, tableName, columnName)
-                if not result:
-                    return result
-            elif not result:
-                return result
-            else:
-                pass
-        encTypeNum = Cryptor.EncType['enc_' + dataType.name].value
+    def encrypt(self, data, columnOrType) -> ResultType:
+        result = self.__getEncryptParams(columnOrType)
+        if not result:
+            return result
+        ctxid, encTypeNum, dk, alg, dataType = result.data
         data = Cryptor.__encodeByDataType(data, dataType)
-        ctxid = result.data['ctxId']
-        dkCipher = result.data['encryptedDek']
-        dkID, dk = mindlakesdk.keyhelper.decryptDKb64(self.__session.mk, dkCipher)
-        alg = result.data['algorithm']
         header = Cryptor.__genCryptoHeader(ctxid, encTypeNum)
         checkCode = Cryptor.__genCheckCode(data, 1)
         data_to_enc = data + checkCode
@@ -67,6 +46,42 @@ def encrypt(self, data, columnOrType: str|DataType) -> ResultType:
         resultHex = '\\x' + result.hex()
         return ResultType(0, None, resultHex)
 
+    def __getEncryptParams(self, columnOrType):
+        result = self.cryptParamsByColumn.get(columnOrType)
+        if not result:
+            if isinstance(columnOrType, DataType):
+                dataType = columnOrType
+                result = mindlakesdk.message.getDKbyName(self.__session)
+                if not result:
+                    return result
+            else:
+                tableName, columnName = columnOrType.split('.')
+                result = mindlakesdk.message.getDataTypeByName(self.__session, tableName, columnName)
+                if not result:
+                    return result
+                dataType = DataType(result.data)
+                result = mindlakesdk.message.getDKbyName(self.__session, self.__session.walletAddress, tableName, columnName)
+                # Temporary solution for MS not returning Error Code
+                if result.code == 40010:
+                    # DK not found, create one
+                    result = mindlakesdk.keyhelper.genDK(self.__session, tableName, columnName)
+                    if not result:
+                        return result
+                elif not result:
+                    return result
+                else:
+                    pass
+            encTypeNum = Cryptor.EncType['enc_' + dataType.name].value
+            ctxid = result.data['ctxId']
+            dkCipher = result.data['encryptedDek']
+            dkID, dk = mindlakesdk.keyhelper.decryptDKb64(self.__session.mk, dkCipher)
+            alg = result.data['algorithm']
+            result = (ctxid, encTypeNum, dk, alg, dataType)
+            self.cryptParamsByColumn[columnOrType] = result
+            self.cryptParamsByID[ctxid] = (dk, alg)
+        return ResultType(0, None, result)
+
+
     def __encodeByDataType(data, dataType: DataType) -> bytes:
         if dataType == DataType.int4:
             result = struct.pack("<i", data)
@@ -91,48 +106,61 @@ def __encodeByDataType(data, dataType: DataType) -> bytes:
             raise Exception("Unsupported encryption type")
         return result
 
-    def decrypt(self, cipher: bytes|str) -> ResultType:
+    def decrypt(self, cipher) -> ResultType:
         if isinstance(cipher, str):
             data = bytes.fromhex(cipher[2:])
         else:
             data = cipher
         header = Cryptor.__extractCryptoHeader(data)
         encTypeNum = Cryptor.__extractEncType(header)
         ctxId = Cryptor.__extractCtxId(header)
-        result = mindlakesdk.message.getDKbyCid(self.__session, ctxId)
+        result = self.__getDecryptParams(ctxId)
         if not result:
             return result
-        dkCipher = result.data['encryptedDek']
-        # TODO: catch decryption error
-        try:
-            dkID, dk = mindlakesdk.keyhelper.decryptDKb64(self.__session.mk, dkCipher)
-        except:
-            return ResultType(60003, "Can't handle DK")
-        alg = result.data['algorithm']
-        if alg is None:
-            raise Exception("Cannot find data key by ctxId")
+        dk, alg = result.data
+        idx = (header[1] & 0x7) + 2
+        if alg == 3:
+            iv = data[idx:idx+16]
+            cipherBlob = data[idx+16:]
+            plainBlob = mindlakesdk.utils.aesDecrypt(dk, iv, cipherBlob)
+        elif alg == 0:
+            iv = data[idx:idx+12]
+            idx += 12
+            mac = data[idx:idx+16]
+            idx += 16
+            cipherBlob = data[idx:]
+            plainBlob = mindlakesdk.utils.aesGCMDecrypt(dk, iv, cipherBlob, mac)
         else:
-            idx = (header[1] & 0x7) + 2
-            if alg == 3:
-                iv = data[idx:idx+16]
-                cipherBlob = data[idx+16:]
-                plainBlob = mindlakesdk.utils.aesDecrypt(dk, iv, cipherBlob)
-            elif alg == 0:
-                iv = data[idx:idx+12]
-                idx += 12
-                mac = data[idx:idx+16]
-                idx += 16
-                cipherBlob = data[idx:]
-                plainBlob = mindlakesdk.utils.aesGCMDecrypt(dk, iv, cipherBlob, mac)
-            else:
-                raise Exception("Unsupported algorithm to decrypt")
-            result = plainBlob[:-1]
-            checkCode = plainBlob[-1:]
-            checkCode2 = Cryptor.__genCheckCode(result, 1)
-            if checkCode != checkCode2:
-                raise Exception("Check code is not correct")
-            result = Cryptor.__decodeByEncType(result, Cryptor.EncType(encTypeNum))
-            return ResultType(0, None, result)
+            return ResultType(60004, "Unsupported algorithm to decrypt")
+        result = plainBlob[:-1]
+        checkCode = plainBlob[-1:]
+        checkCode2 = Cryptor.__genCheckCode(result, 1)
+        if checkCode != checkCode2:
+            return ResultType(60005, "Check code of cipher is not correct")
+        result = Cryptor.__decodeByEncType(result, Cryptor.EncType(encTypeNum))
+        return ResultType(0, None, result)
+
+    def __getDecryptParams(self, ctxId: int):
+        params = self.cryptParamsByID.get(ctxId)
+        if not params:
+            result = mindlakesdk.message.getDKbyCid(self.__session, ctxId)
+            if not result:
+                return result
+            if not result.data:
+                return ResultType(60002, "Can't get data key")
+            dkCipher = result.data['encryptedDek']
+            if not dkCipher:
+                return ResultType(60002, "Can't get data key")
+            try:
+                dkID, dk = mindlakesdk.keyhelper.decryptDKb64(self.__session.mk, dkCipher)
+            except:
+                return ResultType(60003, "Can't handle data key")
+            alg = result.data['algorithm']
+            if alg is None:
+                return ResultType(60002, "Can't get data key")
+            params = (dk, alg)
+            self.cryptParamsByID[ctxId] = params
+        return ResultType(0, None, params)
 
     def __decodeByEncType(data, encType: EncType):
         if encType == Cryptor.EncType.enc_int4:
@@ -187,7 +215,7 @@ def __extractEncType(header):
         type_value = (tmp_value & 0xF8) >> 3
         return type_value
 
-    def __extractCtxId(header):
+    def __extractCtxId(header) -> int:
         ctxIdLen = header[1] & 0x7
         assert len(header) == ctxIdLen + 2
         ctxId = 0
@@ -216,7 +244,7 @@ def __genCryptoHeader(ctxid, encType):
         head[1] = tmp_val
         return bytes(head)
 
-    def __genCheckCode(data, resultSize):
+    def __genCheckCode(data: bytes, resultSize):
         tmpCode = bytearray(resultSize)
         for i in range(len(data)):
             n = i % resultSize

mindlakesdk/message.py

@@ -143,8 +143,16 @@ def sendListGrantee(session: Session) -> ResultType:
     data = {"bizType":126}
     return __requestCommon(session, data)
 
-def sendListGrantedColumn(session: Session, walletAddress: str) -> ResultType:
-    data = {"bizType":127, 'targetWalletAddress': walletAddress}
+def sendListGrantedColumn(session: Session, targetWalletAddress: str) -> ResultType:
+    data = {"bizType":127, 'targetWalletAddress': targetWalletAddress}
+    return __requestCommon(session, data)
+
+def sendListOwner(session: Session) -> ResultType:
+    data = {"bizType":130}
+    return __requestCommon(session, data)
+
+def sendListOwnerColumn(session: Session, targetWalletAddress: str) -> ResultType:
+    data = {"bizType":131, 'targetWalletAddress': targetWalletAddress}
     return __requestCommon(session, data)
 
 def __requestCommon(session: Session, data: str) -> ResultType:

mindlakesdk/permission.py

@@ -139,4 +139,10 @@ def listGrantee(self) -> ResultType:
 
     def listGrantedColumn(self, walletAddress: str) -> ResultType:
         return mindlakesdk.message.sendListGrantedColumn(self.__session, walletAddress) 
+
+    def listOwner(self) -> ResultType:
+        return mindlakesdk.message.sendListOwner(self.__session)
+
+    def listOwnerColumn(self, walletAddress: str) -> ResultType:
+        return mindlakesdk.message.sendListOwnerColumn(self.__session, walletAddress)
 

mindlakesdk/settings.py

@@ -1,4 +1,4 @@
-GATEWAY = 'http://sdk.mindnetwork.xyz/node'
+GATEWAY = 'https://sdk.mindnetwork.xyz/node'
 WEB3API = 'https://goerli.infura.io/v3/744c0ade89464e4b867ca1b002a10231'
 CONTRACT_ADDRESS = '0xF5932e67e84F08965DC6D62C2B67f47a6826E5a7'
 CONTRACT_ABI = [
@@ -70,4 +70,4 @@
                 "type": "function"
         }
 ]
-VERSION = 'v1.0.0'
+VERSION = 'v1.0.2'

mindlakesdk/utils.py

@@ -5,7 +5,6 @@
 from Crypto.Hash import SHA256, HMAC
 from Crypto.Random import get_random_bytes
 from enum import Enum
-import requests
 import json
 import mindlakesdk.settings as settings
 import logging
@@ -41,6 +40,7 @@ def __init__(self) -> None:
         self.pkID = None
         self.appKey = None
         self.gateway = None
+        self.requstSession = None
 
 def genRSAKey():
     rsaKey = RSA.generate(2048)
@@ -117,7 +117,7 @@ def request(data, session: Session):
     headers['app'] = session.appKey
     if session.token:
         headers['token'] = session.token
-    response = requests.post(session.gateway, json=data, headers=headers)
+    response = session.requstSession.post(session.gateway, json=data, headers=headers)
     logging.debug("============== Mind SDK request ==============")
     logging.debug('MindSDKHeaders: %s'%headers)
     logging.debug("MindSDKData: %s"%data)

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "mindlakesdk"
-version = "v1.0.1"
+version = "v1.0.2"
 authors = [
   { name="Mind Labs", email="biz@mindnetwork.xyz" },
 ]
@@ -26,3 +26,6 @@ dependencies = [
 [project.urls]
 "Homepage" = "https://github.com/mind-network/mind-lake-sdk-python"
 "Bug Tracker" = "https://github.com/mind-network/mind-lake-sdk-python/issues"
+
+[tool.hatch.build]
+exclude = ["/examples", "/tests", "/tutorial"]