"""VirusTotal Resolver Tests."""
import datetime
import json
import os
import time
import unittest
from unittest.mock import MagicMock, patch

from libchickadee.resolvers.virustotal import ProResolver

__author__ = "Chapin Bryce"
__date__ = 20200805
__license__ = "MIT Copyright 2020 Chapin Bryce"
__desc__ = """Yet another GeoIP resolution tool."""


class IPAPITestCase(unittest.TestCase):
    """VirusTotal Resolver Tests."""

    def setUp(self):
        """Test config"""
        self.test_data_ips = ["10.0.1.2", "8.8.8.8", "2001:4860:4860::8888"]
        self.expected_result = [
            {"query": "10.0.1.2"},
            {
                "as": "AS15169 Google LLC",
                "country": "United States",
                "org": "Level 3",
                "proxy": False,
                "query": "8.8.8.8",
            },
            {
                "as": "AS15169 Google LLC",
                "country": "United States",
                "org": "Google LLC",
                "proxy": False,
                "query": "2001:4860:4860::8888",
            },
        ]
        self.resolver = ProResolver(api_key="")
        local_dir = os.path.abspath(__file__).rsplit(os.sep, 1)[0]
        resource_file = open(os.path.join(local_dir, "vt_resp_data.json"))
        self.vt_rep_data_list = json.load(resource_file)
        resource_file.close()

    def test_parse_vt_resp(self):
        """Confirm the parsing of VT responses."""
        vt_resp_data = self.vt_rep_data_list["Test0"]["test"]
        query = self.vt_rep_data_list["Test0"]["query"]
        expected = self.vt_rep_data_list["Test0"]["expected"]
        vt_values = self.resolver.parse_vt_resp(query, vt_resp_data)
        self.assertDictEqual(expected, vt_values)

    def test_parse_vt_resp_2(self):
        """Confirm the parsing of VT responses."""
        vt_resp_data = self.vt_rep_data_list["Test1"]["test"]
        query = self.vt_rep_data_list["Test1"]["query"]
        expected = self.vt_rep_data_list["Test1"]["expected"]
        vt_values = self.resolver.parse_vt_resp(query, vt_resp_data)
        self.assertDictEqual(expected, vt_values)

    @patch("libchickadee.resolvers.virustotal.requests.get")
    def test_resolve_single(self, mock_requests):
        """Test the resolution processing of a single item."""

        def mock_json():
            """Return a mocked JSON value from the request"""
            return self.vt_rep_data_list["Test0"]["test"]

        mock_requests.return_value.status_code = 200
        mock_requests.return_value.json = mock_json
        actual = self.resolver.query(self.vt_rep_data_list["Test0"]["query"])
        expected = self.vt_rep_data_list["Test0"]["expected"]
        self.assertEqual(1, len(actual))
        self.assertDictEqual(expected, actual[0])

    @patch("libchickadee.resolvers.virustotal.requests.get")
    def test_resolve_batch(self, mock_requests):
        """Test the resolution processing of a multiple items."""
        # Build 3 requests
        req1 = MagicMock()
        req1.status_code = 200
        req1.json = lambda: self.vt_rep_data_list["Test0"]["test"]

        req2 = MagicMock()
        req2.status_code = 204

        req3 = MagicMock()
        req3.status_code = 200
        req3.json = lambda: self.vt_rep_data_list["Test1"]["test"]

        mock_requests.side_effect = [req1, req2, req3]
        start = datetime.datetime.now()
        actual = self.resolver.query(
            [
                self.vt_rep_data_list["Test0"]["query"],
                self.vt_rep_data_list["Test1"]["query"],
            ]
        )

        self.assertEqual(2, len(actual))
        self.assertGreaterEqual(15, (start - datetime.datetime.now()).total_seconds())
        self.assertDictEqual(self.vt_rep_data_list["Test0"]["expected"], actual[0])
        self.assertDictEqual(self.vt_rep_data_list["Test1"]["expected"], actual[1])

    @patch("libchickadee.resolvers.virustotal.requests.get")
    def test_resolve_errors(self, mock_requests):
        """Testing error handling of the resolver."""
        subtests = {
            400: "Incorrect request. Please check input data",
            403: "Authorization error. Please check API key",
            500: "Unknown error occurred, status code 500, please report",
        }
        for status_code, err_msg in subtests.items():
            mock_requests.return_value.status_code = status_code
            with self.subTest(id=status_code):
                with self.assertLogs(
                    "libchickadee.resolvers.virustotal", level="ERROR"
                ) as mock_log:
                    actual = self.resolver.query(
                        self.vt_rep_data_list["Test0"]["query"]
                    )
                self.assertIsNone(actual)
                self.assertEqual(mock_log.records[0].message, err_msg)

    @patch("libchickadee.resolvers.virustotal.requests.get")
    def test_sleeper(self, mock_requests):
        """Validate that the sleep timer works appropriately for rate limiting"""
        initial_time = datetime.datetime.now()
        self.resolver.last_request = initial_time
        time.sleep(2)
        mock_requests.return_value.status_code = 403

        self.resolver.query(data="1.1.1.1")
        self.assertGreaterEqual(
            self.resolver.last_request, initial_time + datetime.timedelta(seconds=2)
        )


if __name__ == "__main__":
    unittest.main()
