"""远程图片下载的 SSRF 防护回归测试。"""
from __future__ import annotations

import sys
import unittest
from unittest.mock import patch
from pathlib import Path
from io import BytesIO

import httpx
from PIL import Image

SERVER_DIR = Path(__file__).resolve().parents[1]
if str(SERVER_DIR) not in sys.path:
    sys.path.insert(0, str(SERVER_DIR))

from services.storage import (
    ResourceDownloadTimeoutError,
    ImageValidationError,
    UnsafeRemoteUrlError,
    download_bytes,
    normalize_image_bytes,
    validate_remote_url,
)


class RemoteUrlValidationTests(unittest.IsolatedAsyncioTestCase):
    async def test_allows_public_literal_ip(self) -> None:
        await validate_remote_url("https://1.1.1.1/image.jpg")

    async def test_rejects_local_and_private_targets(self) -> None:
        for url in (
            "http://localhost:8800/image.jpg",
            "http://127.0.0.1/image.jpg",
            "http://10.0.0.1/image.jpg",
            "http://169.254.169.254/latest/meta-data/",
            "http://[::1]/image.jpg",
        ):
            with self.subTest(url=url), self.assertRaises(UnsafeRemoteUrlError):
                await validate_remote_url(url)

    async def test_download_rejects_before_opening_a_local_connection(self) -> None:
        with self.assertRaises(UnsafeRemoteUrlError):
            await download_bytes("http://127.0.0.1:1/image.jpg")

    async def test_redirect_target_is_validated_again(self) -> None:
        async def handler(request: httpx.Request) -> httpx.Response:
            self.assertEqual(request.url.host, "1.1.1.1")
            return httpx.Response(
                302,
                headers={"location": "http://127.0.0.1:9000/private.jpg"},
                request=request,
            )

        real_async_client = httpx.AsyncClient
        with patch(
            "services.storage.httpx.AsyncClient",
            side_effect=lambda **kwargs: real_async_client(
                transport=httpx.MockTransport(handler), **kwargs
            ),
        ), self.assertRaises(UnsafeRemoteUrlError):
            await download_bytes("https://1.1.1.1/redirect.jpg")

    async def test_http_client_timeout_is_normalized(self) -> None:
        async def handler(request: httpx.Request) -> httpx.Response:
            raise httpx.ReadTimeout("slow response", request=request)

        real_async_client = httpx.AsyncClient
        with patch(
            "services.storage.httpx.AsyncClient",
            side_effect=lambda **kwargs: real_async_client(
                transport=httpx.MockTransport(handler), **kwargs
            ),
        ), self.assertRaises(ResourceDownloadTimeoutError):
            await download_bytes("https://1.1.1.1/slow.jpg")


class ImageNormalizationTests(unittest.TestCase):
    def test_uses_real_format_and_corrects_exif_orientation(self) -> None:
        source = Image.new("RGB", (30, 10), "red")
        exif = Image.Exif()
        exif[274] = 6  # 顺时针旋转 90°
        raw = BytesIO()
        source.save(raw, format="JPEG", exif=exif)

        image = normalize_image_bytes(raw.getvalue(), max_pixels=1_000)

        self.assertEqual(image.content_type, "image/jpeg")
        self.assertEqual((image.width, image.height), (10, 30))
        self.assertTrue(image.orientation_normalized)
        with Image.open(BytesIO(image.data)) as stored:
            self.assertNotEqual(stored.getexif().get(274), 6)

    def test_rejects_invalid_content_and_pixel_bomb(self) -> None:
        with self.assertRaises(ImageValidationError):
            normalize_image_bytes(b"not an image")

        raw = BytesIO()
        Image.new("RGB", (11, 10), "blue").save(raw, format="PNG")
        with self.assertRaises(ImageValidationError):
            normalize_image_bytes(raw.getvalue(), max_pixels=100)


if __name__ == "__main__":
    unittest.main()