import re

from ..bbot_fixtures import *

from bbot.modules.base import BaseModule
from bbot.modules.output.base import BaseOutputModule
from bbot.modules.report.base import BaseReportModule
from bbot.modules.internal.base import BaseInternalModule


@pytest.mark.asyncio
async def test_modules_basic_checks(events, blasthttp_mock):
    from bbot.scanner import Scanner

    scan = Scanner(config={"omit_event_types": ["URL_UNVERIFIED"]})
    await scan._prep()
    assert "URL_UNVERIFIED" in scan.omitted_event_types

    # output module specific event filtering tests
    base_output_module_1 = BaseOutputModule(scan)
    base_output_module_1.watched_events = ["IP_ADDRESS", "URL_UNVERIFIED"]
    localhost = scan.make_event("127.0.0.1", parent=scan.root_event)
    # ip addresses should be accepted
    result, reason = base_output_module_1._event_precheck(localhost)
    assert result is True
    assert reason == "precheck succeeded"
    # internal events should be rejected
    localhost._internal = True
    result, reason = base_output_module_1._event_precheck(localhost)
    assert result is False
    assert reason == "event is internal and output modules don't accept internal events"
    localhost._internal = False
    result, reason = base_output_module_1._event_precheck(localhost)
    assert result is True
    assert reason == "precheck succeeded"
    # unwatched events should be rejected
    dns_name = scan.make_event("evilcorp.com", parent=scan.root_event)
    result, reason = base_output_module_1._event_precheck(dns_name)
    assert result is False
    assert reason == "its type is not in watched_events"
    # omitted events matching watched types should be accepted
    url_unverified = scan.make_event("http://127.0.0.1", "URL_UNVERIFIED", parent=scan.root_event)
    url_unverified._omit = True
    result, reason = base_output_module_1._event_precheck(url_unverified)
    assert result is True
    assert reason == "its type is explicitly in watched_events"

    base_output_module_2 = BaseOutputModule(scan)
    base_output_module_2.watched_events = ["*"]
    # normal events should be accepted
    localhost = scan.make_event("127.0.0.1", parent=scan.root_event)
    result, reason = base_output_module_2._event_precheck(localhost)
    assert result is True
    assert reason == "precheck succeeded"
    # internal events should be rejected
    localhost._internal = True
    result, reason = base_output_module_2._event_precheck(localhost)
    assert result is False
    assert reason == "event is internal and output modules don't accept internal events"
    localhost._internal = False
    result, reason = base_output_module_2._event_precheck(localhost)
    assert result is True
    assert reason == "precheck succeeded"
    # omitted events should be rejected
    localhost._omit = True
    result, reason = base_output_module_2._event_precheck(localhost)
    assert result is False
    assert reason == "its type is omitted in the config"
    # normal event should be accepted
    url_unverified = scan.make_event("http://127.0.0.1", "URL_UNVERIFIED", parent=scan.root_event)
    result, reason = base_output_module_2._event_precheck(url_unverified)
    assert result is True
    assert reason == "precheck succeeded"
    # omitted event types should be marked during scan egress
    await scan.egress_module.handle_event(url_unverified)
    result, reason = base_output_module_2._event_precheck(url_unverified)
    assert result is False
    assert reason == "its type is omitted in the config"

    egress_module = scan.egress_module
    url = scan.make_event("http://evilcorp.com", "URL_UNVERIFIED", parent=scan.root_event, tags=["target"])
    assert url._omit is False
    # targets should not be omitted
    await egress_module.handle_event(url)
    assert url._omit is False
    # non-targets should be omitted
    url = scan.make_event("http://evilcorp.com", "URL_UNVERIFIED", parent=scan.root_event)
    await egress_module.handle_event(url)
    assert url._omit is True

    # omitted always_emit events should still be omitted
    finding_data = {
        "host": "evilcorp.com",
        "description": "test",
        "severity": "LOW",
        "confidence": "LOW",
        "name": "test",
    }
    finding = scan.make_event(finding_data, "FINDING", parent=scan.root_event)
    assert finding.always_emit is True
    finding._omit = True
    result, reason = base_output_module_2._event_precheck(finding)
    assert result is False
    assert reason == "its type is omitted in the config"

    # always_emit should bypass the internal check
    finding_data2 = {
        "host": "evilcorp.com",
        "description": "test2",
        "severity": "LOW",
        "confidence": "LOW",
        "name": "test2",
    }
    finding2 = scan.make_event(finding_data2, "FINDING", parent=scan.root_event)
    assert finding2.always_emit is True
    finding2._internal = True
    result, reason = base_output_module_2._event_precheck(finding2)
    assert result is True
    assert reason == "event is always emitted"

    # common event filtering tests
    for module_class in (BaseModule, BaseOutputModule, BaseReportModule, BaseInternalModule):
        base_module = module_class(scan)
        localhost2 = scan.make_event("127.0.0.2", parent=events.subdomain)
        localhost2.scope_distance = 0
        # base cases
        base_module._watched_events = None
        base_module.watched_events = ["*"]
        assert base_module._event_precheck(events.emoji)[0] is True
        base_module._watched_events = None
        base_module.watched_events = ["IP_ADDRESS"]
        assert base_module._event_precheck(events.ipv4)[0] is True
        assert base_module._event_precheck(events.domain)[0] is False
        assert base_module._event_precheck(events.localhost)[0] is True
        assert base_module._event_precheck(localhost2)[0] is True
        # target only
        base_module.target_only = True
        assert base_module._event_precheck(localhost2)[0] is False
        localhost2.add_tag("target")
        assert base_module._event_precheck(localhost2)[0] is True
        base_module.target_only = False

        # in scope only
        base_module.in_scope_only = True
        localhost3 = scan.make_event("127.0.0.2", parent=events.subdomain)
        valid, reason = await base_module._event_postcheck(localhost3)
        if base_module._type == "output":
            assert valid
        else:
            assert not valid
            assert reason == "it did not meet in_scope_only filter criteria"
        base_module.in_scope_only = False
        base_module.scope_distance_modifier = 0
        valid, reason = await base_module._event_postcheck(events.localhost)
        assert valid

    # module preloading
    all_preloaded = DEFAULT_PRESET.module_loader.preloaded()
    assert "dnsbrute" in all_preloaded
    assert "DNS_NAME" in all_preloaded["dnsbrute"]["watched_events"]
    assert "DNS_NAME" in all_preloaded["dnsbrute"]["produced_events"]
    assert "subdomain-enum" in all_preloaded["dnsbrute"]["flags"]
    assert "wordlist" in all_preloaded["dnsbrute"]["config"]
    assert type(all_preloaded["dnsbrute"]["config"]["max_depth"]) == int
    assert all_preloaded["baddns"]["deps"]["pip"]
    assert all_preloaded["dnsbrute"]["deps"]["common"]
    assert all_preloaded["gowitness"]["deps"]["ansible"]

    all_flags = set()

    created_date_regex = re.compile(r"2[\d]{3}-[\d]{2}-[\d]{2}")
    for module_name, preloaded in all_preloaded.items():
        # either active or passive and never both
        flags = preloaded.get("flags", [])
        for flag in flags:
            all_flags.add(flag)
        if preloaded["type"] == "scan":
            assert ("active" in flags and "passive" not in flags) or ("active" not in flags and "passive" in flags), (
                f'module "{module_name}" must have either "active" or "passive" flag'
            )
            assert not ("web" in flags and "web-heavy" in flags), (
                f'module "{module_name}" should have either "web" or "web-heavy" flags, not both'
            )
            # every scan module must be classified as safe, loud, and/or invasive
            has_safe = "safe" in flags
            has_loud = "loud" in flags
            has_invasive = "invasive" in flags
            assert has_safe or has_loud or has_invasive, (
                f'module "{module_name}" must have at least one of "safe", "loud", or "invasive" flags'
            )
            assert not (has_safe and (has_loud or has_invasive)), (
                f'module "{module_name}" has "safe" flag but also has "loud" or "invasive" — these are mutually exclusive'
            )
        meta = preloaded.get("meta", {})
        # make sure every module has a description
        assert meta.get("description", ""), f"{module_name} must have a description"
        # make sure every module has an author
        assert meta.get("author", ""), f"{module_name} must have an author"
        # make sure every module has a created date
        created_date = meta.get("created_date", "")
        assert created_date, f"{module_name} must have a created date"
        assert created_date_regex.match(created_date), f"{module_name}'s created_date must match the format YYYY-MM-DD"

        # attribute checks
        watched_events = preloaded.get("watched_events")
        produced_events = preloaded.get("produced_events")

        assert type(watched_events) == list
        assert type(produced_events) == list
        if preloaded.get("type", "") not in ("internal",):
            assert watched_events, f"{module_name}.watched_events must not be empty"
        assert type(watched_events) == list, f"{module_name}.watched_events must be of type list"
        assert type(produced_events) == list, f"{module_name}.produced_events must be of type list"
        assert all(type(t) == str for t in watched_events), (
            f"{module_name}.watched_events entries must be of type string"
        )
        assert all(type(t) == str for t in produced_events), (
            f"{module_name}.produced_events entries must be of type string"
        )

        assert type(preloaded.get("deps_pip", [])) == list, f"{module_name}.deps_pip must be of type list"
        assert type(preloaded.get("deps_pip_constraints", [])) == list, (
            f"{module_name}.deps_pip_constraints must be of type list"
        )
        assert type(preloaded.get("deps_apt", [])) == list, f"{module_name}.deps_apt must be of type list"
        assert type(preloaded.get("deps_shell", [])) == list, f"{module_name}.deps_shell must be of type list"
        assert type(preloaded.get("config", None)) == dict, f"{module_name}.options must be of type list"
        assert type(preloaded.get("options_desc", None)) == dict, f"{module_name}.options_desc must be of type list"
        # options must have descriptions
        assert set(preloaded.get("config", {})) == set(preloaded.get("options_desc", {})), (
            f"{module_name}.options do not match options_desc"
        )
        # descriptions most not be blank
        assert all(o for o in preloaded.get("options_desc", {}).values()), (
            f"{module_name}.options_desc descriptions must not be blank"
        )

    from bbot.core.flags import flag_descriptions

    for flag in all_flags:
        assert flag in flag_descriptions, f'Flag "{flag}" not listed in bbot/core/flags.py'
        description = flag_descriptions.get(flag, "")
        assert description, f'Flag "{flag}" has no description in bbot/core/flags.py'

    await scan._cleanup()


@pytest.mark.asyncio
async def test_modules_basic_perhostonly(bbot_scanner):
    from bbot.modules.base import BaseModule

    class mod_normal(BaseModule):
        _name = "mod_normal"
        watched_events = ["*"]

    class mod_host_only(BaseModule):
        _name = "mod_hostonly"
        watched_events = ["*"]
        per_host_only = True

    class mod_hostport_only(BaseModule):
        _name = "mod_normal"
        watched_events = ["*"]
        per_hostport_only = True

    class mod_domain_only(BaseModule):
        _name = "domain_only"
        watched_events = ["*"]
        per_domain_only = True

    scan = bbot_scanner(
        "evilcorp.com",
        force_start=True,
    )

    await scan._prep()

    scan.modules["mod_normal"] = mod_normal(scan)
    scan.modules["mod_host_only"] = mod_host_only(scan)
    scan.modules["mod_hostport_only"] = mod_hostport_only(scan)
    scan.modules["mod_domain_only"] = mod_domain_only(scan)
    await scan._set_status("RUNNING")

    url_1 = scan.make_event("http://evilcorp.com/1", event_type="URL", parent=scan.root_event, tags=["status-200"])
    url_2 = scan.make_event("http://evilcorp.com/2", event_type="URL", parent=scan.root_event, tags=["status-200"])
    url_3 = scan.make_event("http://evilcorp.com:888/3", event_type="URL", parent=scan.root_event, tags=["status-200"])
    url_4 = scan.make_event("http://www.evilcorp.com/", event_type="URL", parent=scan.root_event, tags=["status-200"])
    url_5 = scan.make_event("http://www.evilcorp.net/", event_type="URL", parent=scan.root_event, tags=["status-200"])

    url_1.scope_distance = 0
    url_2.scope_distance = 0
    url_3.scope_distance = 0
    url_4.scope_distance = 0
    url_5.scope_distance = 0

    for mod_name in ("mod_normal", "mod_host_only", "mod_hostport_only", "mod_domain_only"):
        module = scan.modules[mod_name]

        valid_1, reason_1 = await module._event_postcheck(url_1)
        valid_2, reason_2 = await module._event_postcheck(url_2)
        valid_3, reason_3 = await module._event_postcheck(url_3)
        valid_4, reason_4 = await module._event_postcheck(url_4)
        valid_5, reason_5 = await module._event_postcheck(url_5)

        if mod_name == "mod_normal":
            assert valid_1 is True
            assert valid_2 is True
            assert valid_3 is True
            assert valid_4 is True
            assert valid_5 is True
        elif mod_name == "mod_host_only":
            assert valid_1 is True
            assert valid_2 is False
            assert "per_host_only=True" in reason_2
            assert valid_3 is False
            assert "per_host_only=True" in reason_3
            assert valid_4 is True
            assert valid_5 is True
        elif mod_name == "mod_hostport_only":
            assert valid_1 is True
            assert valid_2 is False
            assert "per_hostport_only=True" in reason_2
            assert valid_3 is True
            assert valid_4 is True
            assert valid_5 is True
        elif mod_name == "mod_domain_only":
            assert valid_1 is True
            assert valid_2 is False
            assert "per_domain_only=True" in reason_2
            assert valid_3 is False
            assert "per_domain_only=True" in reason_3
            assert valid_4 is False
            assert "per_domain_only=True" in reason_4
            assert valid_5 is True

    await scan._cleanup()


@pytest.mark.asyncio
async def test_modules_basic_perdomainonly(bbot_scanner, monkeypatch):
    per_domain_scan = bbot_scanner(
        "evilcorp.com",
        modules=list(available_modules),
        config={i: True for i in available_internal_modules if i != "dnsresolve"},
        force_start=True,
    )

    await per_domain_scan._prep()
    await per_domain_scan.setup_modules()
    await per_domain_scan._set_status("RUNNING")

    # ensure that multiple events to the same "host" (schema + host) are blocked and check the per host tracker

    for module_name, module in sorted(per_domain_scan.modules.items()):
        monkeypatch.setattr(module, "filter_event", BaseModule(per_domain_scan).filter_event)

        if "URL" in module.watched_events:
            url_1 = per_domain_scan.make_event(
                "http://www.evilcorp.com/1", event_type="URL", parent=per_domain_scan.root_event, tags=["status-200"]
            )
            url_1.scope_distance = 0
            url_2 = per_domain_scan.make_event(
                "http://mail.evilcorp.com/2", event_type="URL", parent=per_domain_scan.root_event, tags=["status-200"]
            )
            url_2.scope_distance = 0
            valid_1, reason_1 = await module._event_postcheck(url_1)
            valid_2, reason_2 = await module._event_postcheck(url_2)

            if module.per_domain_only is True:
                assert valid_1 is True
                assert valid_2 is False
                assert hash("evilcorp.com") in module._per_host_tracker
                assert reason_2 == "per_domain_only enabled and already seen domain"

            else:
                assert valid_1 is True
                assert valid_2 is True

    await per_domain_scan._cleanup()


@pytest.mark.asyncio
async def test_modules_basic_setup_deps(bbot_scanner):
    from bbot.modules.base import BaseModule

    class dummy(BaseModule):
        _name = "dummy"
        deps_ran = False
        setup_ran = False

        async def setup_deps(self):
            self.deps_ran = True
            return True

        async def setup(self):
            self.setup_ran = True
            return True

    scan = bbot_scanner()
    scan.modules["dummy"] = dummy(scan)
    await scan.setup_modules(deps_only=True)
    assert scan.modules["dummy"].deps_ran
    assert not scan.modules["dummy"].setup_ran
    await scan._cleanup()


@pytest.mark.asyncio
async def test_modules_basic_stats(helpers, events, bbot_scanner, blasthttp_mock, monkeypatch):
    from bbot.modules.base import BaseModule

    class dummy(BaseModule):
        _name = "dummy"
        watched_events = ["*"]

        async def handle_event(self, event):
            # quick emit events like FINDINGS behave differently than normal ones
            # hosts are not speculated from them
            await self.emit_event(
                {
                    "host": "www.evilcorp.com",
                    "url": "http://www.evilcorp.com",
                    "description": "asdf",
                    "name": "Finding",
                    "severity": "LOW",
                    "confidence": "MEDIUM",
                },
                "FINDING",
                event,
            )
            await self.emit_event("https://asdf.evilcorp.com", "URL", event, tags=["status-200"])

    scan = bbot_scanner(
        "evilcorp.com",
        config={"speculate": True, "dns": {"minimal": False}},
        exclude_output_modules=["csv", "json", "txt"],
        force_start=True,
    )

    await scan._prep()

    await scan.helpers.dns._mock_dns(
        {
            "evilcorp.com": {"A": ["127.0.254.1"]},
            "www.evilcorp.com": {"A": ["127.0.254.2"]},
            "asdf.evilcorp.com": {"A": ["127.0.254.3"]},
        },
    )

    scan.modules["dummy"] = dummy(scan)
    events = [e async for e in scan.async_start()]

    assert len(events) == 11
    assert 2 == len([e for e in events if e.type == "SCAN"])
    assert 4 == len([e for e in events if e.type == "DNS_NAME"])
    # one from target and one from speculate
    assert 2 == len([e for e in events if e.type == "DNS_NAME" and e.data == "evilcorp.com"])
    assert 1 == len([e for e in events if e.type == "DNS_NAME" and e.data == "www.evilcorp.com"])
    assert 1 == len([e for e in events if e.type == "DNS_NAME" and e.data == "asdf.evilcorp.com"])
    assert 1 == len([e for e in events if e.type == "OPEN_TCP_PORT" and e.data == "asdf.evilcorp.com:443"])
    assert 1 == len([e for e in events if e.type == "ORG_STUB" and e.data == "evilcorp"])
    assert 1 == len([e for e in events if e.type == "FINDING"])
    assert 1 == len([e for e in events if e.type == "URL_UNVERIFIED"])

    assert scan.stats.events_emitted_by_type == {
        "SCAN": 1,
        "DNS_NAME": 4,
        "URL": 1,
        "OPEN_TCP_PORT": 1,
        "ORG_STUB": 1,
        "URL_UNVERIFIED": 1,
        "FINDING": 1,
    }

    assert set(scan.stats.module_stats) == {"speculate", "host", "SEED", "python", "dummy", "dnsresolve"}

    target_stats = scan.stats.module_stats["SEED"]
    assert target_stats.produced == {"SCAN": 1, "DNS_NAME": 1}
    assert target_stats.produced_total == 2
    assert target_stats.consumed == {}
    assert target_stats.consumed_total == 0

    dummy_stats = scan.stats.module_stats["dummy"]
    assert dummy_stats.produced == {"FINDING": 1, "URL": 1}
    assert dummy_stats.produced_total == 2
    assert dummy_stats.consumed == {
        "DNS_NAME": 3,
        "FINDING": 1,
        "OPEN_TCP_PORT": 1,
        "ORG_STUB": 1,
        "SCAN": 1,
        "URL": 1,
        "URL_UNVERIFIED": 1,
    }
    assert dummy_stats.consumed_total == 9

    python_stats = scan.stats.module_stats["python"]
    assert python_stats.produced == {}
    assert python_stats.produced_total == 0
    assert python_stats.consumed == {
        "DNS_NAME": 4,
        "FINDING": 1,
        "OPEN_TCP_PORT": 1,
        "ORG_STUB": 1,
        "SCAN": 1,
        "URL": 1,
        "URL_UNVERIFIED": 1,
    }
    assert python_stats.consumed_total == 10

    speculate_stats = scan.stats.module_stats["speculate"]
    assert speculate_stats.produced == {"DNS_NAME": 1, "URL_UNVERIFIED": 1, "ORG_STUB": 1, "OPEN_TCP_PORT": 1}
    assert speculate_stats.produced_total == 4
    assert speculate_stats.consumed == {"URL": 1, "DNS_NAME": 3, "URL_UNVERIFIED": 1, "IP_ADDRESS": 3}
    assert speculate_stats.consumed_total == 8

    await scan._cleanup()


@pytest.mark.asyncio
async def test_module_loading(bbot_scanner):
    scan2 = bbot_scanner(
        modules=list(available_modules),
        output_modules=list(available_output_modules),
        config={i: True for i in available_internal_modules if i != "dnsresolve"},
        force_start=True,
    )
    await scan2._prep()
    await scan2._set_status("RUNNING")

    # attributes, descriptions, etc.
    for module_name, module in sorted(scan2.modules.items()):
        # flags
        assert module._type in ("internal", "output", "scan")
        # async stuff
        not_async = []
        for func_name in ("setup", "ping", "filter_event", "handle_event", "finish", "report", "cleanup"):
            f = getattr(module, func_name)
            if not scan2.helpers.is_async_function(f):
                log.error(f"{f.__qualname__}() is not async")
                not_async.append(f)
    assert not any(not_async)

    await scan2._cleanup()


def test_legacy_options_dict_rejected(tmp_path):
    """A pre-3.0 user module that declares an `options` / `options_desc` dict
    instead of a `class Config(BaseModuleConfig)` schema must be rejected at
    preload time with a clear migration message. Otherwise the loader would
    silently ignore the defaults and the user would never know."""
    from bbot.errors import BBOTError
    from bbot.core.modules import MODULE_LOADER

    legacy_mod = tmp_path / "legacy_user_mod.py"
    legacy_mod.write_text(
        """
from bbot.modules.base import BaseModule


class legacy_user_mod(BaseModule):
    watched_events = ["DNS_NAME"]
    produced_events = ["FINDING"]
    flags = ["passive", "safe"]
    meta = {"description": "x", "created_date": "2024-01-01", "author": "@x"}

    options = {"api_key": "", "max_results": 100}
    options_desc = {"api_key": "API key", "max_results": "Max results"}

    async def handle_event(self, event):
        pass
"""
    )

    with pytest.raises(BBOTError) as excinfo:
        MODULE_LOADER.preload_module(legacy_mod)
    msg = str(excinfo.value)
    assert "legacy_user_mod" in msg
    assert "no longer supported" in msg
    assert "class Config(BaseModuleConfig)" in msg

    # And a module that declares only options_desc (still legacy) is also rejected.
    legacy_mod_desc_only = tmp_path / "legacy_desc_only.py"
    legacy_mod_desc_only.write_text(
        """
from bbot.modules.base import BaseModule


class legacy_desc_only(BaseModule):
    watched_events = ["DNS_NAME"]
    produced_events = ["FINDING"]
    flags = ["passive", "safe"]
    meta = {"description": "x", "created_date": "2024-01-01", "author": "@x"}

    options_desc = {"api_key": "API key"}

    async def handle_event(self, event):
        pass
"""
    )
    with pytest.raises(BBOTError):
        MODULE_LOADER.preload_module(legacy_mod_desc_only)

    # A module with a real `class Config` and NO legacy dict must load cleanly.
    new_mod = tmp_path / "new_style_mod.py"
    new_mod.write_text(
        """
from bbot.modules.base import BaseModule
from bbot.core.config.models import BaseModuleConfig, Field


class new_style_mod(BaseModule):
    watched_events = ["DNS_NAME"]
    produced_events = ["FINDING"]
    flags = ["passive", "safe"]
    meta = {"description": "x", "created_date": "2024-01-01", "author": "@x"}

    class Config(BaseModuleConfig):
        api_key: str = Field("", description="API key", sensitive=True)

    async def handle_event(self, event):
        pass
"""
    )
    preloaded = MODULE_LOADER.preload_module(new_mod)
    assert preloaded is not None
    assert preloaded["config"] == {"api_key": ""}

    # The annotated form `options: dict = {...}` (ast.AnnAssign) must also be rejected,
    # not silently slip the guard.
    annotated_mod = tmp_path / "annotated_legacy_mod.py"
    annotated_mod.write_text(
        """
from bbot.modules.base import BaseModule


class annotated_legacy_mod(BaseModule):
    watched_events = ["DNS_NAME"]
    produced_events = ["FINDING"]
    flags = ["passive", "safe"]
    meta = {"description": "x", "created_date": "2024-01-01", "author": "@x"}

    options: dict = {"api_key": ""}
    options_desc: dict = {"api_key": "API key"}

    async def handle_event(self, event):
        pass
"""
    )
    with pytest.raises(BBOTError) as excinfo:
        MODULE_LOADER.preload_module(annotated_mod)
    assert "no longer supported" in str(excinfo.value)
