127 lines
5.3 KiB
Python
127 lines
5.3 KiB
Python
# requires: fastapi, httpx
|
||
"""Web 安全示例：SQL 注入、XSS、CSRF、CORS、SSRF（Agent 的“抓取网页”工具最容易出这个问题）。
|
||
|
||
每类问题都给出“错误写法”和“正确写法”，test_secure_app.py 用测试证明错误写法可被利用、正确写法可以防住。
|
||
运行：pytest -q test_secure_app.py；直接 `python secure_app.py` 只做一次自检。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import html
|
||
import ipaddress
|
||
import secrets
|
||
import socket
|
||
import sqlite3
|
||
from urllib.parse import urlsplit
|
||
|
||
from fastapi import Cookie, FastAPI, Form, Header, HTTPException, Response
|
||
from fastapi.middleware.cors import CORSMiddleware
|
||
from fastapi.responses import HTMLResponse
|
||
|
||
# ---------- SQL 注入 ----------
|
||
def make_db() -> sqlite3.Connection:
|
||
db = sqlite3.connect(":memory:", check_same_thread=False)
|
||
db.executescript("""
|
||
CREATE TABLE orders (id TEXT, tenant TEXT, amount INTEGER);
|
||
INSERT INTO orders VALUES ('A001', 't1', 100), ('A002', 't1', 200), ('B001', 't2', 999);
|
||
""")
|
||
return db
|
||
|
||
|
||
def find_orders_unsafe(db: sqlite3.Connection, tenant: str, keyword: str) -> list[tuple]:
|
||
# 错误：字符串拼接。keyword = "x' OR 1=1 --" 时会查出其他租户的数据
|
||
return db.execute(f"SELECT id, tenant FROM orders WHERE tenant = '{tenant}' AND id LIKE '%{keyword}%'").fetchall()
|
||
|
||
|
||
def find_orders_safe(db: sqlite3.Connection, tenant: str, keyword: str) -> list[tuple]:
|
||
# 正确：参数化查询，数据永远不会被当成 SQL 执行
|
||
return db.execute("SELECT id, tenant FROM orders WHERE tenant = ? AND id LIKE ?", (tenant, f"%{keyword}%")).fetchall()
|
||
|
||
|
||
# ---------- XSS ----------
|
||
def render_answer_unsafe(answer: str) -> str:
|
||
return f"<div class='answer'>{answer}</div>" # 错误：模型输出 / 用户输入直接拼进 HTML
|
||
|
||
|
||
def render_answer_safe(answer: str) -> str:
|
||
return f"<div class='answer'>{html.escape(answer)}</div>" # 正确：输出到 HTML 时转义
|
||
|
||
|
||
# ---------- SSRF ----------
|
||
ALLOWED_SCHEMES = {"http", "https"}
|
||
|
||
|
||
def check_fetch_url(url: str, resolve=socket.getaddrinfo) -> str:
|
||
"""Agent 抓取网页工具的 URL 校验：只允许 http(s)，解析后的所有 IP 都必须是公网地址。"""
|
||
parts = urlsplit(url)
|
||
if parts.scheme not in ALLOWED_SCHEMES or not parts.hostname:
|
||
raise ValueError("只允许 http / https 地址")
|
||
if parts.username or parts.password:
|
||
raise ValueError("不允许在 URL 中携带凭证")
|
||
infos = resolve(parts.hostname, parts.port or (443 if parts.scheme == "https" else 80))
|
||
for info in infos:
|
||
ip = ipaddress.ip_address(info[4][0])
|
||
if not ip.is_global: # 拒绝内网、回环、链路本地（含云元数据地址 169.254.169.254）
|
||
raise ValueError(f"禁止访问内网地址 {ip}")
|
||
# 注意：真正发请求时还要禁止自动跟随重定向（或对每一跳重新校验），并固定使用这里解析出的 IP，防止 DNS 重绑定
|
||
return url
|
||
|
||
|
||
# ---------- 应用：CORS、CSRF、安全响应头 ----------
|
||
TRUSTED_ORIGINS = ["https://app.example.com"]
|
||
|
||
app = FastAPI()
|
||
app.add_middleware(
|
||
CORSMiddleware,
|
||
allow_origins=TRUSTED_ORIGINS, # 白名单；带凭证时绝不能用 "*"
|
||
allow_credentials=True,
|
||
allow_methods=["GET", "POST"],
|
||
allow_headers=["Content-Type", "X-CSRF-Token"],
|
||
)
|
||
DB = make_db()
|
||
|
||
|
||
@app.middleware("http")
|
||
async def security_headers(request, call_next):
|
||
resp = await call_next(request)
|
||
resp.headers["Content-Security-Policy"] = "default-src 'self'; script-src 'self'; object-src 'none'" # XSS 的第二道防线
|
||
resp.headers["X-Content-Type-Options"] = "nosniff"
|
||
return resp
|
||
|
||
|
||
@app.post("/login")
|
||
def login(response: Response) -> dict:
|
||
csrf = secrets.token_urlsafe(16)
|
||
# 会话 Cookie：HttpOnly（脚本读不到）、Secure、SameSite=Lax（跨站的 POST 不会带上）
|
||
response.set_cookie("session", "user-1", httponly=True, secure=True, samesite="lax")
|
||
# CSRF Token（双重提交）：前端读取后放在请求头中，跨站页面拿不到这个值
|
||
response.set_cookie("csrf_token", csrf, secure=True, samesite="lax")
|
||
return {"ok": True}
|
||
|
||
|
||
@app.post("/transfer")
|
||
def transfer(amount: int = Form(...), session: str | None = Cookie(None), csrf_cookie: str | None = Cookie(None, alias="csrf_token"),
|
||
x_csrf_token: str | None = Header(None)) -> dict:
|
||
if not session:
|
||
raise HTTPException(401, "未登录")
|
||
if not csrf_cookie or not x_csrf_token or not secrets.compare_digest(csrf_cookie, x_csrf_token):
|
||
raise HTTPException(403, "CSRF 校验失败")
|
||
return {"ok": True, "amount": amount}
|
||
|
||
|
||
@app.get("/orders")
|
||
def orders(keyword: str = "") -> list[list[str]]:
|
||
return [list(r) for r in find_orders_safe(DB, "t1", keyword)]
|
||
|
||
|
||
@app.get("/answer", response_class=HTMLResponse)
|
||
def answer(q: str = "") -> str:
|
||
return render_answer_safe(f"你问的是：{q}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
payload = "x' OR 1=1 --"
|
||
assert {r[1] for r in find_orders_unsafe(DB, "t1", payload)} == {"t1", "t2"}
|
||
assert find_orders_safe(DB, "t1", payload) == []
|
||
assert "<script>" not in render_answer_safe("<script>alert(1)</script>")
|
||
print("自检通过：拼接 SQL 被注入读到了 t2 的数据，参数化查询没有；XSS 输出已转义")
|