# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for
# license information.
# --------------------------------------------------------------------------
"""vtlookup test class."""
import unittest
import os
from os import path

import pandas as pd
from ..msticpy.sectools.vtlookup import VTLookup
from ..msticpy.sectools.tiproviders.ti_provider_base import preprocess_observable

_test_data_folders = [
    d for d, _, _ in os.walk(os.getcwd()) if d.endswith("/tests/testdata")
]
if len(_test_data_folders) == 1:
    _TEST_DATA = _test_data_folders[0]
else:
    _TEST_DATA = "./tests/testdata"


class TestVTLookup(unittest.TestCase):

    _TEST_COLS = [
        "Observable",
        "IoCType",
        "Status",
        "ResponseCode",
        "Resource",
        "SourceIndex",
        "VerboseMsg",
        "Resource",
        "ScanId",
        "Permalink",
        "Positives",
        "MD5",
        "SHA1",
        "SHA256",
        "ResolvedDomains",
        "ResolvedIPs",
        "DetectedUrls",
    ]

    def test_ipvalidation(self):

        vtlookup = VTLookup(vtkey="fake", verbosity=2)
        test_ips = [
            ("valid", "90.156.201.27", "90.156.201.27"),
            ("local_ip", "10.0.0.1", None),
            ("mask", "255.255.20.27", None),
            ("loopback", "127.0.0.1", None),
            ("bad_format", "122.1.3", None),
        ]

        for test_case in test_ips:
            result, status = preprocess_observable(test_case[1], "ipv4")
            self.assertEqual(result, test_case[2])
            print(test_case[0], status)

    def test_urlvalidation(self):
        vtlookup = VTLookup(vtkey="fake", verbosity=2)
        v1_url = "http://club-fox.ru/img/www.loginalibaba.com/alibaba/alibaba/login.alibaba.com.php?email=biuro"
        v2_url = "https://www.virustotal.com/en/ip-address/90.156.201.27/information/"
        test_urls = [
            ("valid1", v1_url, v1_url),
            ("valid2", "https://microsoft.com", "https://microsoft.com"),
            ("valid3", "https://python.org", "https://python.org"),
            ("valid3", v2_url, v2_url),
            ("local_ip", "http://10.0.0.1/foo", None),
            ("local_hostname", "https://myhost/path", None),
            ("invalid1", "http//club-fox.ru/foo.html", None),
            ("invalid2", "//club-fox.ru/foo.html", None),
            ("invalid3", "https://123:microsoft.com@user/foo.txt", None),
            ("invalid4", "http//10.0.0.1/foo.txt", None),
        ]

        for test_case in test_urls:
            result, status = preprocess_observable(test_case[1], "url")
            print(test_case[0], status)
            self.assertEqual(
                result,
                test_case[2],
                f"Failed on test case {test_case} ({test_case[1]})",
            )

    def test_parse_file_results(self):
        vtlookup = VTLookup(vtkey="fake", verbosity=2)

        FILE_NAME = path.join(_TEST_DATA, "fileresponse.json")
        with open(FILE_NAME, "r") as file_handle:
            txt = file_handle.read()

        vt_params = vtlookup._VT_API_TYPES["file"]
        vtlookup._parse_vt_results(
            vt_results=txt,
            observable="7657fcb7d772448a6d8504e4b20168b8",
            ioc_type="md5_hash",
            vt_param=vt_params,
        )

        test_df = vtlookup.results[self._TEST_COLS]
        self.assertEqual(len(test_df), 1)
        print(test_df.T)

        vtlookup = VTLookup(vtkey="fake", verbosity=2)
        FILE_NAME2 = path.join(_TEST_DATA, "file-multi_pos.json")
        with open(FILE_NAME2, "r") as file_handle:
            txt = file_handle.read()

        vt_params = vtlookup._VT_API_TYPES["file"]
        vtlookup._parse_vt_results(
            vt_results=txt,
            observable="7657fcb7d772448a6d8504e4b20168b8",
            ioc_type="md5_hash",
            vt_param=vt_params,
        )

        test_df = vtlookup.results[self._TEST_COLS]
        self.assertEqual(len(test_df), 3)
        print(test_df.T)

    def test_parse_url_results(self):
        vtlookup = VTLookup(vtkey="fake", verbosity=2)

        FILE_NAME = path.join(_TEST_DATA, "url_pos.json")
        with open(FILE_NAME, "r") as file_handle:
            txt = file_handle.read()

        vt_params = vtlookup._VT_API_TYPES["url"]
        vtlookup._parse_vt_results(
            vt_results=txt,
            observable="7657fcb7d772448a6d8504e4b20168b8",
            ioc_type="url",
            vt_param=vt_params,
        )
        test_df = vtlookup.results[self._TEST_COLS]
        self.assertEqual(len(test_df), 1)
        print(test_df.T)

        vtlookup = VTLookup(vtkey="fake", verbosity=2)
        FILE_NAME2 = path.join(_TEST_DATA, "url_neg.json")
        with open(FILE_NAME2, "r") as file_handle:
            txt = file_handle.read()

        vt_params = vtlookup._VT_API_TYPES["url"]
        vtlookup._parse_vt_results(
            vt_results=txt,
            observable="7657fcb7d772448a6d8504e4b20168b8",
            ioc_type="url",
            vt_param=vt_params,
        )
        test_df = vtlookup.results[self._TEST_COLS]
        self.assertEqual(len(test_df), 1)
        print(test_df.T)

    def test_parse_domain_results(self):
        vtlookup = VTLookup(vtkey="fake", verbosity=2)

        FILE_NAME = path.join(_TEST_DATA, "domain_pos.json")
        with open(FILE_NAME, "r") as file_handle:
            txt = file_handle.read()

        vt_params = vtlookup._VT_API_TYPES["domain"]
        vtlookup._parse_vt_results(
            vt_results=txt,
            observable="7657fcb7d772448a6d8504e4b20168b8",
            ioc_type="dns",
            vt_param=vt_params,
        )

        test_df = vtlookup.results[self._TEST_COLS]
        self.assertEqual(len(test_df), 1)
        self.assertGreater(len(test_df[["ResolvedIPs"]]), 0)
        self.assertGreater(len(test_df[["DetectedUrls"]].values), 0)
        self.assertGreater(test_df[["Positives"]].values, 0)

        print(test_df.T)

        vtlookup = VTLookup(vtkey="fake", verbosity=2)
        FILE_NAME2 = path.join(_TEST_DATA, "domain_neg.json")
        with open(FILE_NAME2, "r") as file_handle:
            txt = file_handle.read()

        vt_params = vtlookup._VT_API_TYPES["domain"]
        vtlookup._parse_vt_results(
            vt_results=txt,
            observable="7657fcb7d772448a6d8504e4b20168b8",
            ioc_type="dns",
            vt_param=vt_params,
        )
        test_df = vtlookup.results[self._TEST_COLS]
        self.assertEqual(len(test_df), 1)
        self.assertGreater(len(test_df[["ResolvedIPs"]].values), 0)
        self.assertGreater(len(test_df[["DetectedUrls"]].values), 0)
        print(test_df.T)

    def test_parse_ip_results(self):
        vtlookup = VTLookup(vtkey="fake", verbosity=2)

        FILE_NAME = path.join(_TEST_DATA, "ip-address_pos.json")
        with open(FILE_NAME, "r") as file_handle:
            txt = file_handle.read()

        vt_params = vtlookup._VT_API_TYPES["ip-address"]
        vtlookup._parse_vt_results(
            vt_results=txt,
            observable="7657fcb7d772448a6d8504e4b20168b8",
            ioc_type="ipv4",
            vt_param=vt_params,
        )
        test_df = vtlookup.results[self._TEST_COLS]
        self.assertEqual(len(test_df), 1)
        self.assertGreater(len(test_df[["ResolvedDomains"]].values), 0)
        self.assertGreater(len(test_df[["DetectedUrls"]].values), 0)
        self.assertGreater(test_df[["Positives"]].values, 0)
        print(test_df.T)

        vtlookup = VTLookup(vtkey="fake", verbosity=2)
        FILE_NAME2 = path.join(_TEST_DATA, "ip-address_neg.json")
        with open(FILE_NAME2, "r") as file_handle:
            txt = file_handle.read()

        vt_params = vtlookup._VT_API_TYPES["ip-address"]
        vtlookup._parse_vt_results(
            vt_results=txt,
            observable="7657fcb7d772448a6d8504e4b20168b8",
            ioc_type="ipv4",
            vt_param=vt_params,
        )
        test_df = vtlookup.results[self._TEST_COLS]
        self.assertEqual(len(test_df), 1)
        self.assertGreater(len(test_df[["ResolvedDomains"]].values), 0)
        self.assertEqual(test_df[["Positives"]].values, 0)
        print(test_df.T)


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