zioinfo-mail/itsm/test_d2_mfa.py

"""D-2 MFA (TOTP) 테스트"""
import sys, ast, os

os.environ.setdefault("GUARDIA_SECRET_KEY", "test-mfa-secret-key-32bytes-pad!!")
os.environ.setdefault("DATABASE_URL", "sqlite+aiosqlite:///./test_d2.db")

# ── 1. 구문 검사 ─────────────────────────────────────────────────────────────
print("=== 1. 구문 검사 ===")
files = ["core/mfa.py", "routers/auth.py", "models.py"]
ok = True
for f in files:
    try:
        with open(f, encoding="utf-8") as fh:
            src = fh.read()
        ast.parse(src)
        print(f"  OK  {f}")
    except SyntaxError as e:
        print(f"  ERR {f}: {e}")
        ok = False
if not ok:
    sys.exit(1)

# ── 2. AES-GCM 암호화/복호화 라운드트립 ─────────────────────────────────────
print("\n=== 2. AES-GCM 암호화/복호화 테스트 ===")
from core.mfa import (
    generate_totp_secret,
    get_totp_uri,
    verify_totp,
    generate_qr_base64,
    encrypt_totp_secret,
    decrypt_totp_secret,
    create_mfa_pending_token,
    verify_mfa_pending_token,
    _verify_totp_fallback,
)

secret = generate_totp_secret()
print(f"  OK  generate_totp_secret: {secret[:8]}... (len={len(secret)})")

enc = encrypt_totp_secret(secret)
dec = decrypt_totp_secret(enc)
assert dec == secret, f"복호화 실패: {dec!r} != {secret!r}"
print(f"  OK  AES-GCM 암호화/복호화 라운드트립")

# 다른 시크릿과 다른지 확인
enc2 = encrypt_totp_secret(secret)
assert enc != enc2, "동일 평문이 항상 다른 ciphertext 생성 (nonce 랜덤화)"
print(f"  OK  랜덤 nonce로 매번 다른 ciphertext 생성")

# ── 3. TOTP URI 생성 ─────────────────────────────────────────────────────────
print("\n=== 3. TOTP URI 생성 테스트 ===")
uri = get_totp_uri("test_user", secret)
assert "otpauth://totp/" in uri, f"URI 포맷 오류: {uri[:60]}"
assert "GUARDiA" in uri or "guardia" in uri.lower(), f"issuer 없음: {uri[:80]}"
assert secret in uri, "secret 미포함"
print(f"  OK  TOTP URI: {uri[:60]}...")

# ── 4. QR 코드 생성 (옵션) ──────────────────────────────────────────────────
print("\n=== 4. QR 코드 생성 테스트 ===")
qr = generate_qr_base64(uri)
if qr:
    assert len(qr) > 100, "QR base64 너무 짧음"
    print(f"  OK  QR base64 생성 (len={len(qr)})")
else:
    print(f"  INFO  qrcode not installed - QR skip (returns None)")

# ── 5. TOTP 검증 (폴백 함수 사용) ───────────────────────────────────────────
print("\n=== 5. TOTP 검증 테스트 ===")
# 현재 코드 생성 후 검증 (pyotp 있는 경우)
try:
    import pyotp
    totp_obj = pyotp.TOTP(secret)
    current_code = totp_obj.now()
    result = verify_totp(secret, current_code)
    assert result == True, f"현재 코드 검증 실패: {current_code}"
    print(f"  OK  현재 TOTP 코드 검증 성공 (pyotp)")

    # 잘못된 코드
    assert verify_totp(secret, "000000") == False or verify_totp(secret, "000000") == True  # 우연히 맞을 수 있음
    assert verify_totp(secret, "abc123") == False, "비숫자 코드는 무효"
    assert verify_totp(secret, "12345") == False, "5자리는 무효"
    print(f"  OK  잘못된 코드 거부 (형식 검사)")
except ImportError:
    print(f"  INFO  pyotp not installed - using fallback")
    # 폴백은 실제 코드 없이 테스트 불가 (시간 기반)
    assert verify_totp(secret, "abc123") == False
    print(f"  OK  잘못된 코드 거부 (폴백)")

# ── 6. MFA 대기 토큰 발급/검증 ──────────────────────────────────────────────
print("\n=== 6. MFA 대기 토큰 테스트 ===")
mfa_token = create_mfa_pending_token("alice")
assert mfa_token, "토큰 생성 실패"
print(f"  OK  MFA 대기 토큰 생성: {mfa_token[:20]}...")

username = verify_mfa_pending_token(mfa_token)
assert username == "alice", f"사용자 추출 실패: {username}"
print(f"  OK  MFA 대기 토큰 검증: sub={username}")

# 일반 access_token은 mfa_pending 토큰으로 거부
from core.auth import create_access_token
normal_token = create_access_token({"sub": "alice", "role": "ENGINEER"})
assert verify_mfa_pending_token(normal_token) is None, "일반 토큰이 mfa_pending으로 통과됨!"
print(f"  OK  일반 access_token은 mfa_pending 검증 거부")

# 위조된 토큰 거부
assert verify_mfa_pending_token("invalid.token.here") is None
print(f"  OK  위조 토큰 거부")

# ── 7. models.py User 컬럼 확인 ─────────────────────────────────────────────
print("\n=== 7. User 모델 MFA 컬럼 확인 ===")
from models import User, UserOut
cols = {c.key for c in User.__table__.columns}
for col in ["mfa_enabled", "totp_secret_enc"]:
    status = "OK" if col in cols else "ERR"
    print(f"  {status}  User.{col}")
    if status == "ERR":
        ok = False

out_fields = UserOut.model_fields
status = "OK" if "mfa_enabled" in out_fields else "ERR"
print(f"  {status}  UserOut.mfa_enabled")

# ── 8. routers/auth.py 엔드포인트 확인 ─────────────────────────────────────
print("\n=== 8. auth.py MFA 엔드포인트 확인 ===")
with open("routers/auth.py", encoding="utf-8") as f:
    auth_src = f.read()

endpoints = [
    ("/mfa/setup",      "mfa_setup"),
    ("/mfa/enable",     "mfa_enable"),
    ("/mfa/disable",    "mfa_disable"),
    ("/mfa/status",     "mfa_status"),
    ("/login/mfa",      "login_mfa"),
    ("/admin/users",    "admin_mfa_reset"),
]
for path, fn in endpoints:
    status = "OK" if path in auth_src and fn in auth_src else "ERR"
    print(f"  {status}  {path} ({fn})")

print("\n=== 테스트 완료: D-2 MFA (TOTP) ===")
if ok:
    print("모든 검사 통과")
else:
    sys.exit(1)