diff --git a/app.py b/app.py index 60fe918..f5f6e98 100644 --- a/app.py +++ b/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) diff --git a/logo_matching.py b/logo_matching.py index 807f4ed..77590f7 100644 --- a/logo_matching.py +++ b/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()