逻辑优化

wx/zskpay.py

@@ -33,23 +33,20 @@ domain = domainMap[env]
 
 class ZSKPay:
 
-    def __init__(self, appid, taxId, ymId, ymKey, cert):
+    def __init__(self, appid, taxId, ymId, ymKey, public_key, private_key=None):
         self.MAX_ENCRYPTION_SIZE = 117 #214
         self.appid = appid
         self.domain = domain
         self.taxId = taxId
         self.ymId = ymId
         self.ymKey = ymKey
-        self.cert = cert
+        self.public_key = public_key
+        self.private_key = private_key
 
     def encrypt_data(self, data):
         encrypted_chunks = []
-        current_directory = os.getcwd()
-        with open(self.cert + "public_key.txt", "rb") as public_file:
-            public_key_content = public_file.read()
-        
         public_key_loaded = serialization.load_pem_public_key(
-            public_key_content,
+            self.public_key.encode('utf-8'),
             backend=default_backend()
         )
 
@@ -73,14 +70,10 @@ class ZSKPay:
 
 
     def decrypt_data(self, data):
-        current_directory = os.getcwd()
-        with open(current_directory+"/wx/cert/zsk/resp/private_key.txt", "rb") as public_file:
-            private_key_content = public_file.read()
-        
         # 加载私钥
         private_key_loaded = serialization.load_pem_private_key(
-            private_key_content,
-             password=None,
+            self.private_key.encode('utf-8'),
+            password=None,
             backend=default_backend()
         )
         # 解密
@@ -119,9 +112,8 @@ class ZSKPay:
     
     #转账
     def transfer(self,data):
-        if not data['transferNo'] or not data['transferAmount'] or not data['transferRemark'] or not data['openid'] :
+        if 'transferNo' not in data or 'transferAmount' not in data or 'transferRemark' not in data or 'openid' not in data:
             raise ValueError("transferDetailList 必传")
-        company_id = data.get('company_id','')
         batchNo = data.get('transferNo', self.generate_batchNo())
         batchName = data.get('batchName','活动批次')
         batchRemark = data.get('batchRemark','无')