"""Test SQLMap exploitation chain limit.

This test verifies that:
1. Only 3 unique endpoints get exploitation chains (--dbs)
2. Same endpoint with different query params is treated as ONE endpoint
3. Duplicate --dbs jobs for same endpoint are skipped
"""
import pytest
from unittest.mock import patch, MagicMock


class TestSqlmapExploitationLimit:
    """Test the sqlmap exploitation limit logic."""

    def test_base_endpoint_strips_query_params(self):
        """URLs with different query params should normalize to same endpoint."""
        from urllib.parse import urlparse

        def _get_base_endpoint(url):
            if not url or "://" not in url:
                return url
            try:
                parsed = urlparse(url)
                hostname = parsed.netloc.rstrip(".")
                if parsed.scheme == "http" and hostname.endswith(":80"):
                    hostname = hostname[:-3]
                elif parsed.scheme == "https" and hostname.endswith(":443"):
                    hostname = hostname[:-4]
                return f"https://{hostname}{parsed.path}"
            except Exception:
                return url

        # Same endpoint, different query params
        url1 = "http://testasp.vulnweb.com/showforum.asp?id=0"
        url2 = "http://testasp.vulnweb.com/showforum.asp?id=1"
        url3 = "http://testasp.vulnweb.com/showforum.asp?id=2"

        base1 = _get_base_endpoint(url1)
        base2 = _get_base_endpoint(url2)
        base3 = _get_base_endpoint(url3)

        assert base1 == base2 == base3, f"Should be same endpoint: {base1}, {base2}, {base3}"
        assert base1 == "https://testasp.vulnweb.com/showforum.asp"

    def test_different_endpoints_are_different(self):
        """Different paths should be different endpoints."""
        from urllib.parse import urlparse

        def _get_base_endpoint(url):
            if not url or "://" not in url:
                return url
            try:
                parsed = urlparse(url)
                hostname = parsed.netloc.rstrip(".")
                if parsed.scheme == "http" and hostname.endswith(":80"):
                    hostname = hostname[:-3]
                elif parsed.scheme == "https" and hostname.endswith(":443"):
                    hostname = hostname[:-4]
                return f"https://{hostname}{parsed.path}"
            except Exception:
                return url

        url1 = "http://testasp.vulnweb.com/showforum.asp?id=0"
        url2 = "http://testasp.vulnweb.com/login.asp"
        url3 = "http://testphp.vulnweb.com/search.php?q=test"

        base1 = _get_base_endpoint(url1)
        base2 = _get_base_endpoint(url2)
        base3 = _get_base_endpoint(url3)

        assert base1 != base2, "showforum.asp and login.asp should be different"
        assert base1 != base3, "Different hosts should be different"
        assert base2 != base3, "Different hosts should be different"

    def test_exploitation_limit_logic(self):
        """Test the actual limit logic with mock existing jobs."""
        from urllib.parse import urlparse

        def _get_base_endpoint(url):
            if not url or "://" not in url:
                return url
            try:
                parsed = urlparse(url)
                hostname = parsed.netloc.rstrip(".")
                if parsed.scheme == "http" and hostname.endswith(":80"):
                    hostname = hostname[:-3]
                elif parsed.scheme == "https" and hostname.endswith(":443"):
                    hostname = hostname[:-4]
                return f"https://{hostname}{parsed.path}"
            except Exception:
                return url

        MAX_SQLMAP_EXPLOITATION_CHAINS = 3
        exploitation_args = {"--dbs", "--privileges", "--is-dba", "--users", "--passwords"}
        engagement_id = 1

        # Simulate existing jobs - 3 unique endpoints already being exploited
        existing_jobs = [
            {"tool": "sqlmap", "engagement_id": 1, "status": "running",
             "target": "http://test.com/page1.php?id=1",
             "args": ["-u", "http://test.com/page1.php?id=1", "--dbs", "--batch"]},
            {"tool": "sqlmap", "engagement_id": 1, "status": "queued",
             "target": "http://test.com/page2.php?id=1",
             "args": ["-u", "http://test.com/page2.php?id=1", "--dbs", "--batch"]},
            {"tool": "sqlmap", "engagement_id": 1, "status": "done",
             "target": "http://test.com/page3.php?id=1",
             "args": ["-u", "http://test.com/page3.php?id=1", "--dbs", "--batch"]},
        ]

        def should_allow(cmd_target, cmd_args):
            """Replicate the limit check logic."""
            is_exploitation_start = any(arg in exploitation_args for arg in cmd_args)

            if not is_exploitation_start:
                return True  # Not an exploitation job, allow

            # Build set of exploited endpoints
            exploited_endpoints = set()
            for ej in existing_jobs:
                if (
                    ej.get("tool") == "sqlmap"
                    and ej.get("engagement_id") == engagement_id
                    and ej.get("status") not in ["error"]
                ):
                    ej_args = ej.get("args", [])
                    if any(arg in exploitation_args for arg in ej_args):
                        ej_target = ej.get("target", "")
                        if ej_target:
                            exploited_endpoints.add(_get_base_endpoint(ej_target))

            cmd_base_endpoint = _get_base_endpoint(cmd_target)

            # Skip if endpoint already being exploited
            if cmd_base_endpoint in exploited_endpoints:
                return False

            # Skip if limit reached
            if len(exploited_endpoints) >= MAX_SQLMAP_EXPLOITATION_CHAINS:
                return False

            return True

        # Test cases
        # 1. Same endpoint as existing (different param) - should SKIP
        assert should_allow("http://test.com/page1.php?id=99", ["--dbs"]) == False, \
            "Should skip - same endpoint as page1.php"

        # 2. New endpoint (4th unique) - should SKIP (limit reached)
        assert should_allow("http://test.com/page4.php?id=1", ["--dbs"]) == False, \
            "Should skip - 4th unique endpoint, limit is 3"

        # 3. Non-exploitation args - should ALLOW
        assert should_allow("http://test.com/page5.php?id=1", ["--tables"]) == True, \
            "Should allow - not exploitation start args"

    def test_exploitation_limit_allows_first_three(self):
        """Test that first 3 unique endpoints are allowed."""
        from urllib.parse import urlparse

        def _get_base_endpoint(url):
            if not url or "://" not in url:
                return url
            try:
                parsed = urlparse(url)
                hostname = parsed.netloc.rstrip(".")
                return f"https://{hostname}{parsed.path}"
            except Exception:
                return url

        MAX_SQLMAP_EXPLOITATION_CHAINS = 3
        exploitation_args = {"--dbs", "--privileges", "--is-dba", "--users", "--passwords"}
        engagement_id = 1

        # Start with no existing jobs
        existing_jobs = []
        allowed_endpoints = []

        def should_allow_and_track(cmd_target, cmd_args):
            """Check if allowed and track if so."""
            is_exploitation_start = any(arg in exploitation_args for arg in cmd_args)

            if not is_exploitation_start:
                return True

            exploited_endpoints = set()
            for ej in existing_jobs:
                if (
                    ej.get("tool") == "sqlmap"
                    and ej.get("engagement_id") == engagement_id
                    and ej.get("status") not in ["error"]
                ):
                    ej_args = ej.get("args", [])
                    if any(arg in exploitation_args for arg in ej_args):
                        ej_target = ej.get("target", "")
                        if ej_target:
                            exploited_endpoints.add(_get_base_endpoint(ej_target))

            cmd_base_endpoint = _get_base_endpoint(cmd_target)

            if cmd_base_endpoint in exploited_endpoints:
                return False

            if len(exploited_endpoints) >= MAX_SQLMAP_EXPLOITATION_CHAINS:
                return False

            return True

        # Simulate 5 --dbs requests coming in sequence
        test_urls = [
            "http://test.com/page1.php?id=1",  # 1st endpoint - allow
            "http://test.com/page2.php?id=1",  # 2nd endpoint - allow
            "http://test.com/page1.php?id=2",  # same as 1st - skip
            "http://test.com/page3.php?id=1",  # 3rd endpoint - allow
            "http://test.com/page4.php?id=1",  # 4th endpoint - skip (limit)
            "http://test.com/page2.php?id=5",  # same as 2nd - skip
        ]

        results = []
        for url in test_urls:
            allowed = should_allow_and_track(url, ["--dbs"])
            results.append((url, allowed))

            # If allowed, add to existing_jobs (simulating job creation)
            if allowed:
                existing_jobs.append({
                    "tool": "sqlmap",
                    "engagement_id": 1,
                    "status": "queued",
                    "target": url,
                    "args": ["-u", url, "--dbs", "--batch"]
                })

        # Verify results
        assert results[0] == ("http://test.com/page1.php?id=1", True), "1st endpoint should be allowed"
        assert results[1] == ("http://test.com/page2.php?id=1", True), "2nd endpoint should be allowed"
        assert results[2] == ("http://test.com/page1.php?id=2", False), "Duplicate of 1st should be skipped"
        assert results[3] == ("http://test.com/page3.php?id=1", True), "3rd endpoint should be allowed"
        assert results[4] == ("http://test.com/page4.php?id=1", False), "4th endpoint should be skipped (limit)"
        assert results[5] == ("http://test.com/page2.php?id=5", False), "Duplicate of 2nd should be skipped"

        # Should have exactly 3 jobs
        assert len(existing_jobs) == 3, f"Should have 3 jobs, got {len(existing_jobs)}"
