Spaces:
Sleeping
Sleeping
Update main.py
Browse files
main.py
CHANGED
|
@@ -6,11 +6,12 @@ import importlib.util
|
|
| 6 |
import uvicorn # Programmatic bootstrapper
|
| 7 |
|
| 8 |
# =====================================================================
|
| 9 |
-
# 1. CRITICAL PATH INJECTION FIX
|
| 10 |
# =====================================================================
|
| 11 |
CURRENT_BASE_DIR = os.path.dirname(os.path.abspath(__file__))
|
| 12 |
VERSION_1_PATH = os.path.join(CURRENT_BASE_DIR, "Version_1")
|
| 13 |
-
|
|
|
|
| 14 |
VERSION_3_PATH = os.path.join(CURRENT_BASE_DIR, "Version_3")
|
| 15 |
VERSION_4_PATH = os.path.join(CURRENT_BASE_DIR, "Version_4")
|
| 16 |
VERSION_5_PATH = os.path.join(CURRENT_BASE_DIR, "Version_5")
|
|
@@ -70,7 +71,7 @@ except Exception as e:
|
|
| 70 |
print(f"[-] Could not merge Version 1. Error: {e}")
|
| 71 |
|
| 72 |
# =====================================================================
|
| 73 |
-
# [ ENGINE 2 ] VERSION 2: VISUAL RESNET18 PHISHING DETECTOR
|
| 74 |
# =====================================================================
|
| 75 |
v2_router = APIRouter(prefix="/api/v2", tags=["Version 2: Visual ResNet18 Engine"])
|
| 76 |
|
|
@@ -84,7 +85,6 @@ try:
|
|
| 84 |
print("[*] Attempting to load Version 2 visual engine...")
|
| 85 |
v2_main_path = os.path.join(VERSION_2_PATH, "main.py")
|
| 86 |
|
| 87 |
-
# Dynamically load the module to bypass Python's hyphen import restriction
|
| 88 |
spec = importlib.util.spec_from_file_location("v2_main", v2_main_path)
|
| 89 |
v2_main = importlib.util.module_from_spec(spec)
|
| 90 |
sys.modules["v2_main"] = v2_main
|
|
@@ -110,13 +110,11 @@ async def analyze_url_vision(payload: V2VisionRequest):
|
|
| 110 |
temp_img_path = os.path.join(OUTPUT_DIR, f"v2_infer_{int(time.time()*1000)}.png")
|
| 111 |
|
| 112 |
try:
|
| 113 |
-
# Step 1: Capture the visual fingerprint
|
| 114 |
screenshot_file = await v2_capture_screenshot(url, temp_img_path)
|
| 115 |
|
| 116 |
if not screenshot_file or not os.path.exists(screenshot_file):
|
| 117 |
raise HTTPException(status_code=502, detail="Failed to capture screenshot. The target may be offline.")
|
| 118 |
|
| 119 |
-
# Step 2: Run PyTorch Inference
|
| 120 |
result = v2_analyzer.analyze_image(screenshot_file)
|
| 121 |
latency_ms = (time.perf_counter() - start_time) * 1000
|
| 122 |
|
|
@@ -133,7 +131,6 @@ async def analyze_url_vision(payload: V2VisionRequest):
|
|
| 133 |
"latency_ms": round(latency_ms, 2)
|
| 134 |
}
|
| 135 |
finally:
|
| 136 |
-
# Crucial Step 3: Prevent disk bloat
|
| 137 |
if os.path.exists(temp_img_path):
|
| 138 |
os.remove(temp_img_path)
|
| 139 |
|
|
@@ -164,7 +161,7 @@ if os.path.exists(MODEL_PATH):
|
|
| 164 |
bst.load_model(MODEL_PATH)
|
| 165 |
else:
|
| 166 |
bst = None
|
| 167 |
-
print("[-] Warning:
|
| 168 |
|
| 169 |
@v4_router.post("/evaluate")
|
| 170 |
async def evaluate_xgboost_url(payload: XGBoostRequest):
|
|
@@ -347,10 +344,11 @@ if __name__ == "__main__":
|
|
| 347 |
print("[*] Resolving path mapping vectors securely...")
|
| 348 |
print(f"[+] Directing Uvicorn to look inside application anchor: {CURRENT_BASE_DIR}")
|
| 349 |
|
|
|
|
| 350 |
uvicorn.run(
|
| 351 |
"main:app",
|
| 352 |
host="0.0.0.0",
|
| 353 |
-
port=
|
| 354 |
reload=True,
|
| 355 |
-
app_dir=CURRENT_BASE_DIR
|
| 356 |
)
|
|
|
|
| 6 |
import uvicorn # Programmatic bootstrapper
|
| 7 |
|
| 8 |
# =====================================================================
|
| 9 |
+
# 1. CRITICAL PATH INJECTION FIX (Matched to Root Repository Layout)
|
| 10 |
# =====================================================================
|
| 11 |
CURRENT_BASE_DIR = os.path.dirname(os.path.abspath(__file__))
|
| 12 |
VERSION_1_PATH = os.path.join(CURRENT_BASE_DIR, "Version_1")
|
| 13 |
+
# 🔥 FIXED: Added spaces around the hyphen to match your exact directory name
|
| 14 |
+
VERSION_2_PATH = os.path.join(CURRENT_BASE_DIR, "Version_2 - Copy")
|
| 15 |
VERSION_3_PATH = os.path.join(CURRENT_BASE_DIR, "Version_3")
|
| 16 |
VERSION_4_PATH = os.path.join(CURRENT_BASE_DIR, "Version_4")
|
| 17 |
VERSION_5_PATH = os.path.join(CURRENT_BASE_DIR, "Version_5")
|
|
|
|
| 71 |
print(f"[-] Could not merge Version 1. Error: {e}")
|
| 72 |
|
| 73 |
# =====================================================================
|
| 74 |
+
# [ ENGINE 2 ] VERSION 2: VISUAL RESNET18 PHISHING DETECTOR
|
| 75 |
# =====================================================================
|
| 76 |
v2_router = APIRouter(prefix="/api/v2", tags=["Version 2: Visual ResNet18 Engine"])
|
| 77 |
|
|
|
|
| 85 |
print("[*] Attempting to load Version 2 visual engine...")
|
| 86 |
v2_main_path = os.path.join(VERSION_2_PATH, "main.py")
|
| 87 |
|
|
|
|
| 88 |
spec = importlib.util.spec_from_file_location("v2_main", v2_main_path)
|
| 89 |
v2_main = importlib.util.module_from_spec(spec)
|
| 90 |
sys.modules["v2_main"] = v2_main
|
|
|
|
| 110 |
temp_img_path = os.path.join(OUTPUT_DIR, f"v2_infer_{int(time.time()*1000)}.png")
|
| 111 |
|
| 112 |
try:
|
|
|
|
| 113 |
screenshot_file = await v2_capture_screenshot(url, temp_img_path)
|
| 114 |
|
| 115 |
if not screenshot_file or not os.path.exists(screenshot_file):
|
| 116 |
raise HTTPException(status_code=502, detail="Failed to capture screenshot. The target may be offline.")
|
| 117 |
|
|
|
|
| 118 |
result = v2_analyzer.analyze_image(screenshot_file)
|
| 119 |
latency_ms = (time.perf_counter() - start_time) * 1000
|
| 120 |
|
|
|
|
| 131 |
"latency_ms": round(latency_ms, 2)
|
| 132 |
}
|
| 133 |
finally:
|
|
|
|
| 134 |
if os.path.exists(temp_img_path):
|
| 135 |
os.remove(temp_img_path)
|
| 136 |
|
|
|
|
| 161 |
bst.load_model(MODEL_PATH)
|
| 162 |
else:
|
| 163 |
bst = None
|
| 164 |
+
print(f"[-] Warning: {MODEL_PATH} not found. V4 engine will return a configuration error.")
|
| 165 |
|
| 166 |
@v4_router.post("/evaluate")
|
| 167 |
async def evaluate_xgboost_url(payload: XGBoostRequest):
|
|
|
|
| 344 |
print("[*] Resolving path mapping vectors securely...")
|
| 345 |
print(f"[+] Directing Uvicorn to look inside application anchor: {CURRENT_BASE_DIR}")
|
| 346 |
|
| 347 |
+
# 🔥 FIXED: Points directly to port 7860 for Hugging Face standard routing context
|
| 348 |
uvicorn.run(
|
| 349 |
"main:app",
|
| 350 |
host="0.0.0.0",
|
| 351 |
+
port=7860,
|
| 352 |
reload=True,
|
| 353 |
+
app_dir=CURRENT_BASE_DIR
|
| 354 |
)
|