dekun
127 lines
4.4 KiB
Python
127 lines
4.4 KiB
Python
"""单 IP 只允许挂载一台普通节点。白名单单 IP、白名单网段和白名单节点跳过这条限制。"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import ipaddress
|
|
from datetime import datetime
|
|
|
|
from fastapi import Request
|
|
from sqlalchemy import select
|
|
from sqlalchemy.orm import Session
|
|
|
|
from app.clock import utcnow
|
|
from app.errors import AppError
|
|
from app.models import ComputeNode, IPWhitelist
|
|
|
|
IP_LIMIT_MESSAGE = "该 IP 已绑定过设备，非白名单 IP 仅限挂载一台"
|
|
DATACENTER_BATCH_MESSAGE = "机房/批量接入需满足最低 5 台且全系 30/40 显卡门槛，请前往申请机房白名单"
|
|
|
|
|
|
def client_ip(request: Request) -> str:
|
|
"""优先取反向代理写下的真实地址，否则用连接对端。"""
|
|
real = (request.headers.get("x-real-ip") or "").split(",")[0].strip()
|
|
if real:
|
|
return _loose_ip(real)
|
|
forwarded = (request.headers.get("x-forwarded-for") or "").split(",")[0].strip()
|
|
if forwarded:
|
|
return _loose_ip(forwarded)
|
|
host = request.client.host if request.client is not None else ""
|
|
return _loose_ip(host)
|
|
|
|
|
|
def normalize_admin_ip(value: str) -> str:
|
|
"""单 IP 保持地址本身，网段规范成网络地址，例如 192.168.1.10/24 变成 192.168.1.0/24。"""
|
|
text = value.strip()
|
|
if text.lower().startswith("::ffff:"):
|
|
text = text[7:]
|
|
if "/" in text:
|
|
try:
|
|
return str(ipaddress.ip_network(text, strict=False))
|
|
except ValueError as exc:
|
|
raise AppError("IP 或网段格式不正确") from exc
|
|
try:
|
|
return str(ipaddress.ip_address(text))
|
|
except ValueError as exc:
|
|
raise AppError("IP 或网段格式不正确") from exc
|
|
|
|
|
|
def assert_ip_available(session: Session, ip: str, node: ComputeNode | None = None, now: datetime | None = None) -> None:
|
|
"""白名单 IP，或节点已标记 is_whitelist 时，不限制同一地址上的节点数量。"""
|
|
address = _loose_ip(ip)
|
|
if not address:
|
|
return
|
|
if node is not None and node.is_whitelist:
|
|
return
|
|
if matching_whitelist(session, address) is not None:
|
|
return
|
|
moment = now or utcnow()
|
|
query = select(ComputeNode.id).where(
|
|
ComputeNode.client_ip == address,
|
|
ComputeNode.is_whitelist.is_(False),
|
|
ComputeNode.license_expires_at > moment,
|
|
)
|
|
if node is not None:
|
|
query = query.where(ComputeNode.id != node.id)
|
|
if session.scalar(query) is not None:
|
|
raise AppError(DATACENTER_BATCH_MESSAGE, 409)
|
|
|
|
|
|
def list_whitelist(session: Session) -> list[IPWhitelist]:
|
|
return list(session.scalars(select(IPWhitelist).order_by(IPWhitelist.id.desc())).all())
|
|
|
|
|
|
def add_whitelist_ip(session: Session, ip: str, note: str | None = None) -> IPWhitelist:
|
|
address = normalize_admin_ip(ip)
|
|
if _ip_listed(session, address):
|
|
raise AppError("这个 IP 已在白名单中", 409)
|
|
row = IPWhitelist(ip_address=address, note=(note or "").strip() or None, created_at=utcnow())
|
|
session.add(row)
|
|
session.flush()
|
|
return row
|
|
|
|
|
|
def remove_whitelist_ip(session: Session, entry_id: int) -> IPWhitelist:
|
|
row = session.get(IPWhitelist, entry_id)
|
|
if row is None:
|
|
raise AppError("白名单 IP 不存在", 404)
|
|
session.delete(row)
|
|
session.flush()
|
|
return row
|
|
|
|
|
|
def matching_whitelist(session: Session, address: str) -> str | None:
|
|
"""请求 IP 落在某条单 IP 或网段里时，返回那条白名单原文。"""
|
|
entries = list(session.scalars(select(IPWhitelist.ip_address)).all())
|
|
try:
|
|
ip = ipaddress.ip_address(address)
|
|
except ValueError:
|
|
return address if address in entries else None
|
|
subnet_hit = None
|
|
exact_hit = None
|
|
for entry in entries:
|
|
if "/" in entry:
|
|
try:
|
|
if ip in ipaddress.ip_network(entry, strict=False):
|
|
subnet_hit = entry
|
|
except ValueError:
|
|
continue
|
|
elif entry == str(ip):
|
|
exact_hit = entry
|
|
return subnet_hit or exact_hit
|
|
|
|
|
|
def _ip_listed(session: Session, address: str) -> bool:
|
|
return session.scalar(select(IPWhitelist.id).where(IPWhitelist.ip_address == address)) is not None
|
|
|
|
|
|
def _loose_ip(value: str) -> str:
|
|
text = (value or "").strip()
|
|
if text.lower().startswith("::ffff:"):
|
|
text = text[7:]
|
|
if not text:
|
|
return ""
|
|
try:
|
|
return str(ipaddress.ip_address(text))
|
|
except ValueError:
|
|
return text[:64]
|