添加生成密码标签流程及测试接口

watermark_generate/controller/function_test.py

@@ -2,16 +2,15 @@
 功能测试http接口
 """
 
-from flask import Blueprint, request, send_file, jsonify
+from flask import Blueprint, request, jsonify
 
 from watermark_generate.exceptions import BusinessException
 from watermark_generate import logger
-from watermark_generate.tools import sign_verify
+from watermark_generate.tools import sign_verify, secret_label_func
 
 test = Blueprint('test', __name__)
 
 
-
 @test.route('/test/sign', methods=['POST'])
 def test_sign():
     """
@@ -59,4 +58,29 @@ def test_verify():
 
     verify_result = sign_verify.verify_sign(raw_data, sign_data, public_key)
 
-    return jsonify({'verify_result': verify_result}), 200
+    return jsonify({'verify_result': verify_result}), 200
+
+
+@test.route('/test/generate_label', methods=['POST'])
+def test_generate_label():
+    data = request.json
+    logger.info(f'test generate label: {data}')
+    raw_data = data.get('raw_data')
+    if raw_data is None:
+        raise BusinessException(message='原文值不可为空', code=-1)
+    secret_label = secret_label_func.generate_secret_label(raw_data)
+    return jsonify({'secret_label': secret_label}), 200
+
+
+@test.route('/test/verify_label', methods=['POST'])
+def test_verify_label():
+    data = request.json
+    logger.info(f'test verify label: {data}')
+    secret_label = data.get('secret_label')
+    public_key = data.get('public_key')
+    if secret_label is None:
+        raise BusinessException(message='密码标签不可为空', code=-1)
+    if public_key is None:
+        raise BusinessException(message='签名公钥不可为空', code=-1)
+    verify_result = secret_label_func.verify_secret_label(secret_label, public_key)
+    return jsonify({'verify_result': verify_result}), 200

watermark_generate/tools/secret_label_func.py

@@ -0,0 +1,32 @@
+"""
+    密码标签生成与验证功能
+"""
+import json
+
+from watermark_generate.tools import sign_verify
+
+
+def generate_secret_label(raw_data: str):
+    """
+    生成密码标签
+    :param raw_data: 模型版权信息
+    :return: 指定格式密码标签字符串
+    """
+    sign_data, public_key = sign_verify.get_sign(raw_data)
+    # TODO 增加public_key保存公钥的流程
+    secret_label = json.dumps({'data': raw_data, 'sign': sign_data, 'public_key': public_key})
+    return secret_label
+
+
+def verify_secret_label(secret_label: str, public_key: str) -> bool:
+    """
+    验证密码标签
+    :param secret_label: 生成的密码标签
+    :param public_key: 签名公钥
+    :return: 密码标签验证结果
+    """
+    dict_data = json.loads(secret_label)
+    raw_data = dict_data.get("data")
+    sign_data = dict_data.get("sign")
+    verify_result = sign_verify.verify_sign(raw_data, sign_data, public_key)
+    return verify_result