import pytest


def test_strip():
    test_string = "T'hi's >is a t[]es[t t<]ext."
    from PcapMinerV2 import strip

    assert strip(test_string) == "Thisisatesttext."


def test_heirarchy_to_md():
    test_heirarchy = {"First Layer": 50, "Second Layer": 50, "Third Layer": 50, "Fourth Layer": 50}
    from PcapMinerV2 import hierarchy_to_md

    assert (
        hierarchy_to_md(test_heirarchy) == "|Layer| # of Packets|% of Packets|\n|---|----|---|\n"
        "| -> First Layer|50|25.0%|\n| -> Fourth Layer|50|25.0%|\n"
        "| -> Second Layer|50|25.0%|\n| -> Third Layer|50|25.0%|\n"
    )


def test_conversations_to_md():
    converstation_test = {("1.1.1.1", "2.2.2.2"): 15, ("1.1.1.3", "2.2.2.2"): 15, ("1.1.1.3", "8.8.8.8"): 15}
    from PcapMinerV2 import conversations_to_md

    assert (
        conversations_to_md(converstation_test, 4) == "|A|B|# of Packets|\n|---|---|---|\n|1.1.1.1|2.2.2.2|15|\n"
        "|1.1.1.3|2.2.2.2|15|\n|1.1.1.3|8.8.8.8|15|\n"
    )


args_to_test = [
    ({}, {"ID": 1}, None, {}),  # Test that just an ID is not added.
    ({}, {"ID": 2, "Input": "wow"}, None, {2: {"ID": 2, "Input": "wow"}}),  # Test data added correctly
    ({}, {"ID": 2, "Input": "wow"}, 10, {10: {"ID": 2, "Input": "wow"}}),  # Test that a future id is added instead
    ({}, {"ID": 2}, 10, {}),  # Tests that future ID doesn't add empty data.
    (
        {2: {"ID": 2, "Input": "wow"}},
        {"ID": 2, "Input": "amazing", "new_key": "new"},
        None,
        {2: {"ID": 2, "Input": "amazing", "new_key": "new"}},
    ),  # Test that ID changed and new added
    (
        {2: {"ID": 2, "Input": "wow"}},
        {"ID": 2, "Input": "amazing", "new_key": "new"},
        10,
        {10: {"ID": 2, "Input": "amazing", "new_key": "new"}},
    ),  # Test that ID changed and new added and future_id added
    ({}, {"noID": 15}, None, {}),  # Test that data without ID isn't added
    ({}, {"ID": 1, "EntryID": 15}, None, {}),  # Test that just an ID and EntryID is not added.
]


@pytest.mark.parametrize("main_data, data_to_add, future_id, wanted_output", args_to_test)
def test_add_to_data(main_data, data_to_add, future_id, wanted_output):
    from PcapMinerV2 import add_to_data

    add_to_data(main_data, data_to_add, future_id)

    assert main_data == wanted_output


def test_mine_pcap():
    file_path = "./TestData/smb-on-windows-10.pcapng"
    wpa_password = ""
    conversation_number_to_display = 15
    is_flows = True
    is_reg_extract = True
    extracted_protocols = ["DNS", "SMB2"]
    pcap_filter = ""
    homemade_regex = "M-SEARCH * (.+)"
    pcap_filter_new_file_path = ""
    unique_ips = False
    rsa_key_file_path = ""
    from PcapMinerV2 import PCAP

    pcap = PCAP(is_reg_extract, extracted_protocols, homemade_regex, unique_ips, "entry_id")
    pcap.mine(file_path, wpa_password, rsa_key_file_path, is_flows, is_reg_extract, pcap_filter, pcap_filter_new_file_path)
    hr, ec, raw = pcap.get_outputs(conversation_number_to_display, is_flows, is_reg_extract)
    assert raw["EntryID"] == "entry_id"
    assert raw["Packets"] == 1000
    assert len(ec["PCAPResultsDNS"]) == 80
    assert len(ec["PCAPResultsSMB2"]) == 7
    assert raw["URL"][0] == "http://239.255.255.250:1900*"
    assert raw["Regex"] != []