diff --git a/mindlakesdk/__init__.py b/mindlakesdk/__init__.py index d9952b3..c4df229 100644 --- a/mindlakesdk/__init__.py +++ b/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) diff --git a/mindlakesdk/cryptor.py b/mindlakesdk/cryptor.py index f251a0c..fcafc2e 100644 --- a/mindlakesdk/cryptor.py +++ b/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(" 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: @@ -99,40 +114,53 @@ def decrypt(self, cipher: bytes|str) -> ResultType: 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 diff --git a/mindlakesdk/message.py b/mindlakesdk/message.py index 1a8d42e..edde370 100644 --- a/mindlakesdk/message.py +++ b/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: diff --git a/mindlakesdk/permission.py b/mindlakesdk/permission.py index 0b5543d..e73cbc9 100644 --- a/mindlakesdk/permission.py +++ b/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) \ No newline at end of file diff --git a/mindlakesdk/settings.py b/mindlakesdk/settings.py index 6d25894..4cfd7f0 100644 --- a/mindlakesdk/settings.py +++ b/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' \ No newline at end of file +VERSION = 'v1.0.2' \ No newline at end of file diff --git a/mindlakesdk/utils.py b/mindlakesdk/utils.py index f9692c2..9622f5e 100644 --- a/mindlakesdk/utils.py +++ b/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) diff --git a/pyproject.toml b/pyproject.toml index eb19ef5..96289f5 100644 --- a/pyproject.toml +++ b/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"] \ No newline at end of file