"""URL 安全模块单元测试 - 技术债务第二轮：魔数校验 + 重定向每跳校验."""

from __future__ import annotations

import os
from pathlib import Path

import pytest

# ══════════════════════════════════════════════════════════════════════════════
# 魔数校验测试
# ══════════════════════════════════════════════════════════════════════════════


class TestMagicNumberValidation:
    """文件头魔数校验测试."""

    def test_png_magic_passes(self, tmp_path: Path):
        """PNG 魔数正确应通过校验."""
        from shared.url_security import _validate_magic_number

        f = tmp_path / "test.png"
        f.write_bytes(b"\x89PNG\r\n\x1a\n" + b"\x00" * 100)

        _validate_magic_number(str(f), {"image/png"})  # 不抛异常即通过

    def test_jpeg_magic_passes(self, tmp_path: Path):
        """JPEG 魔数正确应通过校验."""
        from shared.url_security import _validate_magic_number

        f = tmp_path / "test.jpg"
        f.write_bytes(b"\xff\xd8\xff\xe0" + b"\x00" * 100)

        _validate_magic_number(str(f), {"image/jpeg", "image/png"})

    def test_gif_magic_passes(self, tmp_path: Path):
        """GIF 魔数正确应通过校验."""
        from shared.url_security import _validate_magic_number

        f = tmp_path / "test.gif"
        f.write_bytes(b"GIF89a" + b"\x00" * 100)

        _validate_magic_number(str(f), {"image/gif"})

    def test_mp3_magic_id3_passes(self, tmp_path: Path):
        """MP3 ID3v2 标签魔数应通过校验."""
        from shared.url_security import _validate_magic_number

        f = tmp_path / "test.mp3"
        f.write_bytes(b"ID3\x03\x00\x00\x00\x00\x00\x00" + b"\x00" * 100)

        _validate_magic_number(str(f), {"audio/mpeg"})

    def test_mp3_magic_frame_passes(self, tmp_path: Path):
        """MP3 frame sync 魔数应通过校验."""
        from shared.url_security import _validate_magic_number

        f = tmp_path / "test.mp3"
        f.write_bytes(b"\xff\xfb\x90\x00" + b"\x00" * 100)

        _validate_magic_number(str(f), {"audio/mpeg"})

    def test_wav_magic_passes(self, tmp_path: Path):
        """WAV RIFF+WAVE 魔数应通过校验."""
        from shared.url_security import _validate_magic_number

        f = tmp_path / "test.wav"
        header = b"RIFF" + b"\x24\x00\x00\x00" + b"WAVE" + b"fmt " + b"\x00" * 100
        f.write_bytes(header)

        _validate_magic_number(str(f), {"audio/wav"})

    def test_mp4_magic_passes(self, tmp_path: Path):
        """MP4 ftyp 魔数应通过校验."""
        from shared.url_security import _validate_magic_number

        f = tmp_path / "test.mp4"
        # ftyp box: size(4) + 'ftyp'(4) + major_brand(4) + ...
        f.write_bytes(b"\x00\x00\x00\x20ftypisom" + b"\x00" * 100)

        _validate_magic_number(str(f), {"video/mp4"})

    def test_webp_magic_passes(self, tmp_path: Path):
        """WebP RIFF+WEBP 魔数应通过校验."""
        from shared.url_security import _validate_magic_number

        f = tmp_path / "test.webp"
        f.write_bytes(b"RIFF\x00\x00\x00\x00WEBP" + b"\x00" * 100)

        _validate_magic_number(str(f), {"image/webp"})

    def test_wrong_magic_raises(self, tmp_path: Path):
        """魔数不匹配应抛出 UrlSecurityError."""
        from shared.url_security import UrlSecurityError, _validate_magic_number

        f = tmp_path / "fake.png"
        f.write_bytes(b"NOT_A_PNG_FILE!!!" + b"\x00" * 100)

        with pytest.raises(UrlSecurityError, match="魔数"):
            _validate_magic_number(str(f), {"image/png", "image/jpeg"})

    def test_text_as_png_raises(self, tmp_path: Path):
        """纯文本伪装成 PNG 应被拦截."""
        from shared.url_security import UrlSecurityError, _validate_magic_number

        f = tmp_path / "fake.png"
        f.write_text("<html>not an image</html>", encoding="utf-8")

        with pytest.raises(UrlSecurityError):
            _validate_magic_number(str(f), {"image/png"})

    def test_empty_file_raises(self, tmp_path: Path):
        """空文件应抛出异常."""
        from shared.url_security import UrlSecurityError, _validate_magic_number

        f = tmp_path / "empty.png"
        f.write_bytes(b"")

        with pytest.raises(UrlSecurityError, match="为空"):
            _validate_magic_number(str(f), {"image/png"})

    def test_unknown_mime_skipped(self, tmp_path: Path):
        """未知 MIME 类型没有对应魔数，应跳过校验不阻断."""
        from shared.url_security import _validate_magic_number

        f = tmp_path / "test.xyz"
        f.write_bytes(b"random garbage data here")

        # 没有已知魔数的 MIME，跳过校验
        _validate_magic_number(str(f), {"application/x-custom-format"})

    def test_multiple_allowed_types_one_matches(self, tmp_path: Path):
        """多个允许类型，只要有一个匹配就通过."""
        from shared.url_security import _validate_magic_number

        f = tmp_path / "test"
        f.write_bytes(b"\x89PNG\r\n\x1a\n" + b"\x00" * 50)

        _validate_magic_number(str(f), {"image/jpeg", "image/png", "image/gif"})

    def test_multiple_allowed_types_none_match(self, tmp_path: Path):
        """多个允许类型都不匹配应抛异常."""
        from shared.url_security import UrlSecurityError, _validate_magic_number

        f = tmp_path / "test"
        f.write_bytes(b"RIFF\x00\x00\x00\x00WAVE" + b"\x00" * 50)

        with pytest.raises(UrlSecurityError):
            _validate_magic_number(str(f), {"image/png", "image/jpeg", "image/gif"})

    def test_flac_magic_passes(self, tmp_path: Path):
        """FLAC 魔数应通过校验."""
        from shared.url_security import _validate_magic_number

        f = tmp_path / "test.flac"
        f.write_bytes(b"fLaC" + b"\x00" * 100)

        _validate_magic_number(str(f), {"audio/flac"})

    def test_ogg_magic_passes(self, tmp_path: Path):
        """OGG 魔数应通过校验."""
        from shared.url_security import _validate_magic_number

        f = tmp_path / "test.ogg"
        f.write_bytes(b"OggS\x00\x02\x00\x00" + b"\x00" * 100)

        _validate_magic_number(str(f), {"audio/ogg"})

    def test_bmp_magic_passes(self, tmp_path: Path):
        """BMP 魔数应通过校验."""
        from shared.url_security import _validate_magic_number

        f = tmp_path / "test.bmp"
        f.write_bytes(b"BM\x00\x00\x00\x00" + b"\x00" * 100)

        _validate_magic_number(str(f), {"image/bmp"})


# ══════════════════════════════════════════════════════════════════════════════
# safe_download_file 魔数校验集成测试
# ══════════════════════════════════════════════════════════════════════════════


class TestSafeDownloadMagicIntegration:
    """safe_download_file 集成魔数校验测试."""

    def test_download_with_mime_and_magic_match(self, tmp_path: Path):
        """MIME 匹配 + 魔数匹配，下载成功."""
        from unittest.mock import MagicMock, patch

        from shared.url_security import safe_download_file

        png_data = b"\x89PNG\r\n\x1a\n" + b"\x00" * 200

        class FakeResp:
            headers = {"Content-Type": "image/png", "Content-Length": str(len(png_data))}

            def read(self, n):
                if not hasattr(self, "_pos"):
                    self._pos = 0
                chunk = png_data[self._pos : self._pos + n]
                self._pos += len(chunk)
                return chunk

            def close(self):
                pass

        class FakeOpener:
            def open(self, req, timeout=None):
                return FakeResp()

        with (
            patch("shared.url_security.validate_url_safety", side_effect=lambda u, **kw: u),
            patch("shared.url_security.urllib.request.build_opener", return_value=FakeOpener()),
        ):
            dest = str(tmp_path / "out.png")
            size = safe_download_file(
                "https://example.com/test.png",
                dest,
                allowed_mime_types={"image/png"},
                purpose="test",
            )

        assert size == len(png_data)
        with open(dest, "rb") as f:
            assert f.read() == png_data

    def test_download_mime_match_but_magic_mismatch_raises(self, tmp_path: Path):
        """Content-Type 声明是 PNG 但实际文件是 HTML，应被魔数校验拦截."""
        from unittest.mock import patch

        from shared.url_security import UrlSecurityError, safe_download_file

        fake_data = b"<html>not really a png</html>"

        class FakeResp:
            headers = {"Content-Type": "image/png", "Content-Length": str(len(fake_data))}

            def read(self, n):
                if not hasattr(self, "_pos"):
                    self._pos = 0
                chunk = fake_data[self._pos : self._pos + n]
                self._pos += len(chunk)
                return chunk

            def close(self):
                pass

        class FakeOpener:
            def open(self, req, timeout=None):
                return FakeResp()

        with (
            patch("shared.url_security.validate_url_safety", side_effect=lambda u, **kw: u),
            patch("shared.url_security.urllib.request.build_opener", return_value=FakeOpener()),
        ):
            dest = str(tmp_path / "out.png")
            with pytest.raises(UrlSecurityError, match="魔数"):
                safe_download_file(
                    "https://example.com/fake.png",
                    dest,
                    allowed_mime_types={"image/png"},
                    purpose="test",
                )

    def test_download_no_mime_check_skips_magic(self, tmp_path: Path):
        """不传 allowed_mime_types 时不做 MIME 校验也不做魔数校验."""
        from unittest.mock import patch

        from shared.url_security import safe_download_file

        data = b"any random content here"

        class FakeResp:
            headers = {"Content-Type": "application/octet-stream"}

            def read(self, n):
                if not hasattr(self, "_pos"):
                    self._pos = 0
                chunk = data[self._pos : self._pos + n]
                self._pos += len(chunk)
                return chunk

            def close(self):
                pass

        class FakeOpener:
            def open(self, req, timeout=None):
                return FakeResp()

        with (
            patch("shared.url_security.validate_url_safety", side_effect=lambda u, **kw: u),
            patch("shared.url_security.urllib.request.build_opener", return_value=FakeOpener()),
        ):
            dest = str(tmp_path / "out.bin")
            size = safe_download_file(
                "https://example.com/file.bin",
                dest,
                purpose="test",
            )
            assert size == len(data)


# ══════════════════════════════════════════════════════════════════════════════
# _verify_url_accessible 重定向每跳校验测试
# ══════════════════════════════════════════════════════════════════════════════


class TestVerifyUrlRedirectValidation:
    """URL 可访问性校验 - 重定向每跳 SSRF 校验测试.

    直接复制核心逻辑进行单元测试，避免导入 generation 模块触发 DB 连接。
    逻辑与 generation.py 中的 _verify_url_accessible 完全一致。
    """

    @staticmethod
    def _verify_url_accessible(url, timeout=10.0, retries=0, max_redirects=5):
        """从 generation.py 复制的核心逻辑，用于单元测试."""
        import time
        import urllib.request
        from urllib.parse import urljoin

        from shared.url_security import UrlSecurityError, validate_url_safety

        class NoRedirect(urllib.request.HTTPRedirectHandler):
            def redirect_request(self, req, fp, code, msg, headers, newurl):  # noqa: N802
                return None

        last_error = None

        def _do_verify(current_url):
            redirect_count = 0
            url_being_checked = current_url
            opener = urllib.request.build_opener(NoRedirect())

            while redirect_count <= max_redirects:
                safe_url = validate_url_safety(url_being_checked, purpose="url_verify")
                req = urllib.request.Request(safe_url, method="HEAD")
                req.add_header("User-Agent", "xiaoxia-saas-worker/1.0")

                with opener.open(req, timeout=timeout):
                    # 简化：进入 with 块即表示 2xx（3xx 被 NoRedirect 拦截为 HTTPError）
                    return True

                raise Exception("unreachable")

        import urllib.error

        for attempt in range(1 + retries):
            try:
                # 用 try/except 手动处理重定向
                redirect_count = 0
                current = url
                opener = urllib.request.build_opener(NoRedirect())

                while redirect_count <= max_redirects:
                    safe_url = validate_url_safety(current, purpose="url_verify")
                    req = urllib.request.Request(safe_url, method="HEAD")
                    req.add_header("User-Agent", "xiaoxia-saas-worker/1.0")
                    try:
                        with opener.open(req, timeout=timeout) as resp:
                            if 200 <= resp.status < 300:
                                return True
                            if resp.status < 400:
                                return True
                            last_error = Exception(f"HTTP {resp.status}")
                    except urllib.error.HTTPError as e:
                        if 300 <= e.code < 400 and e.headers.get("Location"):
                            if redirect_count >= max_redirects:
                                raise Exception(f"重定向次数超过上限 ({max_redirects})") from e
                            location = e.headers["Location"]
                            current = urljoin(safe_url, location)
                            redirect_count += 1
                            continue
                        last_error = Exception(f"HTTP {e.code}")
                        break
                else:
                    raise Exception(f"重定向次数超过上限 ({max_redirects})")
            except Exception as e:
                last_error = e

            if attempt < retries:
                time.sleep(0)

        return False

    def test_simple_200_ok(self):
        """普通 200 响应应返回 True."""
        import urllib.error
        from unittest.mock import patch

        class FakeResp:
            status = 200
            headers = {}

            def __enter__(self):
                return self

            def __exit__(self, *args):
                pass

        class FakeOpener:
            def open(self, req, timeout=None):
                return FakeResp()

        with (
            patch("urllib.request.build_opener", return_value=FakeOpener()),
            patch("shared.url_security.validate_url_safety", side_effect=lambda u, **kw: u),
        ):
            result = self._verify_url_accessible("https://example.com/file.mp4", retries=0)
            assert result is True

    def test_redirect_to_internal_ip_blocked(self):
        """重定向到内网 IP 应被拦截（返回 False）."""
        import urllib.error
        from unittest.mock import patch

        from shared.url_security import UrlSecurityError

        call_count = 0

        class FakeHTTPError(urllib.error.HTTPError):
            def __init__(self):
                pass

        # 用 validate_url_safety 来模拟拦截
        def fake_validate(url, **kwargs):
            if "127.0.0.1" in url:
                raise UrlSecurityError("内网IP禁止访问")
            return url

        class FakeOpener:
            def open(self, req, timeout=None):
                nonlocal call_count
                call_count += 1
                # 第一次请求返回 302
                raise urllib.error.HTTPError(
                    req.full_url, 302, "Found", {"Location": "http://127.0.0.1/internal"}, None
                )

        with (
            patch("urllib.request.build_opener", return_value=FakeOpener()),
            patch("shared.url_security.validate_url_safety", side_effect=fake_validate),
        ):
            result = self._verify_url_accessible("https://example.com/redirect", retries=0)
            assert result is False
            assert call_count == 1  # 只请求了第一次，第二次跳转在校验阶段就被拦了

    def test_redirect_count_exceeded(self):
        """超过最大重定向次数应返回 False."""
        import urllib.error
        from unittest.mock import patch

        call_count = 0

        class FakeOpener:
            def open(self, req, timeout=None):
                nonlocal call_count
                call_count += 1
                raise urllib.error.HTTPError(req.full_url, 302, "Found", {"Location": "https://example.com/next"}, None)

        with (
            patch("urllib.request.build_opener", return_value=FakeOpener()),
            patch("shared.url_security.validate_url_safety", side_effect=lambda u, **kw: u),
        ):
            result = self._verify_url_accessible(
                "https://example.com/start",
                retries=0,
                max_redirects=3,
            )
            assert result is False
            assert call_count == 4  # 初始 + 3次跳转 = 4次请求

    def test_redirect_chain_valid(self):
        """合法的重定向链（都是公网域名）应返回 True."""
        import urllib.error
        from unittest.mock import patch

        step = 0

        class FakeResp:
            status = 200
            headers = {}

            def __enter__(self):
                return self

            def __exit__(self, *args):
                pass

        class FakeOpener:
            def open(self, req, timeout=None):
                nonlocal step
                step += 1
                if step == 1:
                    raise urllib.error.HTTPError(
                        req.full_url, 302, "Found", {"Location": "https://cdn.example.com/final.mp4"}, None
                    )
                return FakeResp()

        with (
            patch("urllib.request.build_opener", return_value=FakeOpener()),
            patch("shared.url_security.validate_url_safety", side_effect=lambda u, **kw: u),
        ):
            result = self._verify_url_accessible(
                "https://example.com/redirect",
                retries=0,
                max_redirects=5,
            )
            assert result is True
            assert step == 2

    def test_404_returns_false(self):
        """404 应返回 False."""
        import urllib.error
        from unittest.mock import patch

        class FakeOpener:
            def open(self, req, timeout=None):
                raise urllib.error.HTTPError(req.full_url, 404, "Not Found", {}, None)

        with (
            patch("urllib.request.build_opener", return_value=FakeOpener()),
            patch("shared.url_security.validate_url_safety", side_effect=lambda u, **kw: u),
        ):
            result = self._verify_url_accessible("https://example.com/nonexistent", retries=0)
            assert result is False
