Merge branch 'lindsey98:main' into main

app.py

@@ -7,11 +7,9 @@
 import os
 from phishpedia import PhishpediaWrapper, result_file_write
 
-
 app = Flask(__name__)
 CORS(app)
 
-
 # 在创建应用时初始化模型
 with app.app_context():
     current_dir = os.path.dirname(os.path.realpath(__file__))
@@ -27,7 +25,7 @@ def analyze():
         data = request.get_json()
         url = data.get('url')
         screenshot_data = data.get('screenshot')
-        
+
         # 解码Base64图片数据
         image_data = base64.b64decode(screenshot_data.split(',')[1])
         image = Image.open(BytesIO(image_data))
@@ -38,51 +36,44 @@ def analyze():
         phish_category, pred_target, matched_domain, \
             plotvis, siamese_conf, pred_boxes, \
             logo_recog_time, logo_match_time = phishpedia_cls.test_orig_phishpedia(url, screenshot_path, None)
-
+
+        # 添加结果处理逻辑
+        result = {
+            "isPhishing": bool(phish_category),
+            "brand": pred_target if pred_target else "unknown",
+            "legitUrl": f"https://{matched_domain[0]}" if matched_domain else "unknown",
+            "confidence": float(siamese_conf) if siamese_conf is not None else 0.0
+        }
+
+        # 记录日志
         today = datetime.now().strftime('%Y%m%d')
         log_file_path = os.path.join(log_dir, f'{today}_results.txt')
 
         try:
             with open(log_file_path, "a+", encoding='ISO-8859-1') as f:
-                result_file_write(f, current_dir, url, phish_category, pred_target, matched_domain, siamese_conf,
+                result_file_write(f, current_dir, url, phish_category, pred_target,
+                                  matched_domain if matched_domain else ["unknown"],
+                                  siamese_conf if siamese_conf is not None else 0.0,
                                   logo_recog_time, logo_match_time)
         except UnicodeError:
             with open(log_file_path, "a+", encoding='utf-8') as f:
-                result_file_write(f, current_dir, url, phish_category, pred_target, matched_domain, siamese_conf,
+                result_file_write(f, current_dir, url, phish_category, pred_target,
+                                  matched_domain if matched_domain else ["unknown"],
+                                  siamese_conf if siamese_conf is not None else 0.0,
                                   logo_recog_time, logo_match_time)
-        # 目前返回示例数据
-        result = {
-            "isPhishing": bool(phish_category),
-            "brand": pred_target,
-            "legitUrl": "https://" + matched_domain[0],
-            "confidence": float(siamese_conf)
-        }
+
         if os.path.exists(screenshot_path):
             os.remove(screenshot_path)
+
         return jsonify(result)
-    
+
     except Exception as e:
-        print(e)
+        print(f"Error in analyze: {str(e)}")
         log_error_path = os.path.join(log_dir, 'log_error.txt')
         with open(log_error_path, "a+", encoding='utf-8') as f:
             f.write(f'{datetime.now().strftime("%Y-%m-%d %H:%M:%S")} - {str(e)}\n')
         return jsonify("ERROR"), 500
 
 
-def is_running_in_docker():
-    """检查是否在 Docker 环境中运行"""
-    # 方法1：检查 /.dockerenv 文件是否存在
-    docker_env = os.path.exists('/.dockerenv')
-    # 方法2：检查 cgroup 中是否包含 docker 字符串
-    try:
-        with open('/proc/1/cgroup', 'r') as f:
-            return docker_env or 'docker' in f.read()
-    except (IOError, OSError):  # 明确指定可能的异常类型
-        return docker_env
-
-
 if __name__ == '__main__':
-    if is_running_in_docker():
-        app.run(host='0.0.0.0', port=5000, debug=False)
-    else:
-        app.run(debug=False)
+    app.run(host='0.0.0.0', port=5000, debug=False)

logo_matching.py

@@ -111,6 +111,8 @@ def cache_reference_list(model, targetlist_path: str, grayscale=False):
                 except OSError:
                     print(f"Error opening image: {os.path.join(targetlist_path, target, logo_path)}")
                     continue
+
+    return logo_feat_list, file_name_list
 
 
 @torch.no_grad()