import json

import MalwareBazaar
import pytest
from CommonServerPython import Common, DBotScoreReliability, FeedIndicatorType

BASE_URL = "https://test.com"


def util_load_json(path):
    """

    Args:
        path:

    Returns:
        json object read from the path given
    """
    with open(path, encoding="utf-8") as f:
        return json.loads(f.read())


def create_client(with_api_key: bool = False):
    client = MalwareBazaar.Client(server_url=BASE_URL, verify=False, proxy=False, headers={}, api_key=None)
    if with_api_key:
        client.api_key = "111"
    return client


def test_dbot_score():
    """

    Given:
        - response from MalwareBazaar on hash file

    When:
        - Running a file command

    Then:
        - Make sure a CommandResult containing DbotScore with a BAD score and a relationship is returned.
    """
    mock_response = util_load_json("test_data/response_data_dbot_score.json")
    result = MalwareBazaar.file_process(
        "094fd325049b8a9cf6d3e5ef2a6d4cc6a567d7d49c35f8bb8dd9e3c6acf3d78d", DBotScoreReliability.A, {}, mock_response
    )
    indicator: Common.File = result.indicator

    assert indicator.dbot_score.score == 3
    assert indicator.relationships[0]._name == "indicator-of"
    assert indicator.relationships[0]._entity_a == "094fd325049b8a9cf6d3e5ef2a6d4cc6a567d7d49c35f8bb8dd9e3c6acf3d78d"
    assert indicator.relationships[0]._entity_b_type == FeedIndicatorType.indicator_type_by_server_version("STIX Malware")


def test_file_command(requests_mock):
    """

    Given:
        - Request file reputation given hash array

    When:
        - Running a file reputation command

    Then:
        - Make sure a file reputation for each file is returned.
    """
    mock_response = util_load_json("test_data/scan_file_results.json")
    requests_mock.post(BASE_URL, json=mock_response)

    client = create_client()

    args = {"file": ["094fd325049b8a9cf6d3e5ef2a6d4cc6a567d7d49c35f8bb8dd9e3c6acf3d78d"]}
    response = MalwareBazaar.file_command(client, args)

    assert response[0].outputs == mock_response.get("data")[0]
    assert response[0].outputs_prefix == "MalwareBazaar.File"
    assert response[0].outputs_key_field == "md5_hash"
    assert response[0].relationships is not None


def test_file_command_no_file_name(requests_mock):
    """

    Given:
        - Request file reputation, given hash array with file hash of a file without name

    When:
        - Running a file reputation command

    Then:
        - Make sure file reputation without file_name is returned (an empty string in outputs and removed from the
        human-readable).
    """
    mock_response = util_load_json("test_data/file_without_name.json")
    requests_mock.post(BASE_URL, json=mock_response)

    client = create_client()

    args = {"file": ["620c496e18e3256af0712541f18f19ed0105b264ce9e1fe40698066480bd7397"]}
    response = MalwareBazaar.file_command(client, args)

    assert response[0].outputs.get("file_name") == ""
    assert response[0].raw_response.get("file_name") is None


def test_does_not_raise_on_yara_rules_none():
    """
    Given:
        - file with empty 'yara_rules' response.

    When:
        - Running file_process method.

    Then:
        - Ensure that an exception is not raised.
    """
    raised = False
    try:
        MalwareBazaar.file_process(
            "094fd325049b8a9cf6d3e5ef2a6d4cc6a567d7d49c35f8bb8dd9e3c6acf3d78d", DBotScoreReliability.A, {}, {"yara_rules": None}
        )
    except TypeError:
        raised = True
    assert not raised


LIST_ARGS = [
    {"sample_type": "tag", "sample_value": "TrickBot", "limit": 2},
    {"sample_type": "tag", "sample_value": "TrickBot", "page": 0, "page_size": 2, "limit": 50},
]


@pytest.mark.parametrize("args", LIST_ARGS)
def test_malwarebazaar_sample_list(args, requests_mock):
    """

    Given:
        - sample_type to search by and sample_value to search for, limit and pagination is
        optional

    When:
        - Running a list command to retreive a list of malware samples

    Then:
        - Make sure a list of malware samples is returned.
    """
    mock_response = util_load_json("test_data/samples_list_results.json")
    requests_mock.post(BASE_URL, json=mock_response)

    client = create_client()

    response = MalwareBazaar.malwarebazaar_samples_list_command(client, args)

    assert response.outputs == mock_response.get("data")
    assert response.outputs_prefix == "MalwareBazaar.MalwarebazaarSamplesList"
    assert response.outputs_key_field == "sha256_hash"


def test_malwarebazaar_comment_add(requests_mock):
    """

    Given:
        - Hash of file and comment to add to MalwareBazaar db about this file

    When:
        - Running a comment add command

    Then:
        - Make sure a success message is returned.
    """
    mock_response = {"query_status": "success"}

    readable_output = (
        "Comment added to 094fd325049b8a9cf6d3e5ef2a6d4cc6a567d7d49c35f8bb8dd9e3c6acf3d78d malware sample successfully"
    )

    outputs = {"sha256_hash": "094fd325049b8a9cf6d3e5ef2a6d4cc6a567d7d49c35f8bb8dd9e3c6acf3d78d", "comment": "test"}

    requests_mock.post(BASE_URL, json=mock_response)

    client = create_client(True)

    args = {"sha256_hash": "094fd325049b8a9cf6d3e5ef2a6d4cc6a567d7d49c35f8bb8dd9e3c6acf3d78d", "comment": "test"}

    response = MalwareBazaar.malwarebazaar_comment_add_command(client, args)

    assert response.readable_output == readable_output
    assert response.outputs_key_field == "sha256_hash"
    assert response.outputs == outputs
    assert response.outputs_prefix == "MalwareBazaar.MalwarebazaarCommentAdd"


def test_malwarebazaar_download_sample(mocker):
    """

    Given:
        - Hash encoding of malware sample to download from MalwareBazaar

    When:
        - Running a download sample command

    Then:
        - Return a file that to download
    """

    class File:
        def __init__(self):
            self.content = b"content"

    file_content = File()
    mocker.patch.object(MalwareBazaar.Client, "malwarebazaar_download_sample_request", return_value=file_content)
    client = create_client()

    args = {"sha256_hash": "094fd325049b8a9cf6d3e5ef2a6d4cc6a567d7d49c35f8bb8dd9e3c6acf3d78d"}
    response = MalwareBazaar.malwarebazaar_download_sample_command(client, args)

    assert response.get("File") == "094fd325049b8a9cf6d3e5ef2a6d4cc6a567d7d49c35f8bb8dd9e3c6acf3d78d.zip"


COMMAND_PACKAGE = [
    (
        {"sha256_hash": "094fd325049b8a9cf6d3e5ef2a6d4cc6a567d7d49c35f8bb8dd9e3c6acf3d78d", "comment": "test"},
        MalwareBazaar.malwarebazaar_comment_add_command,
        True,
    ),
    (
        {"sha256_hash": "094fd325049b8a9cf6d3e5ef2a6d4cc6a567d7d49c35f8bb8dd9e3c6acf3d78d", "comment": "test"},
        MalwareBazaar.malwarebazaar_comment_add_command,
        False,
    ),
    ({"file": ["094fd325049b8a9cf6d3e5ef2a6d4cc6a567d7d49c35f8bb8dd9e3c6acf3d78d"]}, MalwareBazaar.file_command, False),
    ({"sample_type": "tag", "sample_value": "TrickBot", "limit": 2}, MalwareBazaar.malwarebazaar_samples_list_command, False),
]


@pytest.mark.parametrize("args, method_to_run, with_api_key", COMMAND_PACKAGE)
def test_malwarebazaare_exception_raised(requests_mock, args, method_to_run, with_api_key):
    """

    Given:
        - Request with wrong arguments

    When:
        - Running a command

    Then:
        - Make sure an error is raised.
    """
    mock_response = {"query_status": "tag_not_found"}
    requests_mock.post(BASE_URL, json=mock_response)

    client = create_client(with_api_key)

    with pytest.raises(Exception) as e:
        method_to_run(client, args)
    if not e:
        pytest.fail()


def test_hash_not_found(requests_mock):
    """

    Given:
        - Request with file that do not exist in MalwareBazaar

    When:
        - Running a file reputation command

    Then:
        - Make sure an error is raised.
    """
    mock_response = {"query_status": "hash_not_found"}
    requests_mock.post(BASE_URL, json=mock_response)
    client = create_client(False)
    args = {"file": ["094fd325049b8a9cf6d3e5ef2a6d4cc6a567d7d49c35f8bb8dd9e3c6acf3d78d"]}

    result = MalwareBazaar.file_command(client, args)

    assert result[0].indicator.dbot_score.score == Common.DBotScore.NONE