import pytest

from infection_monkey.model.host import VictimHost

DOMAIN_NAME = "domain-name"
IP = "0.0.0.0"
NETBIOS_NAME = "NetBIOS Name"

USERS = ["Administrator", "Bob"]
RIDS = ["500", "1024"]
LM_HASHES = ["abc123", "098zyx"]
NT_HASHES = ["def456", "765vut"]


@pytest.fixture
def zerologon_exploiter_object(monkeypatch):
    from infection_monkey.exploit.zerologon import ZerologonExploiter

    def mock_report_login_attempt(**kwargs):
        return None

    host = VictimHost(IP, DOMAIN_NAME)
    obj = ZerologonExploiter(host)
    monkeypatch.setattr(obj, "dc_name", NETBIOS_NAME, raising=False)
    monkeypatch.setattr(obj, "report_login_attempt", mock_report_login_attempt)
    return obj


@pytest.mark.slow
def test_assess_exploit_attempt_result_no_error(zerologon_exploiter_object):
    dummy_exploit_attempt_result = {"ErrorCode": 0}
    assert zerologon_exploiter_object.assess_exploit_attempt_result(dummy_exploit_attempt_result)


@pytest.mark.slow
def test_assess_exploit_attempt_result_with_error(zerologon_exploiter_object):
    dummy_exploit_attempt_result = {"ErrorCode": 1}
    assert not zerologon_exploiter_object.assess_exploit_attempt_result(
        dummy_exploit_attempt_result
    )


@pytest.mark.slow
def test_assess_restoration_attempt_result_restored(zerologon_exploiter_object):
    dummy_restoration_attempt_result = object()
    assert zerologon_exploiter_object.assess_restoration_attempt_result(
        dummy_restoration_attempt_result
    )


@pytest.mark.slow
def test_assess_restoration_attempt_result_not_restored(zerologon_exploiter_object):
    dummy_restoration_attempt_result = False
    assert not zerologon_exploiter_object.assess_restoration_attempt_result(
        dummy_restoration_attempt_result
    )


@pytest.mark.slow
def test__extract_user_creds_from_secrets_good_data(zerologon_exploiter_object):
    mock_dumped_secrets = [
        f"{USERS[i]}:{RIDS[i]}:{LM_HASHES[i]}:{NT_HASHES[i]}:::" for i in range(len(USERS))
    ]
    expected_extracted_creds = {
        USERS[0]: {
            "RID": int(RIDS[0]),
            "lm_hash": LM_HASHES[0],
            "nt_hash": NT_HASHES[0],
        },
        USERS[1]: {
            "RID": int(RIDS[1]),
            "lm_hash": LM_HASHES[1],
            "nt_hash": NT_HASHES[1],
        },
    }
    assert zerologon_exploiter_object._extract_user_creds_from_secrets(mock_dumped_secrets) is None
    assert zerologon_exploiter_object._extracted_creds == expected_extracted_creds


@pytest.mark.slow
def test__extract_user_creds_from_secrets_bad_data(zerologon_exploiter_object):
    mock_dumped_secrets = [
        f"{USERS[i]}:{RIDS[i]}:::{LM_HASHES[i]}:{NT_HASHES[i]}:::" for i in range(len(USERS))
    ]
    expected_extracted_creds = {
        USERS[0]: {"RID": int(RIDS[0]), "lm_hash": "", "nt_hash": ""},
        USERS[1]: {"RID": int(RIDS[1]), "lm_hash": "", "nt_hash": ""},
    }
    assert zerologon_exploiter_object._extract_user_creds_from_secrets(mock_dumped_secrets) is None
    assert zerologon_exploiter_object._extracted_creds == expected_extracted_creds
