261 lines
7.1 KiB
Text
261 lines
7.1 KiB
Text
from fastapi import APIRouter, HTTPException, Request
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
from pydantic import BaseModel
|
|
from typing import Optional, List
|
|
|
|
from app.rag_service import RAGService
|
|
from app.scanner import Scanner
|
|
from app.blocklist import Blocklist
|
|
from app.wallet_labels import WalletLabels
|
|
from app.entity_intel import EntityIntel
|
|
|
|
router = APIRouter(prefix="/api/v1/protect")
|
|
|
|
# CORS configuration for browser extension access
|
|
origins = [
|
|
"https://chrome-extension://*",
|
|
"https://extension://*",
|
|
"http://localhost:*",
|
|
]
|
|
|
|
router.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=origins,
|
|
allow_credentials=True,
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
class URLCheckRequest(BaseModel):
|
|
url: str
|
|
|
|
class URLCheckResponse(BaseModel):
|
|
safe: bool
|
|
risk: str
|
|
reason: str
|
|
source: str
|
|
|
|
class WalletCheckRequest(BaseModel):
|
|
address: str
|
|
chain: str
|
|
|
|
class WalletCheckResponse(BaseModel):
|
|
safe: bool
|
|
risk_score: int
|
|
flags: List[str]
|
|
|
|
class TokenCheckRequest(BaseModel):
|
|
address: str
|
|
chain: str
|
|
|
|
class TokenCheckResponse(BaseModel):
|
|
safe: bool
|
|
risk_score: int
|
|
flags: List[str]
|
|
details: dict
|
|
|
|
@router.post("/check-url", response_model=URLCheckResponse)
|
|
async def check_url(request: URLCheckRequest):
|
|
"""Check if a URL is safe.
|
|
|
|
Args:
|
|
request: URLCheckRequest containing the URL to check
|
|
|
|
Returns:
|
|
URLCheckResponse with safety assessment
|
|
"""
|
|
url = request.url
|
|
|
|
# Check against RAG known_scams database
|
|
rag_service = RAGService()
|
|
rag_result = await rag_service.check_url(url)
|
|
|
|
if rag_result:
|
|
return URLCheckResponse(
|
|
safe=False,
|
|
risk=rag_result['risk'],
|
|
reason=rag_result['reason'],
|
|
source="rag"
|
|
)
|
|
|
|
# Check against blocklist
|
|
blocklist = Blocklist()
|
|
blocklist_result = await blocklist.check_url(url)
|
|
|
|
if blocklist_result:
|
|
return URLCheckResponse(
|
|
safe=False,
|
|
risk=blocklist_result['risk'],
|
|
reason=blocklist_result['reason'],
|
|
source="blocklist"
|
|
)
|
|
|
|
# If new URL, scan it and feed back to RAG
|
|
scanner = Scanner()
|
|
scan_result = await scanner.scan_url(url)
|
|
|
|
if scan_result['risk'] != "low":
|
|
await rag_service.add_url_to_known_scams(url, scan_result)
|
|
|
|
return URLCheckResponse(
|
|
safe=scan_result['risk'] == "low",
|
|
risk=scan_result['risk'],
|
|
reason=scan_result['reason'],
|
|
source="scanner"
|
|
)
|
|
|
|
@router.post("/check-wallet", response_model=WalletCheckResponse)
|
|
async def check_wallet(request: WalletCheckRequest):
|
|
"""Check the safety of a wallet address.
|
|
|
|
Args:
|
|
request: WalletCheckRequest containing wallet address and chain
|
|
|
|
Returns:
|
|
WalletCheckResponse with safety assessment
|
|
"""
|
|
address = request.address
|
|
chain = request.chain
|
|
|
|
# Check against wallet labels
|
|
wallet_labels = WalletLabels()
|
|
labels_result = await wallet_labels.check_wallet(address, chain)
|
|
|
|
# Check entity intel
|
|
entity_intel = EntityIntel()
|
|
intel_result = await entity_intel.check_wallet(address, chain)
|
|
|
|
# Combine results
|
|
flags = []
|
|
risk_score = 0
|
|
|
|
if labels_result:
|
|
flags.extend(labels_result['flags'])
|
|
risk_score = max(risk_score, labels_result['risk_score'])
|
|
|
|
if intel_result:
|
|
flags.extend(intel_result['flags'])
|
|
risk_score = max(risk_score, intel_result['risk_score'])
|
|
|
|
# Flag known drainers, scammers, sanction lists
|
|
if "drainer" in flags or "scammer" in flags or "sanctioned" in flags:
|
|
risk_score = max(risk_score, 80)
|
|
|
|
return WalletCheckResponse(
|
|
safe=risk_score < 50,
|
|
risk_score=risk_score,
|
|
flags=flags
|
|
)
|
|
|
|
@router.post("/check-token", response_model=TokenCheckResponse)
|
|
async def check_token(request: TokenCheckRequest):
|
|
"""Get a full safety report for a token.
|
|
|
|
Args:
|
|
request: TokenCheckRequest containing token address and chain
|
|
|
|
Returns:
|
|
TokenCheckResponse with full safety report
|
|
"""
|
|
address = request.address
|
|
chain = request.chain
|
|
|
|
# Call existing scanner
|
|
scanner = Scanner()
|
|
scan_result = await scanner.scan_token(address, chain)
|
|
|
|
# Add protection-focused summary
|
|
flags = scan_result.get('flags', [])
|
|
risk_score = scan_result.get('risk_score', 0)
|
|
|
|
# Add additional protection checks
|
|
wallet_labels = WalletLabels()
|
|
labels_result = await wallet_labels.check_token(address, chain)
|
|
|
|
if labels_result:
|
|
flags.extend(labels_result.get('flags', []))
|
|
risk_score = max(risk_score, labels_result.get('risk_score', 0))
|
|
|
|
return TokenCheckResponse(
|
|
safe=risk_score < 50,
|
|
risk_score=risk_score,
|
|
flags=flags,
|
|
details=scan_result
|
|
)
|
|
|
|
@router.get("/domains")
|
|
async def get_domains():
|
|
"""Get the domain blocklist for the browser extension.
|
|
|
|
Returns:
|
|
JSON blocklist containing malicious domains
|
|
"""
|
|
blocklist = Blocklist()
|
|
domains = await blocklist.get_blocklist()
|
|
|
|
# Also include domains from RAG
|
|
rag_service = RAGService()
|
|
rag_domains = await rag_service.get_scams_domains()
|
|
|
|
# Combine and deduplicate
|
|
all_domains = list(set(domains + rag_domains))
|
|
|
|
return {"domains": all_domains}
|
|
|
|
@router.get("/health")
|
|
async def health_check():
|
|
"""Check the health status of all protection modules.
|
|
|
|
Returns:
|
|
JSON object with health status of each module
|
|
"""
|
|
health_status = {
|
|
"status": "ok",
|
|
"modules": {}
|
|
}
|
|
|
|
# Check RAG service
|
|
try:
|
|
rag_service = RAGService()
|
|
rag_health = await rag_service.health_check()
|
|
health_status["modules"]["rag"] = rag_health
|
|
except Exception as e:
|
|
health_status["modules"]["rag"] = {"status": "error", "message": str(e)}
|
|
|
|
# Check Blocklist
|
|
try:
|
|
blocklist = Blocklist()
|
|
blocklist_health = await blocklist.health_check()
|
|
health_status["modules"]["blocklist"] = blocklist_health
|
|
except Exception as e:
|
|
health_status["modules"]["blocklist"] = {"status": "error", "message": str(e)}
|
|
|
|
# Check Scanner
|
|
try:
|
|
scanner = Scanner()
|
|
scanner_health = await scanner.health_check()
|
|
health_status["modules"]["scanner"] = scanner_health
|
|
except Exception as e:
|
|
health_status["modules"]["scanner"] = {"status": "error", "message": str(e)}
|
|
|
|
# Check Wallet Labels
|
|
try:
|
|
wallet_labels = WalletLabels()
|
|
labels_health = await wallet_labels.health_check()
|
|
health_status["modules"]["wallet_labels"] = labels_health
|
|
except Exception as e:
|
|
health_status["modules"]["wallet_labels"] = {"status": "error", "message": str(e)}
|
|
|
|
# Check Entity Intel
|
|
try:
|
|
entity_intel = EntityIntel()
|
|
intel_health = await entity_intel.health_check()
|
|
health_status["modules"]["entity_intel"] = intel_health
|
|
except Exception as e:
|
|
health_status["modules"]["entity_intel"] = {"status": "error", "message": str(e)}
|
|
|
|
# Overall status
|
|
if any(module["status"] != "ok" for module in health_status["modules"].values()):
|
|
health_status["status"] = "degraded"
|
|
|
|
return health_status
|