import os
import sys
import tempfile
import unittest

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.dirname(HERE))

from graylog_ops.asn import AsnDB, build_asn_db, classify_org  # noqa: E402


class AsnTests(unittest.TestCase):
    def test_build_and_lookup(self):
        with tempfile.TemporaryDirectory() as directory:
            csv_path = os.path.join(directory, "asn.csv")
            with open(csv_path, "w", encoding="utf-8") as handle:
                handle.write('1.0.1.0,1.0.3.255,4134,"Chinanet"\n')
                handle.write('8.8.8.0,8.8.8.255,15169,"Google LLC"\n')
                handle.write('47.0.0.0,47.0.0.255,37963,'
                             '"Hangzhou Alibaba Advertising Co.,Ltd."\n')
                handle.write('2001:db8::,2001:db8::ffff,15169,"Google LLC"\n')
            db_path = os.path.join(directory, "asn.db")
            count = build_asn_db(csv_path, db_path)
            self.assertEqual(count, 3)  # IPv6 skipped

            asn = AsnDB(db_path)
            try:
                china = asn.lookup("1.0.2.1")
                self.assertEqual(china["brand"], "中国电信")
                self.assertEqual(china["type"], "broadband")
                aliyun = asn.lookup("47.0.0.1")
                self.assertEqual(aliyun["brand"], "阿里云")
                self.assertEqual(aliyun["type"], "cloud")
                self.assertEqual(aliyun["icon"], "☁️")
                self.assertIsNone(asn.lookup("9.9.9.9"))
            finally:
                asn.close()

    def test_classify_org(self):
        self.assertEqual(classify_org("China Mobile Communications Corporation")[:2],
                         ("中国移动", "mobile"))
        self.assertEqual(classify_org("China Education and Research Network Center")[:2],
                         ("教育网", "education"))
        self.assertEqual(classify_org("Wowrack.com")[:2], ("IDC 机房", "idc"))
        self.assertEqual(classify_org("Some Unknown ISP")[:2], ("未知", "unknown"))


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